Skip to content

[Quantization][GLM-5.3] Let online quantization reach the BF16 attention projections - #64373

Open
suiyoubi wants to merge 1 commit into
vllm-project:mainfrom
CentML:aot/glm53-upstream-fp8-proj
Open

suiyoubi wants to merge 1 commit into
vllm-project:mainfrom
CentML:aot/glm53-upstream-fp8-proj

Conversation

@suiyoubi

Copy link
Copy Markdown
Contributor

Overview

GLM-5.3-Flash builds its BF16 attention projections (KDA, MLA, DSA indexer) with quant_config=None, so online quantization requested with --quantization-config '{"targets": ...}' silently skips them. This PR passes them only that explicitly requested online config, so they can run in MXFP8. Related: #59833, #55741 (see Details).

Claims

  • Online quantization now reaches the 112 BF16 attention projections of GLM-5.3-Flash when --quantization-config names them. On nvidia/GLM-5.3-Flash-NVFP4, the targets under Recommended configuration now quantize 196 layers (attention + shared experts) instead of 84 (shared experts only).
  • No behavior change otherwise: without --quantization-config nothing changes; the checkpoint's own method never applies to these BF16 weights; kv_b_proj stays unquantized; the --quantization mxfp8 / fp8_per_* shorthands keep these projections BF16, as today. No change to shared quantization code, no new flag.
  • Recommended, behavior-preserving configurations (MXFP8 on the 128 non-KDA projections for aggregated serving; MXFP8 on all 196 on prefill instances only for disaggregated serving) keep accuracy and output length unchanged on 7 benchmarks, and give +2% output throughput (aggregated) and up to +12% (disaggregated, prefill-bound) on GB300.
  • MXFP8 on all 196 projections in aggregated serving is faster still (+8% on GB300), but measurably shortens reasoning; this PR describes it as an opt-in with that disclosure.

Validation

Unit tests

pytest tests/models/glm5next

45 passed, 8 skipped (ROCm-only test_sparse_indexer_topk_dispatch cases), on one NVIDIA Blackwell GPU. The branch is on main @ cdd2750cd; its Python tree was installed over the cu134-nightly-edad8b03 image (no compiled code changed in between). CI selects these tests in the glm5next-unit-tests step (pytest -v -s models/glm5next, H200; the MI355 mirror step runs them on ROCm). pre-commit (ruff, mypy, typos, test tethering) passes on the changed files.

The new tests (tests/models/glm5next/test_online_quantization.py, 12 cases, ~1 min on one GPU) build real Glm5NextDecoderLayers (KDA and sparse MLA) with a real VllmConfig. Only the sparse MLA attention backend is stubbed, since it depends on the GPU and is built after every projection. They check:

  • each attention projection's quant method for a block-FP8 checkpoint alone, composed with targets, composed with linear, pure online targets, and the --quantization mxfp8 shorthand, with kv_b_proj, the indexer wk_weights_proj and the KDA conv1d always unquantized;
  • that FP8-checkpoint MLA weights (q_b_proj, and q_a_proj + kv_a_proj_with_mqa → fused_qkv_a_proj) loaded through _try_load_fp8_attn_proj and the real weight loader give exactly the same online-MXFP8 parameters as loading their BF16 dequantization directly.

On main without this change (same test file), the 8 cases where online quantization should reach attention fail, and the 4 no-change cases pass:

$ pytest tests/models/glm5next/test_online_quantization.py   # main cdd2750cd
FAILED test_attention_projections_take_only_online_quantization[checkpoint+targets-deepseek_sparse_attention]
FAILED test_attention_projections_take_only_online_quantization[checkpoint+targets-linear_attention]
FAILED test_attention_projections_take_only_online_quantization[checkpoint+linear-deepseek_sparse_attention]
FAILED test_attention_projections_take_only_online_quantization[checkpoint+linear-linear_attention]
FAILED test_attention_projections_take_only_online_quantization[online-targets-deepseek_sparse_attention]
FAILED test_attention_projections_take_only_online_quantization[online-targets-linear_attention]
FAILED test_fp8_mla_weights_reach_online_target_as_bf16[q_b_proj-checkpoint_shards0]
FAILED test_fp8_mla_weights_reach_online_target_as_bf16[fused_qkv_a_proj-checkpoint_shards1]
8 failed, 4 passed
E   {'q_b_proj': <class '...UnquantizedLinearMethod'>} != {'q_b_proj': <class '...Mxfp8OnlineLinearMethod'>}
E   AssertionError: assert torch.bfloat16 != torch.bfloat16   (q_b_proj weight never quantized)

