chore: import upstream snapshot with attribution

python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
+    NPUW4A16Int4DynamicMoEMethod,
+)
+from sglang.srt.layers.quantization.utils import replace_parameter
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+import torch_npu
+
+
+class AWQAscendLinearKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
+        qweight_tmp = torch.zeros_like(layer.qweight.data)
+        qzeros_tmp = layer.qzeros.data
+        qzeros_list = []
+        shifts = [0, 4, 1, 5, 2, 6, 3, 7]
+
+        for i in range(0, self.quant_config.pack_factor):
+            shift_num = shifts[i] * 4
+            qzeros_list.append((qzeros_tmp.reshape(-1, 1) >> shift_num) & 0xF)
+            qweight_tmp.bitwise_or_(
+                ((layer.qweight.data >> shift_num) & 0xF) << (4 * i)
+            )
+
+        qweight_tmp.bitwise_xor_(0x88888888)
+
+        qzeros_tmp = torch.cat(qzeros_list, dim=-1).reshape(qzeros_tmp.shape[0], -1)
+        qzeros_tmp = -(qzeros_tmp - 8)
+        qzeros_tmp = qzeros_tmp.to(layer.scales.data.dtype)
+
+        layer.zeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False)
+        layer.weight = torch.nn.Parameter(qweight_tmp, requires_grad=False)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        qweight = layer.weight
+        scales = layer.scales
+        qzeros = layer.zeros
+        pack_factor = self.quant_config.pack_factor
+        out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
+        reshaped_x = x.reshape(-1, x.shape[-1])
+
+        if bias is not None and bias.dtype == torch.bfloat16:
+            bias = bias.float()
+
+        out = torch_npu.npu_weight_quant_batchmatmul(
+            reshaped_x,
+            qweight,
+            antiquant_scale=scales,
+            antiquant_offset=qzeros,
+            antiquant_group_size=self.quant_config.group_size,
+            bias=bias,
+        )
+
+        return out.reshape(out_shape)
+
+
+class AWQAscendMoEKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+        self.kernel = NPUW4A16Int4DynamicMoEMethod()
+
+    @staticmethod
+    def _register_or_replace_parameter(
+        layer: torch.nn.Module, name: str, tensor: torch.Tensor
+    ) -> None:
+        if hasattr(layer, name):
+            replace_parameter(layer, name, tensor)
+        else:
+            layer.register_parameter(
+                name, torch.nn.Parameter(tensor, requires_grad=False)
+            )
+
+    def _convert_awq_weight_to_npu_layout(self, qweight: torch.Tensor) -> torch.Tensor:
+        num_experts, input_size, _ = qweight.shape
+        unpacked_weight = (
+            self.kernel._unpack_from_int32(qweight.flatten(0, 1), 4)
+            .view(num_experts, input_size, -1)
+            .transpose(1, 2)
+            .contiguous()
+            .int()
+        )
+        return self.kernel._pack_to_int32(unpacked_weight)
+
+    def _convert_awq_qzeros_to_npu_offset(
+        self, qzeros: torch.Tensor, dtype: torch.dtype
+    ) -> torch.Tensor:
+        num_experts, num_groups, _ = qzeros.shape
+        offset = (
+            -self.kernel._unpack_from_int32(qzeros.flatten(0, 1), 4)
+            .view(num_experts, num_groups, -1)
+            .transpose(1, 2)
+            .contiguous()
+        )
+        return offset.to(dtype)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        self._register_or_replace_parameter(
+            layer,
+            "w13_weight",
+            self._convert_awq_weight_to_npu_layout(layer.w13_qweight.data),
+        )
+        self._register_or_replace_parameter(
+            layer,
+            "w2_weight",
+            self._convert_awq_weight_to_npu_layout(layer.w2_qweight.data),
+        )
+        self._register_or_replace_parameter(
+            layer,
+            "w13_weight_scale",
+            layer.w13_scales.data.transpose(1, 2).contiguous(),
+        )
+        self._register_or_replace_parameter(
+            layer,
+            "w2_weight_scale",
+            layer.w2_scales.data.transpose(1, 2).contiguous(),
+        )
+        self._register_or_replace_parameter(
+            layer,
+            "w13_weight_offset",
+            self._convert_awq_qzeros_to_npu_offset(
+                layer.w13_qzeros.data, layer.w13_scales.data.dtype
+            ),
+        )
+        self._register_or_replace_parameter(
+            layer,
+            "w2_weight_offset",
+            self._convert_awq_qzeros_to_npu_offset(
+                layer.w2_qzeros.data, layer.w2_scales.data.dtype
+            ),
+        )
+
+        self.kernel.process_weights_after_loading(layer)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> torch.Tensor:
+        return self.kernel.apply(layer, dispatch_output)
+
+    def apply_without_routing_weights(
+        self,
+        layer,
+        hidden_states,
+        hidden_states_scale,
+        group_list_type,
+        group_list,
+        output_dtype,
+    ):
+        return self.kernel.apply_without_routing_weights(
+            layer,
+            hidden_states,
+            hidden_states_scale,
+            group_list_type,
+            group_list,
+            output_dtype,
+        )

