🐛 [llm][npu][quant] Fix ModelSlim MXFP4 packed weight loading (#32013)

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
+    The msmodelslim ``W4A4_MXFP4`` checkpoint stores weights as packed ``uint8``
-    container** (``float8_e4m3fn`` [out, in], one FP4 value per byte) plus UE8M0
+    [out, in//2] (two FP4 values per byte) plus UE8M0 block scales (``uint8``
-    block scales (``uint8`` [out, in//32]). The weight is re-packed to
+    [out, in//32]). The weight is transposed and the scale reshaped to 3D; it then
-    ``float4_e2m1fn_x2`` (two FP4 per byte) 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].
+        # The checkpoint is already packed two-FP4-per-byte. Preserve the strided
-        weight_fp4 = torch.ops.npu.npu_dtype_cast(weight, fp4_dtype)
+        # transpose used by vllm-ascend and by the online path.
-        # Transpose to [in//2, out]; no .contiguous() (preserve the strided view so
+        layer.weight = Parameter(weight.transpose(0, 1), requires_grad=False)
-        # the block-scale mapping stays intact).
-        layer.weight = Parameter(weight_fp4.transpose(0, 1), requires_grad=False)
 
         weight_scale = layer.weight_scale.data
         if not weight_scale.is_npu:

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
+``NPUSingleLevelMXFP4OfflineLinearMethod`` (``self.kernel``): the packed weight is
-is re-packed to ``float4_e2m1fn_x2`` (two FP4 per byte) + transposed and the scale
+transposed and the scale reshaped to 3D, then ``npu_quant_matmul`` runs with
-reshaped to 3D, then ``npu_quant_matmul`` runs with ``x1_dtype = x2_dtype =
+``x1_dtype = x2_dtype = float4_e2m1fn_x2`` and
-float4_e2m1fn_x2`` and ``group_sizes=[1, 1, 32]`` — sharing the online
+``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
+This differs from ``W4A8_MXFP`` only in the activation dtype (FP4 vs FP8), and
-from ``W8A8_MXFP8`` (float8_e4m3fn weights of shape [out, in]).
+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
+        # msmodelslim packs two FP4 values per byte along the input dimension.
-        # value per byte (fp8 container). The kernel re-packs it to
-        # float4_e2m1fn_x2 (2 FP4 per byte) in process_weights_after_loading.
         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,