📝 [NPU] Use vendor-neutral wording in quantization comments (#36768)

2026-08-29 23:05:35 +03:00
parent 6afb5e1771
commit 09ecb9aaaa
8 changed files with 34 additions and 45 deletions
@@ -1,4 +1,4 @@
-"""ModelSlim MXFP4 scheme for pre-quantized weight inference on Ascend NPU.
+"""ModelSlim MXFP4 scheme for pre-quantized weight inference on NPU.

 Loads weights pre-quantized by msmodelslim and runs MXFP4 dual-level
 matmul at inference via npu_dual_level_quant_matmul.
@@ -8,9 +8,6 @@ Checkpoint tensor formats (verified from msmodelslim export):
  weight_scale:     [out, in/32]        uint8          (L1 block scales, e8m0+127)
  weight_dual_scale:[out, in/512, 1]    float32        (L0 coarse scales)
  mul_scale:        [in]                float32        (smooth quant activation scale)
-
-Reference: MindIE-SD W4A4MXFP4DualQuantLinear
-(MindIE-SD/mindiesd/quantization/layer.py)
 """

 from typing import List, Optional
@@ -98,7 +95,6 @@ class ModelSlimMXFP4Scheme(ModelSlimLinearScheme):
        # After repack, it becomes `<prefix>.mul_scale`.
        # This is CRITICAL: the offline-quantized weights were calibrated with
        # x * mul_scale applied to the activation. Omitting it causes mosaic output.
-        # Ref: MindIE-SD W4A4MXFP4DualQuantLinear.quant_matmul lines 385-386.
        mul_scale = BasevLLMParameter(
            data=torch.empty(
                (input_size_per_partition,),
@@ -119,7 +115,6 @@ class ModelSlimMXFP4Scheme(ModelSlimLinearScheme):
            weight = weight.to(f"npu:{torch.npu.current_device()}")
        weight = torch_npu.npu_dtype_cast(weight, torch_npu.float4_e2m1fn_x2)
        # npu_dual_level_quant_matmul requires x2 in FRACTAL_NZ format (format 29).
-        # Reference: MindIE-SD W4A4MXFP4DualQuantLinear._init_dynamic_quant_param
        weight = torch_npu.npu_format_cast(
            weight.view(torch.int8), 29, customize_dtype=torch.int8
        )
@@ -170,7 +165,6 @@ class ModelSlimMXFP4Scheme(ModelSlimLinearScheme):
        # Apply smooth quant scale before activation quantization.
        # The offline-quantized weights were calibrated under x * mul_scale,
        # so we MUST apply it here for scale alignment.
-        # Reference: MindIE-SD W4A4MXFP4DualQuantLinear.quant_matmul
        mul_scale = layer.mul_scale
        if getattr(layer, "use_mul_scale", True):
            x_2d = x_2d * mul_scale.to(x_2d.dtype)
@@ -1,4 +1,4 @@
-"""ModelSlim MXFP8 scheme for pre-quantized weight inference on Ascend NPU.
+"""ModelSlim MXFP8 scheme for pre-quantized weight inference on NPU.

 Loads weights pre-quantized by msmodelslim (float8_e4m3fn weights,
 uint8 scales) and runs MXFP8 matmul at inference.
@@ -1,4 +1,4 @@
-"""Online MXFP4 quantization for Diffusion models on Ascend NPU.
+"""Online MXFP4 quantization for Diffusion models on NPU.

 Provides ``NPUMXFP4Config`` (registered as ``"mxfp4_npu"``) and
 ``NPUMXFP4DiffusionLinearMethod`` which quantises FP16/BF16 weights to MXFP4
@@ -10,9 +10,10 @@ The ``"mxfp4_npu"`` key is distinct from upstream's ROCm ``"mxfp4"``
 (``Mxfp4Config`` in ``mxfp4.py``) which targets AMD MI350+ via aiter kernels.

 NOTE: Online weight quantization via ``npu_dynamic_dual_level_mx_quant`` is
-experimental. MindIE-SD only uses an offline (pre-quantized) path for MXFP4
-weights. The online path quantizes FP16/BF16 weights at load time, which may
-produce different numerical results than the offline calibrated path.
+experimental; the established MXFP4 path for these models is offline
+(pre-quantized) only. The online path quantizes FP16/BF16 weights at load
+time, which may produce different numerical results than the offline
+calibrated path.
 """

 from __future__ import annotations
@@ -41,7 +42,7 @@ logger = init_logger(__name__)


 class NPUMXFP4Config(QuantizationConfig):
-    """Config for online MXFP4 quantization on Ascend NPU (Diffusion)."""
+    """Config for online MXFP4 quantization on NPU (Diffusion)."""

    def __init__(self) -> None:
        super().__init__()
@@ -78,13 +79,11 @@ class NPUMXFP4Config(QuantizationConfig):


 class NPUMXFP4DiffusionLinearMethod(LinearMethodBase):
-    """Ascend NPU MXFP4 linear method for Diffusion models (dual-level).
+    """NPU MXFP4 linear method for Diffusion models (dual-level).

    Online mode: loads FP16/BF16 weights → quantises to MXFP4 at load time
    via ``npu_dynamic_dual_level_mx_quant``.
    Inference: dynamic dual-level MXFP4 activation quant + dual-level matmul.
-
-    Reference: MindIE-SD ``W4A4MXFP4DualQuantLinear`` (offline path only).
    """

    def __init__(self, quant_config: NPUMXFP4Config):
@@ -133,9 +132,9 @@ class NPUMXFP4DiffusionLinearMethod(LinearMethodBase):
            weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}")

        # Online dual-level MXFP4 weight quantisation.
-        # NOTE: This is experimental — MindIE-SD only has an offline path for
-        # MXFP4 weights. We assume npu_dynamic_dual_level_mx_quant can also
-        # quantise weights (not just activations).
+        # NOTE: This is experimental — the established MXFP4 path for these
+        # models is offline only. We assume npu_dynamic_dual_level_mx_quant can
+        # also quantise weights (not just activations).
        # Returns: (qw, w_dual_scale, w_scale)
        #   qw          — quantized weight in float4_e2m1fn_x2 (2 FP4 packed/byte)
        #   w_dual_scale — L0-level scale (goes to pos 3 in npu_dual_level_quant_matmul)
@@ -145,14 +144,12 @@ class NPUMXFP4DiffusionLinearMethod(LinearMethodBase):
        )

        # npu_dual_level_quant_matmul requires x2 (weight) in FRACTAL_NZ format.
-        # Reference: MindIE-SD W4A4MXFP4DualQuantLinear._init_dynamic_quant_param
        qw = torch_npu.npu_format_cast(
            qw.view(torch.int8), 29, customize_dtype=torch.int8
        )

        # x2Level0Scale must be [in/level0_block_size, out] — transpose from
        # the [out, in/level0_block_size] shape returned by the quant op.
-        # Reference: MindIE-SD layer.py:409
        w_dual_scale = w_dual_scale.squeeze(-1).transpose(0, 1).contiguous()

        layer.weight = Parameter(qw, requires_grad=False)
@@ -22,8 +22,7 @@ def _normalize_mxfp_scale(scale: torch.Tensor) -> torch.Tensor:

    ``npu_moe_init_routing_v2(quant_mode=3)`` emits the scale flat, while the
    grouped matmul wants the pair-split view. Already-3D scales (what
-    ``npu_dynamic_mx_quant`` returns) pass through untouched. Mirrors
-    vllm-ascend's ``maybe_normalize_mxfp_scale_layout``.
+    ``npu_dynamic_mx_quant`` returns) pass through untouched.
    """
    if scale is None or scale.ndim != 2:
        return scale
@@ -156,7 +156,7 @@ class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase):


 class NPUMXFP8LinearMethod(_NPULinearMethodBase):
