✨ [llm][npu][quant] Add W4A4 MXFP4 quantization support for Qwen3 Dense on Ascend NPU (#23795)
This commit is contained in:
Junlin Wu
@@ -625,3 +625,278 @@ class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase):
|
||||
# Restore original shape (replace last dim with output features).
|
||||
output_shape = list(input_shape[:-1]) + [output.shape[-1]]
|
||||
return output.reshape(output_shape)
|
||||
|
||||
|
||||
class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase):
|
||||
"""Ascend 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
|
||||
activations. All NPU ops go through ``torch.ops.npu.*`` (no top-level
|
||||
``torch_npu``) and the fp4 dtype comes from ``_get_float4_e2m1fn_x2_dtype()``.
|
||||
|
||||
Weight quantization (process_weights_after_loading):
|
||||
BF16/FP16 weight → npu_dynamic_mx_quant(dst=float4_e2m1fn_x2)
|
||||
→ (packed FP4 [out, in//2], UE8M0 block scale) → transpose [in//2, out]
|
||||
|
||||
Inference (apply):
|
||||
BF16/FP16 activation → npu_dynamic_mx_quant(dst=float4_e2m1fn_x2) (A4)
|
||||
→ 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``.
|
||||
"""
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes,
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
"""Register an unquantized (``params_dtype``) weight placeholder.
|
||||
|
||||
The checkpoint still holds full-precision BF16/FP16 weights: the loader
|
||||
fills this buffer, then ``process_weights_after_loading`` quantizes it to
|
||||
MXFP4 in place.
|
||||
"""
|
||||
from sglang.srt.layers.parameter import ModelWeightParameter
|
||||
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
|
||||
layer.logical_widths = output_partition_sizes
|
||||
layer.input_size_per_partition = input_size_per_partition
|
||||
layer.output_size_per_partition = output_size_per_partition
|
||||
layer.orig_dtype = params_dtype
|
||||
|
||||
# Load weights in original dtype; quantise to MXFP4 in
|
||||
# process_weights_after_loading.
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
# Online single-level MXFP4 weight quant. All NPU ops go through
|
||||
# torch.ops.npu.* (no torch_npu); the fp4 dtype comes from the shared
|
||||
# _get_float4_e2m1fn_x2_dtype() helper (the torch_npu int enum).
|
||||
fp4_dtype = _get_float4_e2m1fn_x2_dtype()
|
||||
|
||||
weight_fp = layer.weight.data
|
||||
if weight_fp.dtype not in (torch.float16, torch.bfloat16):
|
||||
weight_fp = weight_fp.to(torch.bfloat16)
|
||||
# Move to NPU if needed (cpu offload may have put it on CPU).
|
||||
if not weight_fp.is_npu:
|
||||
weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}")
|
||||
|
||||
# BF16 -> packed FP4 (float4_e2m1fn_x2, [out, in//2]) + UE8M0 block scale.
|
||||
qw, w_scale = torch.ops.npu.npu_dynamic_mx_quant(
|
||||
weight_fp, dst_type=fp4_dtype, round_mode="round"
|
||||
)
|
||||
# Pre-transpose the weight to [in//2, out] for npu_quant_matmul; use
|
||||
# .data= to preserve the non-contiguous transpose view (npu_quant_matmul
|
||||
# reads strides directly — .contiguous() would reorder data and break
|
||||
# block-scale alignment).
|
||||
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
|
||||
# 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.
|
||||
if w_scale.dim() == 2:
|
||||
n, k = w_scale.shape
|
||||
w_scale = w_scale.reshape(n, k // 2, 2)
|
||||
layer.weight_scale = Parameter(w_scale.transpose(-3, -2), requires_grad=False)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
e8m0_dtype = _get_float8_e8m0fnu_dtype()
|
||||
fp4_dtype = _get_float4_e2m1fn_x2_dtype()
|
||||
|
||||
original_dtype = x.dtype
|
||||
if original_dtype not in (torch.float16, torch.bfloat16):
|
||||
x = x.to(torch.bfloat16)
|
||||
original_dtype = torch.bfloat16
|
||||
|
||||
# Flatten to 2D [tokens, hidden] for npu_dynamic_mx_quant.
|
||||
input_shape = x.shape
|
||||
x_2d = x.reshape(-1, x.shape[-1])
|
||||
|
||||
# Dynamic single-level MXFP4 activation quantisation (A4 — FP4).
|
||||
qx, input_scale = torch.ops.npu.npu_dynamic_mx_quant(
|
||||
x_2d, dst_type=fp4_dtype, round_mode="round"
|
||||
)
|
||||
|
||||
# Single-level MXFP4 matmul (weight & scale already transposed at load
|
||||
# time): x1_dtype = x2_dtype = fp4, group_sizes=[1, 1, block].
|
||||
output = torch.ops.npu.npu_quant_matmul(
|
||||
qx,
|
||||
layer.weight,
|
||||
layer.weight_scale,
|
||||
scale_dtype=e8m0_dtype,
|
||||
pertoken_scale=input_scale,
|
||||
pertoken_scale_dtype=e8m0_dtype,
|
||||
bias=bias.to(torch.float32) if bias is not None else None,
|
||||
output_dtype=original_dtype,
|
||||
x1_dtype=fp4_dtype,
|
||||
x2_dtype=fp4_dtype,
|
||||
group_sizes=[1, 1, MXFP4_BLOCK_SIZE],
|
||||
)
|
||||
|
||||
# Restore original shape (replace last dim with output features).
|
||||
output_shape = list(input_shape[:-1]) + [output.shape[-1]]
|
||||
return output.reshape(output_shape)
|
||||
|
||||
|
||||
class NPUSingleLevelMXFP4OfflineLinearMethod(NPUSingleLevelMXFP4LinearMethod):
|
||||
"""Ascend NPU offline W4A4 (ModelSlim ``W4A4_MXFP4``): fp8-container 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
|
||||
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)
|
||||
|
||||
weight_scale = layer.weight_scale.data
|
||||
if not weight_scale.is_npu:
|
||||
weight_scale = weight_scale.to(f"npu:{torch.npu.current_device()}")
|
||||
# npu_quant_matmul with float4_e2m1fn_x2 requires x2Scale to be 3D:
|
||||
# [out, in/32] -> [out, in/64, 2] -> transpose to [in/64, out, 2].
|
||||
n_dim, k_dim = weight_scale.shape
|
||||
layer.weight_scale = Parameter(
|
||||
weight_scale.reshape(n_dim, k_dim // 2, 2).transpose(0, 1),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
|
||||
class NPUDualLevelMXFP4LinearMethod(NPUSingleLevelMXFP4LinearMethod):
|
||||
"""Ascend 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)
|
||||
block scale plus a coarser L1 scale, so per-block dynamic range is captured far
|
||||
more accurately — this fixed the online-RTN degradation that made single-level
|
||||
decoding loop (never emitting EOS) under greedy sampling. (The single-level
|
||||
:class:`NPUSingleLevelMXFP4LinearMethod` is retained only as the offline path's
|
||||
base — msmodelslim checkpoints ship single-level UE8M0 scales.)
|
||||
|
||||
All NPU ops go through ``torch.ops.npu.*`` (no top-level ``torch_npu``). Only
|
||||
``create_weights`` (the BF16/FP16 placeholder) is shared with the single-level
|
||||
base; weight post-processing and the matmul are fully dual-level.
|
||||
|
||||
Weight quantization (process_weights_after_loading):
|
||||
BF16/FP16 weight → npu_dynamic_dual_level_mx_quant
|
||||
→ (packed FP4 weight, L0 scale, L1 scale); weight cast to FRACTAL_NZ,
|
||||
L0 scale transposed to [in//l0_block, out].
|
||||
|
||||
Inference (apply):
|
||||
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.
|
||||
"""
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
weight_fp = layer.weight.data
|
||||
if weight_fp.dtype not in (torch.float16, torch.bfloat16):
|
||||
weight_fp = weight_fp.to(torch.bfloat16)
|
||||
# Move to NPU if needed (cpu offload may have put it on CPU).
|
||||
if not weight_fp.is_npu:
|
||||
weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}")
|
||||
|
||||
# Dual-level MXFP4 weight quant: packed FP4 weight + L0 (fine, FP8 E4M3)
|
||||
# and L1 (coarse) block scales.
|
||||
qw, w_l0_scale, w_l1_scale = torch.ops.npu.npu_dynamic_dual_level_mx_quant(
|
||||
weight_fp, smooth_scale=None
|
||||
)
|
||||
|
||||
# npu_dual_level_quant_matmul requires the weight (x2) in FRACTAL_NZ.
|
||||
# View the packed FP4 as int8 first (npu_format_cast takes int dtypes).
|
||||
qw_nz = npu_format_cast(
|
||||
qw.view(torch.int8),
|
||||
NPUACLFormat.ACL_FORMAT_FRACTAL_NZ,
|
||||
customize_dtype=torch.int8,
|
||||
)
|
||||
|
||||
# L0 scale -> [in//l0_block, out] (op returns [out, in//l0_block, 1]).
|
||||
w_l0_scale = w_l0_scale.squeeze(-1).transpose(0, 1).contiguous()
|
||||
|
||||
layer.weight = Parameter(qw_nz, requires_grad=False)
|
||||
layer.weight_l0_scale = Parameter(w_l0_scale, requires_grad=False)
|
||||
layer.weight_l1_scale = Parameter(w_l1_scale, requires_grad=False)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
original_dtype = x.dtype
|
||||
if original_dtype not in (torch.float16, torch.bfloat16):
|
||||
x = x.to(torch.bfloat16)
|
||||
original_dtype = torch.bfloat16
|
||||
|
||||
# Flatten to 2D [tokens, hidden] for the quant operators.
|
||||
input_shape = x.shape
|
||||
x_2d = x.reshape(-1, x.shape[-1])
|
||||
|
||||
# Dynamic dual-level MXFP4 activation quant (A4): packed FP4 + L0/L1 scales.
|
||||
qx, act_l0_scale, act_l1_scale = torch.ops.npu.npu_dynamic_dual_level_mx_quant(
|
||||
x_2d, smooth_scale=None
|
||||
)
|
||||
|
||||
# Dual-level matmul. Arg order (act, weight, act_l0, w_l0, act_l1, w_l1);
|
||||
# the weight is NOT transposed here (unlike the single-level path).
|
||||
output = torch.ops.npu.npu_dual_level_quant_matmul(
|
||||
qx,
|
||||
layer.weight,
|
||||
act_l0_scale,
|
||||
layer.weight_l0_scale,
|
||||
act_l1_scale,
|
||||
layer.weight_l1_scale,
|
||||
bias=bias.to(torch.float32) if bias is not None else None,
|
||||
output_dtype=original_dtype,
|
||||
)
|
||||
|
||||
# Restore original shape (replace last dim with output features).
|
||||
output_shape = list(input_shape[:-1]) + [output.shape[-1]]
|
||||
return output.reshape(output_shape)
|
||||
|
||||
@@ -46,6 +46,7 @@ from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
|
||||
from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config
|
||||
from sglang.srt.layers.quantization.mxfp4 import Mxfp4Config
|
||||
from sglang.srt.layers.quantization.npu_mxfp4 import Mxfp4W4A8Config
|
||||
from sglang.srt.layers.quantization.npu_mxfp4_w4a4 import Mxfp4W4A4Config
|
||||
from sglang.srt.layers.quantization.nvfp4_online import NvFp4OnlineConfig
|
||||
from sglang.srt.layers.quantization.petit import PetitNvFp4Config
|
||||
from sglang.srt.layers.quantization.qoq import QoQConfig
|
||||
@@ -117,6 +118,10 @@ if is_npu():
|
||||
BASE_QUANTIZATION_METHODS.update(
|
||||
{
|
||||
"gptq": GPTQAscendConfig,
|
||||
# On NPU, `mxfp4` means single-level W4A4 MXFP4 for dense LLM (the
|
||||
# upstream `Mxfp4Config` OCP-MoE path is only registered on
|
||||
# cpu/cuda/hip above, so there is no collision here).
|
||||
"mxfp4": Mxfp4W4A4Config,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ from sglang.srt.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
)
|
||||
from sglang.srt.layers.quantization.modelslim.schemes import (
|
||||
ModelSlimMXFP4Scheme,
|
||||
ModelSlimMXFP4W4A8Scheme,
|
||||
ModelSlimMXFP8Scheme,
|
||||
ModelSlimW4A4Int4,
|
||||
@@ -208,6 +209,7 @@ class ModelSlimConfig(QuantizationConfig):
|
||||
("W8A8_DYNAMIC", ModelSlimW8A8Int8),
|
||||
("W8A8_MXFP8", ModelSlimMXFP8Scheme),
|
||||
("W4A8_MXFP", ModelSlimMXFP4W4A8Scheme),
|
||||
("W4A4_MXFP4", ModelSlimMXFP4Scheme),
|
||||
]
|
||||
|
||||
quant_schemes = [self.quant_description.get(prefix + ".weight", "")]
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
from .modelslim_scheme import ModelSlimLinearScheme, ModelSlimMoEScheme
|
||||
from .modelslim_mxfp8 import ModelSlimMXFP8Scheme
|
||||
from .modelslim_mxfp4_w4a8 import ModelSlimMXFP4W4A8Scheme
|
||||
from .modelslim_mxfp4 import ModelSlimMXFP4Scheme
|
||||
|
||||
# isort: on
|
||||
from .modelslim_w4a4_int4 import ModelSlimW4A4Int4
|
||||
@@ -20,6 +21,7 @@ __all__ = [
|
||||
"ModelSlimMoEScheme",
|
||||
"ModelSlimMXFP8Scheme",
|
||||
"ModelSlimMXFP4W4A8Scheme",
|
||||
"ModelSlimMXFP4Scheme",
|
||||
"ModelSlimW8A8Int8",
|
||||
"ModelSlimW4A4Int4",
|
||||
"ModelSlimW4A4Int4MoE",
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
"""ModelSlim W4A4_MXFP4 scheme for pre-quantized weight inference on Ascend NPU (SRT).
|
||||
|
||||
The msmodelslim ``W4A4_MXFP4`` checkpoint stores weights in an **fp8 container**:
|
||||
|
||||
weight: float8_e4m3fn shape [out, in] (one FP4 value 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
|
||||
``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]).
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
|
||||
NPUSingleLevelMXFP4OfflineLinearMethod,
|
||||
)
|
||||
from sglang.srt.layers.parameter import GroupQuantScaleParameter, ModelWeightParameter
|
||||
from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimLinearScheme
|
||||
|
||||
# Fixed by the msmodelslim W4A4_MXFP4 export format (group_size=32).
|
||||
MXFP4_BLOCK_SIZE = 32
|
||||
|
||||
|
||||
class ModelSlimMXFP4Scheme(ModelSlimLinearScheme):
|
||||
"""W4A4_MXFP4 offline scheme — fp8-container FP4 weights, MXFP4 activations."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
quant_config: Optional[Dict[str, any]] = None,
|
||||
prefix: Optional[str] = None,
|
||||
):
|
||||
# quant_config / prefix accepted to match ModelSlimConfig.get_linear_scheme's
|
||||
# dispatch signature; W4A4_MXFP4 needs no per-layer config beyond create_weights.
|
||||
del quant_config, prefix
|
||||
self.kernel = NPUSingleLevelMXFP4OfflineLinearMethod()
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: List[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
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.
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
(output_size_per_partition, input_size_per_partition),
|
||||
dtype=torch.float8_e4m3fn,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# UE8M0 block scales: uint8, shape [out, in//32].
|
||||
scale_dim = input_size_per_partition // MXFP4_BLOCK_SIZE
|
||||
weight_scale = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
(output_size_per_partition, scale_dim),
|
||||
dtype=torch.uint8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module):
|
||||
self.kernel.process_weights_after_loading(layer)
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
return self.kernel.apply(layer, x, bias)
|
||||
@@ -0,0 +1,140 @@
|
||||
"""MXFP4 W4A4 online quantization config (dual-level MXFP4 weights + activations).
|
||||
|
||||
Triggered by ``--quantization mxfp4`` on the Ascend NPU backend. On CUDA / AMD /
|
||||
CPU the ``mxfp4`` key resolves to the upstream :class:`Mxfp4Config` (OCP MXFP4
|
||||
MoE) instead; the per-device split is done at registration time in
|
||||
``sglang.srt.layers.quantization.__init__`` (this config is only registered
|
||||
inside the ``is_npu()`` block, mirroring ``GPTQAscendConfig``).
|
||||
|
||||
Online mode: FP16/BF16 weights are quantised to **dual-level** MXFP4 in
|
||||
``process_weights_after_loading`` (a finer FP8 E4M3 L0 block scale plus a coarser
|
||||
L1 scale); activations are dynamically quantised the same way and the matmul runs
|
||||
via ``npu_dual_level_quant_matmul`` (see :class:`NPUDualLevelMXFP4LinearMethod`).
|
||||
Dual-level is the sole online path — it captures per-block dynamic range far more
|
||||
accurately than a single-level UE8M0 scale, avoiding the RTN degradation that made
|
||||
single-level online decoding loop under greedy sampling. Requires Ascend 950 (A5).
|
||||
|
||||
Offline (msmodelslim ``W4A4_MXFP4``) checkpoints are single-level (the checkpoint
|
||||
stores UE8M0 scales) and are handled separately by the ``modelslim`` config
|
||||
(``ModelSlimMXFP4Scheme`` → ``NPUSingleLevelMXFP4OfflineLinearMethod``), not this
|
||||
class.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.quantization.base_config import (
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from sglang.srt.layers.quantization.unquant import (
|
||||
UnquantizedFusedMoEMethod,
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
from sglang.srt.layers.quantization.utils import is_layer_skipped
|
||||
from sglang.srt.utils import is_npu
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Mxfp4W4A4Config(QuantizationConfig):
|
||||
"""Single-level MXFP4 W4A4 online quantization config for Ascend NPU.
|
||||
|
||||
True W4(weight) A4(activation): both weights and activations are quantised
|
||||
to single-level MXFP4 (``float4_e2m1fn_x2``). The device-specific linear
|
||||
method is selected in ``get_quant_method``; only Ascend NPU is wired up
|
||||
today (on other devices ``mxfp4`` maps to the upstream ``Mxfp4Config``).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ignored_layers: Optional[List[str]] = None,
|
||||
packed_modules_mapping: Optional[Dict[str, str]] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.ignored_layers = ignored_layers or []
|
||||
self.packed_modules_mapping = packed_modules_mapping or {}
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> str:
|
||||
return "mxfp4"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> List[torch.dtype]:
|
||||
return [torch.bfloat16, torch.half]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 0 # NPU bypasses CUDA capability checks
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> List[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict) -> Mxfp4W4A4Config:
|
||||
ignored_layers = cls.get_from_keys_or(
|
||||
config, ["ignored_layers", "modules_to_not_convert"], None
|
||||
)
|
||||
if ignored_layers:
|
||||
normalized: List[str] = []
|
||||
for layer in ignored_layers:
|
||||
base = layer.removeprefix("model.")
|
||||
normalized.append(base)
|
||||
normalized.append(f"model.{base}")
|
||||
ignored_layers = normalized
|
||||
packed_modules_mapping = (
|
||||
cls.get_from_keys_or(config, ["packed_modules_mapping"], {}) or {}
|
||||
)
|
||||
return cls(
|
||||
ignored_layers=ignored_layers,
|
||||
packed_modules_mapping=packed_modules_mapping,
|
||||
)
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> Optional[QuantizeMethodBase]:
|
||||
from sglang.srt.layers.linear import LinearBase
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
|
||||
|
||||
if isinstance(layer, LinearBase):
|
||||
if is_layer_skipped(
|
||||
prefix,
|
||||
self.ignored_layers,
|
||||
fused_mapping=self.packed_modules_mapping,
|
||||
):
|
||||
return UnquantizedLinearMethod()
|
||||
if is_npu():
|
||||
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
|
||||
NPUDualLevelMXFP4LinearMethod,
|
||||
)
|
||||
|
||||
# Online W4A4 always uses dual-level MXFP4 (finer FP8 L0 scales):
|
||||
# single-level RTN was too lossy and degenerated under greedy
|
||||
# decoding. Requires Ascend 950 (A5). The single-level kernel is
|
||||
# retained only for the offline msmodelslim path.
|
||||
return NPUDualLevelMXFP4LinearMethod(self)
|
||||
raise NotImplementedError(
|
||||
"mxfp4 W4A4 (single-level MXFP4 weights + activations) is currently "
|
||||
"only implemented for the Ascend NPU backend; no CUDA/other-device "
|
||||
"kernel exists in this config. Add a device branch here when one lands."
|
||||
)
|
||||
elif isinstance(layer, FusedMoE):
|
||||
# MoE single-level MXFP4 W4A4 not yet implemented; fall back to unquantised
|
||||
logger.warning(
|
||||
"MXFP4 W4A4 quantization is not yet supported for FusedMoE layers "
|
||||
"(prefix=%s). Falling back to unquantized MoE — MoE weights will "
|
||||
"run in full precision (BF16/FP16).",
|
||||
prefix,
|
||||
)
|
||||
return UnquantizedFusedMoEMethod(
|
||||
layer.use_triton_kernels, layer.use_flashinfer_trtllm_moe
|
||||
)
|
||||
return None
|
||||
|
||||
def get_scaled_act_names(self) -> List[str]:
|
||||
return []
|
||||
Reference in New Issue
Block a user