[ROCm][MoE] Pad the AITER MoE intermediate size at allocation time, and round the expert-group count to a kernel that exists

[sshlyapn]

Purpose

Two independent startup failures on the ROCm AITER MoE path, both caused by shapes AITER has no compiled kernel for, both focusing on IFM/K2-Horizon-375B-A23B model support.

1. Unaligned per-partition intermediate size - unquantized path only, fp8 AITER alignment is unchanged

AITER's CK 2stages MoE kernel dispatches on inter_dim <= 192: below it both stages use 64-wide tiles, above it at least one stage uses a 128-wide tile. CK's IsSupportedArgument rejects an intermediate not divisible by that width, so moe_intermediate_size = 1792 at TP=8 gives 224 and fails with device_gemm ... does not support this GEMM problem.

Fix: round the intermediate up in maybe_roundup_sizes, the existing hook for this (already overridden by five other quant methods). It runs before create_weights, so the padded size reaches allocation and the loader — which derives shard offsets from the checkpoint, not the parameter shape — fills only the real rows.

Alignment rule: 64 if inter_dim <= 192 else 128, mirroring AITER's own threshold rather than always using 128. A flat 128 mispredicts 2 of the 19 sizes measured against the real kernel (64 and 192 both pass despite not being 128-aligned) — it is a wrong model that happens to be safe. It also costs: inter_dim is a tuned-config lookup key, and across the 148 shipped AITER config files (11279 rows) the threshold rule pads 0 tuned rows while a flat 128 would move 207 off theirs onto the heuristic fallback. Rounding up never crosses the boundary (96->128, 160->192, 192->192 all stay <= 192), so the alignment picked is always the one the dispatcher then applies. Dtype-independent: CK derives both tile widths from sizeof(A0DataType), identical for fp16 and bf16.

Zero-init. The loader narrows to the checkpoint extent and never writes the tail, so w13/w2 are allocated with torch.zeros when either dim is padded — unpadded layers keep torch.empty and pay nothing. Pad lanes are then bit-exactly inert: silu(0) * 0 = 0, and the zero columns in w2 contribute exactly zero to stage 2. With torch.empty the tail holds recycled device memory, which is live weight to the kernel; whether that hurts depends on load-time allocation order, so it presents as sporadic implausible tokens with no error. mxfp4 and quark nvfp4 already allocate zeroed; #55251 is doing the same for fp8.

Scope. Gated on UnquantizedMoeBackend.AITER, only selected on ROCm with AITER MoE enabled; every other unquantized backend falls through the base hook untouched. The quantized AITER paths have their own overrides with different constants and are not touched — fp8 needs 256.

Known interaction. Since intermediate_size_per_partition now differs from ..._unpadded, experts/rocm_aiter_moe.py computes a non-zero intermediate_pad where it previously got 0. Traced: with QuantType.No + bf16 the dispatch falls past every cktile/flydsl branch (all require q_type == per_1x32) to the plain ck_moe_stage1/ck_moe_stage2_fwd partials, which take no pad argument, and intermediate_pad only enters the tuned-config key when config_file is not None, which vLLM never sets. So it is computed and discarded, and the kernel runs the full padded GEMM including the zero rows — up to ~14% extra compute on the MoE GEMMs (448->512 at TP=4, 224->256 at TP=8).

2. AITER biased_grouped_topk expert-group count

_aiter_get_num_expert_group ceil-divides num_experts by the 32-per-group limit then walks to the next divisor, which can produce a count with no compiled kernel: AITER instantiates biased_grouped_topk only for NUM_GRP in {1,2,4,8}, so 96 experts gives 3 and fails to launch.

Grouping is a no-op here (topk_group == num_expert_group), so any divisor within the limit routes identically — the constraint is purely which kernel exists. Round to the largest supported count that divides num_experts and respects the limit. When none fits (320 -> 10, 33 -> 3) the naive value is kept deliberately: it is large enough that the topk >= num_expert_group guard at the call site fails and routing falls back to the generic path rather than reaching AITER.

Test Plan

Both files are CPU-only and gated on current_platform.is_rocm(), so they run in any ROCm CI job without a working AITER runtime.

tests/kernels/moe/test_rocm_aiter_moe.py — new "Weight alignment" section appended to the existing file rather than a standalone module, following that file's conventions (function-local imports, default_vllm_config, rocm_ naming):

• test_aiter_moe_alignment_follows_threshold — the rule, as a table across the boundary.
• test_aiter_moe_padded_size_stays_in_its_dispatch_branch — self-consistency: (intermediate <= 192) == (padded <= 192) and alignment(padded) == alignment(intermediate).
• test_aiter_moe_roundup_pads_intermediate — the hook end-to-end: 96->128, 192->192, 224->256, 448->512, 512->512.
• test_aiter_moe_roundup_is_not_applied_to_other_backends — TRITON and FLASHINFER_CUTLASS keep 224.
• test_aiter_moe_padded_weights_are_zero_initialized — pad lanes are zero, not allocator garbage.
• test_aiter_moe_aligned_weights_are_zero_initialized — zero-init is unconditional, so it holds for an aligned size too.
• test_aiter_moe_bias_matches_padded_weight — w13_bias.shape[1] == w13_weight.shape[1] with has_bias=True; the defect the post-load approach had.
• test_aiter_moe_padding_is_numerically_inert — reproduces the loader's gate/up placement into a zeroed padded parameter, asserts the MoE output is bit-identical to the unpadded reference.

