[mjkvaak-amd]

Overview

Removes _requires_qwen38_tep8_emulation, which forces OCP MX emulation for Qwen3.8 Flash Next's (amd/Qwen3.8-Flash-Next-Quark-MXFP4) routed experts at TEP8 on gfx950 when the fused shared expert is off. The guard was added because AITER padded the MoE width from 640 to 768 for that shape and the native W4A4 kernel lost accuracy. Since #55104 aligns the W4A4 intermediate to 128, Flash Next's width stays at 640 with no pad, and the native kernel is as accurate as emulation and about twice as fast.

Existing work:

• [ROCm][MXFP4] Align the AITER W4A4 MoE intermediate size to 128 #55104 removed the pad that the guard worked around, and on its own fixes Flash Next's accuracy at EP4 and TEP8 (measured below) and at TP1 (0.9651, [ROCm][MXFP4] Align the AITER W4A4 MoE intermediate size to 128 #55104 on nightly-43b4aaea).
• The first revision of this PR also changed the pad rounding at MoE tp_size == 1. With [ROCm][MXFP4] Align the AITER W4A4 MoE intermediate size to 128 #55104 merged, that is a no-op for Flash Next, so it is dropped here.

Claims

• At TEP8 with VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0, the only configuration where the guard fires, the native AITER W4A4 kernel matches emulation's accuracy (GSM8K 0.9651 vs 0.9629) at 2.3x the output throughput at 8k/1k concurrency 4 (362 vs 158 tok/s) and 1.9x at concurrency 32 (1399 vs 741).
• No change for the checkpoint card's recipe: with the fused shared expert each rank holds 65 local experts, the guard requires exactly 64, so it never fired there.
• No change for any other shape: the guard matched only E=512, H=2560, N=640 at EP8 on gfx950.

Validation

Unit tests (MI355X, gfx950, nightly-5555a2d8 with this PR): 5 passed.

pytest -v -s tests/quantization/test_gfx950_moe.py

test_qwen38_tep8_selects_native_aiter replaces the guard's two tests: the TEP8 shape now selects AITER_MXFP4_MXFP4 under both auto and aiter.

Accuracy, GSM8K 5-shot, all 1319 questions, MI355X, the checkpoint card's recipe, vllm/vllm-openai-rocm:nightly-5555a2d8 vs the same image with this PR, one run per cell. Exact match, higher is better:

config	stock	this PR
TEP8, VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0	0.9629 (emulation, guard hit)	0.9651 (native)
TEP8, the card's recipe	0.9644 (native, guard not hit)	unchanged
EP4 (TP4 + --enable-expert-parallel)	0.9636 (native)	unchanged

Flexible-extract; strict-match for the first row is 0.9636 → 0.9659. Standard error is about 0.005.

Throughput, TEP8 with VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0, emulation (stock) vs native (this PR), random dataset, one run per cell. Output tok/s is higher-is-better, mean TPOT lower-is-better:

ISL / OSL	Conc	Output tok/s, emulation → native	Mean TPOT (ms), emulation → native
8192 / 1024	4	158 → 362	24.2 → 9.5
8192 / 1024	32	741 → 1399	38.7 → 20.0
128 / 8	1	40 → 90	21.2 → 7.2

Reproduce:

# plus the checkpoint card's ROCm environment, identical in both arms
VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 \
vllm serve amd/Qwen3.8-Flash-Next-Quark-MXFP4 --tensor-parallel-size 8 --enable-expert-parallel \
  --max-model-len 262144 --no-enable-prefix-caching --trust-remote-code \
  --default-chat-template-kwargs '{"enable_thinking": false}'
lm_eval --model local-completions --tasks gsm8k --num_fewshot 5 --apply_chat_template \
  --model_args model=amd/Qwen3.8-Flash-Next-Quark-MXFP4,base_url=http://localhost:8000/v1/completions,max_gen_toks=2048,num_concurrent=32

Details

• Why the guard is stale. Its docstring says AITER pads N to 768 for this shape. Since [ROCm][MXFP4] Align the AITER W4A4 MoE intermediate size to 128 #55104 the W4A4 intermediate is aligned to 128, so 640 is stored as is, with no pad. The native kernel at 640 is accurate (first accuracy row), so CK-Tile's inter_dim % 256 scale issue does not show up for this shape.
• When the guard fired. It matches num_local_experts == 64. The fused shared expert adds one local expert per rank (ExpertMapManager), so with the card's recipe the count is 65 and the native kernel already ran; only VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 hit the guard.
• Only the ROCm MXFP4 oracle changes; nothing changes on CUDA.

Limitations and CI gaps.