python/sglang/srt/hardware_backend/npu/quantization/gptq_kernels.py

@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+import torch_npu
+
+from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
+    npu_fused_experts,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe import MoeRunnerConfig
+    from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+
+def unpack_from_int32(
+    weight: torch.Tensor,
+    num_bits: int,
+    packed_dim: int = 1,
+) -> torch.Tensor:
+    """
+    Unpacks quantized weights from int32 format back to original bits.
+
+    :param weight: The packed int32 tensor containing quantized weights
+    :param num_bits: The number of bits used for quantization (<= 8)
+    :param packed_dim: Dimension along which weights are packed (0 or 1), defaults to 1
+    :return: Unpacked tensor with int8 dtype after applying offset correction
+    """
+    assert (
+        weight.dtype == torch.int32
+    ), f"Expecting `weight.dtype` is torch.int32 but got {weight.dtype}."
+    assert (
+        num_bits <= 8
+    ), f"Expecting `num_bits` should not be larger than 8 but got {num_bits}."
+
+    pack_factor = 32 // num_bits
+    mask = (1 << num_bits) - 1
+
+    if packed_dim == 1:
+        unpacked_weight = torch.zeros(
+            (weight.shape[0], weight.shape[1] * pack_factor),
+            device=weight.device,
+            dtype=torch.int32,
+        )
+        for i in range(pack_factor):
+            unpacked_weight[:, i::pack_factor] = (weight >> (num_bits * i)) & mask
+    else:
+        unpacked_weight = torch.zeros(
+            (weight.shape[0] * pack_factor, weight.shape[1]),
+            device=weight.device,
+            dtype=torch.int32,
+        )
+        for i in range(pack_factor):
+            unpacked_weight[i::pack_factor, :] = (weight >> (num_bits * i)) & mask
+    offset = pow(2, num_bits) // 2
+    unpacked_weight = (unpacked_weight - offset).to(torch.int8)
+    return unpacked_weight
+
+
+class GPTQLinearAscendKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+        self.use_v2_format = quant_config.checkpoint_format == "gptq_v2"
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+
+        layer.qzeros = torch.nn.Parameter(
+            unpack_from_int32(
+                layer.qzeros.data.contiguous(),
+                self.quant_config.weight_bits,
+                packed_dim=1,
+            ).to(layer.scales.dtype),
+            requires_grad=False,
+        )
+        if not self.use_v2_format:
+            layer.qzeros += 1
+
+        qweight_tmp = unpack_from_int32(
+            layer.qweight.data.contiguous(), self.quant_config.weight_bits, packed_dim=0
+        )
+        # use int8 to store weight by default
+        if self.quant_config.weight_bits != 4:
+            layer.qweight = torch.nn.Parameter(
+                qweight_tmp,
+                requires_grad=False,
+            )
+            return
+
+        # for 4bit case we need to pack 4bit weight to int32 to save memory
+        layer.qweight = torch.nn.Parameter(
+            torch_npu.npu_convert_weight_to_int4pack(qweight_tmp.to(torch.int32)),
+            requires_grad=False,
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        qweight = layer.qweight
+        scales = layer.scales
+        qzeros = layer.qzeros
+
+        reshaped_x = x.reshape(-1, x.shape[-1])
+
+        if bias is not None and bias.dtype == torch.bfloat16:
+            bias = bias.float()
+
+        # 4bit weight is packed to int32(8 x int4)
+        if self.quant_config.weight_bits == 4:
+            out_shape = x.shape[:-1] + (qweight.shape[-1] * 8,)
+        else:
+            out_shape = x.shape[:-1] + (qweight.shape[-1],)
+
+        out = torch_npu.npu_weight_quant_batchmatmul(
+            reshaped_x,
+            qweight,
+            antiquant_scale=scales,
+            antiquant_offset=qzeros,
+            antiquant_group_size=self.quant_config.group_size,
+            bias=bias,
+        )
+
+        return out.reshape(out_shape)
+
+
+class GPTQMoEAscendKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+        self.use_v2_format = quant_config.checkpoint_format == "gptq_v2"
+        self.moe_runner_config: Optional[MoeRunnerConfig] = None
+
+    def create_moe_runner(
+        self,
+        layer: torch.nn.Module,
+        moe_runner_config: MoeRunnerConfig,
+        **extra_weight_attrs,
+    ):
+        self.moe_runner_config = moe_runner_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        w13_qzeros_2d = layer.w13_qzeros.data.contiguous().reshape(
+            -1, layer.w13_qzeros.shape[-1]
+        )
+        layer.w13_qzeros = torch.nn.Parameter(
+            unpack_from_int32(
+                w13_qzeros_2d,
+                self.quant_config.weight_bits,
+                packed_dim=1,
+            )
+            .reshape(layer.w13_qzeros.shape[0], layer.w13_qzeros.shape[1], -1)
+            .to(layer.w13_scales.dtype),
+            requires_grad=False,
+        )
+        if not self.use_v2_format:
+            layer.w13_qzeros += 1
+
+        w2_qzeros_2d = layer.w2_qzeros.data.contiguous().reshape(
+            -1, layer.w2_qzeros.shape[-1]
+        )
+        layer.w2_qzeros = torch.nn.Parameter(
+            unpack_from_int32(
+                w2_qzeros_2d,
+                self.quant_config.weight_bits,
+                packed_dim=1,
+            )
+            .reshape(layer.w2_qzeros.shape[0], layer.w2_qzeros.shape[1], -1)
+            .to(layer.w2_scales.dtype),
+            requires_grad=False,
+        )
+        if not self.use_v2_format:
+            layer.w2_qzeros += 1
+
+        w13_qweight_2d = (
+            layer.w13_qweight.data.transpose(-1, -2)
+            .contiguous()
+            .reshape(-1, layer.w13_qweight.shape[-2])
+        )
+        w13_qweight_tmp = unpack_from_int32(
+            w13_qweight_2d, self.quant_config.weight_bits, packed_dim=1
+        )
+
+        if self.quant_config.weight_bits == 4:
+            group_size = self.quant_config.group_size
+            scale_expanded = layer.w13_scales.data.repeat_interleave(group_size, dim=1)
+
+            neg_mask = scale_expanded < 0
+
+            if neg_mask.any():
+                neg_mask = neg_mask.transpose(-1, -2)
+                neg_mask = neg_mask.contiguous().reshape(w13_qweight_tmp.shape)
+                w13_qweight_tmp[neg_mask] = -w13_qweight_tmp[neg_mask]
+
+                if w13_qweight_tmp.max() > 7:
+                    w13_qweight_tmp.clamp_(max=7)
+
+                layer.w13_scales.data.abs_()
+
+            layer.w13_qweight = torch.nn.Parameter(
+                torch_npu.npu_convert_weight_to_int4pack(
+                    w13_qweight_tmp.reshape(
+                        layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1
+                    )
+                    .transpose(-1, -2)
+                    .contiguous()
+                    .reshape(-1, layer.w13_qweight.shape[2])
+                    .to(torch.int32)
+                )
+                .reshape(layer.w13_qweight.shape[0], layer.w13_qweight.shape[1] * 8, -1)
+                .contiguous(),
+                requires_grad=False,
+            )
+        # use int8 to store weight by default
+        else:
+            layer.w13_qweight = torch.nn.Parameter(
+                w13_qweight_tmp.reshape(
+                    layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1
+                )
+                .transpose(-1, -2)
+                .contiguous(),
+                requires_grad=False,
+            )
+
+        w2_qweight_2d = (
+            layer.w2_qweight.data.transpose(-1, -2)
+            .contiguous()
+            .reshape(-1, layer.w2_qweight.shape[-2])
+        )
+        w2_qweight_tmp = unpack_from_int32(
+            w2_qweight_2d, self.quant_config.weight_bits, packed_dim=1
+        )
+
+        if self.quant_config.weight_bits == 4:
+            group_size = self.quant_config.group_size
+            scale_expanded = layer.w2_scales.data.repeat_interleave(group_size, dim=1)
+
+            neg_mask = scale_expanded < 0
+
+            if neg_mask.any():
+                neg_mask = neg_mask.transpose(-1, -2)
+                neg_mask = neg_mask.contiguous().reshape(w2_qweight_tmp.shape)
+                w2_qweight_tmp[neg_mask] = -w2_qweight_tmp[neg_mask]
+
+                if w2_qweight_tmp.max() > 7:
+                    w2_qweight_tmp.clamp_(max=7)
+
+                layer.w2_scales.data.abs_()
+
+            layer.w2_qweight = torch.nn.Parameter(
+                torch_npu.npu_convert_weight_to_int4pack(
+                    w2_qweight_tmp.reshape(
+                        layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1
+                    )
+                    .transpose(-1, -2)
+                    .contiguous()
+                    .reshape(-1, layer.w2_qweight.shape[2])
+                    .to(torch.int32)
+                )
+                .reshape(layer.w2_qweight.shape[0], layer.w2_qweight.shape[1] * 8, -1)
+                .contiguous(),
+                requires_grad=False,
+            )
+        # use int8 to store weight by default
+        else:
+            layer.w2_qweight = torch.nn.Parameter(
+                w2_qweight_tmp.reshape(
+                    layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1
+                )
+                .transpose(-1, -2)
+                .contiguous(),
+                requires_grad=False,
+            )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> torch.Tensor:
+        from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
+
+        assert self.moe_runner_config is not None, (
+            "moe_runner_config is not set. "
+            "Did you forget to call create_weights/create_moe_runner?"
+        )
+
+        assert self.moe_runner_config.activation in ("silu", "swiglu"), (
+            f"Only SiLU/Swiglu activation is supported, "
+            f"got {self.moe_runner_config.activation!r}."
+        )
+
+        x = dispatch_output.hidden_states
+        topk_output = dispatch_output.topk_output
+        topk_weights, topk_ids, _ = topk_output
+
+        topk_ids = topk_ids.to(torch.int32)
+        topk_weights = topk_weights.to(x.dtype)
+
+        output = npu_fused_experts(
+            hidden_states=x,
+            w13=layer.w13_qweight,
+            w13_scale=layer.w13_scales,
+            w13_offset=layer.w13_qzeros,
+            w2=layer.w2_qweight,
+            w2_scale=layer.w2_scales,
+            w2_offset=layer.w2_qzeros,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            top_k=topk_ids.shape[1],
+            use_wna16=True,
+        )
+
+        return StandardCombineInput(hidden_states=output)

python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py

@@ -0,0 +1,622 @@
+import logging
+from typing import TYPE_CHECKING, Optional
+
+import torch
+from torch.nn.parameter import Parameter
+
+from sglang.srt.hardware_backend.npu.utils import NPUACLFormat, npu_format_cast
+from sglang.srt.layers.quantization.base_config import LinearMethodBase
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+logger = logging.getLogger(__name__)
+
+MXFP8_BLOCK_SIZE = 32
+# W4A8_MXFP block (group) size — fixed at 32 by the msmodelslim export format.
+MXFP4_BLOCK_SIZE = 32
+
+
+# NPU ops are reached via torch.ops.npu.* (registered when torch_npu is imported
+# by the runtime), so this module needs no top-level `import torch_npu` and stays
+# importable on CUDA/CPU/AMD/XPU CI.
+def _get_float8_e8m0fnu_dtype():
+    # Resolve lazily rather than as a module-level constant: this module is
+    # imported early (during quant-scheme registration), so reading the dtype at
+    # call time keeps it correct regardless of import order / platform.
+    return getattr(torch, "float8_e8m0fnu", None)
+
+
+def _get_float4_e2m1fn_x2_dtype():
+    # The packed-FP4 dtype MUST come from torch_npu (an int enum, e.g. 296), not
+    # from torch. The NPU ops that consume it -- npu_dynamic_mx_quant(dst_type=),
+    # npu_quant_matmul(x2_dtype=), npu_format_cast(input_dtype=) -- REJECT the
+    # torch dtype object torch.float4_e2m1fn_x2 in op-plugin on recent torch_npu
+    # builds (it raises, or with None gives "output y must be same shape as input
+    # x"), even though torch.float4_e2m1fn_x2 exists. This is fp4-specific: fp8 /
+    # float8_e8m0fnu is accepted from torch either way. Verified on A5 /
+    # torch_npu 2.10.0.post2.dev20260704 (see llm/probe_fp4_w4a8_chain.py: dst=296
+    # passes the full quant->format_cast->matmul chain, dst=torch dtype fails).
+    #
+    # Lazy import so this NPU-only path keeps the module importable on
+    # CUDA/CPU/AMD/XPU CI (no top-level torch_npu; see AGENTS.md known pitfalls).
+    from sglang.srt.utils import is_npu
+
+    if is_npu():
+        import torch_npu
+
+        npu_dtype = getattr(torch_npu, "float4_e2m1fn_x2", None)
+        if npu_dtype is not None:
+            return npu_dtype
+    return getattr(torch, "float4_e2m1fn_x2", None)
+
+
+class _NPULinearMethodBase(LinearMethodBase):
+
+    def __init__(
+        self,
+        quant_config: Optional["QuantizationConfig"] = None,
+    ):
+        self.quant_config = quant_config
+
+
+class NPUW8A8Int8LinearMethod(_NPULinearMethodBase):
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
+        layer.weight.data = npu_format_cast(layer.weight.data)
+
+        layer.weight_scale.data = layer.weight_scale.data.flatten()
+        # Compressed-tensors format doesn't have this field
+        if hasattr(layer, "weight_offset"):
+            layer.weight_offset.data = layer.weight_offset.data.flatten()
+
+        expanding_factor = layer.weight.data.shape[0]
+        layer.aclnn_input_scale = torch.nn.Parameter(
+            layer.input_scale.data.repeat(expanding_factor).to(device="npu"),
+            requires_grad=False,
+        )
+        layer.aclnn_input_scale_reciprocal = 1 / torch.nn.Parameter(
+            layer.input_scale.data.repeat(expanding_factor).to(device="npu"),
+            requires_grad=False,
+        )
+        layer.aclnn_input_offset = torch.nn.Parameter(
+            layer.input_offset.data.repeat(expanding_factor).to(device="npu"),
+            requires_grad=False,
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        from sglang.srt.layers.linear import RowParallelLinear
+
+        original_dtype = x.dtype
+        if original_dtype != torch.int8:
+            x = torch.ops.npu.npu_quantize(
+                x,
+                layer.aclnn_input_scale_reciprocal,
+                layer.aclnn_input_offset,
+                torch.qint8,
+                -1,
+                False,
+            )
+        # Only fuse bias add into GEMM for rank 0 (this ensures that
+        # bias will not get added more than once in Attention TP>1 case)
+        if isinstance(layer, RowParallelLinear) and layer.tp_rank > 0:
+            quant_bias = None
+        else:
+            quant_bias = layer.quant_bias
+        return torch.ops.npu.npu_quant_matmul(
+            x,
+            layer.weight,
+            layer.deq_scale,
+            bias=quant_bias,
+            output_dtype=original_dtype,
+        )
+
+
+class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase):
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
+        layer.weight.data = npu_format_cast(layer.weight.data)
+
+        layer.weight_scale.data = layer.weight_scale.data.flatten()
+        # Compressed-tensors format doesn't have this field
+        if hasattr(layer, "weight_offset"):
+            layer.weight_offset.data = layer.weight_offset.data.flatten()
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        if isinstance(x, tuple):
+            """dynamic_scale is calculated in malprolog kernel"""
+            original_dtype = torch.bfloat16
+            quant_out, dynamic_scale = x
+        else:
+            original_dtype = x.dtype
+            quant_out, dynamic_scale = torch.ops.npu.npu_dynamic_quant(x)
+        return torch.ops.npu.npu_quant_matmul(
+            quant_out,
+            layer.weight,
+            layer.weight_scale,
+            pertoken_scale=dynamic_scale.flatten(),
+            bias=bias,
+            output_dtype=original_dtype,
+        )
+
+
+class NPUMXFP8LinearMethod(_NPULinearMethodBase):
+    """Ascend 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``).
+    process_weights_after_loading branches on weight dtype: FP16/BF16 weights are
+    quantised to MXFP8 at load time (online); pre-quantised float8_e4m3fn weights
+    are only re-laid-out (offline). Inference: dynamic MXFP8 activation quant +
+    MXFP8 matmul (block_size=32).
+    """
+
+    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,
+    ):
+        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 later 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:
+        weight = layer.weight.data
+        if weight.dtype == torch.float8_e4m3fn:
+            # Offline (ModelSlim) path: weight is already MXFP8-quantised and
+            # layer.weight_scale holds the uint8 block scales [out, in/32]. Only
+            # re-layout to [in, out] / [in//64, out, 2] strided views below.
+            n_dim, k_dim = layer.weight_scale.data.shape
+            scale = layer.weight_scale.data.reshape(n_dim, k_dim // 2, 2)
+            layer.weight = Parameter(weight.transpose(0, 1), requires_grad=False)
+            layer.weight_scale_inv = Parameter(
+                scale.transpose(0, 1), requires_grad=False
+            )
+            # weight_scale is now folded into weight_scale_inv (which keeps the
+            # underlying storage alive via its view); drop the stale parameter so
+            # it doesn't linger in named_parameters() / state_dict().
+            del layer.weight_scale
+        else:
+            # Online path: quantise FP16/BF16 weights to MXFP8 at load time.
+            if weight.dtype not in (torch.float16, torch.bfloat16):
+                logger.warning(
+                    "NPUMXFP8LinearMethod: weight dtype %s is not float16/bfloat16; "
+                    "casting to bfloat16 before MXFP8 quantisation.",
+                    weight.dtype,
+                )
+                weight = weight.to(torch.bfloat16)
+            # Move weight to NPU if needed (cpu offload may move it back to CPU).
+            if not weight.is_npu:
+                weight = weight.to(f"npu:{torch.npu.current_device()}")
+            # Online MXFP8 quantisation of weights (block_size=32).
+            # qw: [out, in] float8_e4m3fn, w_scale: [out, in//64, 2] uint8.
+            qw, w_scale = torch.ops.npu.npu_dynamic_mx_quant(
+                weight, dst_type=torch.float8_e4m3fn
+            )
+            layer.weight = Parameter(qw.transpose(0, 1), requires_grad=False)
+            layer.weight_scale_inv = Parameter(
+                w_scale.transpose(0, 1), requires_grad=False
+            )
+
+        # Both paths produce weight [in, out] and weight_scale_inv [in//64, out,
+        # 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.
+
+        # Cache FP32 bias once to avoid a per-forward dtype conversion + alloc.
+        if (
+            getattr(layer, "bias", None) is not None
+            and layer.bias.dtype != torch.float32
+        ):
+            layer.bias_fp32 = Parameter(
+                layer.bias.data.to(torch.float32), requires_grad=False
+            )
+        else:
+            layer.bias_fp32 = None
+
+    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 npu_dynamic_mx_quant
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Dynamic MXFP8 activation quantisation
+        qx, input_scale = torch.ops.npu.npu_dynamic_mx_quant(
+            x_2d, dst_type=torch.float8_e4m3fn
+        )
+
+        # MXFP8 matmul (weight & scale already transposed at load time)
+        # Use the cached FP32 bias from process_weights_after_loading; fall back
+        # to per-call conversion if the cache was bypassed (e.g. dynamic bias).
+        if bias is None:
+            quant_bias = None
+        elif (
+            bias is getattr(layer, "bias", None)
+            and getattr(layer, "bias_fp32", None) is not None
+        ):
+            quant_bias = layer.bias_fp32
+        else:
+            quant_bias = bias.to(torch.float32)
+
+        e8m0_dtype = _get_float8_e8m0fnu_dtype()
+        output = torch.ops.npu.npu_quant_matmul(
+            qx,
+            layer.weight,
+            layer.weight_scale_inv,
+            scale_dtype=e8m0_dtype,
+            pertoken_scale=input_scale,
+            pertoken_scale_dtype=e8m0_dtype,
+            bias=quant_bias,
+            output_dtype=original_dtype,
+            group_sizes=[1, 1, MXFP8_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 NPU_W4A4DynamicLinearMethod(_NPULinearMethodBase):
+
+    def process_weights_after_loading(self, layer):
+        layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
+        layer.weight_scale.data = layer.weight_scale.data.flatten()
+        layer.weight_scale_fp32 = layer.weight_scale.data.to(torch.float32)
+        layer.weight_offset.data = layer.weight_offset.data.flatten()
+        layer.weight.data = torch.ops.npu.npu_convert_weight_to_int4pack(
+            layer.weight.data.to(torch.int32)
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        tp_rank: Optional[int] = 0,
+    ) -> torch.Tensor:
+        original_dtype = x.dtype
+        quant_out, dynamic_scale = torch.ops.npu.npu_dynamic_quant(
+            x, dst_type=torch.quint4x2
+        )
+        return torch.ops.npu.npu_quant_matmul(
+            quant_out,
+            layer.weight,
+            layer.weight_scale,
+            pertoken_scale=dynamic_scale.flatten(),
+            bias=bias,
+            output_dtype=original_dtype,
+        )
+
+
+class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase):
+    """Ascend 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
+    difference is that the FP4 weights are produced online from BF16/FP16
+    (round-to-nearest, no calibration) instead of being loaded from a msmodelslim
+    checkpoint. An earlier version of this method ran a *dual-level* scheme that
+    also compressed the activation to FP4 (W4A4 compute via
+    ``npu_dual_level_quant_matmul``); that was a large accuracy regression — 4-bit
+    activations — so it was replaced with the single-level FP8-activation path
+    below, aligned with the offline W4A8 implementation.
+
+    Weight quantization (process_weights_after_loading):
+        BF16/FP16 weight → npu_dynamic_mx_quant(dst=float4_e2m1fn_x2) → packed FP4
+        + UE8M0 block scale → npu_format_cast to FRACTAL_NZ → transpose [in//2, out]
+
+    Inference (apply):
+        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
+    (same requirement as the offline W4A8 path — see that class's docstring).
+    """
+
+    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.
+
+        Online quantization needs its own ``create_weights`` because 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. This differs from the offline/int8 methods, whose weights
+        are created by the scheme's own ``create_weights`` to match the
+        already-quantized (FP8 / uint8-packed) layout the checkpoint provides.
+        """
+        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, then lay the weight out exactly
+        # like the offline W4A8 path so the same npu_quant_matmul(x2_dtype=fp4)
+        # kernel accepts it. All NPU ops go through torch.ops.npu.* (no torch_npu).
+        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.
+        # npu_dynamic_mx_quant returns the scale as [out, in//64, 2] (3D); older
+        # builds may return [out, in//32] (2D) — handle both before the transpose.
+        qw, w_scale = torch.ops.npu.npu_dynamic_mx_quant(
+            weight_fp, dst_type=fp4_dtype, round_mode="round"
+        )
+
+        # weight: packed FP4 -> FRACTAL_NZ (float8_e4m3fn view) -> transpose
+        # [in//2, out]. Mirror the offline path (no .contiguous() on the NZ view);
+        # view as uint8 first because npu_format_cast only accepts int-dtype tensors.
+        qw_nz = npu_format_cast(
+            qw.view(torch.uint8),
+            NPUACLFormat.ACL_FORMAT_FRACTAL_NZ,
+            customize_dtype=torch.float8_e4m3fn,
+            input_dtype=fp4_dtype,
+        )
+        layer.weight = Parameter(qw_nz.transpose(-1, -2), requires_grad=False)
+
+        # weight_scale -> [in//64, out, 2] to match npu_quant_matmul.
+        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)
+
+        # Cache FP32 bias once to avoid a per-forward dtype conversion + alloc.
+        if (
+            getattr(layer, "bias", None) is not None
+            and layer.bias.dtype != torch.float32
+        ):
+            layer.bias_fp32 = Parameter(
+                layer.bias.data.to(torch.float32), requires_grad=False
+            )
+        else:
+            layer.bias_fp32 = None
+
+    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 MXFP8 activation quantisation (A8 — FP8, not FP4).
+        quantized_x, dynamic_scale = torch.ops.npu.npu_dynamic_mx_quant(
+            x_2d, dst_type=torch.float8_e4m3fn
+        )
+
+        # Use the cached FP32 bias from process_weights_after_loading; fall back
+        # to per-call conversion if the cache was bypassed (e.g. dynamic bias).
+        if bias is None:
+            quant_bias = None
+        elif (
+            bias is getattr(layer, "bias", None)
+            and getattr(layer, "bias_fp32", None) is not None
+        ):
+            quant_bias = layer.bias_fp32
+        else:
+            quant_bias = bias.to(torch.float32)
+
+        # True W4(weight)A8(activation) matmul, identical to the offline path.
+        output = torch.ops.npu.npu_quant_matmul(
+            quantized_x,
+            layer.weight,
+            layer.weight_scale,
+            scale_dtype=e8m0_dtype,
+            pertoken_scale=dynamic_scale,
+            pertoken_scale_dtype=e8m0_dtype,
+            bias=quant_bias,
+            output_dtype=original_dtype,
+            x2_dtype=fp4_dtype,
+            group_sizes=[0, 0, 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 NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase):
+    """Ascend 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*
+    (``pack_fp4_to_uint8`` → ``uint8`` shape ``[out, in//2]``) plus UE8M0 block
+    scales (``uint8`` shape ``[out, in//group_size]``):
+
+      process_weights_after_loading:
+        weight (uint8 packed FP4 [out, in//2]) → npu_format_cast(29,
+            customize_dtype=float8_e4m3fn, input_dtype=float4_e2m1fn_x2) → FRACTAL_NZ
+            → transpose [in//2, out]
+        weight_scale [out, in/32] → reshape [out, in/64, 2] → transpose → [in/64, out, 2]
+
+      apply:
+        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.
+
+    ⚠️ REQUIRES a recent torch_npu build for the FP4 ``npu_quant_matmul``. On the
+    A5 this device forces ``allow_internal_format=False`` (the NZ cast still produces
+    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
+    switching the weight to ND.
+
+    ⚠️ A ``atb::OperationSetup`` *segfault during decode* (not prefill) is a
+    DIFFERENT, unrelated issue: it is the eager-decode ``ascend`` attention
+    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.
+
+    This is a true W4(weight) A8(activation) single-level matmul. The *online*
+    ``NPUMXFP4W4A8LinearMethod`` now uses this exact apply path — the only
+    difference is that it quantizes BF16/FP16 weights to FP4 at load time instead
+    of loading them from a msmodelslim checkpoint. ``group_size`` is fixed at 32
+    by the ``W4A8_MXFP`` export format.
+    """
+
+    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").
+        fp4_dtype = _get_float4_e2m1fn_x2_dtype()
+
+        # weight: packed-FP4 uint8 [out, in//2] -> FRACTAL_NZ (float8_e4m3fn view)
+        # -> transpose to [in//2, out].
+        layer.weight.data = npu_format_cast(
+            layer.weight.data,
+            NPUACLFormat.ACL_FORMAT_FRACTAL_NZ,
+            customize_dtype=torch.float8_e4m3fn,
+            input_dtype=fp4_dtype,
+        )
+        layer.weight.data = layer.weight.data.transpose(-1, -2)
+        # weight_scale: [out, in/32] uint8 -> [in/64, out, 2].
+        n, k = layer.weight_scale.data.shape
+        layer.weight_scale.data = layer.weight_scale.data.reshape(
+            n, k // 2, 2
+        ).transpose(-3, -2)
+
+    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 MXFP8 activation quantisation (A8).
+        quantized_x, dynamic_scale = torch.ops.npu.npu_dynamic_mx_quant(
+            x_2d, dst_type=torch.float8_e4m3fn
+        )
+
+        if bias is not None and bias.dtype != torch.float32:
+            bias = bias.to(torch.float32)
+
+        # W4(weight)A8(activation) matmul, mirroring vllm-ascend exactly.
+        output = torch.ops.npu.npu_quant_matmul(
+            quantized_x,
+            layer.weight,
+            layer.weight_scale,
+            scale_dtype=e8m0_dtype,
+            pertoken_scale=dynamic_scale,
+            pertoken_scale_dtype=e8m0_dtype,
+            bias=bias,
+            output_dtype=original_dtype,
+            x2_dtype=fp4_dtype,
+            group_sizes=[0, 0, 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)