iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Quantization][Autoround][XPU] Add W4A16(moe) / MXFP4(linear/moe) Support (#47124)

Browse Source 
Signed-off-by: lkk12014402 <kaokao.lv@intel.com>
Signed-off-by: lkk <33276950+lkk12014402@users.noreply.github.com>
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
lkk
2026-07-30 09:52:13 +08:00
committed by GitHub
co-authored by Kunshang Ji
parent 4e582c5b54
commit e5f48dfda2
9 changed files with 816 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/quantization/test_auto_round.py
+391 -11
View File
@@ -8,16 +8,20 @@ Validating the configuration and printing results for manual checking.
Run `pytest tests/quantization/test_auto_round.py`.
"""
from typing import Any, cast
import pytest
import torch
from vllm.model_executor.layers.fused_moe import RoutedExperts
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
from vllm.model_executor.layers.quantization.inc import INCConfig
from vllm.model_executor.layers.quantization.inc.config_parser import INCLayerConfig
from vllm.model_executor.layers.quantization.inc.inc_linear import INCLinearMethod
from vllm.model_executor.layers.quantization.inc.schemes import (
INCMxfp4Scheme,
INCMxfp8Scheme,
INCWna16Scheme,
resolve_scheme,
@@ -71,7 +75,36 @@ MODELS = [
),
]
MODEL_RUNNER_KWARGS = {
QWEN3_AUTOROUND_MODELS = [
pytest.param(
"INCModel/Qwen3-1.7B-AutoRound-MXFP4-W4A4",
marks=pytest.mark.skipif(
not (current_platform.is_cuda() or current_platform.is_xpu()),
reason="Qwen3-1.7B MXFP4 AutoRound model requires CUDA/XPU.",
),
id="auto_round:mxfp4:qwen3-1p7b",
),
pytest.param(
"INCModel/Qwen3-30B-A3B-12L-W4A16-test",
marks=pytest.mark.skipif(
not (current_platform.is_cuda() or current_platform.is_xpu()),
reason="Qwen3-30B-A3B W4A16 AutoRound model requires CUDA/XPU.",
),
id="auto_round:w4a16:qwen3-30b-a3b",
),
pytest.param(
"INCModel/Qwen3-30B-A3B-12L-MXFP4-test",
marks=pytest.mark.skipif(
not (current_platform.is_cuda() or current_platform.is_xpu()),
reason="Qwen3-30B-A3B MXFP4 AutoRound model requires CUDA/XPU.",
),
id="auto_round:mxfp4:qwen3-30b-a3b",
),
]
MODEL_RUNNER_KWARGS: dict[str, dict[str, Any]] = {
"INCModel/Qwen3-1.7B-AutoRound-MXFP4-W4A4": {"enforce_eager": True},
"INCModel/Qwen3-30B-A3B-12L-MXFP4-test": {"enforce_eager": True},
"Intel/Qwen3-8B-w2g64-for-ut": {
"block_size": 64,
"gpu_memory_utilization": 0.8,
@@ -93,7 +126,7 @@ MODEL_RUNNER_KWARGS = {
),
reason="Only supports CPU/XPU/CUDA backend.",
)
@pytest.mark.parametrize("model", MODELS)
@pytest.mark.parametrize("model", MODELS + QWEN3_AUTOROUND_MODELS)
def test_auto_round_model(vllm_runner, model):
with vllm_runner(model, **MODEL_RUNNER_KWARGS.get(model, {})) as llm:
output = llm.generate_greedy(["The capital of France is"], max_tokens=8)
@@ -144,6 +177,65 @@ def make_layer_config(**overrides) -> INCLayerConfig:
return INCLayerConfig(**kwargs)
def make_qwen3_autoround_config(kind: str) -> INCConfig:
configs = {
"qwen3_1p7b_mxfp4": {
"quant_method": "auto-round",
"bits": 4,
"group_size": 32,
"sym": True,
"packing_format": "auto_round:llm_compressor",
"data_type": "mx_fp",
"extra_config": {
"model.layers.0.self_attn.q_proj": {
"bits": 16,
"data_type": "float",
},
},
},
"qwen3_30b_a3b_w4a16": {
"quant_method": "auto-round",
"bits": 4,
"group_size": 32,
"sym": True,
"packing_format": "auto_round:auto_gptq",
"data_type": "int",
"extra_config": {
"model.layers.0.mlp.gate": {
"bits": 16,
"data_type": "float",
},
},
},
"qwen3_30b_a3b_mxfp4": {
"quant_method": "auto-round",
"bits": 4,
"group_size": 32,
"sym": True,
"packing_format": "auto_round:llm_compressor",
"data_type": "mx_fp",
"act_bits": 4,
"act_group_size": 32,
"act_data_type": "mx_fp",
"extra_config": {
"model.layers.0.mlp.gate": {
"bits": 16,
"data_type": "float",
},
"model.layers.0.self_attn.q_proj": {
"bits": 16,
"data_type": "float",
},
},
},
}
try:
config = configs[kind]
except KeyError as err:
raise AssertionError(f"unknown qwen3 autoround config: {kind}") from err
return INCConfig.from_config(config)
def test_inc_config_parser_exact_match() -> None:
config = make_config(
extra_config={
@@ -374,6 +466,301 @@ def test_inc_resolve_scheme_selects_wna16() -> None:
assert isinstance(scheme, INCWna16Scheme)
def test_inc_config_accepts_mxfp_family_llm_compressor() -> None:
config = INCConfig.from_config(
{
"quant_method": "auto-round",
"bits": 4,
"group_size": 32,
"sym": True,
"packing_format": "auto_round:llm_compressor",
"data_type": "mx_fp4e2m1",
}
)
layer_config = config.config_parser.resolve(
DummyLayer(), "model.layers.0.mlp.down_proj"
)
assert config.sym is True
assert layer_config.is_mxfp4 is True
assert isinstance(resolve_scheme(layer_config), INCMxfp4Scheme)
def test_qwen3_1p7b_mxfp4_autoround_uses_mxfp4_linear_scheme(
monkeypatch,
) -> None:
class DummyKernel:
pass
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.inc.schemes."
"inc_mxfp4_linear.init_mxfp4_linear_kernel",
lambda: DummyKernel(),
)
from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_linear import ( # noqa: E501
INCMxfp4LinearMethod,
)
config = make_qwen3_autoround_config("qwen3_1p7b_mxfp4")
assert (
INCConfig.override_quantization_method(
{"quant_method": "auto-round"}, user_quant=None
)
== "inc"
)
ignored_method = config.get_quant_method(
object.__new__(LinearBase), "model.layers.0.self_attn.q_proj"
)
layer_config = config.config_parser.resolve(
DummyLayer(), "model.layers.0.mlp.gate_proj"
)
method = INCMxfp4Scheme().get_linear_method(
config,
object.__new__(LinearBase),
"model.layers.0.mlp.gate_proj",
layer_config,
)
assert isinstance(ignored_method, UnquantizedLinearMethod)
assert layer_config.bits == 4
assert layer_config.group_size == 32
assert layer_config.is_mxfp4 is True
assert isinstance(resolve_scheme(layer_config), INCMxfp4Scheme)
assert isinstance(method, INCLinearMethod)
assert isinstance(method.scheme, INCMxfp4LinearMethod)
assert isinstance(method.scheme.kernel, DummyKernel)
def test_qwen3_30b_a3b_w4a16_autoround_routes_to_gptq_moe(
monkeypatch,
) -> None:
captured = {}
expected_method = object()
class DummyMoeConfig:
pass
def fake_resolve_gptq_moe(layer, layer_config):
captured["layer"] = layer
captured["layer_config"] = layer_config
return expected_method
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.inc.schemes."
"inc_wna16_scheme._resolve_gptq_moe",
fake_resolve_gptq_moe,
)
config = make_qwen3_autoround_config("qwen3_30b_a3b_w4a16")
layer = object.__new__(RoutedExperts)
layer.moe_config = DummyMoeConfig()
method = config.get_quant_method(layer, "model.layers.0.mlp")
assert method is expected_method
assert captured["layer"] is layer
assert captured["layer_config"].bits == 4
assert captured["layer_config"].group_size == 32
assert captured["layer_config"].is_gptq is True
assert captured["layer_config"].is_wna16_int is True
def test_qwen3_30b_a3b_mxfp4_autoround_routes_to_mxfp4_moe(
monkeypatch,
) -> None:
class DummyMoeConfig:
pass
class DummyMxfp4MoEMethod:
def __init__(self, moe_config) -> None:
self.moe_config = moe_config
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe.INCMxfp4MoEMethod",
DummyMxfp4MoEMethod,
)
config = make_qwen3_autoround_config("qwen3_30b_a3b_mxfp4")
layer = object.__new__(RoutedExperts)
layer.moe_config = DummyMoeConfig()
ignored_method = config.get_quant_method(
object.__new__(LinearBase), "model.layers.0.self_attn.q_proj"
)
method = config.get_quant_method(layer, "model.layers.0.mlp")
layer_config = config.config_parser.resolve(DummyLayer(), "model.layers.0.mlp")
assert isinstance(ignored_method, UnquantizedLinearMethod)
assert layer_config.bits == 4
assert layer_config.group_size == 32
assert layer_config.is_mxfp4 is True
assert isinstance(resolve_scheme(layer_config), INCMxfp4Scheme)
assert isinstance(method, DummyMxfp4MoEMethod)
assert method.moe_config is layer.moe_config
def test_inc_mxfp4_linear_method_registers_and_processes_weights(
monkeypatch,
) -> None:
captured = {}
class DummyKernel:
def process_weights_after_loading(self, layer) -> None:
captured["processed_layer"] = layer
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.inc.schemes."
"inc_mxfp4_linear.init_mxfp4_linear_kernel",
lambda: DummyKernel(),
)
monkeypatch.setattr(
"vllm.model_executor.parameter.get_tensor_model_parallel_rank",
lambda: 0,
)
monkeypatch.setattr(
"vllm.model_executor.parameter.get_tensor_model_parallel_world_size",
lambda: 1,
)
from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_linear import ( # noqa: E501
INCMxfp4LinearMethod,
)
layer = torch.nn.Module()
method = INCMxfp4LinearMethod(
make_layer_config(group_size=32, data_type="mx_fp4e2m1")
)
method.create_weights(
layer,
input_size_per_partition=64,
output_partition_sizes=[16, 32],
input_size=64,
output_size=48,
params_dtype=torch.bfloat16,
)
assert layer.weight_packed.shape == (48, 32)
assert layer.weight_packed.dtype is torch.uint8
assert layer.weight_scale.shape == (48, 2)
assert layer.weight_scale.dtype is torch.uint8
assert layer.logical_widths == [16, 32]
assert layer.input_size_per_partition == 64
assert layer.output_size_per_partition == 48
packed_data = layer.weight_packed.data
method.process_weights_after_loading(layer)
assert layer.weight.data.data_ptr() == packed_data.data_ptr()
assert not hasattr(layer, "weight_packed")
assert captured["processed_layer"] is layer
def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel(
monkeypatch,
) -> None:
captured = {}
expected_quant_config = object()
expected_kernel = object()
expected_experts_cls = object()
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe."
"CutlassExpertsMxfp4._supports_current_device",
lambda: False,
)
monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe."
"select_mxfp4_moe_backend",
lambda moe: (Mxfp4MoeBackend.XPU, expected_experts_cls),
)
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe."
"make_mxfp4_moe_quant_config",
lambda **kwargs: captured.update({"quant_config_kwargs": kwargs})
or expected_quant_config,
)
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe."
"make_mxfp4_moe_kernel",
lambda **kwargs: captured.update({"kernel_kwargs": kwargs}) or expected_kernel,
)
from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe import (
INCMxfp4MoEMethod,
)
method = INCMxfp4MoEMethod(moe=cast(Any, "moe-config"))
layer = torch.nn.Module()
layer._expert_routing_tables = lambda: "routing-tables"
method.create_weights(
layer,
num_experts=2,
hidden_size=64,
intermediate_size_per_partition=32,
params_dtype=torch.bfloat16,
)
assert method.experts_cls is expected_experts_cls
assert layer.w13_weight_packed.shape == (2, 64, 32)
assert layer.w2_weight_packed.shape == (2, 64, 16)
assert layer.w13_weight_scale.shape == (2, 64, 2)
assert layer.w2_weight_scale.shape == (2, 64, 1)
w13_packed_data = layer.w13_weight_packed.data
w2_packed_data = layer.w2_weight_packed.data
method.process_weights_after_loading(layer)
assert layer.w13_weight.data.data_ptr() == w13_packed_data.data_ptr()
assert layer.w2_weight.data.data_ptr() == w2_packed_data.data_ptr()
assert not hasattr(layer, "w13_weight_packed")
assert not hasattr(layer, "w2_weight_packed")
assert captured["quant_config_kwargs"]["w1_scale"] is layer.w13_weight_scale
assert captured["quant_config_kwargs"]["w2_scale"] is layer.w2_weight_scale
assert captured["kernel_kwargs"]["moe_quant_config"] is expected_quant_config
assert captured["kernel_kwargs"]["moe_config"] == "moe-config"
assert captured["kernel_kwargs"]["experts_cls"] is expected_experts_cls
assert captured["kernel_kwargs"]["routing_tables"] == "routing-tables"
assert method.moe_kernel is expected_kernel
def test_wna16_xpu_moe_routes_to_gptq_moe(monkeypatch) -> None:
captured = {}
expected_method = object()
class DummyMoeConfig:
pass
monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
monkeypatch.setattr(current_platform, "is_cpu", lambda: False)
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.inc.schemes."
"inc_wna16_scheme._resolve_gptq_moe",
lambda layer, layer_config: captured.update(
{"layer": layer, "layer_config": layer_config}
)
or expected_method,
)
layer = object.__new__(RoutedExperts)
layer.moe_config = DummyMoeConfig()
method = INCWna16Scheme().get_moe_method(
make_config(),
layer,
"model.layers.0.mlp",
make_layer_config(group_size=32),
)
assert method is expected_method
assert captured["layer"] is layer
assert captured["layer_config"].is_gptq is True
def test_inc_resolve_scheme_selects_mxfp8() -> None:
layer_config = INCLayerConfig(
bits=8,
@@ -878,7 +1265,8 @@ def test_resolve_gptq_moe_falls_back_to_moe_wna16(monkeypatch) -> None:
captured["moe"] = moe
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.utils.marlin_utils.check_marlin_supported",
"vllm.model_executor.layers.quantization.utils.marlin_utils."
"check_moe_marlin_supports_layer",
lambda *args, **kwargs: False,
)
monkeypatch.setattr(
@@ -927,10 +1315,6 @@ def test_resolve_gptq_moe_uses_auto_gptq_when_supported(monkeypatch) -> None:
captured["cfg"] = cfg
captured["moe"] = moe
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.utils.marlin_utils.check_marlin_supported",
lambda *args, **kwargs: True,
)
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.utils.marlin_utils."
"check_moe_marlin_supports_layer",
@@ -964,10 +1348,6 @@ def test_resolve_awq_moe_uses_marlin_when_supported(monkeypatch) -> None:
captured["cfg"] = cfg
captured["moe"] = moe
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.utils.marlin_utils.check_marlin_supported",
lambda *args, **kwargs: True,
)
monkeypatch.setattr(
"vllm.model_executor.layers.quantization.utils.marlin_utils.check_moe_marlin_supports_layer",
lambda *args, **kwargs: True,
vllm/model_executor/layers/quantization/inc/config_parser.py
+2 -2
View File
@@ -38,11 +38,11 @@ class INCLayerConfig:
@property
def is_mxfp4(self) -> bool:
return self.data_type == "mx_fp" and self.bits == 4
return "mx_fp" in self.data_type and self.bits == 4
@property
def is_mxfp8(self) -> bool:
return self.data_type == "mx_fp" and self.bits == 8
return "mx_fp" in self.data_type and self.bits == 8
class INCConfigParser:
vllm/model_executor/layers/quantization/inc/inc.py
+19 -8
View File
@@ -72,7 +72,10 @@ class INCConfig(QuantizationConfig):
f"Unsupported weight_bits: {weight_bits}, "
f"currently only support {self.SUPPORTED_BITS}."
)
if data_type not in self.SUPPORTED_DTYPES:
# auto-round mxfp data_type is e.g. "mx_fp4" / "mx_fp4e2m1"; match the
# "mx_fp" family by substring like auto_round.compressors.is_mx_fp.
is_mxfp = "mx_fp" in data_type
if data_type not in self.SUPPORTED_DTYPES and not is_mxfp:
raise ValueError(
f"Unsupported data_type: {data_type}, "
f"currently only support {self.SUPPORTED_DTYPES}."
@@ -111,11 +114,19 @@ class INCConfig(QuantizationConfig):
f"group_size={self.group_size}, sym={self.sym})"
)
@property
def is_mxfp(self) -> bool:
# auto-round mxfp data_type is the "mx_fp" family (e.g. "mx_fp",
# "mx_fp4e2m1"); match by substring like auto_round.compressors.is_mx_fp.
return "mx_fp" in self.data_type
@property
def is_mxfp8(self) -> bool:
# MXFP4 and MXFP8 share data_type "mx_fp" and differ only by bit width.
return self.is_mxfp and self.weight_bits == self.MXFP8_BITS
def _validate_supported_quantization(self) -> None:
if self.data_type == self.MXFP8_DATA_TYPE:
assert self.weight_bits == self.MXFP8_BITS, (
f"INC MXFP8 only supports bits=8, but found bits={self.weight_bits}."
)
if self.is_mxfp8:
assert self.group_size == self.MXFP8_GROUP_SIZE, (
"INC MXFP8 only supports group_size=32, "
f"but found group_size={self.group_size}."
@@ -130,14 +141,14 @@ class INCConfig(QuantizationConfig):
"INC MXFP8 only supports backend='auto', "
f"but found backend={self.backend!r}."
)
elif self.packing_format == self.MXFP8_PACKING_FORMAT:
elif self.packing_format == self.MXFP8_PACKING_FORMAT and not self.is_mxfp:
raise ValueError(
f"packing_format={self.MXFP8_PACKING_FORMAT!r} requires "
f"data_type={self.MXFP8_DATA_TYPE!r}."
f"an {self.MXFP8_DATA_TYPE!r} data_type."
)
def _validate_raw_config(self, config: dict[str, Any]) -> None:
if self.data_type != self.MXFP8_DATA_TYPE:
if not self.is_mxfp8:
return
expected_fields = {
vllm/model_executor/layers/quantization/inc/schemes/__init__.py
+2
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from .factory import resolve_scheme
from .inc_mxfp4_scheme import INCMxfp4Scheme
from .inc_mxfp8_scheme import INCMxfp8Scheme
from .inc_scheme import INCLinearScheme, INCScheme
from .inc_wna16_scheme import INCWna16Scheme
@@ -11,5 +12,6 @@ __all__ = [
"INCLinearScheme",
"INCMxfp8Scheme",
"INCWna16Scheme",
"INCMxfp4Scheme",
"resolve_scheme",
]
vllm/model_executor/layers/quantization/inc/schemes/factory.py
+2
View File
@@ -9,12 +9,14 @@ if TYPE_CHECKING:
def resolve_scheme(layer_config: "INCLayerConfig") -> "INCScheme":
from .inc_mxfp4_scheme import INCMxfp4Scheme
from .inc_mxfp8_scheme import INCMxfp8Scheme
from .inc_wna16_scheme import INCWna16Scheme
scheme_list: list[type[INCScheme]] = [
INCMxfp8Scheme,
INCWna16Scheme,
INCMxfp4Scheme,
]
for scheme_cls in scheme_list:
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py
+91
View File
@@ -0,0 +1,91 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import TYPE_CHECKING, Any
import torch
from torch.nn.parameter import Parameter
from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel
from vllm.model_executor.parameter import (
GroupQuantScaleParameter,
ModelWeightParameter,
)
from .inc_scheme import INCLinearScheme
if TYPE_CHECKING:
from ..config_parser import INCLayerConfig
class INCMxfp4LinearMethod(INCLinearScheme):
"""MXFP4 (W4A4) linear method for AutoRound checkpoints.
E2M1 weights packed two per byte with per-group E8M0 scales
(group_size=32, no global scale). The platform kernel is selected by
``init_mxfp4_linear_kernel`` (FlashInfer / Marlin on CUDA, ``fp4_gemm``
on XPU).
"""
def __init__(self, layer_config: "INCLayerConfig") -> None:
self.group_size = layer_config.group_size or 32
self.kernel = init_mxfp4_linear_kernel()
@classmethod
def get_min_capability(cls) -> int:
return 80
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: Any,
) -> None:
del input_size, output_size
output_size_per_partition = sum(output_partition_sizes)
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.params_dtype = params_dtype
weight_loader = extra_weight_attrs.get("weight_loader")
weight = ModelWeightParameter(
data=torch.empty(
output_size_per_partition,
input_size_per_partition // 2,
dtype=torch.uint8,
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight_packed", weight)
weight_scale = GroupQuantScaleParameter(
data=torch.empty(
output_size_per_partition,
input_size_per_partition // self.group_size,
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) -> None:
layer.weight = Parameter(layer.weight_packed.data, requires_grad=False)
del layer.weight_packed
self.kernel.process_weights_after_loading(layer)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
return self.kernel.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py
+242
View File
@@ -0,0 +1,242 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe import (
FusedMoeWeightScaleSupported,
RoutedExperts,
SharedExperts,
)
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEQuantConfig,
mxfp4_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
CutlassExpertsMxfp4,
)
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
MarlinExperts,
)
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
FusedMoEMethodBase,
)
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
Mxfp4MoeBackend,
make_mxfp4_moe_kernel,
make_mxfp4_moe_quant_config,
select_mxfp4_moe_backend,
)
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
prepare_moe_fp4_layer_for_marlin,
)
from vllm.model_executor.utils import set_weight_attrs
from vllm.platforms import current_platform
logger = init_logger(__name__)
class INCMxfp4MoEMethod(FusedMoEMethodBase):
"""W4A4 MXFP4 group MoE for AutoRound ``auto_round:llm_compressor`` exports.
Registers the packed MXFP4 layout (uint8 ``weight_packed`` + uint8 E8M0
``weight_scale``, ``group_size=32``) and dispatches the fused MoE to the
best backend for the current device: CUTLASS (true W4A4 on supported
GPUs), the native XPU kernel, or Marlin weight-only as a fallback. The
per-expert ``gate_proj`` / ``up_proj`` / ``down_proj`` tensors are folded
into the stacked ``w13`` / ``w2`` parameters by ``make_expert_params_mapping``.
"""
def __init__(self, moe) -> None:
super().__init__(moe)
self.group_size = 32
# Backend selection must stay consistent with the weight preparation in
# process_weights_after_loading / get_fused_moe_quant_config, which only
# implement three layouts: CUTLASS swizzle (true W4A4), the native XPU
# kernel (packed passthrough), and Marlin weight-only. XPU dispatch is
# deferred to the shared oracle; every other non-CUTLASS device falls
# back to Marlin (mirroring CompressedTensorsW4A4Mxfp4MoEMethod).
self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
self.experts_cls: type[mk.FusedMoEExperts] | None = None
if self.use_cutlass_mxfp4:
self.experts_cls = CutlassExpertsMxfp4
logger.info_once("Using CutlassExpertsMxfp4 for AutoRound MXFP4 MoE")
elif current_platform.is_xpu():
self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)
else:
self.experts_cls = MarlinExperts
logger.info_once(
"Using MarlinExperts (weight-only FP4) for AutoRound MXFP4 MoE"
)
def create_weights(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
) -> None:
layer.num_experts = num_experts
layer.params_dtype = params_dtype
# gate + up fused on the output dim; two FP4 packed per input byte.
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // 2,
dtype=torch.uint8,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_packed", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // 2,
dtype=torch.uint8,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_packed", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# Per-group E8M0 block scales (group_size=32), stored as uint8.
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=torch.uint8,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=torch.uint8,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
def get_fused_moe_quant_config(
self, layer: torch.nn.Module
) -> FusedMoEQuantConfig | None:
if self.use_cutlass_mxfp4:
# W4A4: both weights and activations quantized to MXFP4.
return mxfp4_moe_quant_config(
w1_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
)
# Weight-only (Marlin) or native XPU kernel.
return make_mxfp4_moe_quant_config(
mxfp4_backend=self.mxfp4_backend,
w1_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
)
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
layer.w13_weight = torch.nn.Parameter(
layer.w13_weight_packed.data, requires_grad=False
)
delattr(layer, "w13_weight_packed")
layer.w2_weight = torch.nn.Parameter(
layer.w2_weight_packed.data, requires_grad=False
)
delattr(layer, "w2_weight_packed")
if self.use_cutlass_mxfp4:
# Swizzle weight scales from flat checkpoint layout [E, N, K//32]
# to the CUTLASS tiled layout.
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
swizzle_mxfp4_scales,
)
E = layer.w13_weight_scale.shape[0]
w13_N = layer.w13_weight_scale.shape[1]
w13_scale_K = layer.w13_weight_scale.shape[2]
w13_K = w13_scale_K * 32
w2_M = layer.w2_weight_scale.shape[1]
w2_scale_N = layer.w2_weight_scale.shape[2]
w2_N = w2_scale_N * 32
swizzled_w13 = []
swizzled_w2 = []
for e_idx in range(E):
s13 = layer.w13_weight_scale[e_idx]
sw13 = swizzle_mxfp4_scales(s13, w13_N, w13_K)
swizzled_w13.append(sw13.reshape(w13_N, w13_scale_K))
s2 = layer.w2_weight_scale[e_idx]
sw2 = swizzle_mxfp4_scales(s2, w2_M, w2_N)
swizzled_w2.append(sw2.reshape(w2_M, w2_scale_N))
layer.w13_weight_scale = torch.nn.Parameter(
torch.stack(swizzled_w13), requires_grad=False
)
layer.w2_weight_scale = torch.nn.Parameter(
torch.stack(swizzled_w2), requires_grad=False
)
elif current_platform.is_xpu():
# The XPU fused-MoE kernel consumes the packed layout directly; no
# swizzle / repack / transpose is required.
pass
else:
logger.warning_once(
"This device lacks native FP4 compute; using weight-only FP4 "
"via the Marlin kernel, which may reduce performance for "
"compute-heavy workloads."
)
prepare_moe_fp4_layer_for_marlin(layer)
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
if self.moe_quant_config is not None:
assert self.experts_cls is not None
self.moe_kernel = make_mxfp4_moe_kernel(
moe_quant_config=self.moe_quant_config,
moe_config=self.moe,
experts_cls=self.experts_cls,
mxfp4_backend=self.mxfp4_backend,
routing_tables=layer._expert_routing_tables(),
)
def apply(
self,
layer: RoutedExperts,
x: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
shared_experts: SharedExperts | None,
shared_experts_input: torch.Tensor | None,
) -> torch.Tensor:
assert self.moe_kernel is not None
return self.moe_kernel.apply(
x,
layer.w13_weight,
layer.w2_weight,
topk_weights,
topk_ids,
activation=layer.activation,
global_num_experts=layer.global_num_experts,
expert_map=layer.expert_map,
apply_router_weight_on_input=layer.apply_router_weight_on_input,
shared_experts=shared_experts,
shared_experts_input=shared_experts_input,
)
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_scheme.py
+54
View File
@@ -0,0 +1,54 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import TYPE_CHECKING
from vllm.logger import init_logger
from ..inc_linear import INCLinearMethod
from .inc_scheme import INCScheme
if TYPE_CHECKING:
import torch
from ..config_parser import INCLayerConfig
from ..inc import INCConfig
logger = init_logger(__name__)
class INCMxfp4Scheme(INCScheme):
"""MXFP4 (W4A4) scheme for AutoRound checkpoints.
Dispatches to :class:`INCMxfp4LinearMethod` for linear layers and
:class:`INCMxfp4MoEMethod` for fused MoE layers; see those classes for the
per-module weight layout and kernel-selection details.
"""
@staticmethod
def can_handle(layer_config: "INCLayerConfig") -> bool:
return layer_config.is_mxfp4
def get_linear_method(
self,
config: "INCConfig",
layer: "torch.nn.Module",
prefix: str,
layer_config: "INCLayerConfig",
):
del config, layer, prefix
from .inc_mxfp4_linear import INCMxfp4LinearMethod
return INCLinearMethod(INCMxfp4LinearMethod(layer_config))
def get_moe_method(
self,
config: "INCConfig",
layer: "torch.nn.Module",
prefix: str,
layer_config: "INCLayerConfig",
):
del config, prefix, layer_config
from .inc_mxfp4_moe import INCMxfp4MoEMethod
return INCMxfp4MoEMethod(layer.moe_config)
vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py
+13 -27
View File
@@ -7,7 +7,6 @@ from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types
from ..inc_linear import INCLinearMethod
from .inc_scheme import INCScheme
@@ -97,8 +96,8 @@ class INCWna16Scheme(INCScheme):
layer_config: "INCLayerConfig",
):
del config, prefix
# XPU and CPU do not support MoE quantization yet
if current_platform.is_xpu() or current_platform.is_cpu():
# CPU does not support quantized MoE yet.
if current_platform.is_cpu():
from vllm.model_executor.layers.fused_moe import (
UnquantizedFusedMoEMethod,
)
@@ -120,21 +119,17 @@ def _resolve_gptq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"):
MoeWNA16Method,
)
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
check_marlin_supported,
check_moe_marlin_supports_layer,
)
gptq_type_map = {
(4, True): scalar_types.uint4b8,
(8, True): scalar_types.uint8b128,
}
use_marlin = (layer_config.bits, layer_config.sym) in gptq_type_map
if use_marlin:
use_marlin = check_marlin_supported(
gptq_type_map[(layer_config.bits, layer_config.sym)],
layer_config.group_size,
has_zp=not layer_config.sym,
) and check_moe_marlin_supports_layer(layer, layer_config.group_size)
# AutoGPTQMoEMethod selects its fused-MoE backend through the WNA16 oracle
# (Marlin on CUDA, XPUExpertsWNA16 on XPU). Gate only on the layer-shape
# check like compressed-tensors does; the capability-based
# check_marlin_supported is skipped so the XPU path is reachable.
use_marlin = (layer_config.bits, layer_config.sym) in {
(4, True),
(8, True),
} and check_moe_marlin_supports_layer(layer, layer_config.group_size)
if use_marlin:
return AutoGPTQMoEMethod(
@@ -169,21 +164,12 @@ def _resolve_awq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"):
MoeWNA16Method,
)
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
check_marlin_supported,
check_moe_marlin_supports_layer,
)
awq_type_map = {
4: scalar_types.uint4,
8: scalar_types.uint8,
}
use_marlin = layer_config.bits in awq_type_map
if use_marlin:
use_marlin = check_marlin_supported(
awq_type_map[layer_config.bits],
layer_config.group_size,
not layer_config.sym,
) and check_moe_marlin_supports_layer(layer, layer_config.group_size)
use_marlin = layer_config.bits in (4, 8) and check_moe_marlin_supports_layer(
layer, layer_config.group_size
)
if use_marlin:
return AutoAWQMoEMethod(