执行摘要

• 一句话：修复 DP MTP 中 fused top-k v2 的非法内存访问
• 推荐动作：值得精读。作者对 root cause 的定位方法（CUDA-GDB register 分析、instrumented replay）是高质量调试范例；clamp 策略明确、文档契约化的设计值得借鉴；渐进式的 PR 演化（先 gate 再修复）展示了安全的迭代模式。

功能与动机

PR body 明确说明："Removes its TEMP allow_topk_v2 gate and fixes the root cause of the GLM 5.2 MTP illegal memory access"。根因分析指出：在 DP attention 下，draft-extend-v2 CUDA graph replay 中 idle-companion / DP-padded rows 携带的 seq_len fill value (1) 小于 qo_len，导致 expanded seq_lens 为负（如 -4），被 top-k v2 kernel 作为 uint32 读取后变为 ~4e9 token 长度，产生非法地址。此问题导致 DP MTP 在首个请求后立即 crash。

实现拆解

1. Clamp 负值 seq_lens：在 python/sglang/srt/layers/attention/triton_ops/dsa_metadata.py 的 _fused_dsa_draft_extend_metadata_kernel 和 python/sglang/srt/layers/attention/triton_ops/pad.py 的 seqlens_expand_kernel 中，对 expanded seq_lens 添加 tl.maximum(..., 0) 操作，确保任何负值被截断为 0。0 长度会被 kernel 视为无 token，走 trivial all-(-1) 安全路径。
2. 移除 allow_topk_v2 gate：在 python/sglang/srt/layers/attention/dsa_backend.py 的 DSAIndexerMetadata 中删除 allow_topk_v2 字段，并在 get_indexer_metadata 中移除对 is_target_verify / is_draft_extend_v2 的屏蔽；在 python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py 的 topk_transform 中移除 allow_topk_v2 参数和条件判断，使 fused top-k v2 对符合 shape 条件的 PAGED 场景（包括 spec verify / draft-extend）统一生效。
3. 强化文档契约：在 python/sglang/jit_kernel/dsv4/topk.py 的 plan_topk_v2 和 topk_transform_512_v2 函数中添加详细 docstring，明确要求 seq_lens 必须非负，并解释负值的后果。在 _topk_transform_v2_paged 的 docstring 中也补充了非负契约。
4. 修复单元测试回归：在 test/registered/kernels/test_dsa_indexer.py 的 _run_fused_topk_backend_equivalence_test 中，构建 mock DSAMetadata 时调用 plan_topk_v2 预计算 topk_v2_plan，避免 fused v2 dispatch 的 assertion 失败。

关键文件：

• python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py（模块 TopK 调度；类别 source；类型 core-logic；符号 topk_transform, _topk_transform_v2_paged）: 核心调度：移除了 allow_topk_v2 gate，更新注释为 spec verify / draft-extend 也走 v2 路径，并强化了 lengths 非负契约
• python/sglang/srt/layers/attention/dsa_backend.py（模块 索引器元数据；类别 source；类型 core-logic；符号 DSAIndexerMetadata, get_indexer_metadata）: 移除 allow_topk_v2 临时 gate 及相关逻辑
• python/sglang/srt/layers/attention/triton_ops/dsa_metadata.py（模块 Triton 内核；类别 infra；类型 infrastructure；符号 _fused_dsa_draft_extend_metadata_kernel）: 在 fused_dsa_draft_extend_metadata kernel 中添加 clamp，防止负 seq_lens
• python/sglang/srt/layers/attention/triton_ops/pad.py（模块 Triton 内核；类别 infra；类型 infrastructure；符号 seqlens_expand_kernel）: 在 seqlens_expand_kernel 中添加 clamp，确保扩展长度非负
• python/sglang/jit_kernel/dsv4/topk.py（模块 JIT 内核；类别 source；类型 core-logic；符号 plan_topk_v2, topk_transform_512_v2）: 在 Python 接口中添加重要文档，强制要求 seq_lens 非负
• test/registered/kernels/test_dsa_indexer.py（模块 单元测试；类别 test；类型 test-coverage；符号 _run_fused_topk_backend_equivalence_test）: 修复测试回归：在 mock metadata 中预计算 topk_v2_plan

