执行摘要

• 一句话：修复 AITER AR+RMS 核函数的 token 数量计算
• 推荐动作：该 PR 值得合并，修复了一个明确的崩溃问题且改动量小、风险低。代码逻辑清晰，测试覆盖也相应增强。建议维护者关注测试中 Transformers 后端的 fusion 计数调整是否后续需要进一步完善。

功能与动机

Transformers 后端使用 3D 输入张量（batch, sequence, hidden），原 token 数计算 input_.shape[0] 只取了第一个维度，未考虑所有 leading 维度，导致 token_num 过小，错误地选择了 1-stage 核函数并触发 RuntimeError: "Token number is too large for allreduce_fusion_kernel_1stage kernel"。PR 旨在修复该 dispatch 逻辑并增加测试覆盖防止回归。

实现拆解

1. 修复 token 计算逻辑（vllm/_aiter_ops.py）：在 _rocm_aiter_fused_allreduce_rmsnorm_impl 和 _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl 两个函数中，将 token_num = input_.shape[0] 改为 token_num = input_.numel() // hidden_dim。这样无论输入是 2D 还是 3D 张量，都能正确计算 token 总数（所有 leading 维度的乘积），从而确保 1-stage vs 2-stage 核函数选择逻辑正确。
2. 扩展测试覆盖（tests/compile/fusions_e2e/test_tp2_ar_rms.py）：在 test_tp2_ar_rms_fusions 中新增 model_impl 参数，将 llama3_8b 的 transformers 后端加入参数化列表。同时增加 model_impl == "transformers" 时的条件逻辑：若非 ROCm 平台则跳过测试；并调整 matches 中的 fusion 计数（aiter_ar_rms_fusion=1），因为 Transformers 后端的 residual add 和 RMSNorm 尚未融合，需单独处理。

关键文件：

• vllm/_aiter_ops.py（模块 AITER 算子；类别 source；类型 core-logic；符号 _rocm_aiter_fused_allreduce_rmsnorm_impl, _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl）: 核心修复：修改 token 数计算方法，修复 3D 输入时核函数选择错误导致的崩溃。
• tests/compile/fusions_e2e/test_tp2_ar_rms.py（模块 融合测试；类别 test；类型 test-coverage；符号 test_tp2_ar_rms_fusions）: 测试配套：新增 Transformers 后端参数化，确保修复不会回归并覆盖 3D 输入场景。

关键符号：_rocm_aiter_fused_allreduce_rmsnorm_impl, _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl, test_tp2_ar_rms_fusions

关键源码片段

vllm/_aiter_ops.py

核心修复：修改 token 数计算方法，修复 3D 输入时核函数选择错误导致的崩溃。

def _rocm_aiter_fused_allreduce_rmsnorm_impl(
    input_: torch.Tensor,
    residual: torch.Tensor,
    weight: torch.Tensor,
    epsilon: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
    assert aiter_ar is not None, "aiter allreduce must be initialized"
    ca = aiter_ar.aiter_ca
​
    total_bytes = input_.numel() * input_.element_size()
    hidden_dim = input_.shape[-1]
    # 修复前 : token_num = input_.shape[0] # 只取第一个维度，3D 时错误
    # 修复后 : 用 numel() // hidden_dim 计算所有 leading 维度的 token 总数
    token_num = input_.numel() // hidden_dim
    if input_.dtype in (torch.bfloat16, torch.float16):
        pack_size = 16 // input_.element_size()
        hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
    else:
        hidden_ok = False
    token_ok = token_num <= 80
    world_size = ca.world_size
    full_nvlink = ca.fully_connected
​
    if world_size == 2:
        size_ok = True
    elif full_nvlink and world_size <= 4:
        size_ok = total_bytes < 256 * 1024
    elif full_nvlink and world_size <= 8:
        size_ok = total_bytes < 128 * 1024
    else:
        size_ok = False
​
    use_1stage = hidden_ok and token_ok and size_ok
​
    result = ca.custom_fused_ar_rms(
        input_, residual, weight, epsilon, use_1stage=use_1stage,
    )
    assert result is not None
    return result[0], result[1]

评论区精华

Claude bot 自动评论称 PR 来自 fork，需要 maintainer 触发人工 review。此外无其他 review 评论。hmellor 直接批准（APPROVED），无遗留讨论。

• 暂无高价值评论线程

风险与影响

• 风险：高风险较低。变更仅将 shape[0] 改为 numel() // hidden_dim，属于逻辑修复而非架构变动。但需确认所有调用处传入的 input_ 张量维度至少为 2（hidden_dim 是最后一个维度），若存在 1D 张量时 hidden_dim 可能等于 numel() 导致 token_num=1，但该场景在 fused allreduce 中不合理，风险可忽略。测试新增了 Transformers 后端的参数化覆盖，有助于捕获回归。
• 影响：影响范围有限。仅影响 ROCm 平台使用 AITER 且启用 AllReduce+RMSNorm fusion 的 Transformers 后端用户。修复后这些用户可正常使用 fused kernel 而不会崩溃。对 vLLM 原生后端和其他平台无影响。
• 风险标记：暂无

关联脉络

• 暂无明显关联 PR