hyper-parallel CSA/DSv4 融合算子问题汇总（待 kernel 侧定位）

环境

• MindSpore 2.10.0（已更新至最新，含时序同步修复，问题依旧）
• Ascend NPU，PYNATIVE_MODE，8 卡纯 FSDP（dp_shard=8，tp=1，pp=1）
• 模型：DeepSeek-V4 hybrid（dsv4_hybrid），3 层，csa_compress_ratios=[0, 4, 128]，seq=4096
• 涉及算子：npu_lightning_indexer、npu_sparse_flash_mla（及其 grad）
• 对照基线：非融合小算子实现（apply_dsa_kernel_fusion=False）loss 完全正常、可复现

问题 1【严重·阻断】npu_lightning_indexer 在 padding 槽写 inf（污染输出/相邻显存）

根因（已单卡确定性复现，100% 必现）

npu_lightning_indexer 把 top-k 未占满的 padding 槽（即 sparse_indices == -1 的槽，
早期/causal query 行可见压缩块不足时产生）对应的 sparse_values 写成了 inf，
本应写有限哨兵（0 / -inf）。

铁证（见随附最小复现脚本 repro_lightning_indexer_oob.py，单卡每个 iter 必现）：

val_bad_cnt(inf 数量) == pad_cnt(idx==-1 数量) == 524800
bad_all_at_pad == True      # 每一个 inf 都精确落在 -1 padding 槽，正常槽无一 inf
idx 范围 [-1, 1023] 合法，idx_oob == False

与训练 nan 的关联

训练用 return_value=False，这些 value 不返回；但 kernel 仍把 inf 写进 value 缓冲区，
该缓冲区溢出/混叠到相邻的下一个显存分配——实测把某层独立 Parameter attn_sink
(fp32[64]) 写成了 inf，再向下游传播为 nan → 首个 step loss=nan。
出 inf 的卡随机（worker1/worker0 在不同运行/MindSpore 版本间漂移）= 分配布局相关，
与"inf 写在 padding 槽"这一确定性根因一致。

定位证据（已二分确认）

→ 元凶唯一锁定 npu_lightning_indexer：padding 槽的 value 未初始化为有限值。

复现调用参数（BSND，见随附脚本）

q       : bfloat16, [B=1, S=4096, N=64, D=128]
k       : bfloat16, [B=1, Skv=1024, 1, D=128]    # 压缩后 KV，n_compressed=4096/4
weights : float32,  [B=1, S=4096, N=64]
sparse_count = 512        # min(index_topk=512, n_compressed=1024)
layout='BSND', sparse_mode=3, cmp_ratio=4, return_value=True   # True 即可直接看到 inf
cu_seq_lens_q = cu_seq_lens_k = None

sparse_indices 形如 [1,4096,1,512]，范围 [-1,1023]；sparse_values 同形，在 idx==-1 处为 inf。

修复方向

padding（无效 top-k）槽的 sparse_values 应初始化为有限值（0 或 -inf），不要留 inf；
另外 return_value=False 时不应把 value 缓冲区写到会越界/混叠相邻分配的位置。

问题 2【严重·训练发散】npu_sparse_flash_mla 的 cmp_kv=None（SWA）反向对真实输入算错

现象

问题 1、3 修复后，融合路径端到端可跑、前向逐 step 对齐基线，但 grad_norm 从 step1 起异常偏大、
逐步爆炸；loss 因梯度裁剪不 nan，但不再下降（基线在学、融合学不动）。

决定性定位（单卡 gbs=1，python 直跑，真实配置 seq=4096 [0,4,128]）

用梯度 hook 直接读 fused / unfused 两个独立 run 中各中间变量收到的梯度范数（ratio4/128 逐项
1.00 说明非确定性可忽略、对比可信）：

结论：进入 CSA 反向的 cotangent d(csa_out) 完全一致（1.00），但只有 ratio0（cmp_kv=None /
纯 SWA）这层产出的 d_query/d_ori_kv 偏大 3.8~6×；ratio4 / ratio128（有 cmp_kv）反向逐项正确。
排除项全部命中：单卡（world_size=1，sense=1，无卡数/FSDP/sense 影响）、前向匹配、cotangent
匹配、ratio4/128 对照组正确。→ bug 就在 npu_sparse_flash_mla 的 cmp_kv=None 反向分支本身，
把正确的 cotangent 映射成 3.8~6× 错误的 d_query/d_ori_kv。

关键特征：输入结构相关（随机输入测不出）

• 之前用 ms.grad + 随机 q/k/cotangent 的隔离复现（SWA/dense/sparse 各路径）反向都 ≈ 1.000
  对齐——这是误导：随机（平坦 softmax）下碰巧对；
• 用训练 dump 的真实 RoPE'd MLA query/key + inverse-RoPE 来的真实 cotangent，cmp_kv=None 反向
  d_query 偏大 ~3.8×（L2/norm）/ 45×（max 元素）、d_ori_kv ~16×。

最小复现：repro_swa_real_grad.py（+ l0real.npz）

replay 训练 dump 的真实 ratio0 输入 + 真实 cotangent，同输入喂 fused kernel 与 unfused 小算子，
对比 d_query/d_ori_kv/d_sink。实测：

