[AMD][Fix] Restore aiter-padded MoE weight dims for serialized checkpoints

[yichiche]

Motivation

dcd9014f15 ("[AMD][MXFP4] Reland Online MXFP4 quantization 2/N - FP8 to MXFP4 requantization on AMD GPUs", #28291) silently broke accuracy for serialized MXFP4 Quark MoE checkpoints on ROCm whenever the per-rank MoE intermediate size needs aiter padding.

On Qwen3.5-397B-A17B-MXFP4 at TP8 on MI355 this drops GSM8K from 0.99 to 0.54. The failure is not obvious garbage — output stays fluent English, but the model loses knowledge and frequently fails to terminate (runs to max_tokens), so a single smoke prompt does not reliably catch it.

Bisected over the 422 commits between the v0.5.16 branch point (8905cbd42f, good) and current main (fd28242b6, bad); git bisect names dcd9014f15 as the first bad commit, and its parent 696e8f80d1 is good.

Root cause

In QuarkW4A4MXFp4MoE.create_weights, #28291 changed the weight allocation from the aiter-aligned dims returned by get_moe_weight_sizes() to raw dims:

-  w13_weight = torch.empty(num_experts, w13_up_dim,  hidden_size // 2, ...)
-  w2_weight  = torch.empty(num_experts, hidden_size, w2_down_dim,      ...)
+  w13_shape  = (num_experts, 2 * intermediate_size_per_partition, hidden_size // 2)
+  w2_shape   = (num_experts, hidden_size, intermediate_size_per_partition // 2)

but the corresponding scales in the very same function were left on the padded dims:

   w13_weight_scale = torch.ones(num_experts, w13_up_dim,  hidden_size // OCP_MX_BLOCK_SIZE, ...)
   w2_weight_scale  = torch.ones(num_experts, hidden_size, (w2_down_dim * 2) // OCP_MX_BLOCK_SIZE, ...)

get_moe_weight_sizes() is still called and its weight_padded flag is still advertised through extra_weight_attrs (driving FusedMoE.use_padded_loading), so the layer claims padding it no longer has. The weights and their block scales therefore describe different intermediate sizes.

Why only TP >= 8

get_moe_weight_sizes() only pads when w2_down_dim = (moe_intermediate_size / TP) / 2 is not already a multiple of AITER_PADDING_SIZE (128). Below that threshold the new raw expressions happen to equal the padded ones, so the commit is a no-op; at and above it they diverge. For Qwen3.5 (moe_intermediate_size = 1024):

TP	per-rank inter	old (padded) w13 / w2	new (raw) w13 / w2	aiter inter_dim	affected
1	1024	2048 / 512	2048 / 512	1024 -> 1024	no
2	512	1024 / 256	1024 / 256	512 -> 512	no
4	256	512 / 128	512 / 128	256 -> 256	no
8	128	512 / 128	256 / 64	256 -> 128	yes
16	64	512 / 128	128 / 32	256 -> 64	yes

So the trigger condition is moe_intermediate_size / TP < 256. This is not specific to Qwen3.5 or to TP8 — any serialized MXFP4 Quark MoE whose per-rank intermediate is not already 128-aligned is affected.

There is a secondary, silent performance cost: the wrong inter_dim also drops the model off its tuned aiter FlyDSL MoE configs. Before the fix the server logs [fused_moe] no tuned FlyDSL config for (...), using 2stage default for every token bucket; after the fix it resolves tuned kernels again (using 2stage (kernelName1='flydsl_moe1_afp4_wfp4_bf16_...')).

Modifications

python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py — in QuarkW4A4MXFp4MoE.create_weights, use w13_up_dim / w2_down_dim from get_moe_weight_sizes() for the is_checkpoint_mxfp4_serialized branch so the weights match the scales again. The online-quantization branch (else) is untouched and keeps the shapes introduced by #28291.

Two lines of real change plus a comment explaining why the serialized branch must stay on the padded dims.

Accuracy Tests

Qwen3.5-397B-A17B-MXFP4, 8x MI355 (gfx950), TP8, --attention-backend aiter, SGLANG_USE_AITER=1 SGLANG_USE_AITER_UNIFIED_ATTN=1, identical aiter / sgl-kernel / flags across all rows — only the sglang Python tree differs.

GSM8K, 5-shot, --thinking-mode qwen-3, --max-tokens 32768, temperature 0.6, top-p 0.95.

sglang build	GSM8K (n=100)	aiter MoE inter_dim	tuned FlyDSL configs
v0.5.16	0.99	256	resolved
696e8f80d1 (parent of the bad commit)	0.975 (n=40)	256	resolved
dcd9014f15 (first bad commit)	0.525 (n=40)	128	none
main fd28242b6	0.54	128	none
main + this fix	0.98	256	resolved

Qualitative check on the same builds, prompt "What do you know about Kalevipoeg?" (max_tokens 10000):

build	result
v0.5.16	correct ("the national epic of Estonia"), 2177 tokens, finish_reason=stop
main fd28242b6	denies the term exists, rambles to the 10000-token cap, finish_reason=length
main + this fix	correct, 2016 tokens, finish_reason=stop

Full bisect ledger (GSM8K n=40 per point, same harness): 8905cbd42f 0.975 / 66ea79dc2e 0.975 / 696e8f80d1 0.975 — all inter_dim 256; dcd9014f15 0.525 / f369a820d4 0.525 / 1a53f231d6 0.575 / 6f55de0468 0.525 / d708969f68 0.500 / eb242b6c03 0.525 — all inter_dim 128.

Speed Tests and Profiling

No dedicated speed run; this is an accuracy fix. It does restore the tuned aiter FlyDSL MoE kernel selection that the regression disabled, so it should recover MoE throughput lost since #28291 as well. Counting kernel-selection log lines on an identical TP8 startup:

build	tuned FlyDSL selections	no tuned FlyDSL config fallbacks
v0.5.16	112	0
main fd28242b6	0	112
main + this fix	112	0

Happy to add a throughput sweep if reviewers want it.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Note on unit tests: this regression was invisible to the existing suite because the weight and scale shapes are each individually self-consistent — only their relationship is wrong, and only at TP >= 8. A cheap guard would be a shape assertion in create_weights tying the weight dims to the scale dims (deliberately left out of this PR to keep it minimal, happy to add it here or in a follow-up). test/registered/quant/test_quark_mxfp4.py added by #28291 does not cover the serialized + aiter-padded path.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers. cc @fxmarty-amd (author of [AMD][MXFP4] Reland "Online MXFP4 quantization 2/N - FP8 to MXFP4 requantization on AMD GPUs" #28291), @BowenBao, @HaiShaw
3. Trigger CI tests with comments or contact authorized users to do so.
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #30645224321
Latest PR Test (Extra): ❌ Run #30645408755

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[yichiche]

@fxmarty-amd Please help review if this PR fix the issue properly. We observe the accuracy loss only at TP8, and the reason is as below.
Why only TP8?

• Trigger condition: moe_intermediate_size / TP < 256. Qwen3.5 has moe_intermediate_size = 1024, so TP ≤ 4 is safe and TP ≥ 8 breaks — exactly why your TP2 reference script was clean. It generalizes: any serialized-MXFP4 quark MoE with a small enough per-rank intermediate is affected, so "TP8 + unified attention" understates the blast radius on one axis and overstates it on the other.
• Not the attention backend — pure MoE weight-shape arithmetic. get_moe_weight_sizes() only pads when w2_down_dim = (moe_intermediate_size / TP) / 2 isn't already a multiple of AITER_PADDING_SIZE (128). Below that threshold the raw dims dcd9014 introduced happen to equal the padded ones, so the commit is a no-op:

  ┌───────────┬──────────────────┬─────────────────────────────┬──────────────────┬─────────────────┬────────┐
  │    TP     │  per-rank inter  │     old (padded) w13/w2     │ new (raw) w13/w2 │ aiter inter_dim │        │
  ├───────────┼──────────────────┼─────────────────────────────┼──────────────────┼─────────────────┼────────┤
  │ 1 / 2 / 4 │ 1024 / 512 / 256 │ 2048/512, 1024/256, 512/128 │ identical        │ unchanged       │ no-op  │
  ├───────────┼──────────────────┼─────────────────────────────┼──────────────────┼─────────────────┼────────┤
  │ 8         │ 128              │ 512/128                     │ 256/64           │ 256 → 128       │ breaks │
  ├───────────┼──────────────────┼─────────────────────────────┼──────────────────┼─────────────────┼────────┤
  │ 16        │ 64               │ 512/128                     │ 128/32           │ 256 → 64        │ breaks │
  └───────────┴──────────────────┴─────────────────────────────┴──────────────────┴─────────────────┴────────┘

[fxmarty-amd]

LGTM thanks a lot, apologies.

For context, #21097 was merged without unit tests capturing this failure, and its development collided with:

• #18182 / #28291

where I obviously did not properly solve conflicts.

[yichiche]

@amd-bot ci-status

[HaiShaw]

@BowenBao @fxmarty-amd would you please add more UTs to cover ongoing change validations, using following-up PR?