-    """Ascend NPU MXFP8 linear method for LLM (SRT) models.
+    """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``).
@@ -343,7 +343,7 @@ class NPU_W4A4DynamicLinearMethod(_NPULinearMethodBase):


 class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase):
-    """Ascend NPU W4A8 online quantization: MXFP4 weights + MXFP8 activations.
+    """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
@@ -363,7 +363,7 @@ class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase):
        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
+    Hardware: A5 NPU + a recent torch_npu with the FP4 npu_quant_matmul
    (same requirement as the offline W4A8 path — see that class's docstring).
    """

@@ -513,7 +513,7 @@ class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase):


 class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase):
-    """Ascend NPU offline W4A8 (ModelSlim ``W4A8_MXFP``): packed-FP4 weights + MXFP8 activations.
+    """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*
@@ -626,7 +626,7 @@ class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase):


 class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase):
-    """Ascend NPU W4A4 online quantization: single-level MXFP4.
+    """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
@@ -642,8 +642,8 @@ class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase):
        → 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``.
+    Triggered by ``--quantization mxfp4`` on NPU. Hardware: A5 NPU with a recent
+    torch_npu exposing ``float4_e2m1fn_x2``.
    """

    def create_weights(
@@ -765,7 +765,7 @@ class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase):


 class NPUSingleLevelMXFP4OfflineLinearMethod(NPUSingleLevelMXFP4LinearMethod):
-    """Ascend NPU offline W4A4 (ModelSlim ``W4A4_MXFP4``): packed FP4 weights.
+    """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 packed ``uint8``
@@ -796,7 +796,7 @@ class NPUSingleLevelMXFP4OfflineLinearMethod(NPUSingleLevelMXFP4LinearMethod):


 class NPUDualLevelMXFP4LinearMethod(NPUSingleLevelMXFP4LinearMethod):
-    """Ascend NPU W4A4 online quantization: dual-level MXFP4 (higher accuracy).
+    """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)
@@ -819,9 +819,8 @@ class NPUDualLevelMXFP4LinearMethod(NPUSingleLevelMXFP4LinearMethod):
        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.
