@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e5f48dfda2

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

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>

lkk committed Jul 30, 2026 at 09:52 UTC e5f48dfda2b53090f539a2b3988e075e8ac0eba8
9 files changed +816 -48
tests/quantization/test_auto_round.py
+391 -11
@@ -8,16 +8,20 @@ Validating the configuration and printing results for manual checking.
8 Run `pytest tests/quantization/test_auto_round.py`.
9 """
10
11 +from typing import Any, cast
12 +
13 import pytest
14 import torch
15
16 from vllm.model_executor.layers.fused_moe import RoutedExperts
17 +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend
18 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
19 from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
20 from vllm.model_executor.layers.quantization.inc import INCConfig
21 from vllm.model_executor.layers.quantization.inc.config_parser import INCLayerConfig
22 from vllm.model_executor.layers.quantization.inc.inc_linear import INCLinearMethod
23 from vllm.model_executor.layers.quantization.inc.schemes import (
24 + INCMxfp4Scheme,
25 INCMxfp8Scheme,
26 INCWna16Scheme,
27 resolve_scheme,
@@ -71,7 +75,36 @@ MODELS = [
75 ),
76 ]
77
74 -MODEL_RUNNER_KWARGS = {
78 +QWEN3_AUTOROUND_MODELS = [
79 + pytest.param(
80 + "INCModel/Qwen3-1.7B-AutoRound-MXFP4-W4A4",
81 + marks=pytest.mark.skipif(
82 + not (current_platform.is_cuda() or current_platform.is_xpu()),
83 + reason="Qwen3-1.7B MXFP4 AutoRound model requires CUDA/XPU.",
84 + ),
85 + id="auto_round:mxfp4:qwen3-1p7b",
86 + ),
87 + pytest.param(
88 + "INCModel/Qwen3-30B-A3B-12L-W4A16-test",
89 + marks=pytest.mark.skipif(
90 + not (current_platform.is_cuda() or current_platform.is_xpu()),
91 + reason="Qwen3-30B-A3B W4A16 AutoRound model requires CUDA/XPU.",
92 + ),
93 + id="auto_round:w4a16:qwen3-30b-a3b",
94 + ),
95 + pytest.param(
96 + "INCModel/Qwen3-30B-A3B-12L-MXFP4-test",
97 + marks=pytest.mark.skipif(
98 + not (current_platform.is_cuda() or current_platform.is_xpu()),
99 + reason="Qwen3-30B-A3B MXFP4 AutoRound model requires CUDA/XPU.",
100 + ),
101 + id="auto_round:mxfp4:qwen3-30b-a3b",
102 + ),
103 +]
104 +
105 +MODEL_RUNNER_KWARGS: dict[str, dict[str, Any]] = {
106 + "INCModel/Qwen3-1.7B-AutoRound-MXFP4-W4A4": {"enforce_eager": True},
107 + "INCModel/Qwen3-30B-A3B-12L-MXFP4-test": {"enforce_eager": True},
108 "Intel/Qwen3-8B-w2g64-for-ut": {
109 "block_size": 64,
110 "gpu_memory_utilization": 0.8,
@@ -93,7 +126,7 @@ MODEL_RUNNER_KWARGS = {
126 ),
127 reason="Only supports CPU/XPU/CUDA backend.",
128 )
96 -@pytest.mark.parametrize("model", MODELS)
129 +@pytest.mark.parametrize("model", MODELS + QWEN3_AUTOROUND_MODELS)
130 def test_auto_round_model(vllm_runner, model):
131 with vllm_runner(model, **MODEL_RUNNER_KWARGS.get(model, {})) as llm:
132 output = llm.generate_greedy(["The capital of France is"], max_tokens=8)
@@ -144,6 +177,65 @@ def make_layer_config(**overrides) -> INCLayerConfig:
177 return INCLayerConfig(**kwargs)
178
179
180 +def make_qwen3_autoround_config(kind: str) -> INCConfig:
181 + configs = {
182 + "qwen3_1p7b_mxfp4": {
183 + "quant_method": "auto-round",
184 + "bits": 4,
185 + "group_size": 32,
186 + "sym": True,
187 + "packing_format": "auto_round:llm_compressor",
188 + "data_type": "mx_fp",
189 + "extra_config": {
190 + "model.layers.0.self_attn.q_proj": {
191 + "bits": 16,
192 + "data_type": "float",
193 + },
194 + },
195 + },
196 + "qwen3_30b_a3b_w4a16": {
197 + "quant_method": "auto-round",
198 + "bits": 4,
199 + "group_size": 32,
200 + "sym": True,
201 + "packing_format": "auto_round:auto_gptq",
202 + "data_type": "int",
203 + "extra_config": {
204 + "model.layers.0.mlp.gate": {
205 + "bits": 16,
206 + "data_type": "float",
207 + },
208 + },
209 + },
210 + "qwen3_30b_a3b_mxfp4": {
211 + "quant_method": "auto-round",
212 + "bits": 4,
213 + "group_size": 32,
214 + "sym": True,
215 + "packing_format": "auto_round:llm_compressor",
216 + "data_type": "mx_fp",
217 + "act_bits": 4,
218 + "act_group_size": 32,
219 + "act_data_type": "mx_fp",
220 + "extra_config": {
221 + "model.layers.0.mlp.gate": {
222 + "bits": 16,
223 + "data_type": "float",
224 + },
225 + "model.layers.0.self_attn.q_proj": {
226 + "bits": 16,
227 + "data_type": "float",
228 + },
229 + },
230 + },
231 + }
232 + try:
233 + config = configs[kind]
234 + except KeyError as err:
235 + raise AssertionError(f"unknown qwen3 autoround config: {kind}") from err
236 + return INCConfig.from_config(config)
237 +
238 +
239 def test_inc_config_parser_exact_match() -> None:
240 config = make_config(
241 extra_config={
@@ -374,6 +466,301 @@ def test_inc_resolve_scheme_selects_wna16() -> None:
466 assert isinstance(scheme, INCWna16Scheme)
467
468
469 +def test_inc_config_accepts_mxfp_family_llm_compressor() -> None:
470 + config = INCConfig.from_config(
471 + {
472 + "quant_method": "auto-round",
473 + "bits": 4,
474 + "group_size": 32,
475 + "sym": True,
476 + "packing_format": "auto_round:llm_compressor",
477 + "data_type": "mx_fp4e2m1",
478 + }
479 + )
480 +
481 + layer_config = config.config_parser.resolve(
482 + DummyLayer(), "model.layers.0.mlp.down_proj"
483 + )
484 +
485 + assert config.sym is True
486 + assert layer_config.is_mxfp4 is True
487 + assert isinstance(resolve_scheme(layer_config), INCMxfp4Scheme)
488 +
489 +
490 +def test_qwen3_1p7b_mxfp4_autoround_uses_mxfp4_linear_scheme(
491 + monkeypatch,
492 +) -> None:
493 + class DummyKernel:
494 + pass
495 +
496 + monkeypatch.setattr(
497 + "vllm.model_executor.layers.quantization.inc.schemes."
498 + "inc_mxfp4_linear.init_mxfp4_linear_kernel",
499 + lambda: DummyKernel(),
500 + )
501 +
502 + from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_linear import ( # noqa: E501
503 + INCMxfp4LinearMethod,
504 + )
505 +
506 + config = make_qwen3_autoround_config("qwen3_1p7b_mxfp4")
507 +
508 + assert (
509 + INCConfig.override_quantization_method(
510 + {"quant_method": "auto-round"}, user_quant=None
511 + )
512 + == "inc"
513 + )
514 + ignored_method = config.get_quant_method(
515 + object.__new__(LinearBase), "model.layers.0.self_attn.q_proj"
516 + )
517 + layer_config = config.config_parser.resolve(
518 + DummyLayer(), "model.layers.0.mlp.gate_proj"
519 + )
520 + method = INCMxfp4Scheme().get_linear_method(
521 + config,
522 + object.__new__(LinearBase),
523 + "model.layers.0.mlp.gate_proj",
524 + layer_config,
525 + )
526 +
527 + assert isinstance(ignored_method, UnquantizedLinearMethod)
528 + assert layer_config.bits == 4
529 + assert layer_config.group_size == 32
530 + assert layer_config.is_mxfp4 is True
531 + assert isinstance(resolve_scheme(layer_config), INCMxfp4Scheme)
532 + assert isinstance(method, INCLinearMethod)
533 + assert isinstance(method.scheme, INCMxfp4LinearMethod)
534 + assert isinstance(method.scheme.kernel, DummyKernel)
535 +
536 +
537 +def test_qwen3_30b_a3b_w4a16_autoround_routes_to_gptq_moe(
538 + monkeypatch,
539 +) -> None:
540 + captured = {}
541 + expected_method = object()
542 +
543 + class DummyMoeConfig:
544 + pass
545 +
546 + def fake_resolve_gptq_moe(layer, layer_config):
547 + captured["layer"] = layer
548 + captured["layer_config"] = layer_config
549 + return expected_method
550 +
551 + monkeypatch.setattr(
552 + "vllm.model_executor.layers.quantization.inc.schemes."
553 + "inc_wna16_scheme._resolve_gptq_moe",
554 + fake_resolve_gptq_moe,
555 + )
556 +
557 + config = make_qwen3_autoround_config("qwen3_30b_a3b_w4a16")
558 + layer = object.__new__(RoutedExperts)
559 + layer.moe_config = DummyMoeConfig()
560 +
561 + method = config.get_quant_method(layer, "model.layers.0.mlp")
562 +
563 + assert method is expected_method
564 + assert captured["layer"] is layer
565 + assert captured["layer_config"].bits == 4
566 + assert captured["layer_config"].group_size == 32
567 + assert captured["layer_config"].is_gptq is True
568 + assert captured["layer_config"].is_wna16_int is True
569 +
570 +
571 +def test_qwen3_30b_a3b_mxfp4_autoround_routes_to_mxfp4_moe(
572 + monkeypatch,
573 +) -> None:
574 + class DummyMoeConfig:
575 + pass
576 +
577 + class DummyMxfp4MoEMethod:
578 + def __init__(self, moe_config) -> None:
579 + self.moe_config = moe_config
580 +
581 + monkeypatch.setattr(
582 + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe.INCMxfp4MoEMethod",
583 + DummyMxfp4MoEMethod,
584 + )
585 +
586 + config = make_qwen3_autoround_config("qwen3_30b_a3b_mxfp4")
587 + layer = object.__new__(RoutedExperts)
588 + layer.moe_config = DummyMoeConfig()
589 +
590 + ignored_method = config.get_quant_method(
591 + object.__new__(LinearBase), "model.layers.0.self_attn.q_proj"
592 + )
593 + method = config.get_quant_method(layer, "model.layers.0.mlp")
594 + layer_config = config.config_parser.resolve(DummyLayer(), "model.layers.0.mlp")
595 +
596 + assert isinstance(ignored_method, UnquantizedLinearMethod)
597 + assert layer_config.bits == 4
598 + assert layer_config.group_size == 32
599 + assert layer_config.is_mxfp4 is True
600 + assert isinstance(resolve_scheme(layer_config), INCMxfp4Scheme)
601 + assert isinstance(method, DummyMxfp4MoEMethod)
602 + assert method.moe_config is layer.moe_config
603 +
604 +
605 +def test_inc_mxfp4_linear_method_registers_and_processes_weights(
606 + monkeypatch,
607 +) -> None:
608 + captured = {}
609 +
610 + class DummyKernel:
611 + def process_weights_after_loading(self, layer) -> None:
612 + captured["processed_layer"] = layer
613 +
614 + monkeypatch.setattr(
615 + "vllm.model_executor.layers.quantization.inc.schemes."
616 + "inc_mxfp4_linear.init_mxfp4_linear_kernel",
617 + lambda: DummyKernel(),
618 + )
619 + monkeypatch.setattr(
620 + "vllm.model_executor.parameter.get_tensor_model_parallel_rank",
621 + lambda: 0,
622 + )
623 + monkeypatch.setattr(
624 + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size",
625 + lambda: 1,
626 + )
627 +
628 + from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_linear import ( # noqa: E501
629 + INCMxfp4LinearMethod,
630 + )
631 +
632 + layer = torch.nn.Module()
633 + method = INCMxfp4LinearMethod(
634 + make_layer_config(group_size=32, data_type="mx_fp4e2m1")
635 + )
636 +
637 + method.create_weights(
638 + layer,
639 + input_size_per_partition=64,
640 + output_partition_sizes=[16, 32],
641 + input_size=64,
642 + output_size=48,
643 + params_dtype=torch.bfloat16,
644 + )
645 +
646 + assert layer.weight_packed.shape == (48, 32)
647 + assert layer.weight_packed.dtype is torch.uint8
648 + assert layer.weight_scale.shape == (48, 2)
649 + assert layer.weight_scale.dtype is torch.uint8
650 + assert layer.logical_widths == [16, 32]
651 + assert layer.input_size_per_partition == 64
652 + assert layer.output_size_per_partition == 48
653 +
654 + packed_data = layer.weight_packed.data
655 + method.process_weights_after_loading(layer)
656 +
657 + assert layer.weight.data.data_ptr() == packed_data.data_ptr()
658 + assert not hasattr(layer, "weight_packed")
659 + assert captured["processed_layer"] is layer
660 +
661 +
662 +def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel(
663 + monkeypatch,
664 +) -> None:
665 + captured = {}
666 + expected_quant_config = object()
667 + expected_kernel = object()
668 + expected_experts_cls = object()
669 +
670 + monkeypatch.setattr(
671 + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe."
672 + "CutlassExpertsMxfp4._supports_current_device",
673 + lambda: False,
674 + )
675 + monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
676 + monkeypatch.setattr(
677 + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe."
678 + "select_mxfp4_moe_backend",
679 + lambda moe: (Mxfp4MoeBackend.XPU, expected_experts_cls),
680 + )
681 + monkeypatch.setattr(
682 + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe."
683 + "make_mxfp4_moe_quant_config",
684 + lambda **kwargs: captured.update({"quant_config_kwargs": kwargs})
685 + or expected_quant_config,
686 + )
687 + monkeypatch.setattr(
688 + "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe."
689 + "make_mxfp4_moe_kernel",
690 + lambda **kwargs: captured.update({"kernel_kwargs": kwargs}) or expected_kernel,
691 + )
692 +
693 + from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe import (
694 + INCMxfp4MoEMethod,
695 + )
696 +
697 + method = INCMxfp4MoEMethod(moe=cast(Any, "moe-config"))
698 + layer = torch.nn.Module()
699 + layer._expert_routing_tables = lambda: "routing-tables"
700 +
701 + method.create_weights(
702 + layer,
703 + num_experts=2,
704 + hidden_size=64,
705 + intermediate_size_per_partition=32,
706 + params_dtype=torch.bfloat16,
707 + )
708 +
709 + assert method.experts_cls is expected_experts_cls
710 + assert layer.w13_weight_packed.shape == (2, 64, 32)
711 + assert layer.w2_weight_packed.shape == (2, 64, 16)
712 + assert layer.w13_weight_scale.shape == (2, 64, 2)
713 + assert layer.w2_weight_scale.shape == (2, 64, 1)
714 +
715 + w13_packed_data = layer.w13_weight_packed.data
716 + w2_packed_data = layer.w2_weight_packed.data
717 + method.process_weights_after_loading(layer)
718 +
719 + assert layer.w13_weight.data.data_ptr() == w13_packed_data.data_ptr()
720 + assert layer.w2_weight.data.data_ptr() == w2_packed_data.data_ptr()
721 + assert not hasattr(layer, "w13_weight_packed")
722 + assert not hasattr(layer, "w2_weight_packed")
723 + assert captured["quant_config_kwargs"]["w1_scale"] is layer.w13_weight_scale
724 + assert captured["quant_config_kwargs"]["w2_scale"] is layer.w2_weight_scale
725 + assert captured["kernel_kwargs"]["moe_quant_config"] is expected_quant_config
726 + assert captured["kernel_kwargs"]["moe_config"] == "moe-config"
727 + assert captured["kernel_kwargs"]["experts_cls"] is expected_experts_cls
728 + assert captured["kernel_kwargs"]["routing_tables"] == "routing-tables"
729 + assert method.moe_kernel is expected_kernel
730 +
731 +
732 +def test_wna16_xpu_moe_routes_to_gptq_moe(monkeypatch) -> None:
733 + captured = {}
734 + expected_method = object()
735 +
736 + class DummyMoeConfig:
737 + pass
738 +
739 + monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
740 + monkeypatch.setattr(current_platform, "is_cpu", lambda: False)
741 + monkeypatch.setattr(
742 + "vllm.model_executor.layers.quantization.inc.schemes."
743 + "inc_wna16_scheme._resolve_gptq_moe",
744 + lambda layer, layer_config: captured.update(
745 + {"layer": layer, "layer_config": layer_config}
746 + )
747 + or expected_method,
748 + )
749 +
750 + layer = object.__new__(RoutedExperts)
751 + layer.moe_config = DummyMoeConfig()
752 + method = INCWna16Scheme().get_moe_method(
753 + make_config(),
754 + layer,
755 + "model.layers.0.mlp",
756 + make_layer_config(group_size=32),
757 + )
758 +
759 + assert method is expected_method
760 + assert captured["layer"] is layer
761 + assert captured["layer_config"].is_gptq is True
762 +
763 +
764 def test_inc_resolve_scheme_selects_mxfp8() -> None:
765 layer_config = INCLayerConfig(
766 bits=8,
@@ -878,7 +1265,8 @@ def test_resolve_gptq_moe_falls_back_to_moe_wna16(monkeypatch) -> None:
1265 captured["moe"] = moe
1266
1267 monkeypatch.setattr(
881 - "vllm.model_executor.layers.quantization.utils.marlin_utils.check_marlin_supported",
1268 + "vllm.model_executor.layers.quantization.utils.marlin_utils."
1269 + "check_moe_marlin_supports_layer",
1270 lambda *args, **kwargs: False,
1271 )
1272 monkeypatch.setattr(
@@ -927,10 +1315,6 @@ def test_resolve_gptq_moe_uses_auto_gptq_when_supported(monkeypatch) -> None:
1315 captured["cfg"] = cfg
1316 captured["moe"] = moe
1317
930 - monkeypatch.setattr(
931 - "vllm.model_executor.layers.quantization.utils.marlin_utils.check_marlin_supported",
932 - lambda *args, **kwargs: True,
933 - )
1318 monkeypatch.setattr(
1319 "vllm.model_executor.layers.quantization.utils.marlin_utils."
1320 "check_moe_marlin_supports_layer",
@@ -964,10 +1348,6 @@ def test_resolve_awq_moe_uses_marlin_when_supported(monkeypatch) -> None:
1348 captured["cfg"] = cfg
1349 captured["moe"] = moe
1350
967 - monkeypatch.setattr(
968 - "vllm.model_executor.layers.quantization.utils.marlin_utils.check_marlin_supported",
969 - lambda *args, **kwargs: True,
970 - )
1351 monkeypatch.setattr(
1352 "vllm.model_executor.layers.quantization.utils.marlin_utils.check_moe_marlin_supports_layer",
1353 lambda *args, **kwargs: True,
vllm/model_executor/layers/quantization/inc/config_parser.py
+2 -2
@@ -38,11 +38,11 @@ class INCLayerConfig:
38
39 @property
40 def is_mxfp4(self) -> bool:
41 - return self.data_type == "mx_fp" and self.bits == 4
41 + return "mx_fp" in self.data_type and self.bits == 4
42
43 @property
44 def is_mxfp8(self) -> bool:
45 - return self.data_type == "mx_fp" and self.bits == 8
45 + return "mx_fp" in self.data_type and self.bits == 8
46
47
48 class INCConfigParser:
vllm/model_executor/layers/quantization/inc/inc.py
+19 -8
@@ -72,7 +72,10 @@ class INCConfig(QuantizationConfig):
72 f"Unsupported weight_bits: {weight_bits}, "
73 f"currently only support {self.SUPPORTED_BITS}."
74 )
75 - if data_type not in self.SUPPORTED_DTYPES:
75 + # auto-round mxfp data_type is e.g. "mx_fp4" / "mx_fp4e2m1"; match the
76 + # "mx_fp" family by substring like auto_round.compressors.is_mx_fp.
77 + is_mxfp = "mx_fp" in data_type
78 + if data_type not in self.SUPPORTED_DTYPES and not is_mxfp:
79 raise ValueError(
80 f"Unsupported data_type: {data_type}, "
81 f"currently only support {self.SUPPORTED_DTYPES}."
@@ -111,11 +114,19 @@ class INCConfig(QuantizationConfig):
114 f"group_size={self.group_size}, sym={self.sym})"
115 )
116
117 + @property
118 + def is_mxfp(self) -> bool:
119 + # auto-round mxfp data_type is the "mx_fp" family (e.g. "mx_fp",
120 + # "mx_fp4e2m1"); match by substring like auto_round.compressors.is_mx_fp.
121 + return "mx_fp" in self.data_type
122 +
123 + @property
124 + def is_mxfp8(self) -> bool:
125 + # MXFP4 and MXFP8 share data_type "mx_fp" and differ only by bit width.
126 + return self.is_mxfp and self.weight_bits == self.MXFP8_BITS
127 +
128 def _validate_supported_quantization(self) -> None:
115 - if self.data_type == self.MXFP8_DATA_TYPE:
116 - assert self.weight_bits == self.MXFP8_BITS, (
117 - f"INC MXFP8 only supports bits=8, but found bits={self.weight_bits}."
118 - )
129 + if self.is_mxfp8:
130 assert self.group_size == self.MXFP8_GROUP_SIZE, (
131 "INC MXFP8 only supports group_size=32, "
132 f"but found group_size={self.group_size}."
@@ -130,14 +141,14 @@ class INCConfig(QuantizationConfig):
141 "INC MXFP8 only supports backend='auto', "
142 f"but found backend={self.backend!r}."
143 )
133 - elif self.packing_format == self.MXFP8_PACKING_FORMAT:
144 + elif self.packing_format == self.MXFP8_PACKING_FORMAT and not self.is_mxfp:
145 raise ValueError(
146 f"packing_format={self.MXFP8_PACKING_FORMAT!r} requires "
136 - f"data_type={self.MXFP8_DATA_TYPE!r}."
147 + f"an {self.MXFP8_DATA_TYPE!r} data_type."
148 )
149
150 def _validate_raw_config(self, config: dict[str, Any]) -> None:
140 - if self.data_type != self.MXFP8_DATA_TYPE:
151 + if not self.is_mxfp8:
152 return
153
154 expected_fields = {
vllm/model_executor/layers/quantization/inc/schemes/__init__.py
+2
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from .factory import resolve_scheme
5 +from .inc_mxfp4_scheme import INCMxfp4Scheme
6 from .inc_mxfp8_scheme import INCMxfp8Scheme
7 from .inc_scheme import INCLinearScheme, INCScheme
8 from .inc_wna16_scheme import INCWna16Scheme
@@ -11,5 +12,6 @@ __all__ = [
12 "INCLinearScheme",
13 "INCMxfp8Scheme",
14 "INCWna16Scheme",
15 + "INCMxfp4Scheme",
16 "resolve_scheme",
17 ]
vllm/model_executor/layers/quantization/inc/schemes/factory.py
+2
@@ -9,12 +9,14 @@ if TYPE_CHECKING:
9
10
11 def resolve_scheme(layer_config: "INCLayerConfig") -> "INCScheme":
12 + from .inc_mxfp4_scheme import INCMxfp4Scheme
13 from .inc_mxfp8_scheme import INCMxfp8Scheme
14 from .inc_wna16_scheme import INCWna16Scheme
15
16 scheme_list: list[type[INCScheme]] = [
17 INCMxfp8Scheme,
18 INCWna16Scheme,
19 + INCMxfp4Scheme,
20 ]
21
22 for scheme_cls in scheme_list:
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_linear.py new
+91
@@ -0,0 +1,91 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import TYPE_CHECKING, Any
5 +
6 +import torch
7 +from torch.nn.parameter import Parameter
8 +
9 +from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel
10 +from vllm.model_executor.parameter import (
11 + GroupQuantScaleParameter,
12 + ModelWeightParameter,
13 +)
14 +
15 +from .inc_scheme import INCLinearScheme
16 +
17 +if TYPE_CHECKING:
18 + from ..config_parser import INCLayerConfig
19 +
20 +
21 +class INCMxfp4LinearMethod(INCLinearScheme):
22 + """MXFP4 (W4A4) linear method for AutoRound checkpoints.
23 +
24 + E2M1 weights packed two per byte with per-group E8M0 scales
25 + (group_size=32, no global scale). The platform kernel is selected by
26 + ``init_mxfp4_linear_kernel`` (FlashInfer / Marlin on CUDA, ``fp4_gemm``
27 + on XPU).
28 + """
29 +
30 + def __init__(self, layer_config: "INCLayerConfig") -> None:
31 + self.group_size = layer_config.group_size or 32
32 + self.kernel = init_mxfp4_linear_kernel()
33 +
34 + @classmethod
35 + def get_min_capability(cls) -> int:
36 + return 80
37 +
38 + def create_weights(
39 + self,
40 + layer: torch.nn.Module,
41 + input_size_per_partition: int,
42 + output_partition_sizes: list[int],
43 + input_size: int,
44 + output_size: int,
45 + params_dtype: torch.dtype,
46 + **extra_weight_attrs: Any,
47 + ) -> None:
48 + del input_size, output_size
49 + output_size_per_partition = sum(output_partition_sizes)
50 + layer.logical_widths = output_partition_sizes
51 + layer.input_size_per_partition = input_size_per_partition
52 + layer.output_size_per_partition = output_size_per_partition
53 + layer.params_dtype = params_dtype
54 + weight_loader = extra_weight_attrs.get("weight_loader")
55 +
56 + weight = ModelWeightParameter(
57 + data=torch.empty(
58 + output_size_per_partition,
59 + input_size_per_partition // 2,
60 + dtype=torch.uint8,
61 + ),
62 + input_dim=1,
63 + output_dim=0,
64 + weight_loader=weight_loader,
65 + )
66 + layer.register_parameter("weight_packed", weight)
67 +
68 + weight_scale = GroupQuantScaleParameter(
69 + data=torch.empty(
70 + output_size_per_partition,
71 + input_size_per_partition // self.group_size,
72 + dtype=torch.uint8,
73 + ),
74 + input_dim=1,
75 + output_dim=0,
76 + weight_loader=weight_loader,
77 + )
78 + layer.register_parameter("weight_scale", weight_scale)
79 +
80 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
81 + layer.weight = Parameter(layer.weight_packed.data, requires_grad=False)
82 + del layer.weight_packed
83 + self.kernel.process_weights_after_loading(layer)
84 +
85 + def apply_weights(
86 + self,
87 + layer: torch.nn.Module,
88 + x: torch.Tensor,
89 + bias: torch.Tensor | None = None,
90 + ) -> torch.Tensor:
91 + return self.kernel.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py new
+242
@@ -0,0 +1,242 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7 +from vllm.logger import init_logger
8 +from vllm.model_executor.layers.fused_moe import (
9 + FusedMoeWeightScaleSupported,
10 + RoutedExperts,
11 + SharedExperts,
12 +)
13 +from vllm.model_executor.layers.fused_moe.config import (
14 + FusedMoEQuantConfig,
15 + mxfp4_moe_quant_config,
16 +)
17 +from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
18 + CutlassExpertsMxfp4,
19 +)
20 +from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
21 + MarlinExperts,
22 +)
23 +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
24 + FusedMoEMethodBase,
25 +)
26 +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
27 + Mxfp4MoeBackend,
28 + make_mxfp4_moe_kernel,
29 + make_mxfp4_moe_quant_config,
30 + select_mxfp4_moe_backend,
31 +)
32 +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
33 + prepare_moe_fp4_layer_for_marlin,
34 +)
35 +from vllm.model_executor.utils import set_weight_attrs
36 +from vllm.platforms import current_platform
37 +
38 +logger = init_logger(__name__)
39 +
40 +
41 +class INCMxfp4MoEMethod(FusedMoEMethodBase):
42 + """W4A4 MXFP4 group MoE for AutoRound ``auto_round:llm_compressor`` exports.
43 +
44 + Registers the packed MXFP4 layout (uint8 ``weight_packed`` + uint8 E8M0
45 + ``weight_scale``, ``group_size=32``) and dispatches the fused MoE to the
46 + best backend for the current device: CUTLASS (true W4A4 on supported
47 + GPUs), the native XPU kernel, or Marlin weight-only as a fallback. The
48 + per-expert ``gate_proj`` / ``up_proj`` / ``down_proj`` tensors are folded
49 + into the stacked ``w13`` / ``w2`` parameters by ``make_expert_params_mapping``.
50 + """
51 +
52 + def __init__(self, moe) -> None:
53 + super().__init__(moe)
54 + self.group_size = 32
55 + # Backend selection must stay consistent with the weight preparation in
56 + # process_weights_after_loading / get_fused_moe_quant_config, which only
57 + # implement three layouts: CUTLASS swizzle (true W4A4), the native XPU
58 + # kernel (packed passthrough), and Marlin weight-only. XPU dispatch is
59 + # deferred to the shared oracle; every other non-CUTLASS device falls
60 + # back to Marlin (mirroring CompressedTensorsW4A4Mxfp4MoEMethod).
61 + self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
62 + self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
63 + self.experts_cls: type[mk.FusedMoEExperts] | None = None
64 + if self.use_cutlass_mxfp4:
65 + self.experts_cls = CutlassExpertsMxfp4
66 + logger.info_once("Using CutlassExpertsMxfp4 for AutoRound MXFP4 MoE")
67 + elif current_platform.is_xpu():
68 + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)
69 + else:
70 + self.experts_cls = MarlinExperts
71 + logger.info_once(
72 + "Using MarlinExperts (weight-only FP4) for AutoRound MXFP4 MoE"
73 + )
74 +
75 + def create_weights(
76 + self,
77 + layer: torch.nn.Module,
78 + num_experts: int,
79 + hidden_size: int,
80 + intermediate_size_per_partition: int,
81 + params_dtype: torch.dtype,
82 + **extra_weight_attrs,
83 + ) -> None:
84 + layer.num_experts = num_experts
85 + layer.params_dtype = params_dtype
86 +
87 + # gate + up fused on the output dim; two FP4 packed per input byte.
88 + w13_weight = torch.nn.Parameter(
89 + torch.empty(
90 + num_experts,
91 + 2 * intermediate_size_per_partition,
92 + hidden_size // 2,
93 + dtype=torch.uint8,
94 + ),
95 + requires_grad=False,
96 + )
97 + layer.register_parameter("w13_weight_packed", w13_weight)
98 + set_weight_attrs(w13_weight, extra_weight_attrs)
99 +
100 + w2_weight = torch.nn.Parameter(
101 + torch.empty(
102 + num_experts,
103 + hidden_size,
104 + intermediate_size_per_partition // 2,
105 + dtype=torch.uint8,
106 + ),
107 + requires_grad=False,
108 + )
109 + layer.register_parameter("w2_weight_packed", w2_weight)
110 + set_weight_attrs(w2_weight, extra_weight_attrs)
111 +
112 + # Per-group E8M0 block scales (group_size=32), stored as uint8.
113 + w13_weight_scale = torch.nn.Parameter(
114 + torch.empty(
115 + num_experts,
116 + 2 * intermediate_size_per_partition,
117 + hidden_size // self.group_size,
118 + dtype=torch.uint8,
119 + ),
120 + requires_grad=False,
121 + )
122 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
123 + extra_weight_attrs.update(
124 + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
125 + )
126 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
127 +
128 + w2_weight_scale = torch.nn.Parameter(
129 + torch.empty(
130 + num_experts,
131 + hidden_size,
132 + intermediate_size_per_partition // self.group_size,
133 + dtype=torch.uint8,
134 + ),
135 + requires_grad=False,
136 + )
137 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
138 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
139 +
140 + def get_fused_moe_quant_config(
141 + self, layer: torch.nn.Module
142 + ) -> FusedMoEQuantConfig | None:
143 + if self.use_cutlass_mxfp4:
144 + # W4A4: both weights and activations quantized to MXFP4.
145 + return mxfp4_moe_quant_config(
146 + w1_scale=layer.w13_weight_scale,
147 + w2_scale=layer.w2_weight_scale,
148 + )
149 + # Weight-only (Marlin) or native XPU kernel.
150 + return make_mxfp4_moe_quant_config(
151 + mxfp4_backend=self.mxfp4_backend,
152 + w1_scale=layer.w13_weight_scale,
153 + w2_scale=layer.w2_weight_scale,
154 + )
155 +
156 + def process_weights_after_loading(self, layer: RoutedExperts) -> None:
157 + layer.w13_weight = torch.nn.Parameter(
158 + layer.w13_weight_packed.data, requires_grad=False
159 + )
160 + delattr(layer, "w13_weight_packed")
161 + layer.w2_weight = torch.nn.Parameter(
162 + layer.w2_weight_packed.data, requires_grad=False
163 + )
164 + delattr(layer, "w2_weight_packed")
165 +
166 + if self.use_cutlass_mxfp4:
167 + # Swizzle weight scales from flat checkpoint layout [E, N, K//32]
168 + # to the CUTLASS tiled layout.
169 + from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
170 + swizzle_mxfp4_scales,
171 + )
172 +
173 + E = layer.w13_weight_scale.shape[0]
174 + w13_N = layer.w13_weight_scale.shape[1]
175 + w13_scale_K = layer.w13_weight_scale.shape[2]
176 + w13_K = w13_scale_K * 32
177 +
178 + w2_M = layer.w2_weight_scale.shape[1]
179 + w2_scale_N = layer.w2_weight_scale.shape[2]
180 + w2_N = w2_scale_N * 32
181 +
182 + swizzled_w13 = []
183 + swizzled_w2 = []
184 + for e_idx in range(E):
185 + s13 = layer.w13_weight_scale[e_idx]
186 + sw13 = swizzle_mxfp4_scales(s13, w13_N, w13_K)
187 + swizzled_w13.append(sw13.reshape(w13_N, w13_scale_K))
188 + s2 = layer.w2_weight_scale[e_idx]
189 + sw2 = swizzle_mxfp4_scales(s2, w2_M, w2_N)
190 + swizzled_w2.append(sw2.reshape(w2_M, w2_scale_N))
191 + layer.w13_weight_scale = torch.nn.Parameter(
192 + torch.stack(swizzled_w13), requires_grad=False
193 + )
194 + layer.w2_weight_scale = torch.nn.Parameter(
195 + torch.stack(swizzled_w2), requires_grad=False
196 + )
197 + elif current_platform.is_xpu():
198 + # The XPU fused-MoE kernel consumes the packed layout directly; no
199 + # swizzle / repack / transpose is required.
200 + pass
201 + else:
202 + logger.warning_once(
203 + "This device lacks native FP4 compute; using weight-only FP4 "
204 + "via the Marlin kernel, which may reduce performance for "
205 + "compute-heavy workloads."
206 + )
207 + prepare_moe_fp4_layer_for_marlin(layer)
208 +
209 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
210 + if self.moe_quant_config is not None:
211 + assert self.experts_cls is not None
212 + self.moe_kernel = make_mxfp4_moe_kernel(
213 + moe_quant_config=self.moe_quant_config,
214 + moe_config=self.moe,
215 + experts_cls=self.experts_cls,
216 + mxfp4_backend=self.mxfp4_backend,
217 + routing_tables=layer._expert_routing_tables(),
218 + )
219 +
220 + def apply(
221 + self,
222 + layer: RoutedExperts,
223 + x: torch.Tensor,
224 + topk_weights: torch.Tensor,
225 + topk_ids: torch.Tensor,
226 + shared_experts: SharedExperts | None,
227 + shared_experts_input: torch.Tensor | None,
228 + ) -> torch.Tensor:
229 + assert self.moe_kernel is not None
230 + return self.moe_kernel.apply(
231 + x,
232 + layer.w13_weight,
233 + layer.w2_weight,
234 + topk_weights,
235 + topk_ids,
236 + activation=layer.activation,
237 + global_num_experts=layer.global_num_experts,
238 + expert_map=layer.expert_map,
239 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
240 + shared_experts=shared_experts,
241 + shared_experts_input=shared_experts_input,
242 + )
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_scheme.py new
+54
@@ -0,0 +1,54 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import TYPE_CHECKING
5 +
6 +from vllm.logger import init_logger
7 +
8 +from ..inc_linear import INCLinearMethod
9 +from .inc_scheme import INCScheme
10 +
11 +if TYPE_CHECKING:
12 + import torch
13 +
14 + from ..config_parser import INCLayerConfig
15 + from ..inc import INCConfig
16 +
17 +logger = init_logger(__name__)
18 +
19 +
20 +class INCMxfp4Scheme(INCScheme):
21 + """MXFP4 (W4A4) scheme for AutoRound checkpoints.
22 +
23 + Dispatches to :class:`INCMxfp4LinearMethod` for linear layers and
24 + :class:`INCMxfp4MoEMethod` for fused MoE layers; see those classes for the
25 + per-module weight layout and kernel-selection details.
26 + """
27 +
28 + @staticmethod
29 + def can_handle(layer_config: "INCLayerConfig") -> bool:
30 + return layer_config.is_mxfp4
31 +
32 + def get_linear_method(
33 + self,
34 + config: "INCConfig",
35 + layer: "torch.nn.Module",
36 + prefix: str,
37 + layer_config: "INCLayerConfig",
38 + ):
39 + del config, layer, prefix
40 + from .inc_mxfp4_linear import INCMxfp4LinearMethod
41 +
42 + return INCLinearMethod(INCMxfp4LinearMethod(layer_config))
43 +
44 + def get_moe_method(
45 + self,
46 + config: "INCConfig",
47 + layer: "torch.nn.Module",
48 + prefix: str,
49 + layer_config: "INCLayerConfig",
50 + ):
51 + del config, prefix, layer_config
52 + from .inc_mxfp4_moe import INCMxfp4MoEMethod
53 +
54 + return INCMxfp4MoEMethod(layer.moe_config)
vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py
+13 -27
@@ -7,7 +7,6 @@ from vllm.logger import init_logger
7 from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
8 from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
9 from vllm.platforms import current_platform
10 -from vllm.scalar_type import scalar_types
10
11 from ..inc_linear import INCLinearMethod
12 from .inc_scheme import INCScheme
@@ -97,8 +96,8 @@ class INCWna16Scheme(INCScheme):
96 layer_config: "INCLayerConfig",
97 ):
98 del config, prefix
100 - # XPU and CPU do not support MoE quantization yet
101 - if current_platform.is_xpu() or current_platform.is_cpu():
99 + # CPU does not support quantized MoE yet.
100 + if current_platform.is_cpu():
101 from vllm.model_executor.layers.fused_moe import (
102 UnquantizedFusedMoEMethod,
103 )
@@ -120,21 +119,17 @@ def _resolve_gptq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"):
119 MoeWNA16Method,
120 )
121 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
123 - check_marlin_supported,
122 check_moe_marlin_supports_layer,
123 )
124
127 - gptq_type_map = {
128 - (4, True): scalar_types.uint4b8,
129 - (8, True): scalar_types.uint8b128,
130 - }
131 - use_marlin = (layer_config.bits, layer_config.sym) in gptq_type_map
132 - if use_marlin:
133 - use_marlin = check_marlin_supported(
134 - gptq_type_map[(layer_config.bits, layer_config.sym)],
135 - layer_config.group_size,
136 - has_zp=not layer_config.sym,
137 - ) and check_moe_marlin_supports_layer(layer, layer_config.group_size)
125 + # AutoGPTQMoEMethod selects its fused-MoE backend through the WNA16 oracle
126 + # (Marlin on CUDA, XPUExpertsWNA16 on XPU). Gate only on the layer-shape
127 + # check like compressed-tensors does; the capability-based
128 + # check_marlin_supported is skipped so the XPU path is reachable.
129 + use_marlin = (layer_config.bits, layer_config.sym) in {
130 + (4, True),
131 + (8, True),
132 + } and check_moe_marlin_supports_layer(layer, layer_config.group_size)
133
134 if use_marlin:
135 return AutoGPTQMoEMethod(
@@ -169,21 +164,12 @@ def _resolve_awq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"):
164 MoeWNA16Method,
165 )
166 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
172 - check_marlin_supported,
167 check_moe_marlin_supports_layer,
168 )
169
176 - awq_type_map = {
177 - 4: scalar_types.uint4,
178 - 8: scalar_types.uint8,
179 - }
180 - use_marlin = layer_config.bits in awq_type_map
181 - if use_marlin:
182 - use_marlin = check_marlin_supported(
183 - awq_type_map[layer_config.bits],
184 - layer_config.group_size,
185 - not layer_config.sym,
186 - ) and check_moe_marlin_supports_layer(layer, layer_config.group_size)
170 + use_marlin = layer_config.bits in (4, 8) and check_moe_marlin_supports_layer(
171 + layer, layer_config.group_size
172 + )
173
174 if use_marlin:
175 return AutoAWQMoEMethod(