Repository navigation
perf(qwen4-exp): keep 128-aligned MXFP4 expert intermediate unpadded - #2649
Conversation
The MXFP4 MoE method pads expert weights to a multiple of 256, which turns Qwen3.8-Flash-Next's moe_intermediate_size 640 into 768 at TP1: one sixth of every expert is zeros, ~12 GB of HBM that the KV cache could use. Pass pad_align=128 when the per-rank intermediate is already a 128-multiple (as Kimi-K3 / DeepSeek-V4 do) and keep the default otherwise. ATOM_QWEN4_MOE_PAD_ALIGN overrides the choice (e.g. 256 for the old layout). Co-authored-by: Cursor <cursoragent@cursor.com>
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
… 128 pad align The env override is not needed: 128 is used whenever the per-rank expert intermediate is already a multiple of 128 (TP1: 640), and every other shard keeps the MXFP4 default of 256. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Dropped The per-TP check stays instead of a flat Re-validated on MI355X (this branch merged into main
Unit tests: |
|
@valarLip could you take a look at this one? Updated today in Pre Checkin is green. The first run's Ruff job failed on a GitHub API 502, and its unit-test job hit the 3 |
Summary
Mxfp4MoEMethodpads expert weights to a multiple of 256, so Qwen3.8-Flash-Next'smoe_intermediate_size640 becomes 768 at TP1: one sixth of every expert is zeros, ~11.7 GB of HBM. Qwen4-Exp now passespad_align=128when the per-rank intermediate is already a 128-multiple (TP1: 640 stays 640), as Kimi-K3 and DeepSeek-V4 already do, and keeps the default otherwise. The freed memory goes to the KV cache (+25% blocks), which is what AgentX C16 is bound by.Depends on #2606 (to run the model with MTP; this PR does not touch MTP code).
Companion AITER PR (tuned MXMOE table keyed on the unpadded inter 640): ROCm/aiter#6708. This PR works without it (inter 640 then uses AITER's heuristic FlyDSL kernels, same speed as 768 with pad); with it the MoE kernels get faster as well.
Motivation
Qwen3.8-Flash-Next (
amd/Qwen3.8-Flash-Next-Quark-MXFP4) optimization on MI355X, https://github.com/wuhuikx/LLM-Booster/issues/263 (item 3, MoE tuning).Changes
atom/models/qwen4_exp.py:moe_pad_align(intermediate_size, tp_size)returns 128 whenintermediate_size // tp_sizeis a multiple of 128, elseNone(the MXFP4 default 256);Qwen4ExpSparseMoeBlockpasses it toFusedMoE(pad_align=...).pad_alignis only read by the MXFP4 MoE method, so BF16 / FP8 / NVFP4 experts are unaffected.tests/test_qwen4_exp_quantization.py: the choice for TP 1/2/4/8.No env var, no config knob: 2 files, +20 lines.
Per-rank intermediate for 640: TP1 640 (was 768), TP2 320, TP4 160, TP8 80 → not 128-multiples, so they keep the 256 padding (unchanged behavior; only TP1 was validated). The check is per TP on purpose rather than a flat
pad_align=128: TP2 / TP8 would become 384 / 128, shapes with no AITER aux kernel or tuned table, while the Qwen3.8 TP8 aux shape(513, 8192, 256, 11)that AITER main just added (ROCm/aiter#6468) relies on the 256 padding.Accuracy
GSM8K (lm-eval, real MTP verification, i.e. no forced acceptance; same script as the issue's baseline), TP1, MTP 3:
pad_align=1286ea9fabdmerged into main40db24c2)40db24c2, same node and run timeThe shared baseline of the issue measured 0.9757 on main; the differences are run-to-run noise.
Performance
Test config: image
rocm/atom-dev:latest(566aa3dc9d5d), ATOM53e4fb524, AITER948dd267f, MI355X, TP1, MTP 3, bf16 KV, block 16, prefix caching,--gpu-memory-utilization 0.88.Memory (server log):
AgentX C16 (
aiperf --scenario agentx --concurrency 16 --benchmark-duration 3600,semianalysis_cc_traces_weka_062126_256k, golden AL 2.32):Cache = GPU prefix, TP1, concurrency 16 for all rows. Runs were done two at a time on separate GPUs of one node (the "ATOM#2649 only" run and the second "+ aiter#6708" run ran side by side).
beforeruns differ by 7% in tput/GPU: that is the run-to-run noise.beforeruns. The extra 25% of KV blocks lifts the realized prefix hit from 62% to 84% on these long-prefix agentic traces, so far fewer prompt tokens are recomputed.beforeruns differ by only 1.3% there) matches the −10% decode-step kernel time; the tput gain is of the same size as the noise.How to use
Nothing to set: TP1 Qwen3.8-Flash-Next MXFP4 picks
pad_align=128automatically; other TP sizes keep the MXFP4 default.Test plan
black/ruff checkon the changed files (same aspre-checks.yaml).pytest tests/test_qwen4_exp_quantization.py tests/test_envs.pyin a ROCm container with torch + aiter: 137 passed (main: 132; the 5 new cases aremoe_pad_align). (In the CPU-only CI job this file is skipped by itsimportorskip("triton")guard.).github/scripts/run_unit_tests.sh): 3 failed / 9252 passed; the 3 failures (test_graph_marker×2,test_sparse_indexer_side_effect_op) fail on main too.6ea9fabdmerged into main40db24c2, AITER508736740with the ROCm 7.2.4-invalid hipBLASLt rows filtered out of the BF16 GEMM table), TP1, MTP 3, side by side with main on the same node: weight peak 172.44 vs 184.17 GB, available for KV 70.39 vs 58.66 GB (+11.73 GB), KV blocks 126,366 vs 101,062 (+25%); every MXFP4fused_moekey AITER logs (AITER_LOG_LEVEL=INFO) hasinter_dim=640(main: 768); full GSM8K as above (last two rows).pad_align=128GSM8K row were run on a validation tree (ATOM53e4fb524+ fix(qwen4_exp): load unscaled Quark MTP weights as BF16 #2606 + a two-line patch forcingpad_align=128), which gives the same expert layout as this PR's automatic choice (same weight peak and KV block count).Limitations / follow-ups
inter_dim=768for this model would stop matching at TP1; there are none upstream.