执行摘要

• 一句话：迁移通用注意力 Triton kernels 到 sglang.kernels 统一命名空间
• 推荐动作：值得精读，重点观察：内核迁移的模式（字节相同搬迁 + import 重写，确保零回归）；从混合模块中提取内核（如 _build_pa_page_table 从 backend 中抽出）的设计，定义清晰的“kernel 模块”边界；对 review 中安全建议的取舍，因为移动不引入新问题，可以暂缓，但建议作为后续 issue 跟踪。

功能与动机

作为 Phase 2.5 (RFC #29630) 的第 4 个子 PR，旨在将散落在 srt/layers/attention 各处的通用注意力 Triton kernels 集中到 sglang.kernels 规范命名空间下，消除内核分散管理问题，统一内核注册和导入入口，为后续内核性能优化和版本管理奠定基础。

实现拆解

1. 移动 utils.py：将 srt/layers/attention/utils.py 中的 8 个 Triton kernel（MLA fp8 quantize+rope、reshape_and_cache 变体等）原样搬迁至 ops/attention/utils.py。
2. 移动 flash_mla_sm120：将 flash_mla_sm120.py 和 _triton.py 搬迁至 ops/attention/。
3. 移动 NSA decode 包：将 nsa/triton_decode/ 整个目录搬迁至 ops/attention/nsa_triton_decode/（8 个 kernel）。
4. 移动 DCP kernels：将 srt/layers/dcp/kernels.py 移至 ops/attention/dcp_kernels.py，包含 4 个 kernel 和 CPTritonContext。
5. 提取 PA page table builder：从 flashattention_backend.py 中提取内联的 _build_pa_page_table_kernel Triton kernel 和 _build_pa_page_table 启动函数，放入新的 ops/attention/pa_page_table.py，原模块改为导入新路径。
6. 更新全仓库导入：在 aiter_backend.py、dsa_backend.py、flashinfer_backend.py、trtllm_mla_backend.py、flashmla_backend.py、cutlass_mla_backend.py、deepseek_common 等相关文件中，将旧 import sglang.srt.layers.attention.utils 替换为 sglang.kernels.ops.attention.utils，并移除不再使用的局部 import。
7. 注册 KernelSpec：在 sglang.kernels 的注册表中添加所有迁移后的 kernel 入口点，标识 TRITON 后端。

测试验证：test_kernels_namespace.py 和 test_fused_op.py 的 33 个测试全部通过，证明命名空间导入正确。

关键文件：

• python/sglang/srt/layers/attention/flashattention_backend.py（模块 注意力后端；类别 source；类型 core-logic；符号 _build_pa_page_table_kernel, _build_pa_page_table）: 核心文件，删除了 _build_pa_page_table_kernel 和 _build_pa_page_table 的内联定义，改为从 ops/attention/pa_page_table 导入，同时移除了对 triton 的直接依赖。是整个变更中修改量最大的源码文件。
• python/sglang/kernels/ops/attention/pa_page_table.py（模块 内核；类别 infra；类型 infrastructure；符号 _build_pa_page_table_kernel, _build_pa_page_table）: 新增文件，包含从 flashattention_backend 提取的 _build_pa_page_table_kernel Triton kernel 和 _build_pa_page_table 启动函数，构成独立模块。是迁移后新增的最重要的内核文件。
• python/sglang/kernels/ops/attention/utils.py（模块 内核；类别 infra；类型 infrastructure；符号 assert_buffer_fits, create_flashinfer_kv_indices_triton, create_flashmla_kv_indices_triton, get_num_kv_index_blocks_flashmla）: 新位置，包含 8 个 MLA fp8 quantize+rope、reshape_and_cache 等通用 attention kernel，替代了旧位置 srt/layers/attention/utils.py。是 attention 辅助函数的新集中地。
• python/sglang/kernels/ops/attention/dcp_kernels.py（模块 内核；类别 infra；类型 infrastructure；符号 CPTritonContext）: 新位置，包含 4 个 DCP kernel 和 CPTritonContext，替代了旧位置 srt/layers/dcp/kernels.py，是 DCP 通信的内核集合。
• python/sglang/srt/layers/attention/aiter_backend.py（模块 注意力后端；类别 source；类型 dependency-wiring）: 代表所有需要重写 import 的 attention backend 文件：将 sglang.srt.layers.attention.utils 替换为 sglang.kernels.ops.attention.utils。体现 import 重写的典型模式。

关键符号：_build_pa_page_table_kernel, _build_pa_page_table, assert_buffer_fits, create_flashinfer_kv_indices_triton, create_flashmla_kv_indices_triton, get_num_kv_index_blocks_flashmla, concat_mla_absorb_q_general, mla_quantize_and_rope_for_fp8, seqlens_expand_triton, launch_reshape_and_cache_flash, pad_sequence_with_mask

关键源码片段

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

新增文件，包含从 flashattention_backend 提取的 _build_pa_page_table_kernel Triton kernel 和 _build_pa_page_table 启动函数，构成独立模块。是迁移后新增的最重要的内核文件。

"""Paged-attention page-table builder, migrated from
``sglang.srt.layers.attention.flashattention_backend`` (RFC #29630, Phase 2.5).
"""
from typing import Optional
​
import torch
import triton
import triton.language as tl
​
​
@triton.jit
def _build_pa_page_table_kernel(
    req_to_token_ptr,
    req_pool_indices_ptr,
    seq_lens_ptr,
    prefill_lens_ptr,
    dst_page_table_ptr,
    kv_lens_ptr,
    window_size: tl.constexpr,
    req_to_token_stride,
    dst_stride,
    BLOCK_SIZE: tl.constexpr,
):
    # 每个线程块处理一个请求
    bid = tl.program_id(0)
    req_idx = tl.load(req_pool_indices_ptr + bid)
    sl = tl.load(seq_lens_ptr + bid).to(tl.int32)
    pf = tl.load(prefill_lens_ptr + req_idx).to(tl.int32)
​
    # decode 起始位置：取 prefill_len 和 seq_len - window_size 较大者
    decode_start = tl.maximum(pf, sl - window_size)
    gap = tl.where(decode_start > pf, decode_start - pf, 0)
    kv_len = sl - gap
​
    tl.store(kv_lens_ptr + bid, kv_len)
​
    src_base = req_idx * req_to_token_stride
    dst_base = bid * dst_stride
​
    for start in tl.range(0, kv_len, BLOCK_SIZE):
        offs = start + tl.arange(0, BLOCK_SIZE)
        mask = offs < kv_len
        # 构造 page table 中实际 token 位置：小于 prefill_len 的取 offs，否则加上 gap
        pos = tl.where(offs < pf, offs, offs + gap)
        # REVIEW: 当 mask 为 False 时，pos 可能越界，建议加钳制
        kv_loc = tl.load(
            req_to_token_ptr + src_base + pos,
            mask=mask,
            other=0,
        )
        tl.store(dst_page_table_ptr + dst_base + offs, kv_loc.to(tl.int32), mask=mask)
​
​
def _build_pa_page_table(
    req_to_token: torch.Tensor,
    req_pool_indices: torch.Tensor,
    seq_lens: torch.Tensor,
    prefill_lens: torch.Tensor,
    window_size: int,
    bs: int,
    pa_max_len: int,
    device: torch.device,
    dst_page_table: Optional[torch.Tensor] = None,
    dst_kv_lens: Optional[torch.Tensor] = None,
):
    # CUDA-graph 模式下复用已分配缓冲区，否则新建
    if dst_page_table is None:
        dst_page_table = torch.zeros(bs, pa_max_len, dtype=torch.int32, device=device)
    if dst_kv_lens is None:
        dst_kv_lens = torch.empty(bs, dtype=torch.int32, device=device)
    if bs > 0 and pa_max_len > 0:
        _build_pa_page_table_kernel[(bs,)](
            req_to_token,
            req_pool_indices.contiguous(),
            seq_lens.to(torch.int32),
            prefill_lens,
            dst_page_table,
            dst_kv_lens,
            window_size,
            req_to_token.stride(0),
            dst_page_table.stride(0),
            BLOCK_SIZE=256,
        )
    return dst_page_table, dst_kv_lens

评论区精华

Review 中 gemini-code-assist[bot] 指出 pa_page_table.py 中的 _build_pa_page_table_kernel 在 mask 为 false 时，pos 计算可能因越界地址导致 GPU 页面错误，建议添加 pos = tl.where(mask, pos, 0) 钳制。但作者未对此进行修改，且该代码是从原 flashattention_backend 直接搬到新文件，原有实现即如此，可能为误报或已知安全边界。该评论未得到 resolve。

• 潜在 GPU 越界访问风险 (correctness): 未解决；作者未回应或修改；该代码为原搬迁移，原本即有此模式，可能被视为已知安全边界或误报。

风险与影响

• 风险：
  • 回归风险极低：所有迁移均为字节相同搬迁（git rename R100），核心逻辑无变化；_build_pa_page_table 提取虽为剪裁后搬迁，但函数体完全一致。全仓库 import 重写通过测试验证。
  • 潜在内存越界：review 指出的 _build_pa_page_table_kernel 中 mask 未保护越界地址，虽非本次引入，但迁移后该 kernel 被抽取为独立模块，增加了后续独立调用的可能性，若调用不当可能触发 GPU page fault。建议跟进修复。
  • 依赖三方库：部分 backend（如 aiter）仍需 triton 导入，但 flashattention_backend 因提取 kernel 后不再直接依赖 triton，降低了模块耦合。
• 影响：
  • 用户：无感知，推理行为无变化。
  • 开发者：内核位置统一，新增 attention kernel 必须置于 sglang.kernels.ops.attention 下；旧 import 路径失效，需更新所有引用。
  • 系统：启动时内核注册表增加新的 entry，不影响运行时性能。
  • 维护性：长期利于内核的集中管理与优化。
  • 风险标记：字节等价搬迁, import 重写, 内存越界未修复

关联脉络

• PR #30784 [Kernel] Migrate scattered quantization kernels to sglang.kernels (RFC #29630, Phase 2.5, 1/7): 同 Phase 2.5 系列，1/7，迁移量化 kernels，本 PR 依赖其导入路径变更
• PR #30786 [Kernel] Migrate scattered MoE kernels to sglang.kernels (RFC #29630, Phase 2.5, 2/7): 同 Phase 2.5 系列，2/7，迁移 MoE kernels，与本 PR 共同完成内核集中
• PR #30787 [Kernel] Migrate top-level srt/layers stray kernels to sglang.kernels (RFC #29630, Phase 2.5, 3/7): 同 Phase 2.5 系列，3/7，迁移 stray kernels，与本 PR 连续递进