执行摘要
- 一句话:ReplaySSM 环写融合进 GDN MTP verify 内核
- 推荐动作:值得精读,重点学习三处设计:一是条件分支的启用策略(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 侧通道的要求。
功能与动机
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 完全一致的精度结果。
实现拆解
-
新增 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 形态,具体内核实现细节在提供的材料中无法进一步确认。
-
修改 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 选择结果是核心埋点。
-
改造 _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 阶段无感知。
-
新增顶层测试(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 完全不变。
-
清理旧测试的 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(模块 内核实现;类别 source;类型 core-logic;符号 gated_delta_rule_mtp): 本 PR 的核心交付物:4127 行新增的 vendored CuTe DSL GDN MTP verify kernel,提供 gated_delta_rule_mtp 入口并支持 cache_ring 旁路写 ReplaySSM ring buffer(rawv/rawk/g/beta),是让 Blackwell 路径摆脱 Triton verify 的关键。
python/sglang/srt/layers/attention/linear/gdn_backend.py(模块 内核调度;类别 source;类型 core-logic;符号 _replayssm_fold_target_verify, verify_kernel_is_flashinfer): 集成改动落点:dispatcher 新增 verify_kernel_is_flashinfer 标志,_replayssm_fold_target_verify 在 FlashInfer verify + bf16 state + draft_token_num >= 3 时切换 CuTe DSL kernel,否则保持 Triton 回退,是启用新 kernel 的运行时决策点。
test/registered/attention/unittests/gdn/test_gdn_cutedsl_ring_verify.py(模块 验证测试;类别 test;类型 test-coverage;符号 _case, _verify, TestGdnCuteDSLRingVerify, _run): 新增的顶层测试(150 行),完整定义了新 kernel 的验收契约:ring 旁路性(输出 bitwise 不变)、ring 内容忠实性(rawv/rawk 拷贝、g/beta 数值对齐)、fold 一致性(Triton fold 重现 CuTe DSL 自身 committed state,容差 3e-2)。是理解本 PR 数值策略的关键入口。
test/registered/attention/unittests/gdn/test_gdn_replayssm_spec_fold.py(模块 折叠测试;类别 test;类型 test-coverage): 顺带清理:删除 4 行 sys.path 手工注入(pathlib + sys.path.insert),因为新增测试已证明 registered 测试目录可直接导入,降低路径魔法对测试可移植性的影响。
关键符号:gated_delta_rule_mtp, _replayssm_fold_target_verify, _case, _verify, _run
关键源码片段
test/registered/attention/unittests/gdn/test_gdn_cutedsl_ring_verify.py
新增的顶层测试(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]))
评论区精华
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.。
- CuTe DSL 路径数值验证方案确认 (design): 双方认为方案可行,作者确认 we think this is fine 后合入。
风险与影响
-
风险:
- 数值一致性风险:CuTe DSL 与 Triton 的数值路径不同,测试只验证了 fold 结果与 CuTe DSL 自身 committed state 在 3e-2 容差内一致,而非与 Triton 基线 bitwise 对齐。长序列、多轮 verify->commit 链上的累积漂移行为在 CuTe DSL 侧没有像 Triton 侧那样做 256 步漂移测试覆盖。
- 边界条件风险:draft_token_num >= 3 的阈值是硬编码条件,draft 长度为 1-2 时静默走 Triton 路径,同一部署中两种 kernel 混用,若两者数值特性差异在特定输入下放大,可能出现不易察觉的精度分化。
- 大体积 vendored 代码维护风险:cutedsl_gdn_mtp_ring.py 一次性新增 4127 行,属于跨仓库 vendored 性质,后续与上游 CuTe DSL/flashinfer 适配器的同步成本高,且该文件在提供材料中没有可直接审计的内核实现细节。
- CI 状态风险:两个 PR CI run(pr-test 与 pr-test-extra)均为失败状态,且 PR 带有 bypass-fastfail 标签,说明是在 CI 未全绿的情况下合入的,回归检测存在盲区。
- 模块耦合风险: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 侧
关联脉络
- PR #33298 [Spec] Support sampling in the DSPARK graph-folded draft proposal: 同为 speculative decoding 的折叠类机制演进(DSPARK graph-folded draft 与 ReplaySSM fold-every-commit 属于同一大方向),两 PR 都涉及 draft 提议与状态折叠路径的 kernel 侧改造。
- PR #32910 [DeepSeek-V4] Fix nvcc 13 crash building the topk_v2 kernel: 同属 DeepSeek-V4 系 JIT 内核维护线(标签 deepseek + jit-kernel),且本 PR 的验证内核同样面向 SM100+ 上 GDN 模型的 spec decode 路径,存在运行时可组合性。
参与讨论