diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py index f58b8012f..00bb925b9 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py @@ -1,4 +1,4 @@ -"""ModelSlim MXFP4 scheme for pre-quantized weight inference on Ascend NPU. +"""ModelSlim MXFP4 scheme for pre-quantized weight inference on NPU. Loads weights pre-quantized by msmodelslim and runs MXFP4 dual-level matmul at inference via npu_dual_level_quant_matmul. @@ -8,9 +8,6 @@ Checkpoint tensor formats (verified from msmodelslim export): weight_scale: [out, in/32] uint8 (L1 block scales, e8m0+127) weight_dual_scale:[out, in/512, 1] float32 (L0 coarse scales) mul_scale: [in] float32 (smooth quant activation scale) - -Reference: MindIE-SD W4A4MXFP4DualQuantLinear -(MindIE-SD/mindiesd/quantization/layer.py) """ from typing import List, Optional @@ -98,7 +95,6 @@ class ModelSlimMXFP4Scheme(ModelSlimLinearScheme): # After repack, it becomes `.mul_scale`. # This is CRITICAL: the offline-quantized weights were calibrated with # x * mul_scale applied to the activation. Omitting it causes mosaic output. - # Ref: MindIE-SD W4A4MXFP4DualQuantLinear.quant_matmul lines 385-386. mul_scale = BasevLLMParameter( data=torch.empty( (input_size_per_partition,), @@ -119,7 +115,6 @@ class ModelSlimMXFP4Scheme(ModelSlimLinearScheme): weight = weight.to(f"npu:{torch.npu.current_device()}") weight = torch_npu.npu_dtype_cast(weight, torch_npu.float4_e2m1fn_x2) # npu_dual_level_quant_matmul requires x2 in FRACTAL_NZ format (format 29). - # Reference: MindIE-SD W4A4MXFP4DualQuantLinear._init_dynamic_quant_param weight = torch_npu.npu_format_cast( weight.view(torch.int8), 29, customize_dtype=torch.int8 ) @@ -170,7 +165,6 @@ class ModelSlimMXFP4Scheme(ModelSlimLinearScheme): # Apply smooth quant scale before activation quantization. # The offline-quantized weights were calibrated under x * mul_scale, # so we MUST apply it here for scale alignment. - # Reference: MindIE-SD W4A4MXFP4DualQuantLinear.quant_matmul mul_scale = layer.mul_scale if getattr(layer, "use_mul_scale", True): x_2d = x_2d * mul_scale.to(x_2d.dtype) diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py index 79202676e..c4c5dfb7b 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py @@ -1,4 +1,4 @@ -"""ModelSlim MXFP8 scheme for pre-quantized weight inference on Ascend NPU. +"""ModelSlim MXFP8 scheme for pre-quantized weight inference on NPU. Loads weights pre-quantized by msmodelslim (float8_e4m3fn weights, uint8 scales) and runs MXFP8 matmul at inference. diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4_npu.py b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4_npu.py index 09b4d69c7..57128b703 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4_npu.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4_npu.py @@ -1,4 +1,4 @@ -"""Online MXFP4 quantization for Diffusion models on Ascend NPU. +"""Online MXFP4 quantization for Diffusion models on NPU. Provides ``NPUMXFP4Config`` (registered as ``"mxfp4_npu"``) and ``NPUMXFP4DiffusionLinearMethod`` which quantises FP16/BF16 weights to MXFP4 @@ -10,9 +10,10 @@ The ``"mxfp4_npu"`` key is distinct from upstream's ROCm ``"mxfp4"`` (``Mxfp4Config`` in ``mxfp4.py``) which targets AMD MI350+ via aiter kernels. NOTE: Online weight quantization via ``npu_dynamic_dual_level_mx_quant`` is -experimental. MindIE-SD only uses an offline (pre-quantized) path for MXFP4 -weights. The online path quantizes FP16/BF16 weights at load time, which may -produce different numerical results than the offline calibrated path. +experimental; the established MXFP4 path for these models is offline +(pre-quantized) only. The online path quantizes FP16/BF16 weights at load +time, which may produce different numerical results than the offline +calibrated path. """ from __future__ import annotations @@ -41,7 +42,7 @@ logger = init_logger(__name__) class NPUMXFP4Config(QuantizationConfig): - """Config for online MXFP4 quantization on Ascend NPU (Diffusion).""" + """Config for online MXFP4 quantization on NPU (Diffusion).""" def __init__(self) -> None: super().__init__() @@ -78,13 +79,11 @@ class NPUMXFP4Config(QuantizationConfig): class NPUMXFP4DiffusionLinearMethod(LinearMethodBase): - """Ascend NPU MXFP4 linear method for Diffusion models (dual-level). + """NPU MXFP4 linear method for Diffusion models (dual-level). Online mode: loads FP16/BF16 weights → quantises to MXFP4 at load time via ``npu_dynamic_dual_level_mx_quant``. Inference: dynamic dual-level MXFP4 activation quant + dual-level matmul. - - Reference: MindIE-SD ``W4A4MXFP4DualQuantLinear`` (offline path only). """ def __init__(self, quant_config: NPUMXFP4Config): @@ -133,9 +132,9 @@ class NPUMXFP4DiffusionLinearMethod(LinearMethodBase): weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}") # Online dual-level MXFP4 weight quantisation. - # NOTE: This is experimental — MindIE-SD only has an offline path for - # MXFP4 weights. We assume npu_dynamic_dual_level_mx_quant can also - # quantise weights (not just activations). + # NOTE: This is experimental — the established MXFP4 path for these + # models is offline only. We assume npu_dynamic_dual_level_mx_quant can + # also quantise weights (not just activations). # Returns: (qw, w_dual_scale, w_scale) # qw — quantized weight in float4_e2m1fn_x2 (2 FP4 packed/byte) # w_dual_scale — L0-level scale (goes to pos 3 in npu_dual_level_quant_matmul) @@ -145,14 +144,12 @@ class NPUMXFP4DiffusionLinearMethod(LinearMethodBase): ) # npu_dual_level_quant_matmul requires x2 (weight) in FRACTAL_NZ format. - # Reference: MindIE-SD W4A4MXFP4DualQuantLinear._init_dynamic_quant_param qw = torch_npu.npu_format_cast( qw.view(torch.int8), 29, customize_dtype=torch.int8 ) # x2Level0Scale must be [in/level0_block_size, out] — transpose from # the [out, in/level0_block_size] shape returned by the quant op. - # Reference: MindIE-SD layer.py:409 w_dual_scale = w_dual_scale.squeeze(-1).transpose(0, 1).contiguous() layer.weight = Parameter(qw, requires_grad=False) diff --git a/python/sglang/srt/hardware_backend/npu/moe/init_routing.py b/python/sglang/srt/hardware_backend/npu/moe/init_routing.py index e942a369e..f4de357c6 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/init_routing.py +++ b/python/sglang/srt/hardware_backend/npu/moe/init_routing.py @@ -22,8 +22,7 @@ def _normalize_mxfp_scale(scale: torch.Tensor) -> torch.Tensor: ``npu_moe_init_routing_v2(quant_mode=3)`` emits the scale flat, while the grouped matmul wants the pair-split view. Already-3D scales (what - ``npu_dynamic_mx_quant`` returns) pass through untouched. Mirrors - vllm-ascend's ``maybe_normalize_mxfp_scale_layout``. + ``npu_dynamic_mx_quant`` returns) pass through untouched. """ if scale is None or scale.ndim != 2: return scale diff --git a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py index 756a4ca41..aac83c9ca 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py @@ -156,7 +156,7 @@ class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase): class NPUMXFP8LinearMethod(_NPULinearMethodBase): - """Ascend NPU MXFP8 linear method for LLM (SRT) models. + """NPU MXFP8 linear method for LLM (SRT) models. Shared kernel for both the online config path (``--quantization mxfp8``) and the offline ModelSlimMXFP8Scheme (which delegates to this as ``self.kernel``). @@ -343,7 +343,7 @@ class NPU_W4A4DynamicLinearMethod(_NPULinearMethodBase): class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase): - """Ascend NPU W4A8 online quantization: MXFP4 weights + MXFP8 activations. + """NPU W4A8 online quantization: MXFP4 weights + MXFP8 activations. This is a *true* W4(weight) A8(activation) path: it mirrors the offline ``W4A8_MXFP`` kernel (``NPUMXFP4W4A8OfflineLinearMethod``) exactly — the only @@ -363,7 +363,7 @@ class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase): BF16/FP16 activation → npu_dynamic_mx_quant(dst=float8_e4m3fn) (A8, FP8) → npu_quant_matmul(x2_dtype=float4_e2m1fn_x2, group_sizes=[0, 0, block]) - Hardware: Ascend 950 (A5) + a recent torch_npu with the FP4 npu_quant_matmul + Hardware: A5 NPU + a recent torch_npu with the FP4 npu_quant_matmul (same requirement as the offline W4A8 path — see that class's docstring). """ @@ -513,7 +513,7 @@ class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase): class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase): - """Ascend NPU offline W4A8 (ModelSlim ``W4A8_MXFP``): packed-FP4 weights + MXFP8 activations. + """NPU offline W4A8 (ModelSlim ``W4A8_MXFP``): packed-FP4 weights + MXFP8 activations. Kernel for the offline ModelSlimMXFP4W4A8Scheme (delegated as ``self.kernel``). The msmodelslim ``W4A8_MXFP`` checkpoint stores weights as *packed FP4* @@ -626,7 +626,7 @@ class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase): class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase): - """Ascend NPU W4A4 online quantization: single-level MXFP4. + """NPU W4A4 online quantization: single-level MXFP4. True W4(weight) A4(activation): both weights and activations are quantised to single-level MXFP4 (``float4_e2m1fn_x2``), unlike the W4A8 path which keeps FP8 @@ -642,8 +642,8 @@ class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase): → npu_quant_matmul(x1_dtype = x2_dtype = float4_e2m1fn_x2, group_sizes=[1, 1, MXFP4_BLOCK_SIZE]) - Triggered by ``--quantization mxfp4`` on Ascend NPU. Hardware: Ascend 950 (A5) - with a recent torch_npu exposing ``float4_e2m1fn_x2``. + Triggered by ``--quantization mxfp4`` on NPU. Hardware: A5 NPU with a recent + torch_npu exposing ``float4_e2m1fn_x2``. """ def create_weights( @@ -765,7 +765,7 @@ class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase): class NPUSingleLevelMXFP4OfflineLinearMethod(NPUSingleLevelMXFP4LinearMethod): - """Ascend NPU offline W4A4 (ModelSlim ``W4A4_MXFP4``): packed FP4 weights. + """NPU offline W4A4 (ModelSlim ``W4A4_MXFP4``): packed FP4 weights. Kernel for the offline ``ModelSlimMXFP4Scheme`` (delegated as ``self.kernel``). The msmodelslim ``W4A4_MXFP4`` checkpoint stores weights as packed ``uint8`` @@ -796,7 +796,7 @@ class NPUSingleLevelMXFP4OfflineLinearMethod(NPUSingleLevelMXFP4LinearMethod): class NPUDualLevelMXFP4LinearMethod(NPUSingleLevelMXFP4LinearMethod): - """Ascend NPU W4A4 online quantization: dual-level MXFP4 (higher accuracy). + """NPU W4A4 online quantization: dual-level MXFP4 (higher accuracy). This is the sole online ``--quantization mxfp4`` linear path. Instead of a single UE8M0 (power-of-2) block scale, dual-level MX quant produces a finer L0 (FP8 E4M3) @@ -819,9 +819,8 @@ class NPUDualLevelMXFP4LinearMethod(NPUSingleLevelMXFP4LinearMethod): BF16/FP16 activation → npu_dynamic_dual_level_mx_quant (A4, dual-level) → npu_dual_level_quant_matmul(act, weight, act_l0, w_l0, act_l1, w_l1) - Reference: Diffusion ``NPUMXFP4DiffusionLinearMethod`` / MindIE-SD - ``W4A4MXFP4DualQuantLinear``. Hardware: Ascend 950 (A5) only — the - ``DualLevelQuantBatchMatmul`` op is unavailable on A2/A3. + Reference: Diffusion ``NPUMXFP4DiffusionLinearMethod``. Hardware: A5 NPU + only — the ``DualLevelQuantBatchMatmul`` op is unavailable on A2/A3. """ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py b/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py index 6cfe9fb32..df21796f8 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py @@ -223,7 +223,7 @@ class NPUW4A8MXFP4MoEMethod(_NPUMoEMethodBase): ).transpose(1, 2) weight_scale.data = scale - # The refactored Ascend dispatchers currently support BF16 and INT8. + # The refactored NPU dispatchers currently support BF16 and INT8. # Keep dispatch in BF16 and quantize to MXFP8 immediately before GMM. if weight_prefix == "w13": self._set_dispatcher_output_dtype(layer, "bf16") @@ -304,7 +304,7 @@ class NPUW4A4MXFP4MoEMethod(_NPUMoEMethodBase): ).transpose(1, 2) weight_scale.data = scale - # The refactored Ascend dispatchers currently support BF16 and INT8. + # The refactored NPU dispatchers currently support BF16 and INT8. # Keep dispatch in BF16 and quantize immediately before each GMM. if weight_prefix == "w13": self._set_dispatcher_output_dtype(layer, "bf16") @@ -929,7 +929,7 @@ class NPUUnquantMoEMethod(_NPUMoEMethodBase): # NPUMXFP8MoEMethod # --------------------------------------------------------------------------- class NPUMXFP8MoEMethod(_NPUMoEMethodBase): - """MXFP8 MoE on Ascend A5 – float8_e4m3fn weights with e8m0 block scales. + """MXFP8 MoE on A5 NPU – float8_e4m3fn weights with e8m0 block scales. Serves both the online config path (``--quantization mxfp8``, weights quantised at load time) and the offline ModelSlim ``W8A8_MXFP8`` scheme @@ -1016,8 +1016,8 @@ class NPUMXFP8MoEMethod(_NPUMoEMethodBase): # [E, K//64, N, 2] as strided transpose views — DO NOT call # .contiguous(). Beyond breaking the transpose-flag match above, it # measures slower on the same probe: making both sides contiguous costs - # 6.2% on decode. This matches NPUMXFP8LinearMethod, msmodelslim's - # offline layout and vllm-ascend's AscendW8A8MXFP8DynamicFusedMoEMethod. + # 6.2% on decode. This matches NPUMXFP8LinearMethod and msmodelslim's + # offline layout. setattr( layer, f"{weight_prefix}_weight", diff --git a/python/sglang/srt/hardware_backend/npu/quantization/online_moe_methods.py b/python/sglang/srt/hardware_backend/npu/quantization/online_moe_methods.py index 77da812ee..4a9a0ef43 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/online_moe_methods.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/online_moe_methods.py @@ -1,4 +1,4 @@ -"""Online (config-driven) quantized FusedMoE methods for Ascend NPU. +"""Online (config-driven) quantized FusedMoE methods for NPU. These are the ``--quantization `` entry points: the checkpoint holds BF16/FP16 expert weights and the per-gmm kernels quantize them at load time. @@ -28,7 +28,7 @@ class NPUMXFP8OnlineMoEMethod(UnquantizedFusedMoEMethod): """Online MXFP8 FusedMoE entry point (``--quantization mxfp8`` on A5). Weight creation, weight post-processing and the forward pass are identical - to the unquantized Ascend path — the only difference is which per-gmm kernel + to the unquantized NPU path — the only difference is which per-gmm kernel the layer gets, so everything but ``create_moe_runner`` is inherited. ``NPUMXFP8MoEMethod`` then quantizes the BF16 expert weights to MXFP8 in ``process_weights_after_loading``. diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index dac64b987..0937e36e1 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -267,9 +267,9 @@ class Qwen3MoeSparseMoeBlock(nn.Module): routing_method_type=RoutingMethodType.Renormalize, ) - # Router gate: description-driven quant, mirroring vllm-ascend. Only the - # offline ModelSlim path (which carries a per-layer quant_model_description) - # may quantise the gate — if the checkpoint stored it as MXFP8 it is loaded + # Router gate: description-driven quant. Only the offline ModelSlim path + # (which carries a per-layer quant_model_description) may quantise the + # gate — if the checkpoint stored it as MXFP8 it is loaded # and dequantised correctly instead of cast to bf16 without its block scale. # The online Fp8/mxfp8 path keeps the gate in bf16 (unchanged, verified). gate_quant_config = (