[akii96]

Purpose

Replace the hardcoded hidden_dim allowlist used for AITER fused allreduce+RMSNorm 1-stage dispatch with the kernel's actual pack-size constraint.

Today, vLLM only enables the AITER 1-stage AR+RMS path for this fixed set of hidden sizes: 512, 1024, 2048, 4096, 7168.

That allowlist is narrower than the kernel constraint and excludes valid model shapes. One concrete example is GPT-OSS, where the relevant hidden size is 2880.

With bf16/fp16 activations, AITER's 1-stage path uses 16-byte vectorized packs:

The predicate mirrors dispatchFusedAllReduceRMSNorm in AITER v0.1.13, the version pinned by docker/Dockerfile.rocm_base:
https://github.com/ROCm/aiter/blob/cdcfa833bdf554ca75594c90dde4316ea9b50199/csrc/include/custom_all_reduce.cuh#L2553-L2554

• pack_size = 16 / sizeof(T)
• pack_size = 16 / 2 = 8

For GPT-OSS:

• hidden_dim = 2880
• 2880 % 8 == 0
• 2880 / 8 = 360 <= 1024

So hidden_dim=2880 satisfies the 1-stage kernel constraint, but the previous hardcoded allowlist still forced it away from that path.

This PR changes only the hidden_ok predicate. The existing token_ok and size_ok gates are unchanged, so 1-stage dispatch is still limited by the same token-count, world-size, and message-size checks as before.

The intent is not to add a GPT-OSS-specific special case. Instead, this makes the vLLM dispatch condition match the AITER 1-stage layout rule, allowing any hidden size that satisfies the same pack-size constraint to use the fast path.

Note on concurrency: the 1-stage AITER kernel has a hard token_num <= kMaxBlocks (=80) cap built into its launcher (allreduce_fusion_kernel_1stage_launcher); above that, vLLM correctly falls back to 2-stage via the unchanged token_ok gate. So the PR's perf benefit is largest at low-to-moderate concurrency where decode batches stay under that cap. At very high concurrency (e.g. 32+), the 1-stage path is gated out for most decode steps and the patch is effectively neutral — which is the intended behavior, not a regression.

Test Plan

ROCm MI355X serving validation with openai/gpt-oss-120b, TP=2, AITER unified attention, and fp8 KV cache. Compare median TPOT and output throughput before vs after on the same serving recipe (vllm bench serve, random ISL=1000 / OSL=100, concurrency=8, n=5).

for i in 1 2 3 4 5; do
  echo "=== run $i ==="
  vllm bench serve \
    --backend vllm \
    --dataset-name random \
    --ignore-eos \
    --host localhost \
    --port 8022 \
    --model openai/gpt-oss-120b \
    --ready-check-timeout-sec 6000 \
    --percentile-metrics ttft,tpot,itl,e2el \
    --random-input-len 1000 \
    --random-output-len 100 \
    --max-concurrency 8 \
    --seed 10 \
    --num-prompts 1024 \
    --num-warmups 16
done

Test Result

Test image: vllm/vllm-openai-rocm:nightly-28ee78af543c563a2fbf78829a7688120e4e4eb5
Arch: gfx950 (MI355X)
Model: openai/gpt-oss-120b
Serving setup: MI355X, TP=2, ROCM_AITER_UNIFIED_ATTN, --kv-cache-dtype fp8

Before PR

lm-eval GSM8K:

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.8931	±	0.0085
gsm8k	3	strict-match	5	exact_match	↑	0.1895	±	0.0108

vllm bench serve (n=5, conc=8):

Metric	Value
Median TPOT	4.45 ms
Output token throughput	1430 ± 1 tok/s

After PR

lm-eval GSM8K:

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.8939	±	0.0085
gsm8k	3	strict-match	5	exact_match	↑	0.1797	±	0.0106

vllm bench serve (n=5, conc=8):

Metric	Value
Median TPOT	4.30 ms
Output token throughput	1477 ± 8 tok/s

Delta

Metric	Δ
Median TPOT	-3.4%
Output token throughput	+3.3%

[gemini-code-assist]

Code Review

This pull request updates the logic for determining if a hidden dimension is supported in the ROCm AITER fused all-reduce RMS norm implementation within vllm/_aiter_ops.py. The previous hardcoded list of supported dimensions is replaced with a dynamic check that ensures the hidden dimension is compatible with 16-byte vectorized packs and does not exceed 1024 packs per row. I have no feedback to provide as there were no review comments.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[akii96]

@gemini-code-assist-user review please with the updated PR body

[gemini-code-assist]

Code Review

This pull request updates the _rocm_aiter_fused_allreduce_rmsnorm_impl function in vllm/_aiter_ops.py to generalize the hidden_dim check. Instead of using a hardcoded list of supported dimensions, it now calculates a pack_size based on the input element size and verifies that the hidden_dim is compatible with 16-byte vectorized packs and does not exceed 1024 packs per row. I have no feedback to provide.

[tjtanaa]

LGTM. @akii96 Really appreciate for reaching out and also having clear explanation of the situation regarding to technical details of this PR. AITER's internal logic keeps on changing and highlighting his helps expedite the review.

[akii96]

@tjtanaa yeah that needs some work from us still wondering if there is a better longer-term fix but for now this seems to unlock some perf so will proceed with this! Thanks for looking into this!

I think the leftover RTD CI failure just needs a re-run? Can you trigger that please? 😅