执行摘要

• 一句话：修复 ROCm 上 DeepSeek V4 功能与高并发精度问题
• 推荐动作：值得所有 ROCm + DeepSeek V4 用户关注。设计决策（AITER 回退、topk 统一入口）对类似平台适配有参考价值。建议阅读 rocm_aiter_mla_sparse.py 中的重构细节。

功能与动机

PR Body 指出：PR#41263 导致功能崩溃，且高并发时精度退化。AITER MHC kernel 在大 token 数下存在精度问题，prefill topk 索引缓冲区被错误逻辑破坏，因此需要回退 torch 实现并修复 topk 逻辑。

实现拆解

分为四个关键步骤：

1. 回退 AITER MHC 实现（mhc.py）：由于 AITER MHC kernel 在 token 数较大时存在精度问题，将 forward_hip 中条件分支（hidden_size % 256 == 0 走 AITER）永久注释，统一走 torch 实现，并添加 TODO 等待 AITER 修复后重新启用。
2. 修复 topk 索引逻辑（rocm_aiter_mla_sparse.py）：将 _topk_indices_torch 扩展支持 row_starts 参数，对齐 CUDA top_k_per_row_prefill 的语义（索引为行内偏移而非全局列偏移）。删除仅支持 topk_tokens=2048 的 _topk_indices_prefill 和 _topk_indices_decode 包装函数，直接调用 torch.ops._C.top_k_per_row_prefill（支持任意 topk 值）。同时清理 rocm_aiter_sparse_attn_indexer_native 函数接口，移除冗余包装器。
3. 简化 sparse attention indexer（sparse_attn_indexer.py）：移除之前通过 rocm_aiter_sparse_attn_indexer_native 的 fallback 路径，仅保留 AITER 算子（rocm_aiter_sparse_attn_indexer）。如果 AITER 未启用则直接抛出 RuntimeError，要求显式设置 VLLM_ROCM_USE_AITER=1。
4. 移除冗余 ffn_norm（deepseek_v4.py）：在 _forward_rocm 中删除 x = self.ffn_norm(x) 调用。FFN 的归一化已折叠到 self.ffn.norm_gate 中，ffn() 直接接收预归一化激活，此变更对齐了与 CUDA 路径的逻辑。

测试配套：PR 通过 lm-eval 在 GSM8K 上验证了 DeepSeek-V4-Pro 和 Flash，准确率达 95.5%/95.0%，置信度为回归无忧。

关键文件：

• vllm/model_executor/layers/mhc.py（模块 MHC层；类别 source；类型 core-logic；符号 MHCPreOp.forward_hip, MHCPostOp.forward_hip）: 核心修复：禁用有精度问题的 AITER MHC kernel，统一使用 torch 实现，影响 MHC pre/post 操作正确性。
• vllm/v1/attention/ops/rocm_aiter_mla_sparse.py（模块 稀疏索引；类别 infra；类型 refactor；符号 _topk_indices_torch, _topk_indices_prefill, _topk_indices_decode, rocm_aiter_sparse_attn_indexer_native）: 重构热点：修复 topk 索引对齐问题，删除包装函数，简化接口，影响 prefill/decode 稀疏索引正确性。
• vllm/model_executor/layers/sparse_attn_indexer.py（模块 检索器；类别 source；类型 core-logic；符号 SparseAttnIndexer.forward_hip）: 简化控制流：移除 native fallback，强制 AITER 路径，错误信息更明确。
• vllm/model_executor/models/deepseek_v4.py（模块 模型定义；类别 source；类型 core-logic；符号 DeepseekV4DecoderLayer._forward_rocm）: 修复功能崩溃：移除冗余 ffn_norm 调用，对齐 CUDA 路径。

关键符号：_topk_indices_torch, _topk_indices_prefill, _topk_indices_decode, rocm_aiter_sparse_attn_indexer_native, MHCPreOp.forward_hip, MHCPostOp.forward_hip, SparseAttnIndexer.forward_hip, DeepseekV4DecoderLayer._forward_rocm

关键源码片段

vllm/model_executor/layers/mhc.py

核心修复：禁用有精度问题的 AITER MHC kernel，统一使用 torch 实现，影响 MHC pre/post 操作正确性。

# vllm/model_executor/layers/mhc.py 中的 MHCPostOp 类（MHC 后处理操作）
# 修改后的 forward_hip：始终使用 torch kernel，不再根据 hidden_size 分流到 AITER
class MHCPostOp(CustomOp):
    def forward_hip(
        self,
        x: torch.Tensor,
        residual: torch.Tensor,
        post_layer_mix: torch.Tensor,
        comb_res_mix: torch.Tensor,
    ) -> torch.Tensor:
        # TODO: Reenable aiter after we are at the aiter
        # version that has this bugfix
        # https://github.com/ROCm/aiter/commit/b639cb63bcac4672dce33a731fad042a65cb3649
        # It has accuracy problem at large number of tokens.
        # 原条件分支（hidden_size % 256 == 0）被注释，直接走 torch 路径
        return mhc_kernels.mhc_post_torch(
            x,
            residual,
            post_layer_mix,
            comb_res_mix,
        )

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