+    Reference: Diffusion ``NPUMXFP4DiffusionLinearMethod``. Hardware: A5 NPU
+    only — the ``DualLevelQuantBatchMatmul`` op is unavailable on A2/A3.
    """

    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -223,7 +223,7 @@ class NPUW4A8MXFP4MoEMethod(_NPUMoEMethodBase):
        ).transpose(1, 2)
        weight_scale.data = scale

-        # The refactored Ascend dispatchers currently support BF16 and INT8.
+        # The refactored NPU dispatchers currently support BF16 and INT8.
        # Keep dispatch in BF16 and quantize to MXFP8 immediately before GMM.
        if weight_prefix == "w13":
            self._set_dispatcher_output_dtype(layer, "bf16")
@@ -304,7 +304,7 @@ class NPUW4A4MXFP4MoEMethod(_NPUMoEMethodBase):
        ).transpose(1, 2)
        weight_scale.data = scale

-        # The refactored Ascend dispatchers currently support BF16 and INT8.
+        # The refactored NPU dispatchers currently support BF16 and INT8.
        # Keep dispatch in BF16 and quantize immediately before each GMM.
        if weight_prefix == "w13":
            self._set_dispatcher_output_dtype(layer, "bf16")
@@ -929,7 +929,7 @@ class NPUUnquantMoEMethod(_NPUMoEMethodBase):
 #  NPUMXFP8MoEMethod
 # ---------------------------------------------------------------------------
 class NPUMXFP8MoEMethod(_NPUMoEMethodBase):
-    """MXFP8 MoE on Ascend A5 – float8_e4m3fn weights with e8m0 block scales.
+    """MXFP8 MoE on A5 NPU – float8_e4m3fn weights with e8m0 block scales.

    Serves both the online config path (``--quantization mxfp8``, weights
    quantised at load time) and the offline ModelSlim ``W8A8_MXFP8`` scheme
@@ -1016,8 +1016,8 @@ class NPUMXFP8MoEMethod(_NPUMoEMethodBase):
        # [E, K//64, N, 2] as strided transpose views — DO NOT call
        # .contiguous(). Beyond breaking the transpose-flag match above, it
        # measures slower on the same probe: making both sides contiguous costs
-        # 6.2% on decode. This matches NPUMXFP8LinearMethod, msmodelslim's
-        # offline layout and vllm-ascend's AscendW8A8MXFP8DynamicFusedMoEMethod.
+        # 6.2% on decode. This matches NPUMXFP8LinearMethod and msmodelslim's
+        # offline layout.
        setattr(
            layer,
            f"{weight_prefix}_weight",
@@ -1,4 +1,4 @@
-"""Online (config-driven) quantized FusedMoE methods for Ascend NPU.
+"""Online (config-driven) quantized FusedMoE methods for NPU.

 These are the ``--quantization <scheme>`` entry points: the checkpoint holds
 BF16/FP16 expert weights and the per-gmm kernels quantize them at load time.
@@ -28,7 +28,7 @@ class NPUMXFP8OnlineMoEMethod(UnquantizedFusedMoEMethod):
    """Online MXFP8 FusedMoE entry point (``--quantization mxfp8`` on A5).

    Weight creation, weight post-processing and the forward pass are identical
-    to the unquantized Ascend path — the only difference is which per-gmm kernel
+    to the unquantized NPU path — the only difference is which per-gmm kernel
    the layer gets, so everything but ``create_moe_runner`` is inherited.
    ``NPUMXFP8MoEMethod`` then quantizes the BF16 expert weights to MXFP8 in
    ``process_weights_after_loading``.
@@ -267,9 +267,9 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
            routing_method_type=RoutingMethodType.Renormalize,
        )

-        # Router gate: description-driven quant, mirroring vllm-ascend. Only the
-        # offline ModelSlim path (which carries a per-layer quant_model_description)
-        # may quantise the gate — if the checkpoint stored it as MXFP8 it is loaded
+        # Router gate: description-driven quant. Only the offline ModelSlim path
+        # (which carries a per-layer quant_model_description) may quantise the
+        # gate — if the checkpoint stored it as MXFP8 it is loaded
        # and dequantised correctly instead of cast to bf16 without its block scale.
        # The online Fp8/mxfp8 path keeps the gate in bf16 (unchanged, verified).
        gate_quant_config = (