hxhhhlalalaandhyh_hh
9dd36e3afd
Signed-off-by: hyh_hh <huyinghong1@huawei.com> Co-authored-by: hyh_hh <huyinghong1@huawei.com>
710 lines
29 KiB
Python
710 lines
29 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
"""W4A4 MXFP4 (Microscaling FP4) online/offline quantization for diffusion transformers.
|
||
|
||
Architecture mirrors mxfp8_config.py:
|
||
|
||
MXFPLinearMethodBase – platform-agnostic skeleton (imported from mxfp8_config)
|
||
NPUMxfp4LinearMethod – NPU single-scale offline (W4A4 MXFP4)
|
||
NPUMxfp4OnlineLinearMethod – NPU single-scale online (BF16 → FP4)
|
||
NPUMxfp4DualScaleLinearMethod – NPU dual-scale offline (W4A4 MXFP4 DualScale)
|
||
NPUMxfp4DualScaleOnlineLinearMethod – NPU dual-scale online (BF16 → FP4)
|
||
|
||
Quantization configs:
|
||
|
||
DiffusionMXFP4Config – single-scale online/offline (quant_method="mxfp4")
|
||
DiffusionMXFP4DualScaleMixedConfig – dual-scale + per-layer BF16 fallback (quant_method="mxfp4_dualscale")
|
||
Offline: ignored_layers from config.json routes interleaved BF16 layers
|
||
Online: num_bf16_fallback_layers leading blocks stay in BF16 (default 5)
|
||
|
||
Key differences from MXFP8:
|
||
|
||
1. Precision: float4_e2m1fn_x2 (FP4 packed, 2 values per element).
|
||
npu_dynamic_mx_quant(x) without dst_type defaults to float4_e2m1fn_x2.
|
||
|
||
2. Weight layout: stored as (N, K) — NOT pre-transposed.
|
||
FP4 uses a packed format; transposing a packed tensor is not safely contiguous.
|
||
Transpose is done inline in _quant_matmul via layer.weight.transpose(0, 1).
|
||
|
||
3. GEMM signature: npu_quant_matmul requires explicit
|
||
x1_dtype=float4_e2m1fn_x2 and x2_dtype=float4_e2m1fn_x2.
|
||
|
||
Scale layout: (N, S/2, 2) — same reshape as MXFP8, also NOT pre-transposed;
|
||
transposed inline in _quant_matmul.
|
||
|
||
Dual-scale (W4A4_MXFP4_DUALSCALE):
|
||
|
||
Two-level quantization: fine scale (per-32 K) + coarse scale (per-512 K) + per-channel
|
||
activation pre-scale (mul_scale from calibration). Uses npu_dynamic_dual_level_mx_quant
|
||
and npu_dual_level_quant_matmul. Weight stored in NZ hardware format via npu_format_cast.
|
||
|
||
Checkpoint tensor shapes for dual-scale:
|
||
weight : (N, K) float8_e4m3fn – FP4 packed
|
||
weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits)
|
||
weight_dual_scale : (N, K//512, 1) float32 – coarse scale (extra dim avoids shape assert)
|
||
mul_scale : (K,) float32 – per-input-channel activation pre-scale
|
||
|
||
Reference: MindIE-SD W4A4MXFP4QuantLinear / W4A4MXFP4DualQuantLinear (mindiesd/quantization/layer.py).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
from typing import TYPE_CHECKING, Any
|
||
|
||
import torch
|
||
from torch.nn import Module
|
||
from vllm.logger import init_logger
|
||
from vllm.model_executor.layers.linear import (
|
||
LinearBase,
|
||
UnquantizedLinearMethod,
|
||
)
|
||
from vllm.model_executor.layers.quantization import QuantizationMethods
|
||
from vllm.model_executor.layers.quantization.base_config import (
|
||
QuantizationConfig,
|
||
QuantizeMethodBase,
|
||
)
|
||
from vllm.model_executor.layers.quantization.fp8 import _copy_missing_attrs
|
||
from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
|
||
from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight
|
||
from vllm.model_executor.parameter import ModelWeightParameter
|
||
from vllm.model_executor.utils import replace_parameter
|
||
|
||
from vllm_omni.platforms import current_omni_platform
|
||
from vllm_omni.quantization.mxfp8_config import (
|
||
MXFPLinearMethodBase,
|
||
_LazyWeightMixin,
|
||
)
|
||
|
||
if TYPE_CHECKING:
|
||
from vllm.model_executor.models.utils import WeightsMapper
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Config
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class DiffusionMXFP4Config(QuantizationConfig):
|
||
"""W4A4 MXFP4 quantization config for diffusion transformers.
|
||
|
||
Supports both online (BF16 checkpoint → quantize at load time) and offline
|
||
(pre-quantized MXFP4 checkpoint) modes, mirroring DiffusionMXFP8Config.
|
||
|
||
MX (microscaling) format: groups of 32 K-dimension elements share one
|
||
float8_e8m0fnu exponent scale. Weight and activation are float4_e2m1fn_x2.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
is_checkpoint_mxfp4_serialized: bool = False,
|
||
ignored_layers: list[str] | None = None,
|
||
) -> None:
|
||
super().__init__()
|
||
self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized
|
||
self.ignored_layers = ignored_layers or []
|
||
|
||
@classmethod
|
||
def get_name(cls) -> QuantizationMethods:
|
||
return "mxfp4"
|
||
|
||
@classmethod
|
||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||
return [torch.bfloat16, torch.float16]
|
||
|
||
@classmethod
|
||
def get_min_capability(cls) -> int:
|
||
return 80
|
||
|
||
@classmethod
|
||
def get_config_filenames(cls) -> list[str]:
|
||
return []
|
||
|
||
def apply_vllm_mapper(self, hf_to_vllm_mapper: WeightsMapper) -> None:
|
||
if self.ignored_layers:
|
||
self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers)
|
||
|
||
@classmethod
|
||
def from_config(cls, config: dict[str, Any]) -> DiffusionMXFP4Config:
|
||
is_serialized = cls.get_from_keys_or(config, ["is_checkpoint_mxfp4_serialized"], False)
|
||
ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None)
|
||
if not ignored_layers:
|
||
ignored_layers = cls.get_from_keys_or(config, ["modules_to_not_convert"], None)
|
||
return cls(
|
||
is_checkpoint_mxfp4_serialized=is_serialized,
|
||
ignored_layers=ignored_layers,
|
||
)
|
||
|
||
def get_quant_method(
|
||
self,
|
||
layer: torch.nn.Module,
|
||
prefix: str,
|
||
) -> QuantizeMethodBase | None:
|
||
if isinstance(layer, LinearBase):
|
||
if is_layer_skipped(
|
||
prefix=prefix,
|
||
ignored_layers=self.ignored_layers,
|
||
fused_mapping=self.packed_modules_mapping,
|
||
):
|
||
return UnquantizedLinearMethod()
|
||
if current_omni_platform.is_npu():
|
||
if self.is_checkpoint_mxfp4_serialized:
|
||
return NPUMxfp4LinearMethod(self)
|
||
return NPUMxfp4OnlineLinearMethod(self)
|
||
raise NotImplementedError(
|
||
"DiffusionMXFP4Config (W4A4 MXFP4) is currently only supported "
|
||
"on NPU (Ascend) platforms. CUDA support is not yet implemented."
|
||
)
|
||
return None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# NPU MXFP4 single-scale offline method (pre-quantized checkpoint)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class NPUMxfp4LinearMethod(MXFPLinearMethodBase):
|
||
"""NPU W4A4 MXFP4 offline linear method for pre-quantized checkpoints.
|
||
|
||
Weight canonical layout after process_weights_after_loading:
|
||
weight : (N, K) in float4_e2m1fn_x2 — NOT pre-transposed (FP4 packed)
|
||
weight_scale: (N, S/2, 2) in float8_e8m0fnu — NOT pre-transposed
|
||
|
||
Both are transposed inline in _quant_matmul, unlike MXFP8 which pre-transposes.
|
||
NPUMxfp4OnlineLinearMethod normalizes to the same layout so apply() is shared.
|
||
"""
|
||
|
||
def __init__(self, quant_config: DiffusionMXFP4Config) -> None:
|
||
self.quant_config = quant_config
|
||
self.out_dtype = torch.get_default_dtype()
|
||
|
||
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,
|
||
) -> None:
|
||
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
|
||
layer.weight_block_size = None
|
||
|
||
# BF16 placeholder; cast to float4_e2m1fn_x2 in process_weights.
|
||
layer.register_parameter(
|
||
"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,
|
||
),
|
||
)
|
||
|
||
# Scale stored as uint8 in safetensors (float8_e8m0fnu is same bit width).
|
||
# Using uint8 avoids a lossy float32 round-trip when loading the checkpoint.
|
||
num_groups = (input_size_per_partition + 31) // 32
|
||
layer.register_parameter(
|
||
"weight_scale",
|
||
ModelWeightParameter(
|
||
data=torch.empty(output_size_per_partition, num_groups, dtype=torch.uint8),
|
||
input_dim=1,
|
||
output_dim=0,
|
||
weight_loader=weight_loader,
|
||
),
|
||
)
|
||
|
||
def process_weights_after_loading(self, layer: Module) -> None:
|
||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||
return
|
||
|
||
import torch_npu
|
||
|
||
# NPU: cast to float4_e2m1fn_x2. Weight stays (N, K) — no pre-transpose.
|
||
w = layer.weight
|
||
if w.dtype != torch_npu.float4_e2m1fn_x2:
|
||
w = torch_npu.npu_dtype_cast(w.npu(), torch_npu.float4_e2m1fn_x2)
|
||
|
||
# Scale: checkpoint stores uint8 bytes that ARE float8_e8m0fnu bits.
|
||
# Only convert if neither uint8 nor the target NPU dtype already.
|
||
# (N, K_groups) → (N, K_groups/2, 2). Not pre-transposed; done inline in _quant_matmul.
|
||
s = layer.weight_scale.data
|
||
if s.dtype not in (torch.uint8, torch_npu.float8_e8m0fnu):
|
||
s = s.to(torch_npu.float8_e8m0fnu)
|
||
N, K_groups = s.shape
|
||
if K_groups % 2 == 1:
|
||
s = torch.cat([s, torch.zeros(N, 1, dtype=s.dtype, device=s.device)], dim=1)
|
||
s = s.reshape(N, -1, 2).contiguous()
|
||
|
||
replace_parameter(layer, "weight", w)
|
||
replace_parameter(layer, "weight_scale", s)
|
||
layer._already_called_process_weights_after_loading = True
|
||
|
||
# --- NPU MXFP4 ops — shared with online path via inheritance ---
|
||
|
||
def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||
import torch_npu
|
||
|
||
# No dst_type: npu_dynamic_mx_quant defaults to float4_e2m1fn_x2.
|
||
return torch_npu.npu_dynamic_mx_quant(x)
|
||
|
||
def _quant_matmul(
|
||
self,
|
||
x_q: torch.Tensor,
|
||
x_scale: torch.Tensor,
|
||
layer: torch.nn.Module,
|
||
bias: torch.Tensor | None,
|
||
ori_dtype: torch.dtype,
|
||
) -> torch.Tensor:
|
||
import torch_npu
|
||
|
||
if bias is not None and bias.dtype != torch.float32:
|
||
bias = bias.to(torch.float32)
|
||
# FP4 differences vs FP8:
|
||
# weight (N,K) transposed inline → (K,N); scale (N,S/2,2) transposed inline → (S/2,N,2).
|
||
# x1_dtype / x2_dtype required — FP4 dtype not inferred from tensor dtype.
|
||
return torch_npu.npu_quant_matmul(
|
||
x_q,
|
||
layer.weight.transpose(0, 1), # (K, N) inline
|
||
layer.weight_scale.transpose(0, 1), # (S/2, N, 2) inline
|
||
scale_dtype=torch_npu.float8_e8m0fnu,
|
||
x1_dtype=torch_npu.float4_e2m1fn_x2,
|
||
x2_dtype=torch_npu.float4_e2m1fn_x2,
|
||
pertoken_scale=x_scale,
|
||
pertoken_scale_dtype=torch_npu.float8_e8m0fnu,
|
||
bias=bias,
|
||
output_dtype=ori_dtype,
|
||
group_sizes=[1, 1, 32],
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# NPU MXFP4 single-scale online method (BF16 checkpoint → quantize at load time)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod):
|
||
"""NPU W4A4 MXFP4 online linear method.
|
||
|
||
MRO: NPUMxfp4OnlineLinearMethod → _LazyWeightMixin → NPUMxfp4LinearMethod
|
||
→ MXFPLinearMethodBase → LinearMethodBase
|
||
|
||
create_weights : _LazyWeightMixin (meta device + patched loader)
|
||
process_weights : NPUMxfp4OnlineLinearMethod (BF16 → FP4 + normalize)
|
||
apply / ops : NPUMxfp4LinearMethod / MXFPLinearMethodBase (shared)
|
||
"""
|
||
|
||
def process_weights_after_loading(self, layer: Module) -> None:
|
||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||
return
|
||
|
||
import torch_npu
|
||
|
||
if layer.weight.device == torch.device("meta"):
|
||
weight = ModelWeightParameter(
|
||
data=torch.empty_like(layer.weight, device=layer._load_device),
|
||
input_dim=1,
|
||
output_dim=0,
|
||
weight_loader=layer.weight.weight_loader,
|
||
)
|
||
_copy_missing_attrs(layer.weight, weight)
|
||
layer.register_parameter("weight", weight)
|
||
initialize_single_dummy_weight(layer.weight)
|
||
|
||
# NPU: quantize BF16/FP16 (N, K) → FP4. No dst_type → float4_e2m1fn_x2.
|
||
weight_fp4, weight_scale_raw = torch_npu.npu_dynamic_mx_quant(layer.weight)
|
||
|
||
# Weight stays (N, K) — no pre-transpose for FP4 packed format.
|
||
# Scale: (N, S) → (N, S/2, 2). Not pre-transposed; done inline in _quant_matmul.
|
||
weight_scale = weight_scale_raw.reshape(weight_scale_raw.shape[0], -1, 2).contiguous()
|
||
|
||
replace_parameter(layer, "weight", weight_fp4)
|
||
replace_parameter(layer, "weight_scale", weight_scale)
|
||
layer._already_called_process_weights_after_loading = True
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# NPU MXFP4 dual-scale offline method (W4A4_MXFP4_DUALSCALE checkpoint)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class NPUMxfp4DualScaleLinearMethod(MXFPLinearMethodBase):
|
||
"""NPU W4A4 MXFP4 dual-scale offline method for pre-quantized checkpoints.
|
||
|
||
Checkpoint tensors and their canonical post-load shapes:
|
||
weight : (N, K) float8_e4m3fn – FP4 packed (2 values per byte)
|
||
weight_scale : (N, K//32) uint8 – fine scale (float8_e8m0fnu bits); reshaped to (N, K//64, 2)
|
||
weight_dual_scale : (N, K//512, 1) float32 – coarse scale; transposed to (K//512, N)
|
||
mul_scale : (K,) float32 – per-input-channel activation pre-scale (from calibration)
|
||
|
||
Forward pass:
|
||
x_q, l0, l1 = npu_dynamic_dual_level_mx_quant(x, smooth_scale=mul_scale)
|
||
out = npu_dual_level_quant_matmul(x_q, weight, l0, weight_dual_scale, l1, weight_scale)
|
||
|
||
Reference: MindIE-SD W4A4MXFP4DualQuantLinear.
|
||
"""
|
||
|
||
def __init__(self, quant_config: Any) -> None:
|
||
self.quant_config = quant_config
|
||
self.out_dtype = torch.get_default_dtype()
|
||
|
||
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,
|
||
) -> None:
|
||
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
|
||
layer.weight_block_size = None
|
||
|
||
# FP4 packed: 2 values per float8_e4m3fn byte → checkpoint stores as float8_e4m3fn; register as BF16
|
||
layer.register_parameter(
|
||
"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,
|
||
),
|
||
)
|
||
|
||
# Fine scale: one uint8 exponent (float8_e8m0fnu bit pattern) per group of 32 K elements.
|
||
# input_dim=1: RowParallelLinear weight_loader shards along the K-group dimension (dim 1)
|
||
# so each rank receives only its (K/TP)//32 groups; ColumnParallelLinear only uses
|
||
# output_dim=0 for sharding and leaves dim 1 intact.
|
||
num_groups_fine = (input_size_per_partition + 31) // 32
|
||
layer.register_parameter(
|
||
"weight_scale",
|
||
ModelWeightParameter(
|
||
data=torch.empty(output_size_per_partition, num_groups_fine, dtype=torch.uint8),
|
||
input_dim=1,
|
||
output_dim=0,
|
||
weight_loader=weight_loader,
|
||
),
|
||
)
|
||
|
||
# Coarse scale: one float32 per group of 512 K elements.
|
||
# Shape (N, K_coarse, 1) matches checkpoint layout exactly, avoiding the
|
||
# shape-mismatch assert in linear.py:1344.
|
||
# input_dim=1: same TP sharding rationale as weight_scale above.
|
||
num_groups_coarse = (input_size_per_partition + 511) // 512
|
||
layer.register_parameter(
|
||
"weight_dual_scale",
|
||
ModelWeightParameter(
|
||
data=torch.empty(output_size_per_partition, num_groups_coarse, 1, dtype=torch.float32),
|
||
input_dim=1,
|
||
output_dim=0,
|
||
weight_loader=weight_loader,
|
||
),
|
||
)
|
||
|
||
# mul_scale is a float32 calibration tensor; register as float32
|
||
# to avoid precision loss from an implicit BF16 cast during weight loading.
|
||
# input_dim=0: RowParallelLinear shards the 1-D per-input-channel tensor along dim 0
|
||
# so each rank receives only its K/TP channels; ColumnParallelLinear leaves it intact.
|
||
layer.register_parameter(
|
||
"mul_scale",
|
||
ModelWeightParameter(
|
||
data=torch.empty(input_size_per_partition, dtype=torch.float32),
|
||
input_dim=0,
|
||
output_dim=None,
|
||
weight_loader=weight_loader,
|
||
),
|
||
)
|
||
setattr(layer.mul_scale, "ignore_warning", True)
|
||
|
||
def process_weights_after_loading(self, layer: Module) -> None:
|
||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||
return
|
||
|
||
import torch_npu
|
||
|
||
# float8_e4m3fn (FP4 packed) → float4_e2m1fn_x2 → NZ hardware format (format ID 29).
|
||
# NZ layout is required by npu_dual_level_quant_matmul for FP4 weight matrices.
|
||
w = torch_npu.npu_dtype_cast(layer.weight.data.npu(), torch_npu.float4_e2m1fn_x2)
|
||
w = torch_npu.npu_format_cast(w.view(torch.int8), 29, customize_dtype=torch.int8)
|
||
|
||
# Fine scale: (N, K//32) uint8 → cast to float8_e8m0fnu → (N, K//64, 2).
|
||
# The cast reinterprets the stored uint8 bit-patterns as float8_e8m0fnu exponents,
|
||
# which is required for npu_dual_level_quant_matmul to compute correct scale values.
|
||
s = layer.weight_scale.data.npu()
|
||
s = s.reshape(s.shape[0], -1, 2).contiguous()
|
||
|
||
# Coarse scale: (N, K//512, 1) → squeeze → (N, K//512) → transpose → (K//512, N).
|
||
ds = layer.weight_dual_scale.data.squeeze(-1).transpose(0, 1).npu().contiguous()
|
||
|
||
ms = layer.mul_scale.to(torch.bfloat16).data.view(-1).npu().contiguous()
|
||
|
||
replace_parameter(layer, "weight", w)
|
||
replace_parameter(layer, "weight_scale", s)
|
||
replace_parameter(layer, "weight_dual_scale", ds)
|
||
replace_parameter(layer, "mul_scale", ms)
|
||
layer._already_called_process_weights_after_loading = True
|
||
|
||
def _apply_inner(
|
||
self,
|
||
layer: torch.nn.Module,
|
||
x: torch.Tensor,
|
||
bias: torch.Tensor | None,
|
||
ori_dtype: torch.dtype,
|
||
) -> torch.Tensor:
|
||
"""Dual-scale inner loop: dtype guard → 3-tuple quantize (mul_scale as smooth) → matmul.
|
||
|
||
Overrides the default single-scale _apply_inner from MXFPLinearMethodBase.
|
||
apply() in the base class handles reshape/unreshape; this method is not
|
||
responsible for that.
|
||
"""
|
||
if ori_dtype not in (torch.bfloat16, torch.float16):
|
||
x = x.to(torch.bfloat16)
|
||
x_q, l0_scale, l1_scale = self._quantize_activation(x, layer.mul_scale)
|
||
return self._quant_matmul(x_q, l0_scale, l1_scale, layer, bias, ori_dtype)
|
||
|
||
def _quantize_activation( # type: ignore[override]
|
||
self,
|
||
x: torch.Tensor,
|
||
smooth_scale: torch.Tensor,
|
||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||
import torch_npu
|
||
|
||
return torch_npu.npu_dynamic_dual_level_mx_quant(x, smooth_scale=smooth_scale)
|
||
|
||
def _quant_matmul( # type: ignore[override]
|
||
self,
|
||
x_q: torch.Tensor,
|
||
l0_scale: torch.Tensor,
|
||
l1_scale: torch.Tensor,
|
||
layer: torch.nn.Module,
|
||
bias: torch.Tensor | None,
|
||
ori_dtype: torch.dtype,
|
||
) -> torch.Tensor:
|
||
import torch_npu
|
||
|
||
if bias is not None and bias.dtype != torch.float32:
|
||
bias = bias.to(torch.float32)
|
||
# weight_scale is (N, K//64, 2) float8_e8m0fnu — operator expects output-major layout.
|
||
# weight_dual_scale is (K//512, N) float32 — transposed to K-major in process_weights.
|
||
return torch_npu.npu_dual_level_quant_matmul(
|
||
x_q,
|
||
layer.weight,
|
||
l0_scale,
|
||
layer.weight_dual_scale,
|
||
l1_scale,
|
||
layer.weight_scale,
|
||
bias=bias,
|
||
output_dtype=ori_dtype,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# NPU MXFP4 dual-scale online method (BF16 checkpoint → quantize at load time)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class NPUMxfp4DualScaleOnlineLinearMethod(_LazyWeightMixin, NPUMxfp4DualScaleLinearMethod):
|
||
"""NPU W4A4 MXFP4 dual-scale online method: quantises BF16 weights at load time.
|
||
|
||
MRO: NPUMxfp4DualScaleOnlineLinearMethod → _LazyWeightMixin
|
||
→ NPUMxfp4DualScaleLinearMethod → MXFPLinearMethodBase
|
||
|
||
create_weights : _LazyWeightMixin (meta device + patched loader)
|
||
process_weights : NPUMxfp4DualScaleOnlineLinearMethod (BF16 → FP4 + dual scales)
|
||
apply / _quant_matmul : NPUMxfp4DualScaleLinearMethod (shared with offline path)
|
||
"""
|
||
|
||
def process_weights_after_loading(self, layer: Module) -> None:
|
||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||
return
|
||
|
||
import torch_npu
|
||
|
||
if layer.weight.device == torch.device("meta"):
|
||
weight = ModelWeightParameter(
|
||
data=torch.empty_like(layer.weight, device=layer._load_device),
|
||
input_dim=1,
|
||
output_dim=0,
|
||
weight_loader=layer.weight.weight_loader,
|
||
)
|
||
_copy_missing_attrs(layer.weight, weight)
|
||
layer.register_parameter("weight", weight)
|
||
initialize_single_dummy_weight(layer.weight)
|
||
|
||
# Quantize BF16 weight → FP4 + dual-level scales (no smooth pre-scale for online).
|
||
# Returns: (weight_fp4, l0_scale[coarse per-512], l1_scale[fine per-32])
|
||
weight_fp4, w_l0_scale, w_l1_scale = torch_npu.npu_dynamic_dual_level_mx_quant(
|
||
layer.weight.data.npu(), smooth_scale=None
|
||
)
|
||
|
||
# NZ hardware format for the FP4 weight (same as offline path).
|
||
w = torch_npu.npu_format_cast(weight_fp4.view(torch.int8), 29, customize_dtype=torch.int8)
|
||
|
||
# Fine scale (l1): (N, K//32) → (N, K//64, 2). Dtype from op output (float8_e8m0fnu).
|
||
s = w_l1_scale.reshape(w_l1_scale.shape[0], -1, 2).contiguous()
|
||
|
||
# Coarse scale (l0): (N, K_coarse) → (K_coarse, N). Dtype from op output.
|
||
ds = w_l0_scale.reshape(w_l0_scale.shape[0], -1).transpose(0, 1).contiguous()
|
||
|
||
# No calibration available: identity pre-scale (no smooth quantization effect).
|
||
ms = torch.ones(layer.input_size_per_partition, dtype=torch.bfloat16, device="npu")
|
||
|
||
replace_parameter(layer, "weight", w)
|
||
replace_parameter(layer, "weight_scale", s)
|
||
replace_parameter(layer, "weight_dual_scale", ds)
|
||
replace_parameter(layer, "mul_scale", ms)
|
||
layer._already_called_process_weights_after_loading = True
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Block-index helper (shared by DiffusionMXFP4DualScaleMixedConfig)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_BLOCK_IDX_RE = re.compile(r"^blocks\.(\d+)\.")
|
||
|
||
|
||
def _parse_block_idx(prefix: str) -> int | None:
|
||
"""Extract block index from prefix like 'blocks.5.attn1.to_q'."""
|
||
m = _BLOCK_IDX_RE.match(prefix)
|
||
return int(m.group(1)) if m else None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Config: MXFP4 DualScale + per-layer BF16 fallback
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class DiffusionMXFP4DualScaleMixedConfig(QuantizationConfig):
|
||
"""W4A4 MXFP4 DualScale with per-layer BF16 fallback for diffusion transformers.
|
||
|
||
Sensitive layers fall back to BF16 (original weights) while all other linear
|
||
layers use W4A4 MXFP4 DualScale. BF16 fallback layers may be interleaved
|
||
anywhere in the transformer.
|
||
|
||
Offline mode (is_checkpoint_serialized=True):
|
||
Layers whose prefix appears in ignored_layers → UnquantizedLinearMethod (BF16)
|
||
All other linear layers → NPUMxfp4DualScaleLinearMethod
|
||
|
||
ignored_layers is injected into transformer/config.json by the merge script
|
||
and contains the prefixes of all non-MXFP4 linear layers.
|
||
|
||
Online mode (is_checkpoint_serialized=False):
|
||
Layer routing applies two rules in priority order:
|
||
1. ignored_layers (explicit per-layer BF16 override, user-supplied) → BF16
|
||
2. Blocks 0 .. num_bf16_fallback_layers-1 (coarse leading-block rule) → BF16
|
||
3. All other linear layers → NPUMxfp4DualScaleOnlineLinearMethod
|
||
|
||
num_bf16_fallback_layers defaults to 5 when not specified.
|
||
Set ignored_layers to pin arbitrary interleaved layers to BF16 without
|
||
needing an offline checkpoint (useful for accuracy debugging).
|
||
Layers outside "blocks.N.*" (condition_embedder etc.) always use online MXFP4
|
||
unless they appear in ignored_layers.
|
||
|
||
Config injected by merge_mxfp4_dualscale_checkpoint.py:
|
||
{
|
||
"quant_method": "mxfp4_dualscale",
|
||
"is_checkpoint_serialized": true,
|
||
"ignored_layers": ["blocks.0.attn1.to_q", ...]
|
||
}
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
is_checkpoint_serialized: bool = False,
|
||
ignored_layers: list[str] | None = None,
|
||
num_bf16_fallback_layers: int = 5,
|
||
) -> None:
|
||
super().__init__()
|
||
self.is_checkpoint_serialized = is_checkpoint_serialized
|
||
self.ignored_layers = ignored_layers or []
|
||
self.num_bf16_fallback_layers = num_bf16_fallback_layers
|
||
|
||
@classmethod
|
||
def get_name(cls) -> QuantizationMethods:
|
||
return "mxfp4_dualscale"
|
||
|
||
@classmethod
|
||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||
return [torch.bfloat16, torch.float16]
|
||
|
||
@classmethod
|
||
def get_min_capability(cls) -> int:
|
||
return 80
|
||
|
||
@classmethod
|
||
def get_config_filenames(cls) -> list[str]:
|
||
return []
|
||
|
||
def apply_vllm_mapper(self, hf_to_vllm_mapper: WeightsMapper) -> None:
|
||
if self.ignored_layers:
|
||
self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers)
|
||
|
||
@classmethod
|
||
def from_config(cls, config: dict[str, Any]) -> DiffusionMXFP4DualScaleMixedConfig:
|
||
is_serialized = cls.get_from_keys_or(config, ["is_checkpoint_serialized"], False)
|
||
ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None)
|
||
if not ignored_layers:
|
||
ignored_layers = cls.get_from_keys_or(config, ["modules_to_not_convert"], None)
|
||
num_bf16_fallback_layers = cls.get_from_keys_or(config, ["num_bf16_fallback_layers"], 5)
|
||
return cls(
|
||
is_checkpoint_serialized=is_serialized,
|
||
ignored_layers=ignored_layers,
|
||
num_bf16_fallback_layers=num_bf16_fallback_layers,
|
||
)
|
||
|
||
def get_quant_method(
|
||
self,
|
||
layer: torch.nn.Module,
|
||
prefix: str,
|
||
) -> QuantizeMethodBase | None:
|
||
if not isinstance(layer, LinearBase):
|
||
return None
|
||
|
||
if self.is_checkpoint_serialized:
|
||
# Offline: ignored_layers lists interleaved BF16 fallback layer prefixes.
|
||
if is_layer_skipped(
|
||
prefix=prefix,
|
||
ignored_layers=self.ignored_layers,
|
||
fused_mapping=self.packed_modules_mapping,
|
||
):
|
||
return UnquantizedLinearMethod()
|
||
if not current_omni_platform.is_npu():
|
||
raise NotImplementedError(
|
||
"DiffusionMXFP4DualScaleMixedConfig is currently only supported on NPU (Ascend) platforms."
|
||
)
|
||
return NPUMxfp4DualScaleLinearMethod(self)
|
||
|
||
# Online: explicit ignored_layers take priority (user-specified per-layer BF16 override),
|
||
# then fall back to the coarse leading-block rule.
|
||
if is_layer_skipped(
|
||
prefix=prefix,
|
||
ignored_layers=self.ignored_layers,
|
||
fused_mapping=self.packed_modules_mapping,
|
||
):
|
||
return UnquantizedLinearMethod()
|
||
block_idx = _parse_block_idx(prefix)
|
||
if block_idx is not None and block_idx < self.num_bf16_fallback_layers:
|
||
return UnquantizedLinearMethod()
|
||
|
||
if not current_omni_platform.is_npu():
|
||
raise NotImplementedError(
|
||
"DiffusionMXFP4DualScaleMixedConfig is currently only supported on NPU (Ascend) platforms."
|
||
)
|
||
return NPUMxfp4DualScaleOnlineLinearMethod(self)
|