执行摘要

• 一句话：ROCm DSV4默认使用aiter mHC预/后算子
• 推荐动作：值得精读。该 PR 展示了如何通过能力检测优雅地实现多后端 kernel 选择和切换，保留了清晰的 fallback 链，是 ROCm 平台上性能优化的典型案例。评审中的讨论也体现了对代码可维护性和向后兼容性的权衡。

功能与动机

PR #43679引入了tilelang fused mHC内核，并留下钩子准备在aiter修复sqrsum竞争条件后切换回更快的aiter内核。现在aiter >=0.1.14已发布并包含修复，因此恢复aiter作为首选，以获得更好的性能。

实现拆解

1. mhc.py 后端分发重构：添加 HAS_AITER_MHC 模块级标志，基于 is_aiter_found_and_supported() 判断。在 MHCPreOp.forward_hip 和 MHCPostOp.forward_hip 中，将之前注释掉的 aiter 调用恢复为首选，并加入 hidden_size % 256 == 0 约束；若条件不满足则依次 fallback 到 tilelang fused 或 torch/triton 参考实现。

2. amd/model.py 控制流调整：导入 HAS_AITER_MHC，将 self.has_tilelang 替换为 self.use_fused_mhc = HAS_TILELANG_MHC and not (HAS_AITER_MHC and self.hidden_size % 256 == 0)。这样当 aiter 可用且尺寸满足时，use_fused_mhc 为 False，从而选择 unfused 路径（分别调用 hc_pre 和 hc_post）。相应地修改 forward 中的条件，确保在 unfused 路径下正确调用 hc_post。

3. amd/mtp.py 移除直接引用：删除 has_tilelang 属性，改用 self.mtp_block.use_fused_mhc 判断是否调用 hc_post，保持与 model.py 的一致性。

4. Dockerfile 依赖版本提升：将 AITER_BRANCH 从 v0.1.13 提升为 v0.1.14，确保 aiter 中 mHC 预/后算子的 sqrsum 竞争条件修复生效。

5. Review 反馈清理：根据评审意见移除了冗余的自我解释性注释，并保留 tilelang fused 代码路径供未来 aiter fused 内核启用。

关键文件：

• vllm/model_executor/layers/mhc.py（模块 自定义算子；类别 source；类型 core-logic；符号 MHCPreOp.forward_hip, MHCPostOp.forward_hip）: 核心变更文件，在 forward_hip 中启用 aiter mHC 预/后算子，新增 HAS_AITER_MHC 标志和控制流。
• vllm/models/deepseek_v4/amd/model.py（模块 模型定义；类别 source；类型 data-contract；符号 DeepseekV4DecoderLayer.init, DeepseekV4DecoderLayer.forward）: 调整了 fused/unfused 路径选择逻辑，用 use_fused_mhc 替代 has_tilelang，确保 aiter 路径下正确调用 hc_post。
• vllm/models/deepseek_v4/amd/mtp.py（模块 模型定义；类别 source；类型 data-contract；符号 DeepSeekV4MultiTokenPredictorLayer.init, DeepSeekV4MultiTokenPredictorLayer.forward）: 移除对 HAS_TILELANG_MHC 的直接引用，改为通过 mtp_block.use_fused_mhc 判断，与 model.py 保持一致。

关键符号：MHCPreOp.forward_hip, MHCPostOp.forward_hip, DeepseekV4DecoderLayer.init, DeepseekV4DecoderLayer.forward, DeepSeekV4MultiTokenPredictorLayer.init, DeepSeekV4MultiTokenPredictorLayer.forward

关键源码片段

vllm/model_executor/layers/mhc.py

核心变更文件，在 forward_hip 中启用 aiter mHC 预/后算子，新增 HAS_AITER_MHC 标志和控制流。