$ pytest tests/models/glm5next/test_online_quantization.py   # this branch
12 passed

End to end and performance (GB300)

nvidia/GLM-5.3-Flash-NVFP4 (da920bb), upstream nightly edad8b03 with this branch's vllm/models/glm5next mounted over the installed package (the stock nightly for the BF16 rows), Dynamo frontend, sa-bench random prompts (8K in / 1K out, all requests at once, concurrency-capped), --max-num-batched-tokens 32768 (prefill / aggregated), --kernel-config '{"linear_backend_per_quant": {"mxfp8": "torch"}}' for the MXFP8 rows. These runs predate the rebase onto cdd2750cd; the commits in between change no compiled code, and the only one touching GLM-5.3 (#61457) affects AITER fused shared experts on ROCm. Every run completed all requests (0 failed), and the startup log showed the expected Quantizing N layers count (none on the BF16 rows and on the disaggregated decode instance).

Aggregated, 1 node, DP2 + expert parallel (2 GPUs), 5,120 / 10,240 requests:

config layers quantized C512 output tok/s/GPU C1024 output tok/s/GPU C512 / C1024 mean TTFT
BF16 projections (today) 0 2,839 2,441 7.8 s / 52.2 s
MXFP8, non-KDA 128 128 2,896 (+2.0%) 2,488 (+1.9%) 7.3 s / 49.5 s
MXFP8, all 196 (opt-in) 196 3,075 (+8.3%) 2,641 (+8.2%) 6.0 s / 42.2 s

Disaggregated, 2 prefill instances (TP1, 1 GPU each) + 1 decode instance (DP2 + expert parallel), 4 GPUs on 2 nodes, NIXL KV transfer:

config C128 output tok/s/GPU C256 C512 C128 / C256 / C512 mean TTFT
BF16 projections (today) 1,438 2,234 2,695 1.04 s / 1.97 s / 17.5 s
MXFP8 all 196 on prefill only 1,446 (+0.6%) 2,254 (+0.9%) 3,025 (+12.2%) 0.97 s / 1.69 s / 8.7 s

Single runs per point, so the sub-1% C128 / C256 disaggregated differences are not significant. At C512 the 2 prefill GPUs are the bottleneck (17.5 s TTFT), and faster prefill roughly halves TTFT and raises output throughput by 12%; decode step time is unchanged at C128 / C256 (TPOT 20.6 / 25.8 ms in both).

Accuracy

Setup:

  • Checkpoint: nvidia/GLM-5.3-Flash-NVFP4 (revision da920bb), served with --quantization modelopt_fp4.
  • NeMo-Skills 0.7.0 against an OpenAI-compatible endpoint.
  • Model-card sampling: temperature 1.0, top_p 0.95, max_new_tokens 327,680, default chat template (reasoning effort Max); reasoning stripped before scoring.
  • Run on NVIDIA VR200, on a main nightly with this branch's glm5next package. The reference is the same nightly without the branch, so its projections are BF16. Only the --quantization-config targets differ between the arms.
  • Deltas are paired by problem; 95% CIs are a bootstrap over problems. Pass rule: no significant regression, and the pooled CI must exclude a loss of more than 1 point.
benchmark samples BF16 all 196: Δ [95% CI] non-KDA 128: Δ [95% CI]
AIME25 30 x 8 97.08 -1.25 [-4.58, +1.67] +0.83 [-1.67, +3.75]
GPQA-Diamond 198 x 4 90.15 +1.26 [-0.76, +3.41] +1.14 [-0.38, +2.78]
MATH-500 500 x 1 97.60 -0.40 [-1.80, +1.00] +0.60 [-0.60, +2.00]
GSM8K 1319 x 1 96.89 +0.76 [+0.15, +1.44] +0.30 [-0.45, +1.06]
BeyondAIME 100 x 8 87.75 -0.50 [-2.62, +1.50] +0.38 [-1.50, +2.25]
HMMT Feb 2025 30 x 16 97.29 -1.04 [-3.75, +1.46] -0.21 [-2.29, +1.88]
APEX shortlist 47 x 8 86.17 -3.19 [-8.78, +1.60] -2.66 [-7.18, +1.60]
pooled, 3 hard long-output sets -1.58 [-3.69, +0.44] -0.83 [-2.61, +0.84]
pooled macro, 7 benchmarks -0.62 [-1.69, +0.42] +0.05 [-0.85, +0.94]
  • The hard, long-output sets (BeyondAIME, HMMT, APEX shortlist) were added because 8-bit numerics on a linear-attention model's recurrent path have been reported to shorten reasoning with flat accuracy on easier sets but a loss on HLE. HLE itself needs an external LLM judge, so we used these judge-free, symbolically graded sets instead.
  • Non-KDA passes the rule on all 7 benchmarks. Its APEX dip is within the noise of a 47-problem set and is not mirrored on the other two hard sets.
  • All 196 has no significant loss on any benchmark, but it is lower on all three hard sets, and neither pooled CI excludes a loss above 1 point. We disclose it as a possible cost, not a proven one.

