Repository navigation
✨ [llm][npu][quant] Add W8A8 MXFP8 quantization for Qwen3 MoE on Ascend NPU - #30768
Conversation
Add NPUMXFP8FusedMoEMethod (online weight quant via npu_dynamic_mx_quant with 3D input) + npu_fused_experts_mxfp8 kernel (routing -> act quant -> swiglu_quant_v2 -> grouped_matmul -> finalize_routing, mirroring vllm-ascend A5 path). Wire dispatch in Fp8Config.get_quant_method for FusedMoE layers (is_npu + use_mxfp8). Weights kept as strided transpose views — kernel contract verified by A5 probe.
ModelSlimMXFP8MoEScheme for --quantization modelslim with W8A8_MXFP8 pre-quantised checkpoint. NPUMXFP8FusedMoEMethod now detects fp8 dtype and takes offline transpose-only branch. Registered via FusedMoeWeightScaleSupported.BLOCK.
Offline modelslim (W8A8_MXFP8) checkpoints quantise mlp.gate to float8_e4m3fn + weight_scale, but Qwen3MoeSparseMoeBlock hard-coded the gate as ReplicatedLinear(quant_config=None) -> bf16. The fp8 gate was then cast to bf16 WITHOUT its block scale and the scale key dropped, scrambling the router (probe: norm 70195 vs 31.4) -> wrong experts -> garbage. Online (bf16 gate) was unaffected. Pass quant_config to the gate only when it is the offline ModelSlim path (get_name()=='modelslim'); ModelSlimConfig.get_quant_method then dequantises a quantised gate correctly or skips a FLOAT/absent one. The online Fp8/mxfp8 path keeps gate_quant_config=None -> bf16, unchanged. Mirrors vllm-ascend's description-driven dispatch.
ascend_npu_quantization.mdx: mark MXFP8 MoE supported on A5, add an MXFP8-for-LLM-MoE section (online + offline launch commands, kernel path, and the router-gate description-driven caveat). quantization.mdx: extend the mxfp8 Ascend cell to cover LLM MoE (FusedMoE/TP). Mirrors the doc pattern from PR sgl-project#23795.
github-actions
Bot
added
documentation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for MXFP8 quantization (both online and offline) for LLM MoE models (such as Qwen3-30B-A3B) on Ascend NPUs (A5 series). It implements the NPUMXFP8FusedMoEMethod and the corresponding offline ModelSlimMXFP8MoEScheme scheme, utilizing specialized NPU kernels like npu_grouped_matmul_swiglu_quant_v2 and npu_dynamic_mx_quant. Additionally, the router gate in qwen3_moe.py has been updated to support description-driven quantization for offline paths, and the documentation has been updated accordingly. No review comments were provided, so I have no further feedback.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
TallMessiWu
changed the title
Add an opt-in fused path to npu_fused_experts_mxfp8: when SGLANG_ENABLE_NPU_FUSED_MOE_ROUTING_QUANT is on and the op supports it, npu_moe_init_routing_v2(quant_mode=3) does the MXFP8 activation quant inline (e4m3 payload + e8m0 block scale as ret[3]), replacing the separate npu_dynamic_mx_quant pass and saving one kernel launch. Gated by SGLANG_ENABLE_NPU_FUSED_MOE_ROUTING_QUANT (default off) plus a cached runtime capability probe that falls back to the two-step path when quant_mode=3 is unsupported. A5 probe (Q9, torch_npu 2.10.0.post2): quant_mode=3 accepted with x_dtype=None; ret[3] scale is 2D [N,K/32] e8m0 so _normalize_mxfp_scale reshapes it to [N,K/64,2]; fused output is byte-identical to the two-step path (cos=1.0). DeepEP path (apply_without_routing_weights) does not go through init_routing and is unchanged.
Review feedback on sgl-project#30768: the hard rejection at runner construction was wrong. DeepEP works with mxfp MoE even without quantization inside the dispatcher -- it simply returns bf16, and the method can quantize the activations itself. Three pieces: - AscendRunnerCore no longer raises for DeepEP + MXFP8. activation stays None either way, since gmm1 fuses gate/up + swiglu + requant. - process_weights_after_loading sets the dispatcher output dtype per backend. DispatcherOutputDtype has no MXFP8 entry in deepep.py's config_map (only BF16/FP8/INT8/NVFP4), so asking for 'mxfp8' there would KeyError -- hence an explicit bf16 branch rather than letting it through. - apply_fused_gmm1_swiglu takes pertoken_scale as Optional and, when it is None, quantizes the hidden states through the HiddenStatesDynamicQuant wrapper from the previous commit. Both dispatchers therefore reach the fused kernel with the same e4m3 + e8m0 input; only where the activation quant happens differs. ascend_tp is unaffected: is_deepep() is False there, so the dispatcher dtype stays 'mxfp8' and pertoken_scale is never None. Docs updated to match. Only the ascend_tp path is validated end-to-end on A5 -- I have no EP setup, so in particular it is unverified whether low-latency dispatch's 3D hidden states are accepted by npu_grouped_matmul_swiglu_quant_v2.
…oint Review feedback on sgl-project#30768: the online entry point duplicated create_weights, process_weights_after_loading, create_moe_runner and apply from the unquantized method, and the name did not distinguish it from the per-gmm kernel. NPUMXFP8FusedMoEMethod becomes NPUMXFP8OnlineMoEMethod(UnquantizedFusedMoEMethod) and overrides only create_moe_runner, to attach the MXFP8 kernels before MoeRunner is built (AscendRunnerCore.__init__ reads layer.w2_kernel to pick its activation). That removes ~85 lines. Verified the inherited pieces against the current unquant.py rather than assuming: create_weights produces the same [E, 2I, H] / [E, H, I] parameters given is_gated and no bias; the aiter / cpu-amx / deep_gemm / flashinfer branches of process_weights_after_loading are all False on NPU, leaving the _is_npu branch that calls the kernels; apply -> forward_npu -> runner.run(dispatch_output, layer) works because layer carries the same w{13,2}_weight* attribute names AscendQuantInfo did. It has to live in a new module: unquant.py imports moe_methods.py at module scope, so subclassing there would be a circular import. moe_methods.py does not import unquant.py, so the new chain stays acyclic. Two supporting changes: - unquant.py no longer hardcodes dispatcher_output_dtype 'bf16' right after calling the kernels. It duplicated what NPUUnquantMoEMethod already does for w13, and it clobbered the value a quantized subclass kernel had just set. There are no other subclasses, so unquant is unaffected. - create_moe_runner rejects a non-Ascend runner backend loudly. Inheriting UnquantizedFusedMoEMethod puts this method in FusedMoE's w1/w3 shard-swap list (layer.py:942), so an explicit --moe-runner-backend flashinfer_trtllm would load every expert with gate and up exchanged -- silent degenerate output. The automatic override was already fixed in 9ac21f5; this covers the manual one. The other two isinstance sites (layer.py:365 and the weight-update hook at :904) are both gated on is_flashinfer_trtllm_routed() and inert on NPU.
Answers the NZ question raised in review of PR sgl-project#30768. On A5 (torch_npu 2.10.0.post2, Ascend950PR, Qwen3-30B-A3B shapes with 128 experts) this gains 1.4% on decode and 3.8% on prefill against a 0.2-0.3% noise floor, with outputs identical to before. Nowhere near the ~10% the int8 path reported, but consistent across both token shapes. The cast must precede the transpose. gmm1 asserts that weight and weight_scale share a transpose flag (CheckMXTranspose), and the cast returns a physically retiled, non-transposed tensor -- so casting the already-transposed view fails outright. That is why this cannot follow the int8 MoE methods in the same file, which transpose first but carry no MX scale to keep in sync; the dense W4A8 path in linear_method_npu.py already uses this order. Measurements also rule out .contiguous() on either side: making both contiguous costs 6.2% on decode. _is_nz_aligned gains an fp8 branch sharing the int8 column rule, since e4m3 is likewise single-byte. It previously fell through to the catch-all True, which no caller exercised -- the packed-FP4 callers return early via customize_dtype and dense MXFP8 never format-casts -- so this only guards the weights introduced here.
# Conflicts: # python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py
OrangeRedeng
left a comment
OrangeRedeng
left a comment
There was a problem hiding this comment.
Great changes! Thank you!
|
We'll start CI a little later, after we fix it |
ping1jing2
changed the title
|
/tag-and-rerun-ci |
CI investigationThe current CI failures are not caused by this PR's MXFP8 MoE changes. 1. Non-CUDA platforms fail in an upstream DWDP import pathThe MLX, Xeon, and Ascend/NPU jobs all fail while importing This unconditional CUDA-only import was introduced by upstream DWDP PR #29778 (commit The failed NPU case is 2. The Base GPU root failure is runner state, then fast-fail cascadesThe first Base CI root-cause job ( The remaining H100/B200/DeepEP Base failures are the workflow's fast-fail cascade from that root job, rather than independent test or quantization failures. No failing log contains an MXFP8 kernel/quantization assertion or a failure in the files changed by this PR. Lint and XPU checks passed. Conclusion: this PR is blocked by the upstream DWDP non-CUDA import regression plus an independent dirty-GPU-runner failure. The DWDP import needs an upstream platform guard/lazy import; after that, the GPU job should be rerun on a clean runner. |
This known issue was fixed in #31919, I re-ran the CI |
CI investigation — current run (
|
…nd NPU (sgl-project#30768) Co-authored-by: Артем Савкин <58187114+OrangeRedeng@users.noreply.github.com> Co-authored-by: ronnie_zheng <zl19940307@163.com>
…nd NPU (sgl-project#30768) Co-authored-by: Артем Савкин <58187114+OrangeRedeng@users.noreply.github.com> Co-authored-by: ronnie_zheng <zl19940307@163.com>
4 participants
Summary
Adds W8A8 MXFP8 quantization (8-bit weights + 8-bit activations, block_size = 32) for Qwen3 / Qwen3.5 MoE (FusedMoE) LLM models on Ascend NPU, continuing the NPU quantization work tracked in #21584. Both online (
--quantization mxfp8) and offline (msmodelslim, flagless — selected by the checkpoint'squant_model_description.json) paths are supported and share the same fused-experts kernel. Requires Ascend A5 (Ascend 950) — the fused MoE MX kernels (npu_grouped_matmul_swiglu_quant_v2,npu_dynamic_mx_quant) are A5-only.Online quantization (
--quantization mxfp8)Fp8Config.get_quant_method()dispatches FusedMoE layers to the newNPUMXFP8OnlineMoEMethod(theisinstance(layer, FusedMoE)branch, mirroring the existing denseNPUMXFP8LinearMethodlinear branch). It subclassesUnquantizedFusedMoEMethodand overrides onlycreate_moe_runner, to attach the MXFP8 per-gmm kernels — weight creation, weight post-processing and the forward pass are the unquantized Ascend ones. On GPU theFusedMoEbranch is untouched.w13/w2are quantised once at load vianpu_dynamic_mx_quant(dst_type = torch_npu.float8_e4m3fn)— a 3D[E, N, K]input is accepted directly (no per-expert loop) — intofloat8_e4m3fnweights + UE8M0 (float8_e8m0fnu) block scales, kept as non-contiguous transpose views.ascend_tpdispatcher,npu_moe_init_routing_v2withquant_mode=3(countgroup_list) — the per-token MXFP8 activation quant is fused directly into routing, so the same kernel launch emits the sortedfloat8_e4m3fnpayload and itsfloat8_e8m0fnublock scale, saving a separatenpu_dynamic_mx_quantpass — → gmm1npu_grouped_matmul_swiglu_quant_v2(gate_up + swiglu + requant fused, cumulativegroup_list) → gmm2npu_grouped_matmul(countgroup_list+ explicitgroup_list_type). The UE8M0 scale dtypes are passed explicitly; the e4m3x/weightdtypes are left implicit. DeepEP has no MXFP8 dispatch dtype, so on that path the dispatcher keeps sending BF16 and gmm1 quantizes the hidden states itself vianpu_dynamic_mx_quantbefore the fused kernel — both dispatchers reach gmm1 with the same input. Only theascend_tppath is validated end-to-end on A5.quant_mode=3scale comes back flat ([N, K/32]) instead of the pair-split[N, K/64, 2]layout gmm1 expects, so_normalize_mxfp_scalereshapes it before use. Verified via a standalone probe script (npu_moe_init_routing_v2(quant_mode=3)vs. the two-stepquant_mode=-1+npu_dynamic_mx_quantreference) to be byte-identical (cos = 1.0) on Ascend A5 + torch_npu 2.10.0.post2, then confirmed end-to-end on the online and offline launch scripts — so this fused path is the only routing-quant implementation kept (no runtime toggle, no two-step fallback branch).Offline quantization (msmodelslim pre-quantized weights, flagless)
W8A8_MXFP8MoE checkpoints are detected fromquant_model_description.jsonand handled by the newModelSlimMXFP8MoEScheme, which registersfloat8_e4m3fnexpert weights +uint8(E8M0, exponent + 127) block scales and delegates weight post-processing + forward to the sameNPUMXFP8MoEMethodkernel (offline branch: reshape[E, N, K/32] → [E, N, K/64, 2]+ transpose, no re-quant).mlp.gate(W8A8_MXFP8). The gate is aReplicatedLinear, so its quantization is now description-driven — for the offlinemodelslimpath the gate receives the quant config and is dequantised correctly viaModelSlimMXFP8Scheme; the online path keeps it in BF16 (unchanged). Previously the fp8 gate was cast to BF16 without its block scale, scrambling the router and producing garbage output; this is the fix that makes the offline path usable.Benchmark Results
Validated end-to-end on Ascend A5 hardware (Qwen3-30B-A3B). Both W8A8 MXFP8 modes (online and offline msmodelslim) are compared against the BF16 baseline.
Performance
Both paths ship the FRACTAL_NZ expert-weight layout (review item 8). Against BF16, online and offline W8A8 MXFP8 deliver ~+44% throughput and ~-30% TPOT / E2E latency; the two modes are within noise of each other.
FRACTAL_NZ weight layout (review item 8)
ND vs NZ on the same benchmark, everything else equal:
Not the ~10% the int paths saw, but consistent across both modes, so NZ is what this PR
ships. A standalone kernel probe on the same box (gmm1 + gmm2 only, 128 experts) measured
+1.4% on decode-shaped and +3.8% on prefill-shaped inputs against a 0.2-0.3% noise floor,
matching the end-to-end delta.
Two things worth recording, since neither is obvious from the int8 precedent:
npu_grouped_matmul_swiglu_quant_v2assertsthat weight and weight_scale carry the same transpose flag (
CheckMXTranspose), andnpu_format_castreturns a physically retiled, non-transposed tensor — so casting thealready-transposed view fails outright with
transpositions are true/false. The int8 MoEmethods transpose first, but they carry no MX scale to keep in sync; the dense W4A8 path
already uses the order adopted here.
numerically correct but measurably slower (-6.2% decode at 128 experts), so the
transpose views stay. A small-expert-count microbenchmark inverts this conclusion
entirely (E=4 rates contiguous 58% faster), which is why the numbers above come from
the real 128-expert shape.
Accuracy (GSM8K)
Validated on Ascend A5 hardware (Qwen3-30B-A3B), three runs each. Both W8A8 MXFP8 modes stay within run-to-run noise of the BF16 baseline.
Style / consistency with the merged NPU quant PRs (#22352, #23650, #23795)
torch.ops.npu.*; no top-levelimport torch_npu, nocurrent_platformguards (module stays importable on CUDA/CPU/AMD/XPU CI).Fp8Config.get_quant_method(mirrors denseNPUMXFP8LinearMethod). Offline subclassesModelSlimMoESchemeand delegates weight post-processing / forward to aself.kernel(NPUMXFP8MoEMethod), registered inget_moe_scheme'smoe_quant_schemeslist — the same pattern as the int8 / W4A4 / W4A8 MoE schemes..transpose(1, 2)with no.contiguous()), with the expert weight cast to FRACTAL_NZ before that transpose; gmm1 uses the fusednpu_grouped_matmul_swiglu_quant_v2; thegroup_listcount↔cumulative conversion is explicit (gmm1 cumulative, gmm2 count).Files Changed
New
srt/layers/quantization/modelslim/schemes/modelslim_mxfp8_moe.pyModelSlimMXFP8MoESchemeforW8A8_MXFP8MoE checkpointsModified
srt/hardware_backend/npu/quantization/moe_methods.pyNPUMXFP8MoEMethod(online + offlineprocess_weights_after_loading,apply_fused_gmm1_swiglufor w13,applyfor w2); cast the expert weights to FRACTAL_NZ ahead of the transposesrt/hardware_backend/npu/utils.py_is_nz_aligned: add the fp8 branch — e4m3 is single-byte, so it shares int8's column rule. It previously fell through to the catch-allTrue, which no caller exercised (the packed-FP4 callers return early viacustomize_dtype, and dense MXFP8 never format-casts)srt/hardware_backend/npu/quantization/online_moe_methods.py(new)NPUMXFP8OnlineMoEMethod(UnquantizedFusedMoEMethod), the--quantization mxfp8entry pointsrt/hardware_backend/npu/moe/matmul.pyGroupedMatmulSwigluQuant, the reusable wrapper aroundnpu_grouped_matmul_swiglu_quant_v2srt/hardware_backend/npu/moe/hidden_states_quant.py→quant.pyHiddenStatesDynamicQuantthefloat8_e4m3fn(npu_dynamic_mx_quant) dtypesrt/layers/quantization/unquant.pydispatcher_output_dtype: "bf16"that clobbered the value the NPU kernel had just set (it duplicatedNPUUnquantMoEMethod)srt/layers/moe/moe_runner/ascend.pysrt/layers/quantization/fp8.pyNPUMXFP8OnlineMoEMethodunderis_npu() and use_mxfp8srt/arg_groups/overrides.pymxfp8→flashinfer_trtllmrunner-backend override on NPU, so the backend staysautoand resolves to Ascend (it also gates FusedMoE's w1/w3 shard swap)srt/layers/quantization/modelslim/modelslim.py("W8A8_MXFP8", ModelSlimMXFP8MoEScheme)toget_moe_scheme'smoe_quant_schemessrt/layers/quantization/modelslim/schemes/__init__.pyModelSlimMXFP8MoESchemesrt/models/qwen3_moe.pygatequant description-driven (offline modelslim only) so a quantisedmlp.gateloads/dequantises correctly instead of being cast to BF16 without its scaledocs_new/docs/advanced_features/quantization.mdx,docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_quantization.mdxStatus
WIP / Draft. Online (
--quantization mxfp8) and offline (msmodelslim) MoE MXFP8 are both validated end-to-end on Ascend A5 (see Performance above). GSM8K accuracy is within run-to-run noise of the BF16 baseline (see Accuracy above).@OrangeRedeng's review is addressed. It was first applied as a single refactor commit (
6bd7056), which garbled generation on A5 in a way that could not be bisected, so that commit was dropped and the review was redone one item at a time, each verified on A5 before the next:--quantizationfrom the docs9efe9998dhidden_states_quant.py→quant.py, teach it mxfp8efe50dfa3GroupedMatmulSwigluQuantwrapper + pick the matmul byweight_prefix2d06d61a51efd2ca35UnquantizedFusedMoEMethod, rename toNPUMXFP8OnlineMoEMethodea24f248dnpu_format_cast(NZ)d419aa41fRe-validating the baseline first also surfaced three regressions from the preceding
mainmerge, in which upstream restructured the NPU MoE layer (fused_moe_method_npu.pydeleted,moe_methods.py/moe_runner/ascend.py/npu/moe/*.pyadded) and this PR's implementation was re-ported onto the new layout. None are review-related:e9b36bb5dtorchinstead oftorch_npu; the grouped matmuls reject torch's dtype object, so every gmm1 died9ac21f5ff--quantization mxfp8forcedmoe_runner_backendtoflashinfer_trtllmunconditionally, including on NPU. Beyond picking a runner that does not exist there, that flag also drives FusedMoE's w1/w3 shard swap and a 128-alignment round-up — and once item 3 above makes this method inheritUnquantizedFusedMoEMethod, it matches the shard-swap list, loading every expert with gate and up exchanged. Silent degenerate output; this is what garbled the big-bang refactor969087364ModelSlimMoEMethod.applyreads{w13,w2}_weight_offsetunconditionally; MXFP8 is a pure scale format with no zero point, so the offline scheme never created oneRelated Issues
Part of #21584 (MXFP8/MXFP4 support on Ascend NPU for Qwen3 LLM). Builds on merged #23795 (dense W4A4) and #22352 (dense W8A8 MXFP8 — provides the dense linear MXFP8 method +
Fp8ConfigNPU dispatch +ModelSlimMXFP8Schemethis MoE work reuses).CI States
Latest PR Test (Base): 🚫 Run #30416404909
Latest PR Test (Extra): ❌ Run #30416404805