From 308bc1228bc46e11283977e61c1065fbb2abad76 Mon Sep 17 00:00:00 2001 From: Junlin Wu Date: Fri, 21 Aug 2026 03:03:58 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=93=9D=20[NPU]=20Clean=20up=20quantizatio?= =?UTF-8?q?n=20comments=20(#34829)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: ronnie_zheng --- .../npu/quantization/linear_method_npu.py | 33 +++++++++---------- .../modelslim/schemes/modelslim_mxfp4_w4a8.py | 2 +- 2 files changed, 16 insertions(+), 19 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 61cb2961b..756a4ca41 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 @@ -241,10 +241,9 @@ class NPUMXFP8LinearMethod(_NPULinearMethodBase): # 2] as strided transpose views — DO NOT call .contiguous(). The matmul # reduction loop scans the in-dim per output column; the [out, in] # row-major source gives stride-1 access for that scan via the transpose - # view (matches msmodelslim's offline layout and vllm-ascend's - # AscendW8A8MXFP8DynamicLinearMethod). Calling .contiguous() physically - # reorders to [in, out] row-major, making the inner-loop stride = out and - # tanking HBM bandwidth. + # view, matching msmodelslim's offline layout. Calling .contiguous() + # physically reorders to [in, out] row-major, making the inner-loop stride + # equal to out and tanking HBM bandwidth. # Cache FP32 bias once to avoid a per-forward dtype conversion + alloc. if ( @@ -531,7 +530,6 @@ class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase): BF16/FP16 activation → npu_dynamic_mx_quant(dst=float8_e4m3fn) (A8, MXFP8) → npu_quant_matmul(x2_dtype=float4_e2m1fn_x2, group_sizes=[0, 0, block]) - Mirrors vllm-ascend ``AscendW4A8MXFPDynamicLinearMethod`` exactly (Ascend 950/A5). The weight is cast to FRACTAL_NZ then transposed; ``npu_dynamic_mx_quant`` already returns a 3D ``[tokens, in//64, 2]`` block scale so the matmul needs no extra scale-layout normalization. @@ -541,8 +539,8 @@ class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase): a ``FRACTAL_NZ_C0_16`` tensor, which is fine). Older torch_npu (e.g. ``2.10.0.dev20260320``) had a broken FP4 matmul that rejected the NZ weight in *prefill* with ``x2 should be in ... nz format, but it is 2``; - ``2.10.0.post1.dev20260624`` (and later) runs the vllm-aligned NZ path - correctly. If you hit ``it is 2``, update torch_npu — do NOT "fix" it by + ``2.10.0.post1.dev20260624`` (and later) runs the FRACTAL_NZ path correctly. + If you hit ``it is 2``, update torch_npu — do NOT "fix" it by switching the weight to ND. ⚠️ A ``atb::OperationSetup`` *segfault during decode* (not prefill) is a @@ -550,8 +548,8 @@ class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase): backend, NOT this matmul (verified by stage-sync bisection — qkv's matmul syncs clean, the fault surfaces at the entry-sync of the next layer, i.e. the decode attention between qkv and o_proj). Run with the NPU decode graph (do - NOT pass ``--disable-cuda-graph``); graph mode is the NPU default and what - vllm uses. This attention issue is model-agnostic and out of scope for W4A8. + NOT pass ``--disable-cuda-graph``); graph mode is the NPU default. This + attention issue is model-agnostic and out of scope for W4A8. This is a true W4(weight) A8(activation) single-level matmul. The *online* ``NPUMXFP4W4A8LinearMethod`` now uses this exact apply path — the only @@ -561,10 +559,10 @@ class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase): """ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - # Mirror vllm-ascend AscendW4A8MXFPDynamicLinearMethod: cast the packed-FP4 - # weight to FRACTAL_NZ then transpose. All NPU ops go through - # torch.ops.npu.* (no torch_npu). Requires a recent torch_npu build (see - # class docstring): older builds reject the NZ weight ("x2 ... it is 2"). + # Cast the packed-FP4 weight to FRACTAL_NZ then transpose. All NPU ops go + # through torch.ops.npu.* (no torch_npu). A recent torch_npu build is + # required (see class docstring): older builds reject the NZ weight + # ("x2 ... it is 2"). fp4_dtype = _get_float4_e2m1fn_x2_dtype() # weight: packed-FP4 uint8 [out, in//2] -> FRACTAL_NZ (float8_e4m3fn view) @@ -608,7 +606,7 @@ class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase): if bias is not None and bias.dtype != torch.float32: bias = bias.to(torch.float32) - # W4(weight)A8(activation) matmul, mirroring vllm-ascend exactly. + # W4(weight)A8(activation) matmul. output = torch.ops.npu.npu_quant_matmul( quantized_x, layer.weight, @@ -712,8 +710,8 @@ class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase): layer.weight = Parameter(qw, requires_grad=False) layer.weight.data = layer.weight.data.transpose(0, 1) - # weight_scale -> [in//64, out, 2] (3D), matching the offline W4A4 path, - # the W4A8 path and vllm-ascend's W4A4_MXFP4 layout. npu_dynamic_mx_quant + # weight_scale -> [in//64, out, 2] (3D), matching the offline W4A4 and + # W4A8 paths. npu_dynamic_mx_quant # already returns the scale as [out, in//64, 2] (3D) on current builds; # older builds may return [out, in//32] (2D) — reshape those first so the # transpose always yields the 3D layout npu_quant_matmul requires. @@ -775,7 +773,6 @@ class NPUSingleLevelMXFP4OfflineLinearMethod(NPUSingleLevelMXFP4LinearMethod): [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: @@ -783,7 +780,7 @@ class NPUSingleLevelMXFP4OfflineLinearMethod(NPUSingleLevelMXFP4LinearMethod): if not weight.is_npu: weight = weight.to(f"npu:{torch.npu.current_device()}") # The checkpoint is already packed two-FP4-per-byte. Preserve the strided - # transpose used by vllm-ascend and by the online path. + # transpose used by the online path. layer.weight = Parameter(weight.transpose(0, 1), requires_grad=False) weight_scale = layer.weight_scale.data diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp4_w4a8.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp4_w4a8.py index ba841d7a0..104b7c22e 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp4_w4a8.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp4_w4a8.py @@ -12,7 +12,7 @@ scheme: weights are 4-bit FP4, activations are dynamically quantised to MXFP8. This is NOT the same layout as ``W8A8_MXFP8`` (which stores float8_e4m3fn weights of shape [out, in]) — so weight creation and the forward pass differ from MXFP8. Weight post-processing and the matmul are delegated to ``NPUMXFP4W4A8OfflineLinearMethod`` -(``self.kernel``), mirroring vllm-ascend's ``AscendW4A8MXFPDynamicLinearMethod``: +(``self.kernel``): ``npu_format_cast`` the packed FP4 to FRACTAL_NZ + transpose, then ``x2_dtype= float4_e2m1fn_x2`` matmul with ``group_sizes=[0, 0, 32]``. Requires a recent torch_npu for the FP4 matmul on Ascend 950/A5 (older builds reject the NZ weight) —