tests/kernels/moe/test_rocm_aiter_num_expert_group.py (new) — five tests over _aiter_get_num_expert_group: the two invariants the router asserts unconditionally, that a supported count is chosen whenever one fits, a table of known expert counts, and that the fallback is taken deliberately when none fits.

Test Result

gfx950 (MI355X), torch 2.11.0+gitd0c8b1f, on 29af8bd672.

pytest tests/kernels/moe/test_rocm_aiter_num_expert_group.py -q   49 passed, 2 skipped
pytest tests/kernels/moe/test_rocm_aiter_moe.py -q                 1 failed, 63 passed
both files                                                         1 failed, 112 passed, 2 skipped  (72s)

The one failure, test_aiter_fused_moe_mi3xx_fp8_accuracy, is pre-existing and unrelated — verified by running it alone against an unmodified checkout of the parent commit in a separate worktree, same box and image, where it fails identically.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[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.

[tjtanaa]

@sshlyapn will following the MoE tuning guide (online tune AITER_ONLINE_TUNE=1 and offline tuning) in AITER repo resolve this issue? I would like to avoid adding overhead by padding unnecessarily if there are kernels that are actually able to handle the size when there are issues.

Line where ONLINE TUNE feature exists.
https://github.com/ROCm/aiter/blob/eec768a47c1da52e175f9089c640e695f8b572a3/aiter/fused_moe.py#L2440

[sshlyapn]

@sshlyapn will following the MoE tuning guide (online tune AITER_ONLINE_TUNE=1 and offline tuning) in AITER repo resolve this issue? I would like to avoid adding overhead by padding unnecessarily if there are kernels that are actually able to handle the size when there are issues.

Line where ONLINE TUNE feature exists. https://github.com/ROCm/aiter/blob/eec768a47c1da52e175f9089c640e695f8b572a3/aiter/fused_moe.py#L2440

Thank you for comment!

Unfortunately no - tuning can't fix this, because the tile width isn't a tuned parameter. CK instance is selected from inter_dim directly (inter_dim <= 192 then 64-wide, else 128-wide), in both stage-1 and stage-2 dispatch, and CK requires inter_dim to be divisible by whichever N tile it picks, so no tuned config can land inter_dim=224 - IsSupportedArgument rejects it and we get "device_gemm ... does not support this GEMM problem". It's a kernel implementation requirement and it's scoped to the bf16/fp16 data types only.

[sshlyapn]

Hi @tjtanaa @AndreasKaratzas could you please take a look at this PR when you have a moment?

[simondanielsson]

This looks similar to #55251, can we adopt the same type of changes?

This also emphasizes that perhaps we should fix this on the aiter side instead. WDYT?

[ChuanLi1101]

Allocation-time pad + zero tails is the right place to fix the CK shape reject, and it avoids the bias-length bug of a post-load rebuild.

Please address before merge:

1. Rounding num_expert_group to a compiled {1,2,4,8} needs a numerical check vs the Python router (e.g. 96 experts), not only “the kernel exists”. Grouping is a no-op only if topk_group == num_expert_group still selects the same experts.
2. Call out (or fix) that the unquant CK path still GEMMs the padded intermediate (~14% extra on those GEMMs). Prefer AITER-side pad / the same pattern as [ROCm] Add gelu_tanh to the AITER fp8 fused MoE and zero-allocate the padded expert weights #55251 if we can avoid paying that in vLLM.
3. This hook is unquantized-AITER only — please note in the PR that fp8 AITER alignment is unchanged.

[sshlyapn]

@simondanielsson @ChuanLi1101, thanks for the comments!

This looks similar to #55251, can we adopt the same type of changes?

Yeah, the changes are really close, but they address different data types: unquantized BF16/FP16 vs FP8 in the mentioned PR. The logic itself looks the same, I just preferred not to zero out the buffers when the hidden dimension is also padded, as that's not what the MoE triggers/requires - so I'd prefer to keep the scope of the changes narrower

This also emphasizes that perhaps we should fix this on the aiter side instead. WDYT?

Agree, the better fix would definitely be improved unaligned-shape support in AITER. I've created a ticket to track the issue from AITER side: ROCm/aiter#5444. Meanwhile, since that might take some time, I'd suggest keeping this padding to provide functional support for models where such unaligned shapes might appear

Rounding num_expert_group to a compiled {1,2,4,8} needs a numerical check vs the Python router (e.g. 96 experts), not only “the kernel exists”. Grouping is a no-op only if topk_group == num_expert_group still selects the same experts.

Corresponding test has been added, thanks!

Call out (or fix) that the unquant CK path still GEMMs the padded intermediate (~14% extra on those GEMMs). Prefer AITER-side pad / the same pattern

Absolutely agree. Here is the AITER issue for further improvements in this direction: ROCm/aiter#5444

This hook is unquantized-AITER only — please note in the PR that fp8 AITER alignment is unchanged

Done