Output length, paired by (problem, seed), both finished, ties excluded. Under BF16-vs-BF16 replicates, "shorter" stays near 50% (z within about ±1.5):

set all 196: shorter, z non-KDA 128: shorter, z median tokens, BF16 / all 196 / non-KDA
AIME25, GPQA, MATH-500, GSM8K 58.5% of 2,567, +8.6 48.4% of 2,542, -1.6 AIME25 6,642 / 4,546 / 6,833; GPQA 4,237 / 3,378 / 3,965
BeyondAIME, HMMT, APEX 63.6% of 1,649, +11.1 51.2% of 1,648, +1.0 BeyondAIME 20.3K / 15.9K / 20.4K; APEX 69.7K / 59.1K / 71.3K
Broader coverage: all 196 on an earlier build (8 benchmarks + RULER 64K)

Same quantization path (mxfp8_e4m3_quantize plus TorchMxFp8LinearKernel on the same 196 layers, at every M, decode included) on a downstream build of the same model code. The reference is the same build with BF16 projections.

benchmark samples BF16 MXFP8 (all 196) Δ [95% CI]
AIME24 30 x 16 95.83 96.46 +0.62 [-1.25, +2.92]
AIME25 30 x 16 94.17 94.38 +0.21 [-2.71, +3.12]
GPQA-Diamond 198 x 8 91.16 91.04 -0.13 [-1.07, +0.82]
MATH-500 500 x 4 97.70 97.65 -0.05 [-0.80, +0.70]
GSM8K 1319 x 2 97.08 97.31 +0.23 [-0.27, +0.72]
MMLU-Pro 12032 x 1 87.48 87.65 +0.17 [-0.21, +0.53]
IFEval (prompt strict) 541 x 4 90.25 89.97 -0.28 [-1.39, +0.83]
LiveCodeBench v6 (2408-2505) 454 x 4 50.50 51.65 +1.16 [+0.11, +2.26]
RULER 64K (13 tasks) 1300 see note see note no regression
macro average, 8 benchmarks +0.24 [-0.28, +0.77]
  • RULER 64K: NIAH x8 and VT are 100% for both. QA_1 is 92 vs 95 and QA_2 is 74 vs 78. CWE/FWE are noisy in both arms (frequent repetition loops at max effort).
  • Output length: the MXFP8 output is shorter in 59-66% of pairs on AIME/GPQA/LiveCodeBench. Fewer generations end in runaway repetition loops at the token cap (LiveCodeBench 9.7% vs 17.1%), which is where the LiveCodeBench gain comes from.
  • Where the shift comes from: a teacher-forced logprob comparison on identical token sequences, quantizing one projection group at a time, located the shift on the KDA in/out projections. On the hesitation-token measure, each other group's effect was at most about 1/5 of either KDA group's. That localization is what defines the non-KDA config.

Details

Why the projections were unreachable

GLM-5.3-Flash loads its attention projections as BF16 in every checkpoint:

  • KDA in_proj_qkvbfg_a, f_b_proj, g_b_proj, o_proj;
  • MLA fused_qkv_a_proj, q_b_proj, o_proj;
  • DSA indexer wq_b.

The KDA and indexer projections are BF16 in every official checkpoint. The MLA projections are BF16 in the NVFP4 checkpoint and block-FP8 in the Z.ai FP8 checkpoint, which vLLM dequantizes to BF16 at load (_try_load_fp8_attn_proj). The model therefore builds all of them with no quant config: MLA passes quant_config=None, and KDA clears vllm_config.quant_config around its constructor. With --quantization modelopt_fp4 --quantization-config '{"targets": {...}}', composition with the checkpoint method works for the shared-expert linears (84 layers) but silently skips the 112 attention projections. No warning is printed.

What changes

