From d6fcfe02d6258ee8f47b9761a645ee5f842c0e05 Mon Sep 17 00:00:00 2001 From: Junlin Wu Date: Wed, 29 Jul 2026 11:34:41 +0800 Subject: [PATCH] :bug: [llm][npu][quant] Fix ModelSlim MXFP4 packed weight loading (#32013) --- .../npu/quantization/linear_method_npu.py | 22 +++++--------- .../modelslim/schemes/modelslim_mxfp4.py | 30 ++++++++++--------- 2 files changed, 23 insertions(+), 29 deletions(-) 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 166220a97..61cb2961b 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 @@ -767,32 +767,24 @@ class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase): class NPUSingleLevelMXFP4OfflineLinearMethod(NPUSingleLevelMXFP4LinearMethod): - """Ascend NPU offline W4A4 (ModelSlim ``W4A4_MXFP4``): fp8-container FP4 weights. + """Ascend 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 **fp4-in-fp8 - container** (``float8_e4m3fn`` [out, in], one FP4 value per byte) plus UE8M0 - block scales (``uint8`` [out, in//32]). The weight is re-packed to - ``float4_e2m1fn_x2`` (two FP4 per byte) and the scale reshaped to 3D; it then + The msmodelslim ``W4A4_MXFP4`` checkpoint stores weights as packed ``uint8`` + [out, in//2] (two FP4 values per byte) plus UE8M0 block scales (``uint8`` + [out, in//32]). The weight is transposed and the scale reshaped to 3D; it then shares the online :class:`NPUSingleLevelMXFP4LinearMethod` matmul (``apply``) exactly — only the weight source differs (msmodelslim checkpoint vs online RTN). Mirrors vllm-ascend's single-level W4A4 MXFP4 layout. """ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - # Re-pack fp8-container FP4 to float4_e2m1fn_x2 and pre-transpose to match - # the online path's layout. All NPU ops go through torch.ops.npu.* (no - # torch_npu); the fp4 dtype must be the torch_npu enum (helper). - fp4_dtype = _get_float4_e2m1fn_x2_dtype() - weight = layer.weight.data if not weight.is_npu: weight = weight.to(f"npu:{torch.npu.current_device()}") - # fp8 container -> float4_e2m1fn_x2 (2 FP4 per byte): [out, in] -> [out, in//2]. - weight_fp4 = torch.ops.npu.npu_dtype_cast(weight, fp4_dtype) - # Transpose to [in//2, out]; no .contiguous() (preserve the strided view so - # the block-scale mapping stays intact). - layer.weight = Parameter(weight_fp4.transpose(0, 1), requires_grad=False) + # The checkpoint is already packed two-FP4-per-byte. Preserve the strided + # transpose used by vllm-ascend and by the online path. + layer.weight = Parameter(weight.transpose(0, 1), requires_grad=False) weight_scale = layer.weight_scale.data if not weight_scale.is_npu: diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp4.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp4.py index 1dcb4d5e5..c9c03fbbb 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp4.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp4.py @@ -1,20 +1,20 @@ """ModelSlim W4A4_MXFP4 scheme for pre-quantized weight inference on Ascend NPU (SRT). -The msmodelslim ``W4A4_MXFP4`` checkpoint stores weights in an **fp8 container**: +The msmodelslim ``W4A4_MXFP4`` checkpoint stores weights as **packed FP4**: - weight: float8_e4m3fn shape [out, in] (one FP4 value per byte) + weight: uint8 shape [out, in//2] (two FP4 values per byte) weight_scale: uint8 (UE8M0) shape [out, in//32] (block scales, group_size=32) This is a true W4(weight) A4(activation) scheme: both weights and activations are single-level MXFP4. Weight post-processing and the matmul are delegated to -``NPUSingleLevelMXFP4OfflineLinearMethod`` (``self.kernel``): the fp8-container FP4 -is re-packed to ``float4_e2m1fn_x2`` (two FP4 per byte) + transposed and the scale -reshaped to 3D, then ``npu_quant_matmul`` runs with ``x1_dtype = x2_dtype = -float4_e2m1fn_x2`` and ``group_sizes=[1, 1, 32]`` — sharing the online +``NPUSingleLevelMXFP4OfflineLinearMethod`` (``self.kernel``): the packed weight is +transposed and the scale reshaped to 3D, then ``npu_quant_matmul`` runs with +``x1_dtype = x2_dtype = float4_e2m1fn_x2`` and +``group_sizes=[1, 1, 32]`` — sharing the online ``NPUSingleLevelMXFP4LinearMethod`` matmul exactly (only the weight source differs). -This differs from ``W4A8_MXFP`` (packed-uint8 FP4 weights + FP8 activations) and -from ``W8A8_MXFP8`` (float8_e4m3fn weights of shape [out, in]). +This differs from ``W4A8_MXFP`` only in the activation dtype (FP4 vs FP8), and +from ``W8A8_MXFP8`` in both weight packing and activation dtype. """ from typing import Dict, List, Optional @@ -29,10 +29,11 @@ from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimLinearSche # Fixed by the msmodelslim W4A4_MXFP4 export format (group_size=32). MXFP4_BLOCK_SIZE = 32 +MXFP4_PACK_FACTOR = 2 class ModelSlimMXFP4Scheme(ModelSlimLinearScheme): - """W4A4_MXFP4 offline scheme — fp8-container FP4 weights, MXFP4 activations.""" + """W4A4_MXFP4 offline scheme — packed-FP4 weights, MXFP4 activations.""" def __init__( self, @@ -57,13 +58,14 @@ class ModelSlimMXFP4Scheme(ModelSlimLinearScheme): weight_loader = extra_weight_attrs.get("weight_loader") output_size_per_partition = sum(output_partition_sizes) - # msmodelslim exports weight as float8_e4m3fn, shape [out, in] — one FP4 - # value per byte (fp8 container). The kernel re-packs it to - # float4_e2m1fn_x2 (2 FP4 per byte) in process_weights_after_loading. + # msmodelslim packs two FP4 values per byte along the input dimension. weight = ModelWeightParameter( data=torch.empty( - (output_size_per_partition, input_size_per_partition), - dtype=torch.float8_e4m3fn, + ( + output_size_per_partition, + input_size_per_partition // MXFP4_PACK_FACTOR, + ), + dtype=torch.uint8, ), input_dim=1, output_dim=0,