执行摘要

• 一句话：迁移 srt/layers 散落内核至 sglang.kernels
• 推荐动作：作为 Phase 2.5 的重要一环，适合架构组和维护者精读，了解内核迁移的目录约定和注册模板。review 中指出的 missing import 和防御性断言可作为后续 PR 改进的参考。建议在后续 PR 中补充这些断言并修复遗漏的导入。

功能与动机

作为 RFC #29630 Phase 2.5 清除计划的第 3/7 步，目标是将 srt/layers 下剩余的内核统一搬迁至 sglang.kernels，解决内核散落各处、难以维护和集中注册的问题。PR body 指出此次移动的 11 个内核全部为 byte-identical 搬迁（git rename R100），只改写 import 路径。

实现拆解

1. 目录规划：在 sglang.kernels.ops 下创建或扩展子目录 layernorm、attention、memory、sampling，分别收容对应功能的内核。
2. 整体搬迁：将 elementwise.py、gemma4_fused_ops.py、mhc.py、mhc_head.py、deepseek_v4_rope.py、fused_qk_norm_rope_store.py、fused_qk_rmsnorm_rope_gate.py、fused_qk_norm.py、rotary_embedding/triton_kernels.py、multimodal.py、layers/utils/hash.py 这些文件完整移动到目标子目录（字节不变，git rename 检测为 100%）。
3. 内核提取：从 mrope.py 中剥离 apply_interleaved_rope_kernel 及其 Python 启动函数到 ops/attention/mrope.py；从 dp_attention.py 中剥离 memcpy_triton_kernel、prod、memcpy_triton 到 ops/memory/memcpy_triton.py。原模块保留 MRotaryEmbedding 等类，不再定义内核。
4. 导入重写：使用 sed 或类似工具在 python/sglang/、test/、benchmark/ 等目录的所有文件中将旧导入路径替换为新路径，确保编译时和运行时符号正确。
5. 注册与验证：在目标子包的 init.py 中通过 register_kernel 和 KernelSpec 注册迁移的内核；运行 test_kernels_namespace.py 和 test_fused_op.py（共 33 个测试）全部通过，现有正确性测试通过重写的 import 覆盖。

关键文件：

• python/sglang/srt/layers/rotary_embedding/mrope.py（模块 RoPE 层；类别 source；类型 core-logic；符号 apply_interleaved_rope_kernel, apply_interleaved_rope_triton）: 迁移的核心源文件之一：移除了 apply_interleaved_rope_kernel Triton 内核和 apply_interleaved_rope_triton 包装函数，改为从 sglang.kernels.ops.attention.mrope 导入；同时将 triton_mrope_fused 和 triton_ernie45_rope_fused_inplace 的导入源更新为 sglang.kernels.ops.attention.rotary_triton。
• python/sglang/srt/layers/dp_attention.py（模块 DP 注意力；类别 source；类型 core-logic；符号 memcpy_triton_kernel, prod, memcpy_triton）: 迁移核心源文件之二：移除了 memcpy_triton_kernel、prod 和 memcpy_triton 函数，改为从 sglang.kernels.ops.memory.memcpy_triton 导入；同时移除了 functools、triton 等不再需要的导入。
• python/sglang/kernels/ops/attention/mrope.py（模块 Attention 内核；类别 infra；类型 infrastructure；符号 apply_interleaved_rope_kernel, apply_interleaved_rope_triton）: 新增的目标文件：包含从 mrope.py 提取的 apply_interleaved_rope_kernel Triton 内核和 apply_interleaved_rope_triton 函数，是 M-RoPE 交互式旋转位置编码的核心实现。
• python/sglang/kernels/ops/memory/memcpy_triton.py（模块 Memory 内核；类别 infra；类型 infrastructure；符号 memcpy_triton_kernel, prod, memcpy_triton）: 新增的目标文件：包含从 dp_attention.py 提取的 memcpy_triton_kernel、prod 和 memcpy_triton 函数，用于 DP 数据并行中带有偏移量和大小的设备间内存拷贝。
• python/sglang/srt/models/deepseek_v4.py（模块 DeepSeek V4；类别 source；类型 data-contract）: 大量 import 重写：将 deepseek_v4_rope、dp_attention、mhc、mhc_head、fused_qk_norm_rope_store 等模块的导入路径从 srt/layers 更新为 sglang/kernels/ops 下的新位置，同时保持了符号名称不变。
• python/sglang/srt/models/grok.py（模块 Grok 模型；类别 source；类型 data-contract）: 将 elementwise 内核（fused_dual_residual_rmsnorm、fused_rmsnorm、gelu_and_mul_triton）的导入从 srt/layers/elementwise 改为 sglang/kernels/ops/layernorm/elementwise。

关键符号：apply_interleaved_rope_kernel, apply_interleaved_rope_triton, memcpy_triton_kernel, prod, memcpy_triton

