[ROCm][Kernel] W4A16 MoE prefill: M-aware moe_align block size (1.6x TTFT)

[stoivonen-amd]

Purpose

Closes the short-prefill TTFT regression on the hybrid_triton W4A16 MoE path (AIESW-36448). At 128 input tokens the Triton fallback was 69% slower than the exllama HIP kernel it replaced; after this series it is within 4%, and it is already 1.65x / 2.10x faster than exllama at 512 / 1024 tokens.

The cost is moe_align_block_size padding, not tiling. moe_align_block_size pads every touched expert up to a full alignment block, and fused_moe_kernel_gptq_awq's only early exit is

if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
    return

so an all-padding block that lands inside num_tokens_post_padded runs the entire K-contraction — it loads and dequantizes its full slice of INT4 weights and only discards the result at the store. (gemm2 does not even discard: it is invoked with num_valid_tokens = num_slots, so its padding rows are stored.) Kernel work is therefore proportional to num_tokens_post_padded, not to M_routed, and

npp = sum over touched experts of ceil(tokens_e / A) * A

With E=128 a short prefill touches nearly every expert with a handful of tokens each, so almost all of npp is padding. Measured on routing captured from live prefills of cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit (E=128, top_k=8, group_size=32), averaged over MoE layers 0–3:

tokens	experts touched	npp @ A=64	npp @ A=128	npp@128 / routed rows
6	33	2112	4224	88.0x
128	110	7120	14080	13.8x
512	122	9456	16576	4.0x
1024	122	12864	18752	2.3x
2048	124	20816	25984	1.6x

A 128-token prefill does 13.8x the GEMM work it needs to. _select_block_size_m returned a single alignment for every prefill batch with no M-awareness, so this was unavoidable.

Two commits, which must land together:

1. Retune the small-group-size Triton tiles. The group_size <= 64 branch of _triton_config (added in f47a359c96 for Qwen3-Omni) was derived from a fixed-M sweep at one captured problem size. num_stages=1 was inherited from the group_size=128 branch, whose "pipeliner regresses" note was measured at much narrower tiles; and gemm2's (BLOCK_SIZE_N=64, num_warps=4) was not the best tile at any M swept. Four fields change. No change to which kernel runs, to the alignment, or to the group_size > 64 path.

2. Make the alignment M-aware. _select_block_size_m gains a second tier: alignment 64 below one full padding block of routed rows per expert, 128 at or above it. This commit also carries the gemm1 tile that pairs with alignment 64 — that tile is worse than commit 1's at the old alignment, which is why the two cannot be split.

64 is a floor, not an arbitrary halving: with BLOCK_SIZE_M = A the block count is pinned to the number of touched experts regardless of A, so the per-block INT4 weight load + dequant is roughly constant and below 64 the M-tile can no longer amortize it. A=16 measures worse than A=32 everywhere.

The crossover sits where the mechanism says it should — one full block per expert — so the tier is selected on num_tokens * topk / E rather than a hard-coded token count, and it tracks E and top_k on other models.

Scope and safety

• The group_size > 64 path is untouched.
• Batches of 2048+ tokens take the new _LARGE tier and get byte-identical config to before, so f47a359c96's Qwen3-Omni tune — validated at --input-len 4096 — is unaffected. Only sub-2048 batches change, which is the regime that tune never measured.
• Alignment is not a numerical parameter: it regroups routed rows into blocks, and each row does the same full-K contraction against the same expert weights in the same order. Output is expected to be bit-identical, and this is verified rather than assumed (below).
• Workspace shrinks as a side effect: num_slots at 128 tokens 17280 → 9088, at 1024 tokens 24448 → 16256.

Test Plan

All on Strix Halo (gfx1151, Ryzen AI MAX+ 395), ROCm, cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit, float16.

1. Existing kernel suites — nothing in them is expected to change, since the touched path is MoE-only and the perf golden covers the linear path:

pytest tests/kernels/moe/test_hybrid_w4a16_moe.py
pytest tests/kernels/quantization/test_hybrid_w4a16_triton.py \
       tests/kernels/quantization/test_moe_gemm_w4a16.py
pytest tests/kernels/quantization/test_hybrid_w4a16_perf.py

