Prhub

#33102 [gdn] fused replayssm ring write into flashinfer gdn mtp verify kernel

原始 PR 作者 Qiaolin-Yu 合并时间 2026-08-03 12:31 文件变更 4 提交数 7 评论 3 代码增减 +4316 / -7

执行摘要

ReplaySSM 环写融合进 GDN MTP verify 内核

PR body 的 Motivation 章节为空,但从代码可推导出明确动机:base 版本 gdn_backend.py 的注释写着 the ring-write exists only in the Triton kernel,这意味着此前在 SM100+ 上选择 FlashInfer CuTe DSL verify kernel 时,ReplaySSM fold-every-commit 因拿不到 ring 数据通道而无法工作,只能回退到更慢的 Triton recurrent kernel。本 PR 把 ring write 融合进 CuTe DSL MTP verify kernel,让 Blackwell 用户既获得 CuTe DSL verify 的性能优势,又保留 ReplaySSM 折叠能力。PR body 的 Accuracy Tests 记录了 aime26 在 qwen 3.5 上与 triton replayssm 及非 replayssm 完全一致的精度结果。

值得精读,重点学习三处设计:一是条件分支的启用策略(backend x dtype x draft 长度三重条件),它限定了新 kernel 只作用于已验证的部署形态;二是 mixed-numerics 验证锚点设计——当两个 kernel 无法 bitwise 对齐时,改为让 Triton fold 消费新 kernel 写的 ring 并与其自身 committed state 对比,既验证了 ring 通道忠实性又避免了跨实现数值争议;三是 ring 作为纯旁路通道的契约(cache_ring 不影响输出,bitwise 断言保障)。阅读时可对照 test_gdn_cutedsl_ring_verify.py 的三条锚点与 gdn_backend.py 的分支逻辑,理解 spec decode 折叠机制对 kernel 侧通道的要求。

讨论亮点

review 阶段没有公开的 review comment,PR 内 3 条评论均为流程性质:Gemini Code Assist bot 的停机通知、作者触发的 /tag-and-rerun-ci、以及作者对与维护者离线讨论的结论陈述。核心设计确认过程发生在 PR 之外:作者与 zcnrex 离线同步,确定了新 kernel 的数值验证策略——由于 CuTe DSL 与 Triton 数值路径不同,测试不再强求 triton-vs-triton 的 bitwise 锚点,而是采用 mixed-numerics 约束(fold 结果对照 kernel 自身 committed state,容差放宽到 bf16-ulp 量级)。作者原话:offline sync with @zcnrex, and we think this is fine.。