bf16_projection_quant_config() in a new hardware-neutral vllm/models/glm5next/quant_config.py (following the model-level quant_config.py of deepseek_v4 / deepseek_v41) returns the config these projections get:

  • with a quantized checkpoint, the online config composed through --quantization-config (any form);
  • with a BF16 checkpoint (--quantization online), that config only when it names targets;
  • otherwise None, as today.

Layers that match no online target get UnquantizedLinearMethod (existing OnlineQuantizationConfig behavior), so the user picks the projections with the targets regex. kv_b_proj is pinned to quant_config=None, because it is absorbed into the W_UK / W_UV bmm weights after loading. Loading is unchanged; the change is construction-time only, with no effect on the forward path when no target matches.

The MLA attention module (self_attn.attn) and the indexer now also receive the online config instead of None. The online config returns no method for non-linear layers, so the KV-cache scales stay at their defaults, as before. A targets pattern that also matches self_attn.attn is now rejected at startup ("online quantization is not supported for MLAAttention"), as for DeepSeek-V2, whose MLA attention already receives the quant config; before, it was silently ignored. The recommended patterns below name the projections, so they do not match it.

file change
vllm/models/glm5next/quant_config.py new: bf16_projection_quant_config() (32 lines)
vllm/models/glm5next/common/model.py MLA gets bf16_projection_quant_config(quant_config) instead of None
vllm/models/glm5next/common/kda.py KDA projections get the same; the fused conv is unchanged
vllm/models/glm5next/common/attention.py kv_b_proj pinned to None
tests/models/glm5next/test_online_quantization.py new, 2 tests (12 parametrized cases)

No docs change: the projections now follow the existing contract in docs/features/quantization/online.md ("Online quantization on unquantized layers from partially-quantized checkpoints"). The GLM-5.3-Flash-specific recommendation stays in this description.

Recommended configuration

All measurements use NVIDIA's NVFP4 checkpoint, nvidia/GLM-5.3-Flash-NVFP4 (revision da920bb). Its routed experts and dense MLPs are NVFP4. The 196 projections studied are BF16:

group projections layers
KDA (34 linear-attention layers) in_proj_qkvbfg_a, o_proj 68
MLA / DSA (11 layers: 3, 7, ..., 43) fused_qkv_a_proj, q_b_proj, o_proj, indexer.wq_b 44
shared expert (42 MoE layers) gate_up_proj, down_proj 84

We tested MXFP8 (E4M3 values, one E8M0 scale per 32 elements, for weights and dynamic activations). Accuracy holds in every configuration. However, MXFP8 on the 68 KDA in/out projections at decode measurably changes reasoning behavior: traces get 20-30% shorter on AIME/GPQA, with fewer hesitation and self-correction tokens, and accuracy on hard long-output problems trends lower (not significant). The other 128 projections leave behavior unchanged.

deployment recommended config behavior vs BF16
aggregated (or any decode-serving instance) MXFP8 on the 128 non-KDA projections (MLA, indexer wq_b, shared expert); KDA in/out stay BF16 unchanged: accuracy and output length match BF16 on 7 benchmarks
disaggregated MXFP8 on all 196 projections on the prefill instances only; decode instances unchanged decode, which produces nearly all of the output tokens, runs exactly as today
opt-in, aggregated MXFP8 on all 196 projections changed: shorter reasoning, and a possible 1-2 point cost on hard long-output problems
  • Where the throughput comes from: the prefill-heavy, large-M GEMMs. The largest projection is KDA in_proj_qkvbfg_a (4096 → 24,896), so excluding the KDA projections keeps only part of the gain in aggregated serving (+2% instead of +8% on GB300).
  • Why prefill-only is behavior-preserving: the behavior change comes from MXFP8 numerics at decode, which pick each next token. With MXFP8 only on prefill instances, prefill computes the KV cache and KDA state in MXFP8, while every generated token after the first comes from BF16 decode. An aggregated variant with exactly that split (MXFP8 for GEMMs of ≥ 2048 tokens, BF16 weights kept for decode; not part of this PR, it needs a second weight copy) showed no length change and no accuracy change: macro +0.09 pt [-0.43, +0.61] over 8 benchmarks. The disaggregated deployment itself was not accuracy-tested.
  • Use the cuBLASLt MXFP8 kernel (--kernel-config '{"linear_backend_per_quant": {"mxfp8": "torch"}}'). On SM10x, the automatic choice (FlashInfer CuTe-DSL) was markedly slower for these shapes at large M. The kernel priority is a separate follow-up.

