Repository navigation
[Perf] [AMD] qwen3_5: Quantize the BF16 GDN projections to MXFP4 online - #43839
Merged
Merged
Conversation
The Quark MXFP4 Qwen3.5 / Qwen3.8 checkpoints keep the GDN in_proj_qkvz and out_proj in BF16. in_proj_qkvz is the largest dense weight in each decode step, and no BF16 kernel beats hipBLASLt on its shape. Add SGLANG_ROCM_GDN_ONLINE_MXFP4. The value is a list of the BF16 GDN projections (in_proj_qkvz, out_proj) to quantize after load. The projections use the AITER preshuffled MXFP4 GEMM for 64 rows or less when AITER has a tuned config. Other inputs use the AITER ASM a4w4 GEMM. The default is empty, which keeps the BF16 path. Co-authored-by: Cursor <cursoragent@cursor.com>
github-actions
Bot
added
the
documentation
The test skips on any GPU that is not gfx95. The stage-b-test-1-gpu-small-amd suite runs on MI300, so the test never ran in CI. Move it to the MI35x suite. Raise ValueError for an unknown SGLANG_ROCM_GDN_ONLINE_MXFP4 value. An assert does not run under python -O and gives no useful message. Co-authored-by: Cursor <cursoragent@cursor.com>
omirosh
changed the title
Rewrite the comments, the docstring, the error message and the docs row for SGLANG_ROCM_GDN_ONLINE_MXFP4 as full sentences in the active voice. The code does not change. Co-authored-by: Cursor <cursoragent@cursor.com>
omirosh
marked this pull request as ready for review
October 17, 2026 12:29
omirosh
requested review from
1am9trash,
BowenBao,
HaiShaw,
JustinTong0323,
hubertlu-tw,
kkHuang-amd,
sogalin,
wisclmy0611,
yichiche and
zijiexia
as code owners
October 17, 2026 12:29
Quantize the online MXFP4 GDN projections in process_weights_after_loading, not in prepare_before_cuda_graph_capture. A monokernel spec then sees the MXFP4 weights when install_monokernel_forward runs at the end of load_model. The KV pool also gets the memory that the BF16 weights free. Co-authored-by: Cursor <cursoragent@cursor.com>
5 tasks done
Contributor
Preview approval requiredThis pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit |
omirosh
force-pushed
the
rocm/qwen38-gdn-online-mxfp4
branch
from
October 22, 2026 11:26
488c832 to
48a6dcc
Compare
| @@ -1877,14 +1901,24 @@ def get_input_embeddings(self): | |||
|
|
|||
| def prepare_before_cuda_graph_capture(self, model_runner) -> None: | |||
| if _use_aiter and self.config.model_type in _QWEN3_5_ROCM_PACKED_MODEL_TYPES: | |||
| packed = 0 | |||
| packed = mxfp4 = 0 | |||
There was a problem hiding this comment.
I don't like this syntax, but I am told it is correct.
Contributor
Author
There was a problem hiding this comment.
Agree that 2 lines are more clear to read. Changed it in c6f7fb9
| # Comma-separated BF16 Qwen3.5/3.8 GDN projections ("in_proj_qkvz", | ||
| # "out_proj") to quantize to MXFP4 after loading. Needs gfx95x and | ||
| # SGLANG_USE_AITER=1. | ||
| SGLANG_ROCM_GDN_ONLINE_MXFP4 = EnvStr("") |
There was a problem hiding this comment.
This is not very usable. If agent can find this env variable and use it, that's ok.
For human the list of strings is more complicated than "in", "out" or "all".
Also, if results are always better, should there rather be an optional parameter for disabling it?
Contributor
Author
There was a problem hiding this comment.
Agreed, made this variable a boolean and on by default for Qwen3.5/3.8 and for N, K multiples of 256 in 24b5622
Co-authored-by: Cursor <cursoragent@cursor.com>
SGLANG_ROCM_GDN_ONLINE_MXFP4=1 quantizes in_proj_qkvz and out_proj to MXFP4. The comma-separated list is removed. All PR results use both projections. An old list value now raises ValueError from EnvBool. Co-authored-by: Cursor <cursoragent@cursor.com>
SGLANG_ROCM_GDN_ONLINE_MXFP4 is now true by default. Set it to 0 to keep the BF16 projections. The variable has an effect only on gfx95x with SGLANG_USE_AITER=1, so the warning for other platforms is removed. Only the Qwen3.5 / Qwen3.8 model types (qwen3_5_text, qwen3_5_moe_text) get the MXFP4 projections. Models that reuse Qwen3_5GatedDeltaNet, such as qwen4_exp and interns2_mobius, keep BF16. A projection with N or K not a multiple of 256 also keeps BF16, so a default run does not stop on the gemm_a4w4 shape assert. Co-authored-by: Cursor <cursoragent@cursor.com>
5 tasks done
HaiShaw
approved these changes
Oct 27, 2026
prepare_before_cuda_graph_capture counts the MXFP4 GDN projections with getattr(module, name). The test_rocm_prepare_hook_is_limited_to_qwen35 case of test_qwen35_gdn_packed_in_proj.py passes a Mock with the Qwen3_5GatedDeltaNet spec. That Mock has no in_proj_qkvz, so the call raises AttributeError on MI35x CI, where the variable is on by default. Use getattr(module, name, None). A missing projection counts as BF16. Co-authored-by: Cursor <cursoragent@cursor.com>
This was referenced Oct 28, 2026
This was referenced Oct 29, 2026
Open
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
This ROCm performance change quantizes the BF16 GDN
in_proj_qkvzandout_projof Qwen3.5 / Qwen3.8 to MXFP4 once after load and runs them on AITER MXFP4 GEMMs instead of BF16 hipBLASLt, because the Quark MXFP4 checkpoints keep both projections in BF16, and cuts median TPOT by 5.1% with NEXTN MTP and 4.0% without MTP for Qwen3.8-2.4T-A95B MXFP4 on 8x MI355X, TP8, geomean over 12 serving shapes on the SGLang main nightly with only this PR added. Output tok/s goes up by 5.9% and 4.2%. The change is on by default on gfx95x withSGLANG_USE_AITER=1.SGLANG_ROCM_GDN_ONLINE_MXFP4=0turns it off.At TP8,
in_proj_qkvzis a 4608x8192 weight per GPU and the largest dense weight that every decode and verify step reads. A retuned BF16 GEMM table for these shapes gave no end-to-end gain (99.8% of the current table), and no BF16 kernel we tried beat hipBLASLt on 4608x8192. An MXFP4 weight with its scales is about a quarter of the bytes. At 16 rows, the MXFP4in_proj_qkvzand the BF16in_proj_batake 8.1 + 7.9 us, and the packed BF16 GEMM of both takes 23.1 us.Modifications
environ.pyadds the booleanSGLANG_ROCM_GDN_ONLINE_MXFP4. The default is true.0keeps the BF16 projections.quark_w4a4_mxfp4.pyaddsAiterMxfp4OnlineLinearMethod. A model sets it on a BF16 linear layer aftercreate_weights. The loader then callsprocess_weights_after_loading(), which quantizes the loaded BF16 weight with the existingdynamic_mxfp4_quant, preshuffles the weight (16x16) and the E8M0 scales, and records whether AITER has a tuned preshuffled config for (N, K). A second call does nothing.apply()quantizes the activation withper_1x32_f4_quant_hip. Up to 64 rows with a tuned config, it runs the Tritongemm_afp4wfp4_preshuffle. All other batches run the ASMgemm_a4w4(bpreshuffle=True).qwen3_5.pyreads the variable at import. It has an effect only on gfx95x withSGLANG_USE_AITER=1, and only for the Qwen3.5 / Qwen3.8 model types (qwen3_5_text,qwen3_5_moe_text). There, each GDN layer setsAiterMxfp4OnlineLinearMethodonin_proj_qkvzandout_projin__init__, so the model loader quantizes them inload_model, before SGLang sizes the KV cache and captures the CUDA graphs. It skips projections that the checkpoint already quantized, and projections with N or K not a multiple of 256. Models that reuseQwen3_5GatedDeltaNet, such asqwen4_expandinterns2_mobius, keep BF16. An MXFP4in_proj_qkvzis no longer anUnquantizedLinearMethod, sofinalize_fused_in_proj()inprepare_before_cuda_graph_capturedoes not build the packed BF16 qkvz+ba GEMM, andin_proj_ba(N=32) runs alone in BF16.test/registered/amd/test_qwen35_gdn_online_mxfp4.pychecks both TP8 shapes at 1, 4, 16, 31, 32, 64, 65 and 256 rows against a dequantized FP32 reference. The rows cover both sides of the 32-row scale-layout split and of the 64-row kernel cut.register_amd_ciputs it in the MI35x suite, because it skips on anything other than gfx95x.environment_variables.mdxadds one row for the new variable.With
SGLANG_ROCM_GDN_ONLINE_MXFP4=0, nothing changes. No layer is quantized, andfinalize_fused_in_proj()builds the packed BF16 path as before. With the default, only the GDN projections of Qwen3.5 / Qwen3.8 on gfx95x with AITER change. Other platforms and other models do not change. The MI35x nightly accuracy tests of Qwen3.5-397B (BF16 and MXFP4) and Qwen3.8 MXFP4 now run the MXFP4 projections.The Triton path for
in_proj_qkvzneeds the tuned config from Rocm/aiter5998 (merged). The AITER pin on main,2bb549102b(#44379), includes it. With an older AITER, such ase7d2453f,in_proj_qkvztakes the ASM path at every row count. The output is still correct there, but we did not measure the speed.The MXFP4 projections do not support online weight updates, because quantization changes the weight shape and dtype at load. Set
SGLANG_ROCM_GDN_ONLINE_MXFP4=0on servers that receive weight updates. The memory that the BF16 weights free (about 5.5 GB per GPU on Qwen3.8 TP8) goes to the KV cache, because quantization runs before SGLang sizes the pool. The pool grows by 6.4%, from 6.65M to 7.07M tokens with NEXTN and from 7.39M to 7.86M without MTP. The speed tests below ran on the first commit, which quantized after the pool was sized. There, both servers of each pair have the same KV pool, and peak usage was 63% and 56%. The speedup does not come from memory.Accuracy Tests
The unit test passes on MI355X, with a relative error below 1e-2 for every shape and row count. It also checks that a second
process_weights_after_loading()call does not quantize the MXFP4 weight again. It ran in the benchmark job before the servers started.The end-to-end evals ran on the same servers as "Speed Tests and Profiling" below. BF16 is the same image with the variable off.
With plain MXFP4 on these projections, the model thinks for longer. We repeated the thinking-length check of vllm#62750 without MTP: thinking-on GSM8K, temperature 1.0, 8 reps of 1319 questions per arm. The MXFP4 answers are 10.7% longer than BF16 (95% CI 6.1% to 15.7%), and the accuracy goes from 0.9641 to 0.9597 (McNemar p = 0.0017). Where long reasoning matters more than the per-token speed, set
SGLANG_ROCM_GDN_ONLINE_MXFP4=0. A Hadamard-32 rotation in the MXFP4 activation quant removes most of the increase (Rocm/aiter6734, open). It is a follow-up PR after the AITER pin moves.Speed Tests and Profiling
8x MI355X, TP8, one node, all four servers in one job. The image is
rocm/sgl-dev:v0.5.22-rocm10-mi35x-20261016(SGLang main34c4ba14c8, AITER2bb549102b) plus the first commit of this PR,ad94c3c8f2, and no other change. The later commits make the variable a boolean that is on by default, limit it to the Qwen3.5 / Qwen3.8 model types and to shapes with N and K multiples of 256, change the CI suite and comments, and move the quantization into the load path. "Load path check" below compares the load path with the first commit. Page size 16,SGLANG_USE_AITER_UNIFIED_ATTN=1, fp8 KV cache,--mamba-ssm-dtype bfloat16, chunked prefill 16384, radix cache off, at most 256 running requests. NEXTN runs use 3 steps, top-k 1, 4 draft tokens,--enable-linear-replayssm-specand--speculative-draft-model-pathset to the model path.bench_serverandom dataset, range ratio 0.8, one run per cell.NEXTN MTP:
No MTP:
Median TTFT is lower in every cell too, by 8.6% (geomean) with NEXTN and 6.5% without MTP. The NEXTN accept length does not drop.
bench_servereports a lifetime average per server, and at the end of the run it was 3.14 for BF16 and 3.18 for MXFP4.The MXFP4 projections run the Triton kernel up to 64 rows and the ASM kernel above that. Prefill always runs ASM. A NEXTN verify step has 4 rows per request, so NEXTN decode runs ASM from concurrency 32 up, and decode without MTP from concurrency 128 up.
Per-call time on one MI355X GPU in the same image, CUDA-graph replays, in us. The MXFP4 times include the activation quant. The BF16
in_projcolumn is what main runs: one packed qkvz+ba GEMM up to 64 rows, and separate qkvz and ba GEMMs above 64 rows. With MXFP4,in_proj_bastays BF16 and runs alone.Load path check
The same nightly image plus this PR at
0542071e02(quantization inprepare_before_cuda_graph_capture) or at2c6ceb33c0(quantization in the load path). Both setSGLANG_ROCM_GDN_ONLINE_MXFP4=in_proj_qkvz,out_proj, the list form of those commits, and the flags above. The servers ran on one node in the order capture, load, load, capture, with and without NEXTN. 8k / 1k, onebench_serverun per cell and round. The table gives the mean of the two median TPOTs.Between two rounds of the same server, the median TPOT moves by up to 2.8%. Every server logs 138 MXFP4 projections. GSM8K (first 50 questions, 5-shot, chat, thinking, greedy, flexible-extract) gives 0.98 in all four runs without MTP. With NEXTN, it gives 0.98 and 0.96 for each commit. All eight runs miss the same question, and one round of each NEXTN server misses one more, a different question each time.
We ran only Qwen3.8-2.4T-A95B MXFP4 at TP8. Qwen3.5 checkpoints and other TP sizes are untested.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ✅ Run #40492362617
Latest PR Test (Extra): ❌ Run #40492362515
Latest PR Test (AMD ROCm 10): ❌ Run #40492362624