Prhub

#39553 [CI] Run test_unified_radix_cache_kl_dcp on cutedsl_mla with bf16 KV cache

原始 PR 作者 kpham-sgl 合并时间 2026-09-15 15:35 文件变更 1 提交数 1 评论 4 代码增减 +2 / -4

执行摘要

DCP KL 测试切换 MLA 后端并回退 bf16 KV

PR body 明确写出:test_unified_radix_cache_kl_dcp.py 自 09/11 起在 scheduled extra-b-test-4-gpu-b200 上约 50% 的运行失败(两次 retry 都失败),avg_kl_div 为 0.010-0.013 而阈值为 0.01。作者把回归 bisect 到 #35770:resolve_attn_backend() 现在会 unwrap HybridLinearAttnBackend.full_attn_backend,于是 Kimi Linear prefill 落到 TokenspeedMLABackend.prepare_prefill_qkv,在 prefill kernel 之前把 Q/K/V 量化成 FP8(#38152 后来让这个钩子对 NoPE 安全,之前曾在 09/05 直接崩溃)。量化误差的影响有明确的时间线证据:09/04 之前的调度运行测得 0.003-0.006,之后为 0.007-0.013。作者同时强调这不是 runner 问题、也没有依赖包版本变化,因此选择把测试切到更贴近 DCP 语义的 cutedsl_mla + 默认 bf16 KV cache,与其余注册 KL 测试口径一致,而不是放宽阈值。

值得精读,但重点不在代码(+2/-4 的测试参数切换,几分钟即可看完),而在 motivation 里的因果链:一次 attention backend 解析逻辑的改动(#35770 unwrap HybridLinearAttnBackend.full_attn_backend)如何让 Kimi Linear prefill 意外落入 FP8 量化钩子,并把 KL 从 0.003-0.006 抬到 0.007-0.013。建议关注三个设计决策:(1) 为什么选择换后端而不是放宽阈值——保持了测试的判别力,是正确的取舍;(2) 为什么 cutedsl_mla 更合适——它是 SM100 上的 DCP 原生 MLA 内核,且已被 Kimi Linear DCP 测试采用,避免为单个测试保留一条非主路径;(3) 遗留问题——是否应把“tokenspeed_mla + fp8 KV + DCP”组合的精度损失记录为已知限制,或单独加一个小规模的 FP8 KL 守护测试,并考虑把阈值从 0.01 收紧到约 0.008 以保留回归灵敏度。

讨论亮点

本 PR 没有任何 review comment(review_comments_count = 0),也缺少 reviewer 的设计讨论,4 条评论全部是 CI 流程噪音:

  1. 作者两次发起 /rerun-test test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py,第一次被 github-actions[bot] 拦下:“A major update has landed on main. Your PR is diverged relative to required base commit ddd4600197a4. Re-run was not dispatched.”,并提示可对 CI-fix PR 申请 bypass-maintenance 标签跳过该检查;作者没有走例外通道,而是推进分支后重新触发。

  2. 第二次重跑成功:github-actions[bot] 回报 4-gpu-b200(1 test)✅,并附上调度命令 cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py。

  3. 真正有价值的技术论证全部写在 PR body 而非评论区:bisect 结论(#35770 解包 HybridLinearAttnBackend.full_attn_backend)、量化钩子时间线(#38152)、以及“阈值不变”这一明确取舍。由于没有 reviewer 追问,Kimi Linear + tokenspeed_mla + FP8 KV 组合的精度损失是否可接受、是否需要单独的工作项,在讨论中并未被确认。

实现拆解

  1. 变更入口:唯一改动文件是 test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py(+2/-4),改动集中在模块 docstring 与 TestUnifiedKimiLinearDcpHiCache.setUpClass() 的 popen_launch_server(other_args=[...]) 列表里,属于纯粹的测试配置调整,没有触碰任何 python/sglang/srt/ 生产代码。

  2. 后端切换(核心改动):--attention-backend 的值由 tokenspeed_mla 改为 cutedsl_mla。原因是 tokenspeed_mla 在这条 Kimi Linear + DCP4 路径上会先经过 TokenspeedMLABackend.prepare_prefill_qkv 做 FP8 量化,从而引入额外数值误差;cutedsl_mla 被 PR 描述为 Blackwell(SM100)上的 DCP 原生 MLA 解码内核,并已在 test_kimi_linear_unified_memory_dcp_blackwell.py 中承担 Kimi Linear DCP 的验证职责,因此测试覆盖的仍是 DCP MLA 主路径,而不是退化路径。

  3. KV cache 精度回退:删除 --kv-cache-dtype 与 fp8_e4m3 两个列表项,KV cache 回落到默认 bfloat16,与 --dtype bfloat16 保持一致,也和“其余所有注册 KL 测试”的口径对齐。这一步同时消掉了 FP8 KV 与 FP8 QKV 量化误差的叠加。

  4. 文档同步:模块 docstring 中 Blackwell-only 的说明从“MLA DCP decode path needs tokenspeed_mla (SM100/12x)”改为“cutedsl_mla is the DCP-native MLA decode kernel on SM100”,避免注释与实现相互矛盾——语义上由“唯一可用后端”变为“DCP 原生后端”。

  5. 未改动项(刻意保持):kl_threshold = 0.01、gsm8k_threshold = 0.85、register_cuda_ci(est_time=280, stage="extra-b", runner_config="4-gpu-b200")、DCP 相关参数(--tp-size 4、--dcp-size 4、--dcp-comm-backend a2a、--dcp-replicate-q-proj)、HiCache 配置(--enable-hierarchical-cache、--hicache-size 10、--hicache-write-policy write_through)以及 SGLANG_ENABLE_UNIFIED_RADIX_TREE / SGLANG_ENABLE_RANK_CONSENSUS_CHECKER 两个环境变量均保持不变。这说明测试意图(守护 DCP 下 widened logical index 到 per-rank shard 的翻译正确性)与 CI 成本都没有被削弱,只有数值路径变干净。

  6. 验证配套:PR body 附验证表格,10 次连续运行(main @ 2fca6d69,约 4.5 min/次)10/10 通过:logprobs_match 0.0052-0.0057、prefill_cache_hit 0.0029-0.0037、decode_cache_hit 0.0026-0.0042,GSM8K 0.930-0.945;对照组是同一配置下 09/06 起的 tokenspeed_mla 调度运行 0.007-0.013。合并前 /rerun-test 在 4-gpu-b200 上单测通过。

文件 模块 状态 重要度
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py 统一缓存测试 modified 4.11

关键符号

TestUnifiedKimiLinearDcpHiCache.setUpClass TestUnifiedKimiLinearDcpHiCache

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

后端切换后的 KL 余量是否足够 测试

PR body 给出 10 次连续运行(main @ 2fca6d69,约 4.5 min/ 次)全部通过:logprobs_match 0.0052-0.0057、prefill_cache_hit 0.0029-0.0037、decode_cache_hit 0.0026-0.0042,GSM8K 0.930-0.945;对照组是同配置下 09/06 以后调度 CI 上的 tokenspeed_mla,读数为 0.007-0.013。

结论:bf16 KV + cutedsl_mla 后三项 KL 指标相对 0.01 阈值有接近一倍以上余量,判定可接受,阈值维持 0.01 不变;合并前 4-gpu-b200 重跑亦通过。 · 已解决

分支落后导致 /rerun-test 未被派发 测试

作者两次执行 /rerun-test,第一次被 github-actions[bot] 拒绝:PR 相对 required base commit ddd4600197a4 处于 diverged 状态,重跑未派发,并提示 CI-fix PR 可申请 bypass-maintenance 标签跳过该检查。

结论:作者未使用例外标签,而是在推进分支后重新触发,最终 4-gpu-b200 单测 ✅ 通过,流程闭环。 · 已解决

回归根因(#35770 引入的 FP8 QKV 量化)归属 正确性

PR body 明确 bisect 到 #35770:resolve_attn_backend() 现在会 unwrap HybridLinearAttnBackend.full_attn_backend,使 Kimi Linear prefill 命中 TokenspeedMLABackend.prepare_prefill_qkv 并对 Q/K/V 做 FP8 量化;#38152 让该钩子对 NoPE 安全。时间线显示 KL 从 09/04 前的 0.003-0.006 变为 09/06 后的 0.007-0.013。

结论:本 PR 只在测试侧规避,未修改生产代码;Kimi Linear + tokenspeed_mla + fp8_e4m3 组合的精度损失是否需要在实现层处理,讨论中没有结论,属于遗留项。 · 待处理

风险与影响

  1. 根因未修,只做规避:KL 翻倍(0.003-0.006 → 0.007-0.013)的直接触发点是 #35770 之后 Kimi Linear prefill 走 TokenspeedMLABackend.prepare_prefill_qkv 做 Q/K/V FP8 量化,本 PR 没有修改生产代码,只是让测试绕开该组合。若真实用户使用 --attention-backend tokenspeed_mla + --kv-cache-dtype fp8_e4m3 + DCP 部署 Kimi Linear,精度损失仍然存在且不再有 CI 守护。

  2. 测试覆盖形态发生变化:被守护路径由“tokenspeed_mla + fp8_e4m3 KV”变为“cutedsl_mla + bf16 KV”。prepare_prefill_qkv 钩子、FP8 KV cache 与 DCP widened index 的交互(例如 #39487 修的 widened KV id 本地化场景)在 FP8 组合下失去端到端覆盖;DCP MLA 的两条解码内核路径也因此只保留一条被 KL 测试覆盖。

  3. 阈值未收紧,回归灵敏度被放弃:新配置实测上界为 0.0057,而阈值仍为 0.01,余量接近一倍。未来若再出现类似量级的 KV 索引翻译退化(本测试存在的意义正是“missing translation 会读到别的 rank 的 KV”),有可能被 0.01 的宽松阈值掩盖。

  4. 验证样本量有限:10 次运行对 ~50% 的失败率有足够统计分辨力,但对低频 flake(例如 5% 级别)置信度不足;PR body 未给出跨机器/跨时间窗口的复现数据,只有作者本地 4x B200 与一次合并前调度重跑。

  5. 无运行时影响:改动只在 test 目录,不涉及服务启动参数默认值、显存占用或线上行为,不存在功能回归风险。

对最终用户零直接影响——本 PR 不修改任何推理、调度或缓存实现。对开发团队影响明确且正面:恢复 extra-b / 4-gpu-b200 排班约 50% 失败率的 CI 噪音,避免该 job 长期占用 retry 预算、也避免 maintainer 对 KL 测试产生“狼来了”的钝感。对测试资产的影响是覆盖面重排:UnifiedRadixCache + DCP4 + HiCache L2 的 KL/GSM8K 端到端守护依然存在(阈值、断言、缓存命中校验都没动),但守护的精度组合从 FP8 收敛到 bf16,代价是 FP8 量化组合失去同等强度的守护。对后续 PR 的隐含影响是:任何人在 Kimi Linear DCP 场景引入 FP8 QKV 量化时,将不再有 CI 报警提示 KL 抬升。

仅规避未修根因 FP8 量化路径失去 CI 覆盖 KL 阈值未收紧 验证样本仅 10 次

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论