def forward_hip(
    self,
    residual: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    n_splits: int = 1,
    norm_weight: torch.Tensor | None = None,
    norm_eps: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    hidden_size = residual.shape[-1]
    # 优先选择 aiter mHC pre kernel，要求 hidden_size 能被 256 整除
    if HAS_AITER_MHC and hidden_size % 256 == 0:
        return torch.ops.vllm.mhc_pre_aiter(
            residual, fn, hc_scale, hc_base, rms_eps, hc_pre_eps,
            hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat,
        )
    # 其次选择 tilelang fused kernel（CUDA 或 ROCm 且非 gfx942）
    elif HAS_TILELANG_MHC:
        return torch.ops.vllm.mhc_pre_tilelang(
            residual, fn, hc_scale, hc_base, rms_eps, hc_pre_eps,
            hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat,
            n_splits, norm_weight, norm_eps,
        )
    # 最后使用 torch/triton 参考实现
    else:
        return self.forward_native(
            residual, fn, hc_scale, hc_base, rms_eps, hc_pre_eps,
            hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat,
            n_splits, norm_weight, norm_eps,
        )

评论区精华

Review 中主要讨论了几点：

• Dockerfile 版本提升协调：AndreasKaratzas 指出 ROCm 大版本提升有发布节奏，tjtanaa 回复称 PR 需要等待 aiter 升级到 v0.1.14 后继续，随后 tuukkjs 确认主线 aiter 已为 v0.1.16.post2，可继续。
• 保留 fused 代码路径：tjtanaa 建议不要删除 fused post+pre 代码路径，以便将来可以轻松启用 aiter fused 内核。该建议被采纳，保留了逻辑。
• 移除冗余注释：tjtanaa 建议移除 mhc.py 和 model.py 中的详细注释，因为代码自解释，且不保证向后兼容。最终按照建议清理了注释。

• MTP 路径验证：tjtanaa 要求对 MTP（多 token 预测）进行最终检查，作者额外验证了 MTP 的 gsm8k 精度，通过。

  • aiter 版本提升与依赖管理 (other): PR 在 aiter 升级到所需版本后继续，最终在 main 上 aiter 已满足条件。
  • 保留 fused post+pre 代码路径 (design): 保留 fused 路径，但当前使用 aiter 非 fused 路径（分别调用 pre/post）。
  • 移除冗余注释 (style): 移除了所有提议的注释。

风险与影响

• 风险：
  1. aiter 版本依赖：若用户使用的 aiter 版本低于 v0.1.14，aiter mHC 路径将不可用（HAS_AITER_MHC 返回 False），自动 fallback 到 tilelang 或 torch/triton，不会崩溃，但性能可能不达预期。
  2. ROCm 专用变更：CUDA 路径不受影响，所有修改均在 HIP 条件内或通过 current_platform.is_rocm() 隔离。
  3. Hidden size 约束：aiter mHC 内核要求 hidden_size % 256 == 0，若不满足会自动回退，不影响正确性。
  4. MTP 路径回归：MTP 层中移除 has_tilelang 改为依赖 use_fused_mhc，需确保所有控制流正确。作者验证了 MTP 精度和性能，风险较低。
  5. 注释删除不影响功能：仅删除注释，不影响代码执行。
     - 影响：对用户：ROCm 上运行 DeepSeek V4 的用户将在不修改任何配置的情况下获得 4.7%-7.2% 的吞吐量提升。对系统：无需额外环境变量，基于能力自动选择最佳路径，降低了运维复杂度。对团队：维护两套 mHC 后端路径（aiter 与 tilelang），但代码结构清晰，fallback 机制完善。对测试：本次未新增测试文件，但 PR 作者在 MI3xx 上进行了准确性（gsm8k）和性能对比测试。
     - 风险标记：依赖外部库 aiter >=0.1.14, ROCm 专用变更, hidden size 256 约束, MTP 控制流影响

关联脉络

• PR #43679 [Perf][Kernel] Add tilelang fused mHC post+pre kernel for DSV4: 引入了 tilelang fused kernel 并留下切换钩子，本 PR 是该钩子的后续实现，将默认路径切换到 aiter。
• PR #42735 [Perf][Kernel] Use bf16 shared staging in mHC pre TileLang kernel: 对 tilelang mHC pre kernel 的优化，与本 PR 的 tilelang 路径相关。