关键符号：topk_transform, _topk_transform_v2_paged, _fused_dsa_draft_extend_metadata_kernel, seqlens_expand_kernel, plan_topk_v2, topk_transform_512_v2, get_indexer_metadata, _run_fused_topk_backend_equivalence_test

关键源码片段

python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py

核心调度：移除了 allow_topk_v2 gate，更新注释为 spec verify / draft-extend 也走 v2 路径，并强化了 lengths 非负契约

# dsa_topk_backend.py 中 topk_transform 方法的 dispatch 逻辑
# 现在 fused top-k v2 对所有符合条件的 PAGED 场景（包括 spec verify / draft-extend）开放
def topk_transform(
    self,
    logits: torch.Tensor,
    lengths: torch.Tensor,
    topk: int,
    topk_transform_method: TopkTransformMethod,
    attn_metadata,
    ...
) -> torch.Tensor:
    if not envs.SGLANG_DSA_FUSE_TOPK.get() or force_unfused_topk:
        return self.topk_func(logits, lengths, topk, row_starts=row_starts)
​
    # 关键分支：fused top-k v2 路径（移除 allow_topk_v2 gate）
    if (
        envs.SGLANG_OPT_USE_TOPK_V2.get()
        and topk_transform_method == TopkTransformMethod.PAGED
        and row_starts is None
        and batch_idx_list is None
        and 0 < topk <= 2048
        and lengths.shape[0] == logits.shape[0] == attn_metadata.real_page_table.shape[0]
    ):
        return _topk_transform_v2_paged(logits, lengths, topk, attn_metadata)
​
    # 回退到 legacy 路径（需要 page_table_1）
    assert attn_metadata.page_table_1 is not None
    # ... ( 后续代码省略 )

python/sglang/srt/layers/attention/triton_ops/dsa_metadata.py

在 fused_dsa_draft_extend_metadata kernel 中添加 clamp，防止负 seq_lens

# dsa_metadata.py 中 _fused_dsa_draft_extend_metadata_kernel 的 clamp
# 防止负 seq_lens 被下游 uint32 kernel 识别为超大值
expanded_seq = base_seq - qo_len_for_row + local_off + 1
expanded_seq = tl.maximum(expanded_seq, 0) # Clamp to >= 0: DP-padded rows 的安全兜底
expanded_seq = tl.where(mask_e, expanded_seq, 0)
dsa_seq = tl.minimum(expanded_seq, dsa_index_topk)
dsa_cu = tl.cumsum(dsa_seq, 0)

评论区精华

PR body 中作者对根因进行了详尽的 CUDA-GDB 调试分析，但 review 评论无实质性讨论。审核者 Fridge003 直接批准。

• 无实质性讨论 (other): 直接批准，无需修改。

风险与影响

• 风险：风险较低，因修复经过本地（4x B200 DP 配置）和 CI rerun 验证（GLM5.2 MTP 测试通过）。主要风险点：clamp 到 0 可能掩盖其他错误的 seq_lens 负值来源（但逻辑上合理的 padding 应该得到 0）；另外重新启用 fused top-k v2 可能在其他未测试的模型或配置中触发类似问题，但相同 shape 条件的 DP MTP 已覆盖。单元测试的修复确保合并后主分支测试通过。
• 影响：影响范围集中于使用 DP attention 且启用 MTP（multi-token prediction）的场景，主要为 GLM 5.2 和 DeepSeek-V3.2 等模型。受益用户无需再遭遇 illegal memory access 崩溃，且 spec decode 性能恢复至 fused top-k v2 水平（avg_spec_accept_length ~4.0）。对非 DP 或非 MTP 场景无影响。团队维护工作量降低（移除临时 gate）。
• 风险标记：CUDA graph 填充值依赖, DP padding 假设

关联脉络

• PR #30274 Add TEMP allow_topk_v2 gate for MTP: 前一个 PR 引入 allow_topk_v2 临时 gate，本 PR 移除它并修复根因。