Commands for the NVFP4 checkpoint. Each regex covers layers 0-44, so the MTP layer (45) is excluded; the K=128 KDA gate projections, kv_b_proj, the indexer wk_weights_proj and the dense MLPs never match.

Non-KDA, 128 layers (aggregated; behavior-preserving). The MLA layer indices are those of the released 45-layer config:

vllm serve nvidia/GLM-5.3-Flash-NVFP4 --quantization modelopt_fp4 \
  --quantization-config '{"targets": {"re:.*[.]layers[.]((3|7|11|15|19|23|27|31|35|39|43)[.]self_attn[.](o_proj|fused_qkv_a_proj|q_b_proj|indexer[.]wq_b)|([0-9]|[1-3][0-9]|4[0-4])[.]mlp[.]shared_experts[.](gate_proj|up_proj|gate_up_proj|down_proj))": "mxfp8"}}' \
  --kernel-config '{"linear_backend_per_quant": {"mxfp8": "torch"}}'

All 196 (prefill instances of a disaggregated deployment; or the disclosed aggregated opt-in):

vllm serve nvidia/GLM-5.3-Flash-NVFP4 --quantization modelopt_fp4 \
  --quantization-config '{"targets": {"re:.*[.]layers[.]([0-9]|[1-3][0-9]|4[0-4])[.](self_attn[.](in_proj_qkvbfg_a|o_proj|fused_qkv_a_proj|q_b_proj|indexer[.]wq_b)|mlp[.]shared_experts[.](gate_proj|up_proj|gate_up_proj|down_proj))": "mxfp8"}}' \
  --kernel-config '{"linear_backend_per_quant": {"mxfp8": "torch"}}'

The startup log confirms the coverage, e.g. Quantizing 196 layers of types: ... self_attn.in_proj_qkvbfg_a: 34, self_attn.o_proj: 45, ... (128 layers for the non-KDA regex).

Other checkpoints and other FP8 schemes

Other checkpoints (static analysis; only the NVFP4 checkpoint was run):

checkpoint how to use it
shared experts BF16 (NVIDIA NVFP4, RedHat / LibertAI NVFP4, W4A16 quants) the regexes above
shared experts already quantized (Z.ai FP8, RadixArk NVFP4, AMD Quark) drop the mlp.shared_experts alternative (44 or 112 attention layers). With it, those checkpoint methods cannot re-quantize (no dequantize_weight) and loading fails, as on main. For the FP8 checkpoint, the MLA weights go FP8 → BF16 (load-time dequant) → MXFP8.
BF16 (zai-org/GLM-5.3-Flash-BF16) add --quantization online; --quantization-config alone is ignored (existing behavior)

--quantization-config does not reach the MTP draft today (existing behavior), so the MTP layer keeps the checkpoint method. If #49900 lands, the regexes above still exclude layer 45.

  • Do not use a per-tensor activation scale on these layers at decode. One scale shared by every row of a decode batch couples all requests. FP8 per-tensor W+A at every batch size was broken in our runs: AIME 94-97% → 53-60%, MATH-500 98% → 65%, with 20-25% runaway repetition loops. Upstream's online fp8_per_tensor uses per-token activation scales on CUTLASS-capable GPUs and does not have this coupling; we did not evaluate it.
  • Per-channel FP8 (fp8_per_channel) on all 196 mostly avoids the behavior change but also keeps only a small part of the throughput gain, so it offers nothing over the non-KDA MXFP8 config.

Limitations

  • Hardware-neutral, but validated on NVIDIA only; ROCm uses the same code path (MXFP8 via RocmDotScaledMxfp8LinearKernel) and is untested.
  • CI covers the wiring and the helper (unit tests); no CI job serves GLM-5.3-Flash (191 GB checkpoint). End-to-end validation was done outside CI as listed above.
  • The disaggregated prefill-only configuration was measured for throughput, not for accuracy (see the aggregated equivalent above).

AI assistance

This change was prepared with AI assistance (Claude Code). The submitter reviewed every changed line and ran the tests above.


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.

🤖 Generated with Claude Code

…ion projections

GLM-5.3-Flash builds its KDA, MLA and DSA-indexer projections with
quant_config=None because they load as BF16 in every checkpoint, so online
quantization requested with --quantization-config silently skipped them.
Give them only the explicitly requested online config
(bf16_projection_quant_config), keep kv_b_proj unquantized, and leave the
--quantization shorthands on BF16 checkpoints unchanged. Add tests on real
decoder layers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Ao Tang <aot@nvidia.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant