执行摘要

• 一句话：禁用 DSV4 稀疏 MLA 后向 aggressive smem merge，修复 NaN 梯度
• 推荐动作：值得精读，重点看 PR body 的根因分析与验证设计，以及 Issue #1571 的 bisect 过程。设计决策上，“用 2% 内核性能换取确定性正确性”是合理权衡；更值得借鉴的是其基于实验排除干扰项（fp8、DSA-indexer、坏节点）的方法。建议团队为所有 sparse-MLA 内核（glm5、DSV4 等）建立共享的 pass_configs 基线，并增加编译配置的回归测试，避免同类 bug 在不同内核副本间重复出现。

功能与动机

DSV4-Flash RL 训练在第一个 grad-norm 检查即崩溃，报错“found NaN in local grad norm for bucket #0 in backward pass before data-parallel communication collective”。PR body 指出 forward/rollout 均有限，首次后向即出现 NaN；根因是 bwd 内核 JIT 编译时开启 TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE 引发共享内存别名 miscompile，导致 dQ/dKV 为 NaN。同一缺陷此前由 Issue #1571 针对 glm5 内核修复，但该修复只覆盖 glm5 且限制在 Blackwell（sm103）上；DSV4 的 kernel 变体在 Hopper（H200, cc 9.0）上也复现，因此需要无条件禁用该 pass。

实现拆解

1. 定位变更入口：miles_plugins/models/deepseek_v4/ops/kernel/tilelang_sparse_mla_bwd.py 中 bwd 内核的 tilelang JIT pass_configs。
2. 核心变更：将 tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE 由 True 改为 False，禁止 JIT 对共享内存做激进合并。这是全 PR 唯一改动（+1/-1），内核主体逻辑不变。
3. 验证方式：在 H200 上，dev 镜像 6 节点验证 grad_norm 为 0.075/0.056，:latest + 其他修复 7 节点验证为 0.0753；并实验排除了 fp8 vs bf16、DSA-indexer 配置、坏节点等干扰因素。
4. 成本评估：按 Issue #1571 的 ablation，该内核性能损失约 2%，端到端影响小于 1%。PR 未附带自动化测试，回归保障依赖后续补充或上游 tilelang 修复。

关键文件：

• miles_plugins/models/deepseek_v4/ops/kernel/tilelang_sparse_mla_bwd.py（模块 模型内核；类别 source；类型 configuration；符号 bwd）: 唯一的变更文件；在 DSV4 sparse-MLA 后向内核的 tilelang pass_configs 中关闭 aggressive shared-memory merge，修复共享内存别名导致的 NaN dQ/dKV，解除 DSV4-Flash RL 首次 grad-norm 检查即崩溃的训练阻断。

关键符号：bwd

关键源码片段

miles_plugins/models/deepseek_v4/ops/kernel/tilelang_sparse_mla_bwd.py

唯一的变更文件；在 DSV4 sparse-MLA 后向内核的 tilelang pass_configs 中关闭 aggressive shared-memory merge，修复共享内存别名导致的 NaN dQ/dKV，解除 DSV4-Flash RL 首次 grad-norm 检查即崩溃的训练阻断。

文件：miles_plugins/models/deepseek_v4/ops/kernel/tilelang_sparse_mla_bwd.py

片段：DSV4 sparse-MLA 后向内核 bwd 的 tilelang JIT 装饰器 pass_configs 配置

# 该配置块位于 bwd 内核的 tilelang JIT 装饰器中，控制编译 pass。
pass_configs={
    tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
    tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
    # 修复点：aggressive smem merge 会让仍存活的 shared buffer 互相别名（miscompile），
    # 产生 NaN dQ/dKV → NaN 梯度 → DSV4-Flash RL 首次 grad-norm 检查崩溃。
    # 与 PR #1571 修复的 glm5 内核属同一 tilelang 0.1.9 缺陷；
    # DSV4 变体在 Hopper（H200, cc 9.0）上同样复现，因此不能像 #1571 那样只对
    # Blackwell 生效，必须无条件关闭。代价约 2% 内核性能（来自 #1571 的 ablation）。
    tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: False,
}

变更仅一行：将 True 改为 False，内核主体逻辑不变。

文件：miles_plugins/models/deepseek_v4/ops/kernel/tilelang_sparse_mla_bwd.py

片段：DSV4 sparse-MLA 后向内核 bwd 的 tilelang JIT 装饰器 pass_configs 配置

# 该配置块位于 bwd 内核的 tilelang JIT 装饰器中，控制编译 pass。
pass_configs={
    tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
    tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
    # 修复点：aggressive smem merge 会让仍存活的 shared buffer 互相别名（miscompile），
    # 产生 NaN dQ/dKV → NaN 梯度 → DSV4-Flash RL 首次 grad-norm 检查崩溃。
    # 与 PR #1571 修复的 glm5 内核属同一 tilelang 0.1.9 缺陷；
    # DSV4 变体在 Hopper（H200, cc 9.0）上同样复现，因此不能像 #1571 那样只对
    # Blackwell 生效，必须无条件关闭。代价约 2% 内核性能（来自 #1571 的 ablation）。
    tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: False,
}

变更仅一行：将 True 改为 False，内核主体逻辑不变。

评论区精华

本 PR 没有实质性的评审讨论：Gemini Code Assist bot 仅给出“no feedback”，维护者 yueming-yuan 直接 approve。技术依据主要来自 PR body 中的根因分析和多节点验证表格，以及 Issue #1571 中对该 tilelang 缺陷的 bisect 结论。

• 暂无高价值评论线程

风险与影响

• 风险：正确性风险：关闭 aggressive smem merge 后，需确认没有其他 pass 组合导致新的 miscompile；#1571 曾以 50/50 压力迭代和 rel_err ~0.003 验证过类似修复，本 PR 在 H200 多节点上也确认 grad_norm 有限，风险较低。性能风险：该内核约有 2% 性能回退（#1571 的 ablation），端到端影响 <1%。兼容性：无条件关闭意味着 Blackwell 上的 DSV4 也会承担该性能成本；但这是避免崩溃的必要代价。测试风险：PR 未新增自动化回归测试，若未来有人重新打开该 flag 或升级 tilelang 改变默认行为，问题可能回归，建议补充针对 pass_configs 的编译级回归测试。
• 影响：影响范围：所有使用 DSV4-Flash RL 训练（fp8 与 bf16 均受影响）的用户，修复首次训练步崩溃，使训练可以持续推进；对 H200 与 Blackwell 均生效。团队影响：暴露了 tilelang PassConfig 在不同架构上的正确性差异，以及“同一缺陷在不同内核副本间漂移”的维护问题；建议统一管理 sparse-MLA 家族内核的编译配置，并跟进上游 tilelang 修复。
• 风险标记：核心路径变更, 缺少自动化回归测试, 小幅性能回退, 同类 miscompile 或波及其他内核

关联脉络

• PR #1571 GLM-5.2 744B on GB300: sm103 kernel fix, training config: 同一 tilelang aggressive smem merge miscompile 的原始修复，但只覆盖 glm5 且限制 Blackwell；本 PR 将其推广到 DSV4 并解除硬件限制，是其直接延续和补充。