关键源码片段

python/sglang/kernels/ops/attention/mrope.py

新增的目标文件：包含从 mrope.py 提取的 apply_interleaved_rope_kernel Triton 内核和 apply_interleaved_rope_triton 函数，是 M-RoPE 交互式旋转位置编码的核心实现。

# Interleaved M-RoPE Triton kernel, migrated from
# sglang.srt.layers.rotary_embedding.mrope (RFC #29630, Phase 2.5).
​
import torch
import triton
import triton.language as tl
​
​
@triton.jit
def apply_interleaved_rope_kernel(
    x_ptr,
    out_ptr,
    S: tl.constexpr,
    D: tl.constexpr,
    stride_x_m,
    stride_x_s,
    stride_out_s,
    section_1_end,
    section_2_end,
    BLOCK_S: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    # 计算当前 block 的序列位置偏移和维度偏移
    start_s = tl.program_id(0) * BLOCK_S
    s_offsets = start_s + tl.arange(0, BLOCK_S)
​
    dim_offset = tl.program_id(1) * BLOCK_SIZE
    dim_indices = dim_offset + tl.arange(0, BLOCK_SIZE)
​
    mask_s = s_offsets < S
    mask_d = dim_indices < D
    mask = mask_s[:, None] & mask_d[None, :]
​
    # 从 x_ptr 的 m=0 平面加载基础值（对应原始 x[0]）
    val_ptr = (
        x_ptr + 0 * stride_x_m + s_offsets[:, None] * stride_x_s + dim_indices[None, :]
    )
    val = tl.load(val_ptr, mask=mask, other=0.0)
​
    # 条件覆盖：当维度索引模 3 为 1 且小于 section_1_end * 3 时，
    # 从 x_ptr 的 m=1 平面读取覆盖值（对应 x[1]）
    cond_a = (dim_indices[None, :] % 3 == 1) & (
        dim_indices[None, :] < section_1_end * 3
    )
    val_a_ptr = (
        x_ptr + 1 * stride_x_m + s_offsets[:, None] * stride_x_s + dim_indices[None, :]
    )
    val_a = tl.load(val_a_ptr, mask=mask & cond_a, other=0.0)
​
    # 条件覆盖：当维度索引模 3 为 2 且小于 section_2_end * 3 时，
    # 从 x_ptr 的 m=2 平面读取覆盖值（对应 x[2]）
    cond_b = (dim_indices[None, :] % 3 == 2) & (
        dim_indices[None, :] < section_2_end * 3
    )
    val_b_ptr = (
        x_ptr + 2 * stride_x_m + s_offsets[:, None] * stride_x_s + dim_indices[None, :]
    )
    val_b = tl.load(val_b_ptr, mask=mask & cond_b, other=0.0)
​
    val = tl.where(cond_a, val_a, val)
    val = tl.where(cond_b, val_b, val)
​
    out_ptr = out_ptr + s_offsets[:, None] * stride_out_s + dim_indices[None, :]
    tl.store(out_ptr, val, mask=mask)
​
​
def apply_interleaved_rope_triton(x: torch.Tensor, mrope_section: list) -> torch.Tensor:
    # 对 3D 输入 x 应用 interleaved M-RoPE 变换。
    # x 的形状为 (M, S, D)，其中 M 必须为 3（分别对应 x[0], x[1], x[2]）。
    # mrope_section 是一个包含 3 个整数的列表，分别控制三个维度的动态 RoPE 边界。
    x = x.contiguous()
    M, S, D = x.shape
​
    out = torch.empty((S, D), dtype=x.dtype, device=x.device)
​
    BLOCK_S = 64
    BLOCK_SIZE = 128
​
    grid = (triton.cdiv(S, BLOCK_S), triton.cdiv(D, BLOCK_SIZE))
​
    section_1_end = mrope_section[1]
    section_2_end = mrope_section[2]
​
    apply_interleaved_rope_kernel[grid](
        x, out,
        S, D,
        x.stride(0), x.stride(1), out.stride(0),
        section_1_end, section_2_end,
        BLOCK_S=BLOCK_S, BLOCK_SIZE=BLOCK_SIZE,
    )
    return out

python/sglang/kernels/ops/memory/memcpy_triton.py

新增的目标文件：包含从 dp_attention.py 提取的 memcpy_triton_kernel、prod 和 memcpy_triton 函数，用于 DP 数据并行中带有偏移量和大小的设备间内存拷贝。

# Offset/size-driven device memcpy kernel, migrated from
# sglang.srt.layers.dp_attention (RFC #29630, Phase 2.5).
​
import functools
​
import triton
import triton.language as tl
​
​
@triton.jit
def memcpy_triton_kernel(
    dst_ptr,
    src_ptr,
    offset_ptr,
    sz_ptr,
    offset_src: tl.constexpr,
    chunk_size, # multiplied for offset and sz
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0).to(tl.int64)
    # 从指针加载偏移量（单位是元素个数）和拷贝大小，再乘以 chunk_size 以适配元素
    offset = tl.load(offset_ptr).to(tl.int64) * chunk_size
    sz = tl.load(sz_ptr).to(tl.int64) * chunk_size
​
    start_index = pid * BLOCK_SIZE
    offs = tl.arange(0, BLOCK_SIZE)
    mask = start_index + offs < sz
​
    if offset_src:
        # 从源地址的 offset 处读取，写入目标地址的起始处
        data = tl.load(src_ptr + offset + start_index + offs, mask=mask)
        tl.store(dst_ptr + start_index + offs, data, mask=mask)
    else:
        # 从源地址的起始处读取，写入目标地址的 offset 处
        data = tl.load(src_ptr + start_index + offs, mask=mask)
        tl.store(dst_ptr + offset + start_index + offs, data, mask=mask)
​
​
def prod(x):
    return functools.reduce(lambda a, b: a * b, x, 1)
​
​
def memcpy_triton(dst, src, dim, offset, sz, offset_src):
    # 带动态偏移量和大小的设备内存拷贝，用于 DP 数据并行中的 token 分发。
    # dim 必须为 0（仅支持第一维度的 slice）；dst 和 src 的尾部维度必须相同。
    # offset 和 sz 应为 torch.Tensor（包含整数），位于与 dst 相同的设备上。
    max_size = min(src.numel(), dst.numel())
    assert dim == 0, 'dim != 0 unsupported'
    assert src.shape[1:] == dst.shape[1:], 'src and dst must have same shape'
    chunk_size = prod(src.shape[1:])
    BLOCK_SIZE = 8192
    grid = (triton.cdiv(max_size, BLOCK_SIZE),)
​
    memcpy_triton_kernel[grid](dst, src, offset, sz, offset_src, chunk_size, BLOCK_SIZE)

评论区精华

Gemini Code Assist 机器人提出了三点审查意见：

• 在 python/sglang/kernels/ops/layernorm/__init__.py 中缺少 KernelBackend 导入，可能导致注册时 NameError（高优先级）。
• 在 apply_interleaved_rope_triton 中建议添加输入形状和 mrope_section 长度的防御性断言（中优先级）。

• 在 memcpy_triton 中建议确认 offset 和 sz 为 torch.Tensor 且与目标张量位于同一设备（中优先级）。
  截至 PR 合并，这些意见未被作者回复或解决，但 PR 仍被合并。

• 缺少 KernelBackend 导入及防御性断言 (correctness): 未解决。这些意见未被作者采纳或修复，PR 仍被合并。

风险与影响

• 风险：
  1. import 遗漏风险：虽然进行了仓库级别的导入重写，但动态导入（如 __getattr__、延迟导入）或外部脚本可能仍引用旧路径，导致运行时 ImportError。审查中发现的 KernelBackend 缺失即为典型遗漏。
  2. 缺少防御性断言：新内核文件（mrope.py、memcpy_triton.py）在异常输入下可能产生难以调试的越界访问或设备不匹配，生成不正确的计算结果或 CUDA 崩溃。
  3. 配置与性能无回归：所有移动内核字节相同，且未移动调优配置，故性能回归风险极低。
     - 影响：对终端用户无功能或性能影响。对开发团队而言，内核组织结构更加清晰：layernorm 相关内核集中至 ops/layernorm，attention 相关内核集中至 ops/attention，内存操作内核集中至 ops/memory，采样相关内核集中至 ops/sampling。后续新增内核可直接在规范位置添加，无需在 srt/layers 中随意放置。但短期内开发者需适应新的导入路径。
     - 风险标记：import 重写可能遗漏动态引用, 缺少防御性断言, KernelBackend 导入缺失

关联脉络

• PR #30784 [Kernel] Migrate scattered quantization kernels to sglang.kernels (RFC #29630, Phase 2.5, 1/7): 同为 Phase 2.5 系列，第 1/7 步，迁移量化内核。与本 PR 共享迁移计划，部分文件（如 ops/moe/init.py、fp8_kernel）可能有重叠的修改，需按顺序合并。
• PR #30786 [Kernel] Migrate scattered MoE kernels to sglang.kernels (RFC #29630, Phase 2.5, 2/7): 同为 Phase 2.5 系列，第 2/7 步，迁移 MoE 内核。与本 PR 共享迁移计划，也需要按任意顺序合并（PR body 说明独立于彼此）。
• PR #29630 RFC: Kernel migration plan (Phase 2.5): 定义整体迁移计划的 RFC 提案，本 PR 是该 RFC 的一个具体执行步骤。