[fwd] amax fused=3.7188e+00 unfused=3.7031e+00  L2rel=1.06e-01
[bwd] d_query   amax unfused=2.18e-05 fused=9.99e-04  ratio(f/u)=45.8  L2rel=3.60   <-- 错
[bwd] d_ori_kv  amax unfused=7.39e-05 fused=1.23e-03  ratio(f/u)=16.6  L2rel=5.73   <-- 错
[bwd] d_sink    amax unfused=4.64e-05 fused=4.73e-05  ratio(f/u)=1.02              <-- 基本对

（l0real.npz 由 csa.py 在 CSA_DUMP_L0=1 下单卡跑一步自动 dump：第一个 ratio0 层的
query/key/sink/d_csa_out/scale/win_left，约 1GB，可单独提供。）

对照：随机输入下反向正确（证明是结构相关）

repro_flash_mla_swa_grad.py（SWA，随机，N=16/64，sink 0~50 全扫）、repro_flash_mla_cmp_grad.py
（dense cmp）、repro_flash_mla_sparse_grad.py（sparse cmp ratio4 topk512）、repro_l0_threeway.py
（ratio0 三方）——随机输入下 forward+backward 全部 ratio≈1.000。只有真实结构才暴露 cmp_kv=None
反向 bug。

修复方向

1. kernel 侧：修 npu_sparse_flash_mla cmp_kv=None（SWA / 纯 band-mask）分支的反向，使其对真实
   RoPE'd MLA 输入也正确（cmp_kv 非空分支是好的，可对比定位）；
2. mindformers 侧临时 workaround（可立即解锁训练）：ratio0 层走非融合小算子
   unfused_compressed_sparse_attn（仅滑窗、无压缩/无 indexer、fp32 精确且不贵），ratio4/128 保持融合。

问题 3【缺陷·阻断反向】custom_op_impl.py 中 NpuSparseFlashMlaDFunction 的 wrapper 缺陷

npu_sparse_flash_mla 的反向链路明显未经过端到端训练验证（ST 仅覆盖 forward）。为让反向跑通，临时在 hyper_parallel/platform/mindspore/custom_ops/custom_op_impl.py 打了 3 处补丁，建议合入源码：

1. metadata 应传 None
   backward 把 forward 算好的 metadata 传给 npu_sparse_flash_mla_grad，但 grad kernel 断言失败：
   metadata should be nullptr now（sparse_flash_mla_grad_tiling_general.cpp:401）。
   → 改为传 None，让 grad kernel 内部自算 metadata。

2. cmp_residual_kv 未透传到反向
   forward 未把 cmp_residual_kv 存入 ctx，backward 给 grad kernel 传了 None，导致：
   cmpResidual cannot be null when CFA/SCFA and cmpMaskMode=3（同文件:565）。
   → forward 存入 ctx、backward 透传。

3. forward 内部 return_softmax_lse 硬编码 False，导致反向 LSE 无效 → 梯度爆炸
   forward kernel 调用末尾固定 ..., 1, False（注释 "kernel currently only supports False"），
   于是 softmax_lse = result[1] 是无效/零值，却被存起来喂给 grad kernel；
   结果反向梯度 ≈ 9e9，首 step 即 nan。
   → 把 forward 内部改为 return_softmax_lse=True（DFunction 反向恒需有效 LSE，与外部是否返回 LSE 无关）。改后问题 2 的即时 nan 消失（但 SWA 反向仍残留发散，见问题 2）。

另：SparseFlashMlaGrad not support deterministic yet

开启 deterministic=True 时 grad kernel 直接报不支持（sparse_flash_mla_grad_tiling_general.cpp:414）。
当前只能 deterministic=False 训练。若需对齐确定性精度，反向需支持 deterministic。

状态与优先级

1. 问题 1（npu_lightning_indexer padding 槽写 inf + return_value=False 时 0-size buffer 越界）
   —— 已在 hp 绑定层 lightning_indexer_v2.cc 修复（始终分配完整 value buffer）；
   底层 aclnn kernel 仍在 padding 槽写 inf，属质量问题但已被遏制，建议 kernel 侧一并修。
2. 问题 3（flash_mla wrapper 三处缺陷）—— 已合入 hp 源码 custom_op_impl.py 并重编译安装。
3. 问题 2（npu_sparse_flash_mla cmp_kv=None 反向对真实输入算错）—— kernel 侧 bug，待修；
   已用 repro_swa_real_grad.py 单算子最小复现（真实 dump 输入，d_query 3.8~45×）。mindformers
   侧可先 workaround（ratio0→小算子）解锁训练。（可选：grad kernel 支持 deterministic。）

结论修正：之前误判 cmp_kv=None 单算子反向"对齐"（基于随机输入 + amax 比值），现已证明
随机输入碰巧对、真实 RoPE'd MLA 结构下 cmp_kv=None 反向偏大 3.8~45×。cmp_kv 非空（ratio4/128）
反向正确。问题 1/3 已修；问题 2 待 kernel 修 cmp_kv=None 反向分支。

随附复现脚本（服务器 /home/w00906328/yaml/csa/）：

• 问题 1：repro_lightning_indexer_oob.py（单卡确定性复现 inf）
• 问题 2：repro_swa_real_grad.py + l0real.npz（真实输入触发 cmp_kv=None 反向 bug）；
  对照组（随机输入，反向正确）：repro_flash_mla_swa_grad.py / repro_flash_mla_cmp_grad.py /
  repro_flash_mla_sparse_grad.py / repro_l0_threeway.py
  运行需设 ASCEND_CUSTOM_OPP_PATH（指向 Ascend_official/vendors/custom_transformer 等）+
  LD_LIBRARY_PATH，见各脚本头注释 / 本仓 README 片段。