变更摘要

本 PR 为 DeepSeek-V4 模型实现了 TND（变长序列）布局下的 Context Parallelism + 前缀 KV 特性。核心改动包括：引入 LayerCompressMode 枚举统一压缩模式判断；在 kvallgather CP 路径上支持 TND 布局的前缀 KV gather、截取与重排；将多个融合算子按 TND/BSND 布局拆分以消除不必要的 chunk 操作；修正 TND 下 indexer loss 反向传播的 num_seqs 计算；以及新增 --ring-degree 参数和 tnd 参数传递链。

主要改动

• 引入 LayerCompressMode 枚举并重构压缩逻辑：在 g2_attention.py 中新增 LayerCompressMode 枚举（NO_COMPRESS / COMPRESSOR_ONLY / INDEXER），替代原有 compress_ratio 与魔数（1、4）的直接比较；同时支持 compressor 返回 None 时将 cmp_ratio 降级为 1 并跳过稀疏 attention。

• TND 布局下前缀 KV 的 gather 与重排：在 deepseek_utils.py 中新增 slice_tnd_kv_prefix、_compute_prefix_kv_cu_seqlens、_rearrange_prefix_kv 三个辅助函数；g2_attention.py 的 forward 中，TND 布局下先 gather 所有 rank 的 KV，再按本 rank 偏移截取前缀段并重排，同时派生本地 cu_seqlens_q / cu_seqlens_kv，使前缀 KV 的 cu_seqlens 与注意力算子期望匹配。

• kvallgather CP 路径按布局拆分为 TND/BSND 分支：在 dsa_indexer.py 中将 fused_sparse_attn_shared_kv_kvallgather、fused_lightning_indexer_with_compress_kvallgather、fused_sparse_flash_mla_with_indexer_loss_kvallgather 拆分为 _tnd 和 _bsnd 内部实现：TND 下直接透传不 chunk，BSND 下保留原有双段 chunk 逻辑。

• TND 布局下 indexer loss 的 num_seqs 修正：在 npu_sparse_flash_mla_with_indexer_loss.py 的 backward 中，TND 布局下通过 all_gather_into_tensor 收集各 rank 本地 token 数求和得到全局 num_seqs，替代原先错误的 query_index.shape[0]；同时将 grad_scale 存入 ctx 以避免多层反向时的类变量竞态。

• 新增 --ring-degree 参数与 tnd 参数传递链：在 context_parallel_feature.py 中新增 --ring-degree 参数（默认 1）；在 gather_from_sp_cp 中增加 tnd 参数，tnd=True 时跳过 permute_cp_shard 重排；在 all_gather_qk_weight_kvallgather 中增加 tnd 参数控制 TND 下跳过 key/weight 的 all-gather。

openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=MindSpeed-LLM&number=4842&platform=gitcode&codeHostingPlatformFlag=gitcode

代码审查

Now I have completed the review. Let me provide the closing summary.

审查总结

审查范围：共审查了 8 个变更文件，逐一核实了每个文件的变更内容及上下文。

发现问题统计：

• P2：1 个
• P3：1 个
• P0/P1：0 个

各文件审查结论：

整体风险评估：中等风险。P2 问题影响 COMPRESSOR_ONLY 模式（compress_ratio > 1 且 != 4）下的 compressor 训练，若该模式实际被使用则会导致 compressor 无法学习。建议在合入前确认 COMPRESSOR_ONLY 模式的使用场景，并根据结论调整 detach() 的作用范围。

⛔ 需要修改

/approve