2. Bit-identity. Greedy decode (temperature=0.0, enforce_eager=True, enable_prefix_caching=False) over 3 fixed prompts x 48 output tokens, with prompt lengths 128/128/512 chosen to land on the changed tier. Token ids dumped to JSON on each build and diffed.

3. End-to-end TTFT, three builds x three prefill sizes:

vllm-bench.py --model cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit \
  --num-prompts 10 --max-num-seqs 1 \
  --input-len {128,512,1024} --output-len 8 \
  --max-model-len $((inlen + 64)) --dtype float16 --no-cudagraph

--output-len 8 so the decode path (num_tokens <= MAX_SKINNY_BATCH_SIZE, untouched here) is exercised too. Each run's log is asserted to have loaded the vLLM install it was supposed to, since vllm-bench.py spawns vllm serve and console-script shebangs are absolute.

4. Kernel-level sweeps behind the constants: alignment x tile x M, joint gemm1+gemm2 (they share one alignment, so the pair is what matters), across 6..2048 tokens and three routing distributions — a uniform 128-expert spread plus topk_ids captured from real prefills.

Test Result

1. Existing kernel suites — all passed:

suite	result
tests/kernels/moe/test_hybrid_w4a16_moe.py	115 passed
+ test_hybrid_w4a16_triton.py + test_moe_gemm_w4a16.py	160 passed
tests/kernels/quantization/test_hybrid_w4a16_perf.py	140 passed

(Those runs also included 45 GPU-free unit tests pinning the alignment staircase and the BLOCK_SIZE_M | alignment invariant, which are not part of this PR; the counts above include them.)

test_hybrid_w4a16_perf.py failed one case (i2560-o19456-g32-hybrid-w4a16-bf16) on the first run. That case is the linear path (wvsplitk_int4 / hybrid_triton_w4a16, in_features/out_features, no experts) and is unreachable from an MoE-only change; it passed 3/3 in isolation and 140/140 on a repeat full run. Treated as a thermal/ordering flake.

2. Bit-identity — confirmed. The two JSON dumps (3 prompts x 48 greedy tokens = 144 tokens) are byte-identical before and after.

3. End-to-end TTFT, median ms, 10 prompts:

input-len	before	after 1/2	after 2/2	total	exllama HIP
128	205.11	184.84	126.39	1.62x	121.56
512	288.80	270.43	217.29	1.33x	357.90
1024	408.45	393.62	348.11	1.17x	731.35

Improvement is consistent across mean/median/P99. Decode is unaffected — median TPOT 17.80 → 17.16 at 128 tokens and 22.97 → 22.76 at 1024.

4. Kernel-level sweeps.

Joint gemm1+gemm2 vs alignment (us), real captured routing:

tokens	A=32	A=64	A=128	best
6	1509.2	876.1	1537.3	64 (1.75x vs A=128)
128	3485.3	2619.7	4655.7	64 (1.78x)
512	4391.1	3407.5	5453.1	64 (1.60x)
1024	6264.1	4754.0	6248.9	64 (1.31x)
2048	9838.3	8256.8	8355.5	64 (1.01x)

A=64 wins at every point from 6 to 1024 tokens under all three routing distributions (1.14–1.85x) and is a wash at 2048, which is where the _LARGE tier takes over.

This is the alignment, not the tile. Decomposed at 128 tokens (gemm1+gemm2, us):

config	alignment	us
pre-retune tile	128	5237.2
commit 1 tile	128	4755.6
commit 2 tile	128	4882.9
commit 2 tile	64	2619.7

The new gemm1 tile is worse than commit 1's at A=128 — it only pays once the alignment moves. Keeping commit 1's tile at A=64 would give back 1.10–1.76x of the win. BLOCK_SIZE_N=32, num_warps=2 is best at all 27 (tokens x routing) points at A=64, by 1.10x to 2.13x over the runner-up; gemm2's tile is unchanged between tiers (geomean 1.000, worst 1.004 at A=64).

Caveat: tuned on one GPU and one model family. The alignment tier is selected on routed rows per expert so it tracks E and top_k, but the gemm1 tile at the low alignment is a fixed choice.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors.

You ask your reviewers to trigger select CI tests on top of fastcheck CI.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