执行摘要

• 一句话：放宽 AITER 1-stage AR+RMS kernel 准入条件
• 推荐动作：推荐合并。PR 逻辑清晰、影响局部、收益明确，且经过维护者批准。无需深入精读，但可作为 ROCm 上 AITER 集成中与内核约束对齐的简洁示例。

功能与动机

AITER 1-stage 内核的 pack-size 约束（16 字节向量化）比 vLLM 硬编码的白名单更宽松，导致多个有效模型尺寸（如 GPT-OSS 的 hidden_dim=2880）被排除在快速路径之外。此 PR 旨在消除这一不必要的限制，使 vLLM 的 dispatch 条件与 AITER 内核的实际布局规则对齐。

实现拆解

1. 替换 hidden_dim 检查：在 vllm/_aiter_ops.py 的 _rocm_aiter_fused_allreduce_rmsnorm_impl 中，将 hidden_ok = hidden_dim in (512, 1024, 2048, 4096, 7168) 替换为动态检查：若数据类型为 bf16 或 fp16，则计算 pack_size = 16 // element_size()，并检查 hidden_dim % pack_size == 0 且 hidden_dim // pack_size <= 1024；否则 hidden_ok = False。
2. 保持其他門控不变：token_ok（<=80 tokens）、world_size 和 size_ok 门控保持不变，确保仅在低并发（decode 阶段）启用 1-stage 路径。
3. 测试验证：在 ROCm MI355X 上使用 GPT-OSS 120B 模型进行 serving benchmark，结果显示 TPOT 降低 3.4%，吞吐量提高 3.3%。

关键文件：

• vllm/_aiter_ops.py（模块 ROCm 内核；类别 source；类型 core-logic；符号 _rocm_aiter_fused_allreduce_rmsnorm_impl）: 唯一变更文件，修改了 AITER 融合 allreduce+RMSNorm 的 hidden_dim 校验逻辑，从硬编码白名单改为动态 pack-size 约束。

关键符号：_rocm_aiter_fused_allreduce_rmsnorm_impl

关键源码片段

vllm/_aiter_ops.py

唯一变更文件，修改了 AITER 融合 allreduce+RMSNorm 的 hidden_dim 校验逻辑，从硬编码白名单改为动态 pack-size 约束。

# vllm/_aiter_ops.py
​
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"
​
    total_bytes = input_.numel() * input_.element_size()
    hidden_dim = input_.shape[-1]
    token_num = input_.shape[0]
​
    # PR#42409: 不再使用硬编码白名单，而是根据 AITER kernel 的 pack-size 约束动态判断。
    # AITER 1-stage 内核要求 hidden_dim 能被 16 字节向量化并 <= 1024 个 pack。
    # 仅在 bf16/fp16 下启用，其他数据类型保持关闭。
    if input_.dtype in (torch.bfloat16, torch.float16):
        pack_size = 16 // input_.element_size() # 16 字节 / 每个元素字节数
        hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
    else:
        hidden_ok = False
​
    token_ok = token_num <= 80
    # ... 后续 world_size 和 size_ok 检查不变 ...
    use_1stage = hidden_ok and token_ok and size_ok
    result = aiter_ar.fused_ar_rms(
        input_,
        residual,
        w=weight,
        eps=epsilon,
        registered=torch.cuda.is_current_stream_capturing(),
        use_1stage=use_1stage,
    )
    assert result is not None
    return result[0], result[1]

评论区精华

review 由自动化 bot（gemini-code-assist[bot]）完成，未提出修改意见。维护者 tjtanaa 认可 PR 的技术细节和清晰说明，并予以批准。未发现争议点。

• 暂无高价值评论线程

风险与影响

• 风险：风险较低。变更仅影响 hidden_dim 的判断逻辑，且已通过 type guard（仅 bf16/fp16）和 pack-size 约束限制。token_ok 和 size_ok 门控保持不变，不会在高并发下意外启用 1-stage 路径。可能的风险是未来 AITER 内核 pack-size 变更时需同步更新此逻辑，但该风险可通过定期依赖更新管理。

• 影响：对用户：使用 ROCm GPU 且模型 hidden_dim 不在旧白名单内的用户（如 GPT-OSS）将获得 3%+ 的 decode 性能提升。
  对系统：单文件、6 行变更，无外部依赖，无兼容性问题。
  对团队：维护成本极低，但需注意 AITER 内核版本升级时 pack-size 约束可能变化。

• 风险标记：暂无

关联脉络

• PR #42072 [ROCm] Restore fast top_k_per_row kernels for sparse MLA when topk_tokens=2048: 同为 ROCm kernel 优化，涉及 AITER 后端性能提升。
• PR #42509 [ROCm][MLA] FP8 ASM prefill for AITER dense MLA backend on gfx950: 同为 ROCm AITER kernel 性能优化，与本 PR 在 ROCm 内核优化目标上相似。