实现拆解

  1. 新增 vendored CuTe DSL GDN MTP kernel(python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py,+4127 行):提供 gated_delta_rule_mtp 入口,在 MTP verify 计算的同时以纯旁路(side channel)方式写入 ReplaySSM ring buffer(rawv、rawk、g、beta),使 CuTe DSL 路径首次具备 ReplaySSM fold-every-commit 所需的数据通道。该文件为 vendored 形态,具体内核实现细节在提供的材料中无法进一步确认。

  2. 修改 GDN kernel dispatcher(python/sglang/srt/layers/attention/linear/gdn_backend.py):在 init 的 verify kernel 选择逻辑中新增 verify_kernel_is_flashinfer 标志:当 decode/prefill 后端为 FlashInfer 且 flashinfer_kernel.supports_target_verify 为真时置 True,否则置 False。该标志供后续 _replayssm_fold_target_verify 运行时分支判断,同时 rank0_log 打印 dispatcher 选择结果是核心埋点。

  3. 改造 _replayssm_fold_target_verify 的路径选择:在原有 assert(仅支持线性 draft 链 topk <= 1)之后,计算 seq_len、batch_size、draft_token_num(seq_len // batch_size),当 verify_kernel_is_flashinfer、ssm_states.dtype == bfloat16、draft_token_num >= 3 三个条件同时满足时,将 q/k/v/a/b 按 (batch_size, draft_token_num, ...) 重排后调用 gated_delta_rule_mtp,以 cache_ring=True 写 ring 并返回 verify 输出;否则走原有 Triton fused_sigmoid_gating_delta_rule_update(cache_ring=True)。两个 kernel 写入相同格式的 raw window,保证 fold 阶段无感知。

  4. 新增顶层测试(test/registered/attention/unittests/gdn/test_gdn_cutedsl_ring_verify.py,+150 行):注册到 base-b CI stage(1-gpu-large),覆盖 test_ilp4_small_batch(B=1)与 test_wide_vec_batch(B=8)两种 kernel 变体。验证三条锚点:cache_ring 开关注入下 verify 输出 bitwise 不变(纯旁路契约);rawv/rawk 是 kernel 输入的 bitwise 拷贝、g/beta 分别匹配 Triton gating 与 fp32 sigmoid(容差 5e-5);用 Triton fold 内核消费 CuTe DSL 写的 ring,其结果与 CuTe DSL 自身 committed state(disable_state_update=False 运行)在 3e-2 容差内一致,且未触碰的 slot 完全不变。

  5. 清理旧测试的 sys.path hack(test/registered/attention/unittests/gdn/test_gdn_replayssm_spec_fold.py,-4 行):删除手工 sys.path.insert 与 pathlib.Path 导入,因为新增测试文件已处于可被正常导入的 registered 测试目录结构中,减少路径魔法对可移植性的影响。

文件 模块 状态 重要度
python/sglang/kernels/ops/attention/cutedsl_gdn_mtp_ring.py 内核实现 added 6.22
python/sglang/srt/layers/attention/linear/gdn_backend.py 内核调度 modified 6.55
test/registered/attention/unittests/gdn/test_gdn_cutedsl_ring_verify.py 验证测试 added 7.3
test/registered/attention/unittests/gdn/test_gdn_replayssm_spec_fold.py 折叠测试 modified 3.59

关键符号

gated_delta_rule_mtp _replayssm_fold_target_verify _case _verify _run

关键源码片段

test/registered/attention/unittests/gdn/test_gdn_cutedsl_ring_verify.py test-coverage

新增的顶层测试(150 行),完整定义了新 kernel 的验收契约:ring 旁路性(输出 bitwise 不变)、ring 内容忠实性(rawv/rawk 拷贝、g/beta 数值对齐)、fold 一致性(Triton fold 重现 CuTe DSL 自身 committed state,容差 3e-2)。是理解本 PR 数值策略的关键入口。

class TestGdnCuteDSLRingVerify(CustomTestCase):
    def _run(self, B):
        gating, inputs, state0, slots, rings = _case(B)
​
        # 锚点 1:ring-write 必须是纯旁路(side channel)——
        # cache_ring 开关不应改变 verify 输出(bitwise 相等)。
        out_ref = _verify(gating, inputs, state0.clone(), slots)
        out_ring = _verify(gating, inputs, state0.clone(), slots, rings=rings)
        self.assertTrue(torch.equal(out_ref, out_ring), f"{B=}")
​
        # 锚点 2:ring 内容必须忠实反映 kernel 输入窗口,
        # 否则后续 fold 阶段会基于错误的 raw window 提交状态。
        for i, s in enumerate(slots.tolist()):
            self.assertTrue(
                torch.equal(rings["rawv"][0, s], inputs["v"][i].transpose(0, 1))
            )
            self.assertTrue(
                torch.equal(rings["rawk"][0, s], inputs["k"][i].transpose(0, 1))
            )
​
        # 锚点 3:g / beta 与 Triton gating 及 fp32 sigmoid 对齐(fastmath 容差)。
        g_ref, _ = fused_gdn_gating(
            gating["A_log"],
            inputs["a"].view(B * T, HV),
            inputs["b"].view(B * T, HV),
            gating["dt_bias"],
        )
        g_ref = g_ref.view(B, T, HV).transpose(1, 2).float()
        beta_ref = torch.sigmoid(inputs["b"].float()).transpose(1, 2)
        self.assertLess((rings["g"][0, slots.long()] - g_ref).abs().max().item(), 5e-5)
        self.assertLess(
            (rings["beta"][0, slots.long()] - beta_ref).abs().max().item(), 5e-5
        )
​
        # 锚点 4:fold 一致性——用 Triton fold 内核消费 CuTe DSL 写的 ring,
        # 结果应与 CuTe DSL 自身 committed state(disable_state_update=False)
        # 在 3e-2 容差内一致;未触碰的 slot 必须保持完全不变。
        state_ref = state0.clone()
        _verify(gating, inputs, state_ref, slots, disable_state_update=False)
        fold_state = state0.clone().unsqueeze(0)
        commit_gdn_replayssm_fold_all_layers(
            checkpoint_state=fold_state,
            rawv_cache=rings["rawv"],
            rawk_cache=rings["rawk"],
            g_cache=rings["g"],
            beta_cache=rings["beta"],
            ssm_state_indices=slots,
            accept_lens=torch.full((B,), T, device=DEVICE, dtype=torch.int32),
            max_cache_len=T,
            num_k_heads=H,
        )
        touched = slots.long()
        err = (
            (fold_state[0, touched].float() - state_ref[touched].float())
            .abs()
            .max()
            .item()
        )
        self.assertLess(err, 3e-2, f"{B=} fold vs own update: {err}")
        untouched = [s for s in range(SLOTS) if s not in slots.tolist()]
        self.assertTrue(torch.equal(fold_state[0, untouched], state0[untouched]))

评论区精华

CuTe DSL 路径数值验证方案确认 设计

作者与 zcnrex 离线同步讨论新 kernel 的验证策略。由于 CuTe DSL 与 Triton 数值路径不同,无法沿用 triton-vs-triton 的 bitwise 锚点,改为 mixed-numerics 约束:fold 结果对照 kernel 自身 committed state,容差放宽到 bf16-ulp 量级(3e-2),同时保留输出 bitwise 不变的旁路性断言。

结论:双方认为方案可行,作者确认 we think this is fine 后合入。 · 已解决

风险与影响

  1. 数值一致性风险:CuTe DSL 与 Triton 的数值路径不同,测试只验证了 fold 结果与 CuTe DSL 自身 committed state 在 3e-2 容差内一致,而非与 Triton 基线 bitwise 对齐。长序列、多轮 verify->commit 链上的累积漂移行为在 CuTe DSL 侧没有像 Triton 侧那样做 256 步漂移测试覆盖。
  2. 边界条件风险:draft_token_num >= 3 的阈值是硬编码条件,draft 长度为 1-2 时静默走 Triton 路径,同一部署中两种 kernel 混用,若两者数值特性差异在特定输入下放大,可能出现不易察觉的精度分化。
  3. 大体积 vendored 代码维护风险:cutedsl_gdn_mtp_ring.py 一次性新增 4127 行,属于跨仓库 vendored 性质,后续与上游 CuTe DSL/flashinfer 适配器的同步成本高,且该文件在提供材料中没有可直接审计的内核实现细节。
  4. CI 状态风险:两个 PR CI run(pr-test 与 pr-test-extra)均为失败状态,且 PR 带有 bypass-fastfail 标签,说明是在 CI 未全绿的情况下合入的,回归检测存在盲区。
  5. 模块耦合风险:gdn_backend.py 中新增的 verify_kernel_is_flashinfer 标志与 _replayssm_fold_target_verify 的运行时分支耦合了 dispatcher 选择结果,后续若 verify kernel 选择逻辑调整,需要同步维护该分支条件。

用户侧:SM100+(Blackwell)GPU 上使用 FlashInfer GDN verify 的部署(GDN 系模型如 qwen 3.5 等的 MTP spec decode)将在 ReplaySSM 折叠场景下从 Triton verify 切换到 CuTe DSL verify,预期降低 verify 阶段延迟;SM90 及其他 backend 完全不受影响(回退路径保持原样)。
系统侧:新增一个 4k+ 行的 kernel 文件、dispatcher 新增一个布尔状态、热路径 _replayssm_fold_target_verify 增加条件分支,默认行为不变,改变集中在启用条件全部满足时。
团队侧:kernel 维护负担增加,vended CuTe DSL 代码需要与 FlashInfer bf16-state 适配器(zcnrex 维护线)保持同步;测试注册到 base-b CI(est_time=120s,1-gpu-large),CI 资源成本上升。

新增 4k+ 行 vendored kernel CuTe DSL 与 Triton 数值非 bitwise 一致 条件分支启用新路径 CI 失败且 bypass-fastfail 合入 长序列漂移测试未覆盖 CuTe DSL 侧

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论