From e5f48dfda2b53090f539a2b3988e075e8ac0eba8 Mon Sep 17 00:00:00 2001 From: lkk <33276950+lkk12014402@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:52:13 +0800 Subject: [PATCH] [Quantization][Autoround][XPU] Add W4A16(moe) / MXFP4(linear/moe) Support (#47124) Signed-off-by: lkk12014402 Signed-off-by: lkk <33276950+lkk12014402@users.noreply.github.com> Signed-off-by: Kunshang Ji Co-authored-by: Kunshang Ji --- tests/quantization/test_auto_round.py | 402 +++++++++++++++++- .../layers/quantization/inc/config_parser.py | 4 +- .../layers/quantization/inc/inc.py | 27 +- .../quantization/inc/schemes/__init__.py | 2 + .../quantization/inc/schemes/factory.py | 2 + .../inc/schemes/inc_mxfp4_linear.py | 91 ++++ .../quantization/inc/schemes/inc_mxfp4_moe.py | 242 +++++++++++ .../inc/schemes/inc_mxfp4_scheme.py | 54 +++ .../inc/schemes/inc_wna16_scheme.py | 40 +- 9 files changed, 816 insertions(+), 48 deletions(-) create mode 100644 vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py create mode 100644 vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py create mode 100644 vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_scheme.py diff --git a/tests/quantization/test_auto_round.py b/tests/quantization/test_auto_round.py index 59c6a1e326..838ca44956 100644 --- a/tests/quantization/test_auto_round.py +++ b/tests/quantization/test_auto_round.py @@ -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, diff --git a/vllm/model_executor/layers/quantization/inc/config_parser.py b/vllm/model_executor/layers/quantization/inc/config_parser.py index 6e94cad2cc..7a0d8496b8 100644 --- a/vllm/model_executor/layers/quantization/inc/config_parser.py +++ b/vllm/model_executor/layers/quantization/inc/config_parser.py @@ -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: diff --git a/vllm/model_executor/layers/quantization/inc/inc.py b/vllm/model_executor/layers/quantization/inc/inc.py index ad8389a8b9..219fc2b046 100644 --- a/vllm/model_executor/layers/quantization/inc/inc.py +++ b/vllm/model_executor/layers/quantization/inc/inc.py @@ -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 = { diff --git a/vllm/model_executor/layers/quantization/inc/schemes/__init__.py b/vllm/model_executor/layers/quantization/inc/schemes/__init__.py index e4dc5cb5dc..1f2d5f4401 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/__init__.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/__init__.py @@ -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", ] diff --git a/vllm/model_executor/layers/quantization/inc/schemes/factory.py b/vllm/model_executor/layers/quantization/inc/schemes/factory.py index 30ce420b50..61335e602f 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/factory.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/factory.py @@ -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: diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py new file mode 100644 index 0000000000..8fbbc65a0b --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py @@ -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) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py new file mode 100644 index 0000000000..7cb3adce09 --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py @@ -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, + ) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_scheme.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_scheme.py new file mode 100644 index 0000000000..e22bb900ce --- /dev/null +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_scheme.py @@ -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) diff --git a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py index 46ad24ea5b..f83cc22e8e 100644 --- a/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py +++ b/vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py @@ -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(