重构热点：修复 topk 索引对齐问题，删除包装函数，简化接口，影响 prefill/decode 稀疏索引正确性。

# vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
# 修改后的 _topk_indices_torch：支持 row_starts 参数，输出行内本地索引
​
def _topk_indices_torch(
    logits: torch.Tensor,
    topk_tokens: int,
    row_starts: torch.Tensor | None = None,
) -> torch.Tensor:
    """Compute top-k indices using torch.topk, with optional row offset correction.
​
    When `row_starts` is provided, indices are local within each row's valid
    range (matching CUDA `top_k_per_row_prefill` contract).
    """
    k = min(topk_tokens, logits.shape[-1])
    values, indices = torch.topk(logits, k=k, dim=-1)
    indices = indices.to(torch.int32)
    if k < topk_tokens:
        # 填充无效列索引为 -1
        padding = torch.full_like(indices, -1, dtype=torch.int32)
        indices = torch.cat([indices, padding], dim=-1)
    if row_starts is not None:
        starts = row_starts.to(dtype=torch.int32).view(-1, 1)
        # 将绝对列索引转换为行内本地索引
        indices = torch.where(indices < 0, indices, indices - starts)
    if k == topk_tokens:
        return indices
    # 如果实际 topk 小于请求的 topk_tokens，则填充到指定大小
    padded = torch.full(
        (logits.shape[0], topk_tokens), -1, dtype=torch.int32, device=logits.device
    )
    padded[:, :k] = indices
    return padded
​
# 原有的 _topk_indices_prefill 和 _topk_indices_decode 函数已被移除，
# 因为它们仅针对 topk_tokens=2048 做了专用路径优化，而新版 C++ kernel 支持任意值。

评论区精华

1. Critical：topk buffer 切片大小不匹配（@gemini-code-assist[bot]）
   - 问题：解码路径用 num_decode_tokens 而非 num_padded_tokens 切片 topk_indices_buffer，当 padding 开启时导致 OOB 写入和后续 reshape 失败。
   - 结论：作者在后续 commit（"num_padded_tokens"）中修复了此问题，切片改为 num_padded_tokens。

2. 禁用 AITER MHC 的性能考量（@tjtanaa）
   - 作者说明：虽然回退到 torch 实现，但在 CUDA graph 和 torch.compile 模式下，该操作已被优化，性能优于常规 torch，并非瓶颈。

3. 清理包装函数的原因（@tjtanaa）
   - 自评：不再需要另一层包装，直接调用 C++ kernel 更简洁且支持所有 topk_tokens 值。

• topk buffer 切片大小不匹配 (correctness): 作者在后续 commit "num_padded_tokens" 中修复了该问题，最终代码使用 num_padded_tokens 切片。
• 禁用 AITER MHC 的性能考量 (performance): 团队接受该变通方案，等待 AITER 修复后重新启用。
• 包装函数清理 (design): 同意清理，减少间接层。

风险与影响

• 风险：
  • AITER MHC 回退风险：禁用 AITER 后，MHC pre/post 操作使用纯 torch 实现。尽管作者称 torch.compile 已优化，但在某些场景（如极端长度）仍可能比专用 AITER kernel 慢。需后续跟踪性能回归。
  • Sparse attn indexer 强制依赖 AITER：若用户未设置 VLLM_ROCM_USE_AITER=1，模型会直接抛 RuntimeError。环境依赖增强，但明确性提升。
  • Topk buffer 越界风险（已修复）：gemini 指出的 decode 路径切片问题在后续 commit 中已解决，审查时应确认最终代码使用 num_padded_tokens。
  • 无新增测试：本次改动未添加对应单元测试，回归覆盖依赖 manual lmeval。
• 影响：
  • 用户影响：DeepSeek V4/V4 Flash 在 ROCm 上功能和精度恢复，但需确保 AITER 已安装并设置环境变量。
  • 系统影响：删除 ~177 行代码（净减 89 行），降低维护成本；sparse_attn_indexer 路径简化，错误信息更明确。
  • 团队影响：与近期 PR#41710（清理 norm）方向一致，体现持续简化 DeepSeek V4 代码的风格。
  • 风险标记：AITER MHC 精度风险, Sparse attn indexer 强制依赖 AITER, Topk buffer size 潜在越界(已修复)

关联脉络

• PR #41263 Unavailable (PR#41263): 此 PR 引入了功能崩溃，本 PR 的第一项修复就是回退其影响。
• PR #41710 fix: remove unused norm for dpskv4: 同样涉及 DeepSeek V4 的 norm 清理，体现了团队持续简化代码的方向。