执行摘要

• 一句话：融合 mHC 后处理与前归一化 GEMM 内核，提升 DeepSeek-V4 推理性能
• 推荐动作：值得精读，尤其是 FMA 代替 tensor core 的融合策略和 TileLang 内核编写方法。review 中关于阈值和 UnboundLocalError 的讨论也值得关注，可作为代码审查的 checklist。

功能与动机

This PR adds a new mHC kernel, which fuses the hc_post operation with the prenorm_gemm portion of hc_pre. The approach is adapted from TRTLLM, and performs the GEMM using FMA (rather than tensor cores), improving speed at low concurrency.

实现拆解

1. 新增 TileLang 融合内核（vllm/model_executor/layers/mhc.py）：定义 mhc_fused_tilelang 函数，使用 TileLang JIT 编译。该内核在一个 grid 内完成：从 post_mix 和 comb_res_mix 计算新的残差流；利用 weight 矩阵对 new_r 做 GEMM 得到 yp_out；同时计算残差平方和（rp_out）并写回 residual_out。内核支持 split-k 和 tile-n 分块，通过 FMA 指令逐线程累加，最后做 warp reduce 和 cross-warp reduce。

2. Python 调度函数（mhc.py）：mhc_fused_post_pre 封装了内核调用，负责：根据 num_tokens 选择是否使用融合路径（阈值 <16 走原有分步路径）；对输入 tensor 做 flatten（合并 batch 和 sequence 维度）；根据 n_splits 选择 split-k 参数，调用 TileLang 内核；将输出拆包为 4 个 tensor 返回。同时提供 _mhc_fused_post_pre_fake 作为 fallback。

3. 模型集成（vllm/model_executor/models/deepseek_v4.py）：DeepseekV4DecoderLayer.forward 签名新增 post_mix, res_mix, residual 三个可选参数。当上层传入残差状态时，直接调用 mhc_fused_post_pre 替代原来的 hc_post + hc_pre 序列；否则（首层）执行独立的 hc_pre。DeepseekV4Model.forward 循环前初始化 residual, post_mix, res_mix = None, None, None，逐层传递状态，循环后调用最后一层（或首层）的 hc_post 完成最终输出。注意：当 start_layer == end_layer 时，for 循环为空，else 块中的 layer 变量存在 UnboundLocalError 风险（参见评论区精华）。

4. 测试覆盖（tests/kernels/test_mhc_kernels.py）：新增测试文件，包含 sinkhorn_normalize_ref, mhc_pre_ref, mhc_post_ref 三个参考实现（源自 TileLang 示例），使用纯 PyTorch 实现作为 ground truth。test_mhc_fused_post_pre 参数化测试（num_tokens∈[1,4,8,128], hidden_size∈[4096,7168], hc_mult=4），比较融合内核输出与参考实现的数值一致性，atol/rtol=1e-2。

关键文件：

• vllm/model_executor/layers/mhc.py（模块 mHC 内核；类别 source；类型 core-logic；符号 mhc_fused_tilelang, mhc_fused_post_pre, _mhc_fused_post_pre_fake）: 核心 kernel 实现，新增 mhc_fused_tilelang 和 mhc_fused_post_pre
• tests/kernels/test_mhc_kernels.py（模块 测试用例；类别 test；类型 test-coverage；符号 sinkhorn_normalize_ref, mhc_pre_ref, mhc_post_ref, test_mhc_fused_post_pre）: 全面测试融合 kernel 的正确性，包含参考实现和参数化测试
• vllm/model_executor/models/deepseek_v4.py（模块 模型定义；类别 source；类型 core-logic）: 修改 DeepseekV4DecoderLayer 和 DeepseekV4Model 以使用融合 kernel，影响模型推理流程

关键符号：mhc_fused_tilelang, mhc_fused_post_pre, _mhc_fused_post_pre_fake, test_mhc_fused_post_pre, mhc_pre_ref, mhc_post_ref

评论区精华

• 缺少 CUDA 绑定文件：gemini-code-assist[bot] 指出 csrc/torch_bindings.cpp 中引用了 mhc_fused_hc 但对应的 .cu 文件未包含。作者回复已修复。
• 阈值不一致：路径选择条件 num_tokens < 16 与另一处的 num_tokens <= 16 不一致，可能导致 num_tokens == 16 时走错路径。作者回复已修复。
• 非扁平张量传递：融合 kernel 内部假设张量首维为 num_tokens，但传入的可能是未扁平化的 4D/3D 张量，可能导致索引错误。作者回复已修复。

• UnboundLocalError：在 DeepseekV4Model.forward 中，如果层循环为空（如 start_layer == end_layer），else 块中引用的 layer 变量未定义。评论建议添加 if residual is not None 保护。截至合并，此问题在代码中仍可见，可能存在风险。

  • mhc_fused_hc CUDA 绑定缺失 (correctness): 作者回复 'fixed the comments here'，表明已移除相关绑定或添加文件。
  • 路径选择阈值不一致 (correctness): 作者回复 'fixed this'，表明已将条件统一为 <=16 或类似。
  • 非平坦张量传递给 kernel 可能导致索引错误 (correctness): 作者回复 'fixed this'，表明已修正为使用扁平化张量。
  • UnboundLocalError when loop empty (correctness): 评论建议添加 if residual is not None 检查，但 head 版本中仍使用 else 块，可能存在风险。截至合并未看到作者回复。

风险与影响

• 风险：
  • 核心路径变更：模型 forward 逻辑重写，可能影响流水线并行、MTP 等其他特性。
  • 新 kernel 编译依赖：TileLang kernel 依赖编译器运行时，在不支持的平台（如 ROCM）上可能无法编译或退化到 Python fallback。
  • 缺少 autotuning：tile_n 和 split_k 参数为硬编码默认值，未做自动调优，可能在不同模型配置下不是最优。
  • UnboundLocalError 潜在风险：当 start_layer == end_layer 时，else 块中 layer 变量未定义，在流水线并行等场景下可能触发运行时错误。
• 影响：
  • 用户影响：仅影响 DeepSeek-V4 模型推理，性能提升在低并发（concurrency=4）下显著，输出吞吐提升约 6%，TTFT 和 TPOT 均改善约 5-6%。用户无需修改配置或启动参数。
  • 系统影响：新增 TileLang JIT 编译开销，首次推理可能因编译略微延迟。
  • 团队影响：需维护 TileLang kernel 和 Python 调度代码，确保跨平台兼容。后续可能引入 autotuning 和预热逻辑。
  • 风险标记：核心路径变更, 新 kernel 编译依赖, 缺少 autotuning, UnboundLocalError 潜在风险

关联脉络

• PR #41694 [DSV4] Add PP support for deepseek-v4: 修改了同一个模型文件 deepseek_v4.py，两个 PR 都是对 DeepSeek-V4 的支持和优化，存在合并冲突风险。