• gfx950 only. The W4A4 MXFP4 MoE path does not run on gfx942.
• The end-to-end image applied the previous revision of this PR, which also changed the pad rounding at MoE tp_size == 1. With [ROCm][MXFP4] Align the AITER W4A4 MoE intermediate size to 128 #55104, Flash Next's pad there is 0, so that part did nothing in these runs; the guard removal is the same code.
• The unit tests run in the AMD kernels CI job. The accuracy and throughput runs happened outside CI.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

---

Co-authored-by: Cursor cursoragent@cursor.com

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

[mjkvaak-amd]

@simondanielsson: could you please review this? 🙏 Pinging also @eppaneamd for visibility.

[maeehart]

Looks good to me. In the pinned AITER v0.1.24.post1, the CK-Tile stage 1 already doubles the pad for the fused gate/up weight (n_pad_zeros = intermediate_pad // 64 * 64 * (2 if use_g1u1 else 1)), so the extra doubling at MoE tp_size == 1 masked live rows. The paired GSM8K results are convincing, with 596 to 648 problems fixed against 5 to 12 broken, and TP2, which this PR does not change, stays within noise. Removing _requires_qwen38_tep8_emulation is also justified. Native TEP8 matches the emulation accuracy (0.9666 against 0.9674) at 1.9x to 2.6x the output throughput, and the guard did not fire for the checkpoint card's recipe with fused shared experts.

I have one request. The old comment linked the AITER lines that this rounding follows. Can the aiter_moe_pads docstring point to the AITER code that does the doubling, and say whether the FlyDSL stage 1, which takes inter_dim_pad unchanged, expects the same undoubled value? That makes the next AITER bump easy to check. @tjtanaa, can you take the final look?

[mjkvaak-amd]

Looks good to me. In the pinned AITER v0.1.24.post1, the CK-Tile stage 1 already doubles the pad for the fused gate/up weight (n_pad_zeros = intermediate_pad // 64 * 64 * (2 if use_g1u1 else 1)), so the extra doubling at MoE tp_size == 1 masked live rows. The paired GSM8K results are convincing, with 596 to 648 problems fixed against 5 to 12 broken, and TP2, which this PR does not change, stays within noise. Removing _requires_qwen38_tep8_emulation is also justified. Native TEP8 matches the emulation accuracy (0.9666 against 0.9674) at 1.9x to 2.6x the output throughput, and the guard did not fire for the checkpoint card's recipe with fused shared experts.

I have one request. The old comment linked the AITER lines that this rounding follows. Can the aiter_moe_pads docstring point to the AITER code that does the doubling, and say whether the FlyDSL stage 1, which takes inter_dim_pad unchanged, expects the same undoubled value? That makes the next AITER bump easy to check. @tjtanaa, can you take the final look?

Thanks for the review 🙏 . I pushed c41620b, which extends the aiter_moe_pads docstring. It points at the AITER lines and says what each stage 1 kernel expects, so the next AITER bump has something to check.

[simondanielsson]

Thanks for the fix!

Suggestion: can we add a regression test for this? I.e. a test that checks that no weights are being overriden after calling aiter's moe op. The test should fail without the fix

[simondanielsson]

CC claims this is fixed by #55104 already, can you please check?

[mjkvaak-amd]

CC claims this is fixed by #55104 already, can you please check?

Thanks for the review! You were right about the accuracy. #55104 covers most of the original scope of this PR. What it left is _requires_qwen38_tep8_emulation. It still forces OCP MX emulation at TEP8 without fused shared experts, although the 768 pad it was written for is gone. Native at 640 matches emulation's accuracy at 1.9x to 2.3x the output throughput, so I've rescoped the PR to removing the guard. The body has the new numbers.

[simondanielsson]

Thanks for the simplification!

Suggestion: Can we add full niah and aime25 results to check that the non-emulated kernel indeed is numerically accurate?

Suggestion: can we check if we can add a test for this? I.e. verify that the qwen38 tp8ep moe shapes now match a reference implementation

[github-actions]

✅ @mjkvaak-amd, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[dllehr-amd]

/ci run

[simondanielsson]

Thanks for attaching the results.

Just following up on this, then we should be good to go

Suggestion: can we add a regression test for this? I.e. a test that checks that no weights are being overriden after calling aiter's moe op. The test should fail without the fix

[mjkvaak-amd]

Thanks for attaching the results.

Just following up on this, then we should be good to go

Suggestion: can we add a regression test for this? I.e. a test that checks that no weights are being overriden after calling aiter's moe op. The test should fail without the fix

Added the test. Could you have a look? 🙏

[mjkvaak-amd]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95345 for commit d1de800a4017.

[tjtanaa]

LGTM

[mjkvaak-amd]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #95345.

[simondanielsson]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95419 for commit 24ff8ee1116c.

[simondanielsson]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #95419.

[mjkvaak-amd]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95733 for commit 79b33c9a4c19.

[mjkvaak-amd]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95856 for commit 844f04bcfe73.