Repository navigation
🚧 [WIP][llm][npu][quant] Fix offline W8A8 MXFP8 quantization for Qwen3.5 on Ascend NPU - #32266
Closed
TallMessiWu wants to merge 40 commits into
Closed
TallMessiWu wants to merge 40 commits into
TallMessiWu wants to merge 40 commits into
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.
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.
A5 online+offline e2e verified, so flip SGLANG_ENABLE_NPU_FUSED_MOE_ROUTING_QUANT default from False to True. Behavior change: the MXFP8 MoE TP path now fuses activation quant into npu_moe_init_routing_v2 (quant_mode=3) by default; the runtime capability probe still falls back to the two-step path if unsupported, and SGLANG_ENABLE_NPU_FUSED_MOE_ROUTING_QUANT=0 forces the two-step path.
A5 e2e validated the fused quant_mode=3 path (npu_fused_experts_mxfp8), so keep it as the only path -- remove the two-step (quant_mode=-1 + npu_dynamic_mx_quant) fallback branch, the _supports_fused_mxfp8_routing capability probe, and the SGLANG_ENABLE_NPU_FUSED_MOE_ROUTING_QUANT env switch.
Port MXFP8 MoE onto the refactored Ascend MoE stack. Upstream sgl-project#25663 (8ed82af) deleted fused_moe_method_npu.py -- the file this branch extended -- and split the monolithic NPU MoE path into four layers: routing/act-quant in AscendTPDispatcher, orchestration in AscendRunnerCore, per-gmm kernels in quantization/moe_methods.py, and weight creation in the schemes. NPUMXFP8FusedMoEMethod and npu_fused_experts_mxfp8 had no place in that shape, so the conflict is resolved by accepting the deletion and porting the implementation rather than restoring the file. - moe/init_routing.py: NPUMoEInitRouting_v2 already threads quant_mode, so quant_mode=3 needs only the e8m0 scale reshape ([N,K/32] -> [N,K/64,2]), which _normalize_mxfp_scale now does there. - moe/utils.py + token_dispatcher/ascend_tp.py: new DispatcherOutputDtype.MXFP8 selects NPUMoEInitRouting_v2(quant_mode=3); the selector raises on unknown dtypes, so a missing branch fails loudly. - quantization/moe_methods.py: NPUMXFP8MoEMethod is the per-gmm kernel (online and offline branch by weight dtype, strided transpose views preserved); NPUMXFP8FusedMoEMethod is the online --quantization mxfp8 entry point. - moe_runner/ascend.py: gmm1 uses npu_grouped_matmul_swiglu_quant_v2, which fuses gate/up + swiglu + requant and does not decompose into the runner's gmm1 -> activation -> gmm2 split. Keep the fusion via a dedicated branch; the activation is None for MXFP8 and DeepEP is rejected explicitly (the activation quant comes from init_routing, which DeepEP bypasses). - modelslim_mxfp8_moe.py: rewritten per weight group to match the new scheme contract; drops apply_weights/create_moe_runner (forward is the runner's now). - The DeepEP branch (apply_without_routing_weights) is dropped: its upstream counterpart is gone, and MoE was documented TP-only and never validated. Reuses upstream's _get_float8_e8m0fnu_dtype from linear_method_npu behind a non-null assert, since a None scale dtype would silently change how the MXFP8 ops read the scales.
Port regression from the main merge (e983b8f). The pre-merge implementation in fused_moe_method_npu.py resolved the block-scale dtype as getattr(torch_npu, "float8_e8m0fnu", getattr(torch, ...)), torch_npu first. The port to moe_methods.py routed it through the dense helper _get_float8_e8m0fnu_dtype, which only reads torch. npu_grouped_matmul_swiglu_quant_v2 / npu_grouped_matmul validate their scale dtype arguments against torch_npu's own enum (float8_e8m0fnu == 293 on A5) and reject torch's dtype object: "weight_scale_dtype only supports float8_e8m0fnu or None, but the actual value is Float8_e8m0fnu" -- so every MXFP8 MoE forward dies on the first gmm1, online and offline alike. Restore torch_npu-first resolution with a lazy import and a module-level cache, keeping torch as fallback and the loud raise when neither exposes it. The dense helper stays as is: npu_quant_matmul accepts either source.
--quantization mxfp8 makes ServerArgs override moe_runner_backend from 'auto' to
'flashinfer_trtllm' unconditionally, including on NPU where flashinfer does not
exist. NPUMXFP8FusedMoEMethod.create_moe_runner only substitutes ASCEND when the
backend is still 'auto', so it built a flashinfer runner and every decode died
with 'Unexpected quant_info type for flashinfer_trtllm: AscendQuantInfo' during
cuda-graph capture.
The value is not confined to runner selection. FusedMoE keys two more behaviours
off flashinfer_trtllm: the w1/w3 shard swap ('flashinfer assumes w31 format for
w13_weight', layer.py:939) for a listed set of methods, and a 128-alignment
round-up of intermediate_size_per_partition. The current method subclasses
FusedMoEMethodBase so it escapes the shard-swap list today, but it is one
inheritance change away from loading every expert with gate and up exchanged --
silent degenerate output rather than a loud TypeError.
Skip the override on NPU so the backend stays 'auto' and resolves to ASCEND.
ModelSlimMoEMethod.apply builds AscendQuantInfo by reading
layer.{w13,w2}_weight_offset unconditionally, which the int8 and int4 MoE
schemes all register. MXFP8 is a pure scale format -- the e8m0 block exponent
carries everything and there is no zero point -- so the offline scheme never
created one and every decode raised AttributeError: 'FusedMoE' object has no
attribute 'w13_weight_offset' during cuda-graph capture.
Before the main merge the scheme delegated apply to the fused method directly
and never went through the shared builder, so the assumption held.
AscendQuantInfo already types the field Optional, so register it as None -- the
attribute then exists and resolves to 'no offset'. Follows the register_parameter
(name, None) convention used for bias in linear.py and input_scale in fp8.py; a
None parameter is skipped by named_parameters(), so no weight loader looks for
it.
Review feedback on sgl-project#30768: --quantization modelslim does not need to be passed for a msmodelslim checkpoint. Confirmed in the loader. ModelConfig._find_quant_modelslim_config() reads quant_model_description.json and injects quant_method='modelslim', and _verify_quantization() adopts it whenever --quantization was not given -- the comment there calls the field 'required for flagless model loading'. Drop the flag from the offline MoE example and say the description file selects both the ModelSlim path and the scheme. This was the only --quantization modelslim left in the docs.
…type Review feedback on sgl-project#30768: make the hidden-states quant a reusable kernel wrapper like the matmuls, so other quantization schemes can pick it up. Rename npu/moe/hidden_states_quant.py to quant.py (moe_methods.py was its only importer) and route HiddenStatesDynamicQuant by dtype: float8_e4m3fn goes to npu_dynamic_mx_quant, int8/quint4x2 keep npu_dynamic_quant, anything else raises. The two existing call sites use quint4x2 and int8, so behaviour is unchanged; the mxfp8 branch has no consumer yet and is wired up in a later commit. Written fresh rather than cherry-picked from sgl-project#29762: that version does a module-level import torch_npu, which breaks the CUDA/CPU/AMD/XPU CI jobs since the module is imported on every platform. Going through torch.ops.npu.* keeps it importable everywhere, and torch.float8_e4m3fn is a plain torch dtype so no torch_npu reference is needed here at all.
…y prefix Review feedback on sgl-project#30768, two items. Move npu_grouped_matmul_swiglu_quant_v2 out of NPUMXFP8MoEMethod into a GroupedMatmulSwigluQuant wrapper in npu/moe/matmul.py, next to GroupedMatmul, so other block-scaled schemes can reuse it. Two deviations from GroupedMatmul worth knowing: it returns (quantized_activations, block_scale) rather than a single tensor, because the requant is part of the kernel, and it takes no output_dtype (that comes from quant_dtype in scale_args). The count -> cumulative group_list conversion moves inside the wrapper, since this op wants cumulative while npu_grouped_matmul keeps the count form the dispatcher produces. NPUMXFP8MoEMethod.__init__ now takes weight_prefix and picks the matmul there -- GroupedMatmulSwigluQuant for w13, GroupedMatmul for w2 -- instead of gmm1 reaching for the op directly. The instances were already per-prefix (layer.w13_kernel / layer.w2_kernel, one ModelSlimMXFP8MoEScheme per weight group), so only the two construction sites had to pass the prefix. apply() and process_weights_after_loading() keep their weight_prefix parameter, since the runner calls every kernel through that same interface. The resulting op call is argument-for-argument identical to the previous one.
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
npu_gemma_rms_norm has no kernel registered for socVersion ascend950 (A5), causing graph capture to fail with "does not support opType [GemmaRmsNorm]". Gemma RMSNorm equals normalize(x) * (1 + weight), so a plain npu_rms_norm with gamma = 1 + weight is mathematically equivalent and is implemented across A2/A3/A5.
Only Ascend A5 (soc_version 260 / ascend950) lacks the npu_gemma_rms_norm kernel; A2/A3 (910) support it. Branch on is_npu_a5() so A5 uses the equivalent npu_rms_norm(x, 1 + weight) fallback while A2/A3 keep the native fused op, mirroring vllm-ascend's per-device dispatch.
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
github-actions
Bot
added
the
documentation
4 of 5 tasks
1. Guard the new provider import in srt/layers/layernorm.py and fall back to torch_npu, so NPU deployments on an sgl-kernel-npu wheel predating the staged provider keep importing; a hard import would break every model, not just Gemma, at module import time 2. Restore Gemma3RMSNorm.forward_npu's residual path to forward_native and drop the added env switch; both were outside the scope of replacing direct torch_npu calls and would shift numerics on all Ascend SoCs 3. Extract the duplicated lazy-import guard into _sgl_kernel_npu_gemma, which fixes the fused-add error message naming the wrong module 4. Keep GemmaFusedAddRMSNormOp on the SGL_KERNEL_NPU provenance alongside GemmaRMSNormOp 5. Pin each op's error to its own module, drop the reverted-behaviour case, and replace defensive hasattr with explicit parametrization 6. Tell A5 users in the Ascend install guide that released kernel wheels are 910-only
TallMessiWu
changed the title
TallMessiWu
changed the title
10 tasks
Resolve conflicts against the current fused-op kernel registry: - kernels/fused_op.py, kernels/ops/layernorm, srt/layers/layernorm.py: take the up-to-date PR sgl-project#32745 Gemma RMSNorm state, which already uses main's forward_torch_npu / _op_label naming and keeps the Ascend Gemma provider selected at sgl-kernel-npu build time (RFC sgl-project#35709). - models/qwen3_vl.py: keep main's language_model_only gate and apply the ModelSlim quant_config propagation inside its else branch. - tests: follow main's runtime_context.get_mm and the fused-op auto_selected_backend API.
Contributor
Author
|
Closing: everything still unique to this PR has moved elsewhere.
Nothing unique is left here, and the branch is ~1200 commits behind main, so a fresh PR would be cheaper than reviving this one. |
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
1 participant
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.
Summary
Fixes Qwen3.5-specific ModelSlim loading issues for offline W8A8 MXFP8 checkpoints on Ascend NPU, covering the language-model GDN projections and the Qwen3-VL visual tower.
Inherited Ascend Gemma RMSNorm support
sgl-kernel-npudependency are owned by 🐛 [NPU] Fix Qwen3.5 dense serving on Ascend 950 #32745.Qwen3.5 offline ModelSlim packed-module resolution
in_proj_qkvz -> [in_proj_qkv, in_proj_z]andin_proj_ba -> [in_proj_b, in_proj_a]at the top level.model/visualmappings. Previously, selecting only the scoped mapping dropped the Qwen3.5 mappings, so those projections silently fell back to unquantized Linear and their checkpointweight_scaletensors were ignored.Qwen3-VL offline visual-tower quantization
Qwen3VLForConditionalGenerationnow passes the checkpointquant_configintoQwen3VLMoeVisionModelinstead of hard-codingNone.visual.*.weight_scaleparameters.Offline partial MXFP8 block-scale layout
ceil(K / 32)so a final partial block is retained.K=4304: the checkpoint has 135 scales, while the previous floor division allocated 134 and caused the row-parallel loader to truncate the final scale.135 -> 136 -> 68pairs, matching the kernel requirementceil(4304 / 64)=68and the vllm-ascend implementation.Key NPU APIs used
sgl-kernel-npuGemma RMSNorm APIs from #32745torch.ops.npu.npu_quant_matmul(..., group_sizes=[1, 1, 32])Files Changed
srt/layers/quantization/modelslim/modelslim.pysrt/models/qwen3_vl.pyquant_configinto the Qwen3-VL visual towersrt/layers/quantization/modelslim/schemes/modelslim_mxfp8.pyceil(K / 32)srt/hardware_backend/npu/quantization/linear_method_npu.pytest/registered/unit/layers/quantization/test_modelslim_config.pytest/registered/unit/models/test_qwen3_vl_quant_config.pytest/registered/unit/layers/quantization/test_modelslim_mxfp8.pyImplementation Notes
Offline-only scope
These changes are exercised when SGLang detects
quant_model_description.jsonand loads a pre-quantized ModelSlim checkpoint. The online--quantization mxfp8path continues to use the existing implementation from #30768 and is unchanged.Dependency on #32745
The Gemma RMSNorm support required for Ascend 950 is inherited from #32745. To keep that implementation single-sourced and preserve the intended branch history, merge #32745 first and then merge this PR.
Packed mapping precedence
Top-level list-valued mappings are model-specific and remain valid alongside the generic scoped mappings added by the NPU loader. The resolver merges the two levels instead of choosing one and discarding the other; a scoped entry wins only when the same packed-module key exists at both levels.
Offline scale padding
The zero padding is applied only to the synthetic partner of the final partial 32-element block. It does not replace or alter any checkpoint scale. After transpose, the kernel-facing layout is
[ceil(K / 64), out, 2].Validation
Benchmark Results
Performance
Online W8A8 results are intentionally omitted because this PR changes only the offline ModelSlim path.
For throughput metrics, improvement is
(Offline - BF16) / BF16. For latency metrics, improvement is(BF16 - Offline) / BF16; positive values are better.Accuracy
Online W8A8 results are intentionally omitted because this PR changes only the offline ModelSlim path.
Status
WIP / Draft. The offline ModelSlim fixes and focused regression cases are in place. Online MXFP8 is already supported and is unchanged by this PR. The remaining work is to run the five CPU tests, complete offline text + image end-to-end validation on Ascend 950 (A5), and merge #32745 first.
Related Issues
Part of #21584 (MXFP8 / MXFP4 support on Ascend NPU for Qwen3.5 LLM). Builds on the merged MoE W8A8 MXFP8 implementation from #30768 and dense MXFP8 infrastructure from #22352. Must be merged after #32745.
CI States
Latest PR Test (Base): ❌ Run #32805800005
Latest PR Test (Extra): ❌ Run #32805799860
Latest PR Test (AMD ROCm 7.2): ❌ Run #32805799970