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

[ROCm][Quark][7/N] Use MXFP4 linear kernel abstraction for `emulation` backend (#48949)

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

fxmarty-amd committed Jul 31, 2026 at 19:07 UTC e67a2e0a56b7ba535bae8bfd98c1c4e4fa742e63
17 files changed +825 -91
csrc/core/scalar_type.hpp
+3
@@ -327,6 +327,8 @@ static inline constexpr auto kFE2M1f =
327 ScalarType::float_(2, 1, true, ScalarType::NAN_NONE);
328 static inline constexpr auto kFE3M2f =
329 ScalarType::float_(3, 2, true, ScalarType::NAN_NONE);
330 +static inline constexpr auto kFE2M3f =
331 + ScalarType::float_(2, 3, true, ScalarType::NAN_NONE);
332 static inline constexpr auto kFE4M3fn =
333 ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN);
334 static inline constexpr auto kFE8M0fnu =
@@ -346,6 +348,7 @@ static inline constexpr auto kUint8b128 = kU8B128;
348
349 static inline constexpr auto kFloat4_e2m1f = kFE2M1f;
350 static inline constexpr auto kFloat6_e3m2f = kFE3M2f;
351 +static inline constexpr auto kFloat6_e2m3f = kFE2M3f;
352 static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn;
353 static inline constexpr auto kFloat8_e5m2 = kFE5M2;
354 static inline constexpr auto kFloat16_e8m7 = kFE8M7;
tests/kernels/quantization/test_mxfp4_kernel_selection.py
+118 -22
@@ -12,15 +12,40 @@ import torch
12
13 from vllm.model_executor.kernels.linear import (
14 AiterMxfp4LinearKernel,
15 + EmulationMxfp4LinearKernel,
16 + FlashInferMxFp4LinearKernel,
17 + HummingMxFp4LinearKernel,
18 + MarlinMxFp4LinearKernel,
19 MxFp4LinearKernel,
20 MxFp4LinearLayerConfig,
21 + XPUMxFp4LinearKernel,
22 init_mxfp4_linear_kernel,
23 register_linear_kernel,
24 )
25 +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
26 + quant_dequant_mxfp4,
27 +)
28 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
29 + kMxfp4Dynamic,
30 + kMxfp6E2M3Dynamic,
31 + kMxfp6E3M2Dynamic,
32 +)
33 from vllm.platforms import PlatformEnum
34
35 pytestmark = pytest.mark.cpu_test
36
37 +# Kernels that quantize activations themselves (true W4A4): they require an
38 +# explicit MXFP4-dynamic activation key.
39 +_TRUE_W4A4_KERNELS = [
40 + FlashInferMxFp4LinearKernel,
41 + XPUMxFp4LinearKernel,
42 + AiterMxfp4LinearKernel,
43 +]
44 +
45 +# Weight-only (A16) kernels: they never quantize activations. They still accept
46 +# MXFP4 activation keys as an intentional compatibility fallback.
47 +_WEIGHT_ONLY_KERNELS = [MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel]
48 +
49
50 def test_can_implement_is_abstract():
51 """Test that can_implement()/is_supported() are properly defined."""
@@ -28,6 +53,97 @@ def test_can_implement_is_abstract():
53 assert hasattr(MxFp4LinearKernel, "is_supported")
54
55
56 +@pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS)
57 +def test_true_w4a4_kernels_accept_dynamic_mxfp4_activation(kernel_cls):
58 + config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic)
59 + can_implement, reason = kernel_cls.can_implement(config)
60 + assert can_implement, reason
61 +
62 +
63 +@pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS)
64 +def test_true_w4a4_kernels_reject_unset_activation(kernel_cls):
65 + """None means weight-only/unquantized activations, not dynamic MXFP4."""
66 + config = MxFp4LinearLayerConfig()
67 + can_implement, reason = kernel_cls.can_implement(config)
68 + assert not can_implement
69 + assert reason
70 +
71 +
72 +@pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS)
73 +def test_true_w4a4_kernels_reject_explicit_non_mxfp4_activation(kernel_cls):
74 + """FlashInfer/XPU/Aiter quantize activations to MXFP4 internally, so an
75 + explicit request for a different activation format must be rejected."""
76 + config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp6E3M2Dynamic)
77 + can_implement, reason = kernel_cls.can_implement(config)
78 + assert not can_implement
79 + assert reason
80 +
81 +
82 +@pytest.mark.parametrize("kernel_cls", _WEIGHT_ONLY_KERNELS)
83 +@pytest.mark.parametrize("activation_quant_key", [None, kMxfp4Dynamic])
84 +def test_weight_only_kernels_accept_unquantized_or_mxfp4_activation(
85 + kernel_cls, activation_quant_key
86 +):
87 + """Marlin/Humming never quantize activations, so an unset activation key,
88 + or one that already describes MXFP4-shaped data, is tolerated. When an
89 + activation key is explicitly set, a warning must be logged noting that it
90 + is ignored, since these kernels are weight-only (A16)."""
91 + config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key)
92 + with patch(f"{kernel_cls.__module__}.logger.warning_once") as warning_once:
93 + can_implement, reason = kernel_cls.can_implement(config)
94 + assert can_implement, reason
95 +
96 + if activation_quant_key is None:
97 + warning_once.assert_not_called()
98 + else:
99 + warning_once.assert_called_once()
100 + message = warning_once.call_args.args[0]
101 + assert "the requested activation quantization" in message
102 + assert "is ignored" in message
103 +
104 +
105 +@pytest.mark.parametrize("kernel_cls", _WEIGHT_ONLY_KERNELS)
106 +def test_weight_only_kernels_reject_non_mxfp4_activation(kernel_cls):
107 + config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp6E3M2Dynamic)
108 + can_implement, reason = kernel_cls.can_implement(config)
109 + assert not can_implement
110 + assert reason
111 +
112 +
113 +@pytest.mark.parametrize(
114 + "activation_quant_key",
115 + [None, kMxfp4Dynamic, kMxfp6E3M2Dynamic, kMxfp6E2M3Dynamic],
116 +)
117 +def test_emulation_kernel_accepts_any_config(activation_quant_key):
118 + """EmulationMxfp4LinearKernel is the universal fallback: it must accept
119 + every supported activation format."""
120 + config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key)
121 + with patch(
122 + "vllm.model_executor.kernels.linear._get_linear_backend",
123 + return_value="emulation",
124 + ):
125 + can_implement, reason = EmulationMxfp4LinearKernel.can_implement(config)
126 + assert can_implement, reason
127 +
128 +
129 +def test_emulation_kernel_derives_quant_dequant_func_from_config():
130 + """quant_dequant_func must be derived purely from the config's activation
131 + QuantKey, not set externally."""
132 + with patch(
133 + "vllm.model_executor.kernels.linear._get_linear_backend",
134 + return_value="emulation",
135 + ):
136 + weight_only_config = MxFp4LinearLayerConfig()
137 + kernel = EmulationMxfp4LinearKernel(weight_only_config)
138 + x = torch.randn(4)
139 + # identity for weight-only
140 + assert torch.equal(kernel.quant_dequant_func(x), x)
141 +
142 + w4a4_config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic)
143 + kernel = EmulationMxfp4LinearKernel(w4a4_config)
144 + assert kernel.quant_dequant_func is quant_dequant_mxfp4
145 +
146 +
147 def test_aiter_kernel_is_supported_requires_native_mx_support():
148 """AiterMxfp4LinearKernel must not be selected on platforms without
149 native MX compute, even if AITER itself is importable."""
@@ -70,10 +186,10 @@ def test_init_mxfp4_linear_kernel_dispatches_to_registered_kernel(platform_mock)
186 platform_mock._enum = PlatformEnum.OOT
187 register_linear_kernel(OOTMxFp4LinearKernel, PlatformEnum.OOT, "mxfp4")
188
73 - kernel = init_mxfp4_linear_kernel()
189 + kernel = init_mxfp4_linear_kernel(activation_quant_key=kMxfp4Dynamic)
190
191 assert isinstance(kernel, OOTMxFp4LinearKernel)
76 - assert kernel.config == MxFp4LinearLayerConfig()
192 + assert kernel.config == MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic)
193
194
195 class UnsupportedMxFp4LinearKernel(MxFp4LinearKernel):
@@ -108,23 +224,3 @@ def test_init_mxfp4_linear_kernel_raises_when_no_kernel_matches(platform_mock):
224
225 with pytest.raises(ValueError, match="Failed to find a kernel"):
226 init_mxfp4_linear_kernel()
111 -
112 -
113 -@patch("vllm.model_executor.kernels.linear.mxfp4.aiter.is_aiter_found_and_supported")
114 -@patch("vllm.model_executor.kernels.linear.mxfp4.aiter.current_platform")
115 -@patch("vllm.model_executor.kernels.linear.current_platform")
116 -def test_init_mxfp4_linear_kernel_raises_on_rocm_without_aiter(
117 - linear_platform_mock, aiter_platform_mock, is_aiter_found_and_supported_mock
118 -):
119 - """On ROCm, the only registered MXFP4 linear kernel is AITER-based.
120 - If AITER is not found/supported, no kernel should be selected."""
121 - linear_platform_mock._enum = PlatformEnum.ROCM
122 - aiter_platform_mock.supports_mx.return_value = True
123 - is_aiter_found_and_supported_mock.return_value = False
124 -
125 - with pytest.raises(
126 - ValueError,
127 - match="(?s)Failed to find a kernel.*"
128 - "AITER not found or not supported on the current platform",
129 - ):
130 - init_mxfp4_linear_kernel()
tests/kernels/quantization/test_mxfp6_kernel_selection.py new
+150
@@ -0,0 +1,150 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for MXFP6 linear kernel selection logic (CPU-only)
4 +
5 +Run `pytest tests/kernels/quantization/test_mxfp6_kernel_selection.py`.
6 +"""
7 +
8 +from unittest.mock import patch
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm.model_executor.kernels.linear import (
14 + EmulationMxfp6LinearKernel,
15 + MxFp6LinearKernel,
16 + MxFp6LinearLayerConfig,
17 + init_mxfp6_linear_kernel,
18 + register_linear_kernel,
19 +)
20 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
21 + kMxfp4Dynamic,
22 + kMxfp4Static,
23 + kMxfp6E2M3Dynamic,
24 + kMxfp6E2M3Static,
25 + kMxfp6E3M2Dynamic,
26 + kMxfp6E3M2Static,
27 +)
28 +from vllm.platforms import PlatformEnum
29 +
30 +pytestmark = pytest.mark.cpu_test
31 +
32 +# The only implementation available at the moment is software emulation.
33 +_WEIGHT_QUANT_KEYS = [kMxfp6E3M2Static, kMxfp6E2M3Static]
34 +
35 +
36 +def test_can_implement_is_abstract():
37 + """Test that can_implement()/is_supported() are properly defined."""
38 + assert hasattr(MxFp6LinearKernel, "can_implement")
39 + assert hasattr(MxFp6LinearKernel, "is_supported")
40 +
41 +
42 +def test_emulation_kernel_rejects_non_mxfp6_weights():
43 + """EmulationMxfp6LinearKernel must not implement a non-MXFP6 weight
44 + format."""
45 + config = MxFp6LinearLayerConfig(weight_quant_key=kMxfp4Static)
46 + can_implement, reason = EmulationMxfp6LinearKernel.can_implement(config)
47 + assert not can_implement
48 + assert reason
49 +
50 +
51 +@pytest.mark.parametrize("weight_quant_key", _WEIGHT_QUANT_KEYS)
52 +@pytest.mark.parametrize(
53 + "activation_quant_key",
54 + [None, kMxfp4Dynamic, kMxfp6E3M2Dynamic, kMxfp6E2M3Dynamic],
55 +)
56 +def test_emulation_kernel_accepts_any_supported_config(
57 + weight_quant_key, activation_quant_key
58 +):
59 + """EmulationMxfp6LinearKernel is the only backend today: it must accept
60 + every supported weight/activation format combination."""
61 + config = MxFp6LinearLayerConfig(
62 + weight_quant_key=weight_quant_key, activation_quant_key=activation_quant_key
63 + )
64 + can_implement, reason = EmulationMxfp6LinearKernel.can_implement(config)
65 + assert can_implement, reason
66 +
67 +
68 +@pytest.mark.parametrize("weight_quant_key", _WEIGHT_QUANT_KEYS)
69 +def test_emulation_kernel_rejects_non_mxfp4_or_mxfp6_activation(weight_quant_key):
70 + config = MxFp6LinearLayerConfig(
71 + weight_quant_key=weight_quant_key, activation_quant_key=kMxfp4Static
72 + )
73 + can_implement, reason = EmulationMxfp6LinearKernel.can_implement(config)
74 + assert not can_implement
75 + assert reason
76 +
77 +
78 +class OOTMxFp6LinearKernel(MxFp6LinearKernel):
79 + @classmethod
80 + def is_supported(
81 + cls, compute_capability: int | None = None
82 + ) -> tuple[bool, str | None]:
83 + return True, None
84 +
85 + @classmethod
86 + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]:
87 + return True, None
88 +
89 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
90 + pass
91 +
92 + def apply_weights(
93 + self,
94 + layer: torch.nn.Module,
95 + x: torch.Tensor,
96 + bias: torch.Tensor | None = None,
97 + ) -> torch.Tensor:
98 + pass
99 +
100 +
101 +@patch("vllm.model_executor.kernels.linear.current_platform")
102 +def test_init_mxfp6_linear_kernel_dispatches_to_registered_kernel(platform_mock):
103 + """init_mxfp6_linear_kernel should select a registered kernel that
104 + reports itself as supported/able to implement the given config, and
105 + construct it with that exact config."""
106 + platform_mock._enum = PlatformEnum.OOT
107 + register_linear_kernel(OOTMxFp6LinearKernel, PlatformEnum.OOT, "mxfp6")
108 +
109 + kernel = init_mxfp6_linear_kernel(
110 + weight_quant_key=kMxfp6E3M2Static, activation_quant_key=kMxfp6E3M2Dynamic
111 + )
112 +
113 + assert isinstance(kernel, OOTMxFp6LinearKernel)
114 + assert kernel.config == MxFp6LinearLayerConfig(
115 + weight_quant_key=kMxfp6E3M2Static, activation_quant_key=kMxfp6E3M2Dynamic
116 + )
117 +
118 +
119 +class UnsupportedMxFp6LinearKernel(MxFp6LinearKernel):
120 + @classmethod
121 + def is_supported(
122 + cls, compute_capability: int | None = None
123 + ) -> tuple[bool, str | None]:
124 + return False, "never supported"
125 +
126 + @classmethod
127 + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]:
128 + return True, None
129 +
130 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
131 + pass
132 +
133 + def apply_weights(
134 + self,
135 + layer: torch.nn.Module,
136 + x: torch.Tensor,
137 + bias: torch.Tensor | None = None,
138 + ) -> torch.Tensor:
139 + pass
140 +
141 +
142 +@patch("vllm.model_executor.kernels.linear.current_platform")
143 +def test_init_mxfp6_linear_kernel_raises_when_no_kernel_matches(platform_mock):
144 + platform_mock._enum = PlatformEnum.UNSPECIFIED
145 + register_linear_kernel(
146 + UnsupportedMxFp6LinearKernel, PlatformEnum.UNSPECIFIED, "mxfp6"
147 + )
148 +
149 + with pytest.raises(ValueError, match="Failed to find a kernel"):
150 + init_mxfp6_linear_kernel(weight_quant_key=kMxfp6E3M2Static)
vllm/model_executor/kernels/linear/__init__.py
+98 -2
@@ -77,6 +77,9 @@ from vllm.model_executor.kernels.linear.mxfp4 import (
77 from vllm.model_executor.kernels.linear.mxfp4.aiter import (
78 AiterMxfp4LinearKernel,
79 )
80 +from vllm.model_executor.kernels.linear.mxfp4.emulation import (
81 + EmulationMxfp4LinearKernel,
82 +)
83 from vllm.model_executor.kernels.linear.mxfp4.flashinfer import (
84 FlashInferMxFp4LinearKernel,
85 )
@@ -89,6 +92,13 @@ from vllm.model_executor.kernels.linear.mxfp4.marlin import (
92 from vllm.model_executor.kernels.linear.mxfp4.xpu import (
93 XPUMxFp4LinearKernel,
94 )
95 +from vllm.model_executor.kernels.linear.mxfp6 import (
96 + MxFp6LinearKernel,
97 + MxFp6LinearLayerConfig,
98 +)
99 +from vllm.model_executor.kernels.linear.mxfp6.emulation import (
100 + EmulationMxfp6LinearKernel,
101 +)
102 from vllm.model_executor.kernels.linear.mxfp8 import (
103 Mxfp8LinearKernel,
104 Mxfp8LinearLayerConfig,
@@ -294,6 +304,8 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
304 "emulation": {
305 EmulationMxfp8LinearKernel,
306 EmulationNvFp4LinearKernel,
307 + EmulationMxfp6LinearKernel,
308 + EmulationMxfp4LinearKernel,
309 },
310 "xpu": {
311 XPUW8A8FP8LinearKernel,
@@ -467,14 +479,25 @@ _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
479 ],
480 }
481
482 +_POSSIBLE_MXFP6_KERNELS: dict[PlatformEnum, list[type[MxFp6LinearKernel]]] = {
483 + PlatformEnum.CUDA: [
484 + EmulationMxfp6LinearKernel,
485 + ],
486 + PlatformEnum.ROCM: [
487 + EmulationMxfp6LinearKernel,
488 + ],
489 +}
490 +
491 _POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = {
492 PlatformEnum.CUDA: [
493 FlashInferMxFp4LinearKernel,
494 MarlinMxFp4LinearKernel,
495 HummingMxFp4LinearKernel,
496 + EmulationMxfp4LinearKernel,
497 ],
498 PlatformEnum.ROCM: [
499 AiterMxfp4LinearKernel,
500 + EmulationMxfp4LinearKernel,
501 ],
502 PlatformEnum.XPU: [
503 XPUMxFp4LinearKernel,
@@ -812,9 +835,15 @@ def init_mxfp8_linear_kernel() -> Mxfp8LinearKernel:
835 )
836
837
815 -def init_mxfp4_linear_kernel() -> MxFp4LinearKernel:
838 +def init_mxfp4_linear_kernel(
839 + activation_quant_key: QuantKey | None = None,
840 +) -> MxFp4LinearKernel:
841 """Select and instantiate the best MXFP4 linear kernel for the
842 current platform."""
843 + config = MxFp4LinearLayerConfig(
844 + activation_quant_key=activation_quant_key,
845 + )
846 +
847 linear_backend = _get_linear_backend()
848
849 platform = current_platform._enum
@@ -843,8 +872,13 @@ def init_mxfp4_linear_kernel() -> MxFp4LinearKernel:
872 failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
873 continue
874
875 + can_implement, reason = kernel_cls.can_implement(config)
876 + if not can_implement:
877 + failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
878 + continue
879 +
880 logger.info_once("Using %s for MXFP4 GEMM", kernel_cls.__name__)
847 - return kernel_cls(MxFp4LinearLayerConfig())
881 + return kernel_cls(config)
882
883 raise ValueError(
884 "Failed to find a kernel that can implement the "
@@ -852,6 +886,59 @@ def init_mxfp4_linear_kernel() -> MxFp4LinearKernel:
886 )
887
888
889 +def init_mxfp6_linear_kernel(
890 + weight_quant_key: QuantKey,
891 + activation_quant_key: QuantKey | None = None,
892 +) -> MxFp6LinearKernel:
893 + """Select and instantiate the best MXFP6 linear kernel for the
894 + current platform."""
895 + config = MxFp6LinearLayerConfig(
896 + weight_quant_key=weight_quant_key,
897 + activation_quant_key=activation_quant_key,
898 + )
899 +
900 + linear_backend = _get_linear_backend()
901 +
902 + platform = current_platform._enum
903 + possible = list(_POSSIBLE_MXFP6_KERNELS.get(platform, []))
904 +
905 + # Apply --linear-backend filtering when set.
906 + if linear_backend != "auto":
907 + filtered = _filter_kernels_by_backend(linear_backend, possible)
908 + if not filtered:
909 + raise ValueError(
910 + f"--linear-backend={linear_backend} was requested but no "
911 + f"'{linear_backend}' kernel exists for MXFP6 layers."
912 + )
913 + possible = filtered
914 +
915 + failure_reasons = []
916 + for kernel_cls in possible:
917 + if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS:
918 + failure_reasons.append(
919 + f" {kernel_cls.__name__} disabled by environment variable"
920 + )
921 + continue
922 +
923 + is_supported, reason = kernel_cls.is_supported()
924 + if not is_supported:
925 + failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
926 + continue
927 +
928 + can_implement, reason = kernel_cls.can_implement(config)
929 + if not can_implement:
930 + failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
931 + continue
932 +
933 + logger.info_once("Using %s for MXFP6 GEMM", kernel_cls.__name__)
934 + return kernel_cls(config)
935 +
936 + raise ValueError(
937 + "Failed to find a kernel that can implement the "
938 + "MXFP6 linear layer. Reasons: \n" + "\n".join(failure_reasons)
939 + )
940 +
941 +
942 def init_wfp8_a16_linear_kernel(
943 weight_quant_key: QuantKey,
944 activation_quant_key: QuantKey,
@@ -1037,6 +1124,10 @@ def register_linear_kernel(
1124 if platform not in _POSSIBLE_MXFP4_KERNELS:
1125 _POSSIBLE_MXFP4_KERNELS[platform] = []
1126 _POSSIBLE_MXFP4_KERNELS[platform].append(kernel_class)
1127 + elif kernel_type == "mxfp6":
1128 + if platform not in _POSSIBLE_MXFP6_KERNELS:
1129 + _POSSIBLE_MXFP6_KERNELS[platform] = []
1130 + _POSSIBLE_MXFP6_KERNELS[platform].append(kernel_class)
1131 else:
1132 raise ValueError(f"Unrecognized kernel type: {kernel_type}")
1133
@@ -1091,7 +1182,12 @@ __all__ = [
1182 "init_mxfp4_linear_kernel",
1183 "MxFp4LinearKernel",
1184 "MxFp4LinearLayerConfig",
1185 + "MxFp6LinearKernel",
1186 + "MxFp6LinearLayerConfig",
1187 + "init_mxfp6_linear_kernel",
1188 + "EmulationMxfp6LinearKernel",
1189 "AiterMxfp4LinearKernel",
1190 + "EmulationMxfp4LinearKernel",
1191 "FlashInferMxFp4LinearKernel",
1192 "MarlinMxFp4LinearKernel",
1193 "FlashInferCutedslMxfp8LinearKernel",
vllm/model_executor/kernels/linear/mxfp4/aiter.py
+30 -1
@@ -4,11 +4,18 @@
4 import torch
5 from torch.nn.parameter import Parameter
6
7 +import vllm.envs as envs
8 from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
9 +from vllm.logger import init_logger
10 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
11 + kMxfp4Dynamic,
12 +)
13 from vllm.platforms import current_platform
14
15 from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
16
17 +logger = init_logger(__name__)
18 +
19 # NOTE: Do not import aiter at module scope. Importing aiter eagerly initializes HIP
20 # which can force the engine core to spawn instead of fork.
21 # is_aiter_found_and_supported() checks platform + arch + library availability via
@@ -131,12 +138,34 @@ class AiterMxfp4LinearKernel(MxFp4LinearKernel):
138 ) -> tuple[bool, str | None]:
139 if not current_platform.supports_mx():
140 return False, "current platform does not support native MXFP4 computation"
141 +
142 + from vllm._aiter_ops import is_aiter_found_and_supported
143 + from vllm.model_executor.kernels.linear import _get_linear_backend
144 +
145 + linear_backend = _get_linear_backend()
146 +
147 + if (
148 + current_platform.is_rocm()
149 + and current_platform.supports_mx()
150 + and "AiterMxfp4LinearKernel" not in envs.VLLM_DISABLED_KERNELS
151 + and linear_backend == "auto"
152 + and not is_aiter_found_and_supported()
153 + ):
154 + logger.warning_once(
155 + "This platform supports native MXFP4 W4A4 MOE "
156 + "computation via AITER MOE backend, but AITER is not "
157 + "found or not supported. Consider installing AITER: "
158 + "https://github.com/ROCm/aiter."
159 + )
160 +
161 if is_aiter_found_and_supported():
162 return True, None
163 return False, "AITER not found or not supported on the current platform"
164
165 @classmethod
139 - def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
166 + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
167 + if config.activation_quant_key != kMxfp4Dynamic:
168 + return False, "only supports MXFP4 dynamic activation"
169 return True, None
170
171 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
vllm/model_executor/kernels/linear/mxfp4/base.py
+7 -1
@@ -6,6 +6,8 @@ from dataclasses import dataclass
6
7 import torch
8
9 +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
10 +
11
12 @dataclass
13 class MxFp4LinearLayerConfig:
@@ -13,9 +15,13 @@ class MxFp4LinearLayerConfig:
15
16 All MXFP4 layers share the same structure: packed uint8 weights (2 FP4 values per
17 byte) and per-block weight scales (group size 32).
18 +
19 + Attributes:
20 + activation_quant_key: Identifies the activation quantization format,
21 + or `None` when activations must not be quantized.
22 """
23
18 - pass
24 + activation_quant_key: QuantKey | None = None
25
26
27 class MxFp4LinearKernel(ABC):
vllm/model_executor/kernels/linear/mxfp4/emulation.py new
+106
@@ -0,0 +1,106 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Callable
5 +from functools import partial
6 +
7 +import torch
8 +import torch.nn.functional as F
9 +from torch.nn.parameter import Parameter
10 +
11 +from vllm.logger import init_logger
12 +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
13 + dequant_mxfp4,
14 + quant_dequant_mxfp4,
15 +)
16 +from vllm.model_executor.layers.quantization.utils.mxfp6_utils import (
17 + quant_dequant_mxfp6,
18 +)
19 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
20 + QuantKey,
21 + kMxfp4Dynamic,
22 + kMxfp6E2M3Dynamic,
23 + kMxfp6E3M2Dynamic,
24 +)
25 +from vllm.platforms import current_platform
26 +
27 +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
28 +
29 +logger = init_logger(__name__)
30 +
31 +_ACTIVATION_QUANT_DEQUANT_FUNCS: dict[
32 + QuantKey, Callable[[torch.Tensor], torch.Tensor]
33 +] = {
34 + kMxfp4Dynamic: quant_dequant_mxfp4,
35 + kMxfp6E3M2Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e3m2"),
36 + kMxfp6E2M3Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e2m3"),
37 +}
38 +
39 +
40 +class EmulationMxfp4LinearKernel(MxFp4LinearKernel):
41 + """Software emulation fallback for OCP MXFP4/MXFP6 (dequant + F.linear)."""
42 +
43 + def __init__(self, config: MxFp4LinearLayerConfig) -> None:
44 + super().__init__(config)
45 + if config.activation_quant_key is None:
46 + # no input Q/DQ for weight-only
47 + self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = (
48 + lambda x: x
49 + )
50 + else:
51 + self.quant_dequant_func = _ACTIVATION_QUANT_DEQUANT_FUNCS[
52 + config.activation_quant_key
53 + ]
54 +
55 + @classmethod
56 + def is_supported(
57 + cls, compute_capability: int | None = None
58 + ) -> tuple[bool, str | None]:
59 + return True, None
60 +
61 + @classmethod
62 + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
63 + if config.activation_quant_key not in (
64 + None,
65 + kMxfp4Dynamic,
66 + kMxfp6E3M2Dynamic,
67 + kMxfp6E2M3Dynamic,
68 + ):
69 + return False, "only supports MXFP4 or MXFP6 or unquantized activations"
70 +
71 + if (
72 + current_platform.is_rocm()
73 + and current_platform.supports_mx()
74 + and config.activation_quant_key != kMxfp4Dynamic
75 + ):
76 + logger.warning_once(
77 + "The current platform supports native MXFP4/MXFP6 computation, "
78 + f"but kernels for activation_quant_key={config.activation_quant_key} "
79 + f"are not yet integrated in vLLM. Using EmulationMxfp4LinearKernel, "
80 + "with simulated weight dequantization and activation "
81 + "QDQ (quantize and dequantize), with the linear "
82 + "layers computed in high precision."
83 + )
84 +
85 + if not current_platform.supports_mx():
86 + logger.warning_once(
87 + "The current platform does not support native MXFP4 "
88 + "computation. Using EmulationMxfp4LinearKernel, with simulated weight "
89 + "dequantization and activation QDQ (quantize and dequantize), with "
90 + "the linear layers computed in high precision."
91 + )
92 +
93 + return True, None
94 +
95 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
96 + layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
97 +
98 + def apply_weights(
99 + self,
100 + layer: torch.nn.Module,
101 + x: torch.Tensor,
102 + bias: torch.Tensor | None = None,
103 + ) -> torch.Tensor:
104 + dq_w = dequant_mxfp4(layer.weight, layer.weight_scale, x.dtype)
105 + qdq_x = self.quant_dequant_func(x)
106 + return F.linear(qdq_x, dq_w, bias)
vllm/model_executor/kernels/linear/mxfp4/flashinfer.py
+5
@@ -7,6 +7,9 @@ from torch.nn.parameter import Parameter
7 from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
8 swizzle_mxfp4_scales,
9 )
10 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
11 + kMxfp4Dynamic,
12 +)
13 from vllm.platforms import current_platform
14 from vllm.utils.flashinfer import has_flashinfer_cutedsl
15
@@ -28,6 +31,8 @@ class FlashInferMxFp4LinearKernel(MxFp4LinearKernel):
31
32 @classmethod
33 def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
34 + if config.activation_quant_key != kMxfp4Dynamic:
35 + return False, "only supports MXFP4 dynamic activation"
36 return True, None
37
38 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
vllm/model_executor/kernels/linear/mxfp4/humming.py
+13 -1
@@ -3,14 +3,18 @@
3
4 import torch
5
6 +from vllm.logger import init_logger
7 from vllm.model_executor.layers.quantization.utils.humming_utils import (
8 convert_linear_layer_to_humming_standard,
9 prepare_humming_layer,
10 )
11 +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic
12 from vllm.platforms import current_platform
13
14 from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
15
16 +logger = init_logger(__name__)
17 +
18
19 class HummingMxFp4LinearKernel(MxFp4LinearKernel):
20 """Humming GEMM Kernel for MXFP4."""
@@ -28,7 +32,15 @@ class HummingMxFp4LinearKernel(MxFp4LinearKernel):
32 return True, None
33
34 @classmethod
31 - def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
35 + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
36 + if config.activation_quant_key not in (None, kMxfp4Dynamic):
37 + return False, "only supports MXFP4 dynamic or unquantized activations"
38 + if config.activation_quant_key is not None:
39 + logger.warning_once(
40 + "HummingMxFp4LinearKernel is a weight-only (A16) kernel; "
41 + "the requested activation quantization (%s) is ignored.",
42 + config.activation_quant_key,
43 + )
44 return True, None
45
46 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
vllm/model_executor/kernels/linear/mxfp4/marlin.py
+14 -1
@@ -3,8 +3,13 @@
3
4 import torch
5
6 +from vllm.logger import init_logger
7 +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic
8 +
9 from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
10
11 +logger = init_logger(__name__)
12 +
13
14 class MarlinMxFp4LinearKernel(MxFp4LinearKernel):
15 @classmethod
@@ -20,7 +25,15 @@ class MarlinMxFp4LinearKernel(MxFp4LinearKernel):
25 return False, "Marlin FP4 not available"
26
27 @classmethod
23 - def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
28 + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
29 + if config.activation_quant_key not in (None, kMxfp4Dynamic):
30 + return False, "only supports MXFP4 dynamic or unquantized activations"
31 + if config.activation_quant_key is not None:
32 + logger.warning_once(
33 + "MarlinMxFp4LinearKernel is a weight-only (A16) kernel; "
34 + "the requested activation quantization (%s) is ignored.",
35 + config.activation_quant_key,
36 + )
37 return True, None
38
39 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
vllm/model_executor/kernels/linear/mxfp4/xpu.py
+6 -1
@@ -6,6 +6,9 @@ import torch
6 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
7 xpu_mxfp4_quantize as quant_mxfp4,
8 )
9 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
10 + kMxfp4Dynamic,
11 +)
12 from vllm.model_executor.utils import replace_parameter
13 from vllm.platforms import current_platform
14
@@ -24,7 +27,9 @@ class XPUMxFp4LinearKernel(MxFp4LinearKernel):
27 return True, None
28
29 @classmethod
27 - def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
30 + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
31 + if config.activation_quant_key != kMxfp4Dynamic:
32 + return False, "only supports MXFP4 dynamic activation"
33 return True, None
34
35 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
vllm/model_executor/kernels/linear/mxfp6/__init__.py new
+12
@@ -0,0 +1,12 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.model_executor.kernels.linear.mxfp6.base import (
5 + MxFp6LinearKernel,
6 + MxFp6LinearLayerConfig,
7 +)
8 +
9 +__all__ = [
10 + "MxFp6LinearKernel",
11 + "MxFp6LinearLayerConfig",
12 +]
vllm/model_executor/kernels/linear/mxfp6/base.py new
+76
@@ -0,0 +1,76 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from abc import ABC, abstractmethod
5 +from dataclasses import dataclass
6 +
7 +import torch
8 +
9 +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
10 +
11 +
12 +@dataclass
13 +class MxFp6LinearLayerConfig:
14 + """Configuration for an MXFP6 linear layer.
15 +
16 + All MXFP6 layers share the same structure: packed uint8 weights (2 FP4 values per
17 + byte) and per-block weight scales (group size 32).
18 +
19 + Attributes:
20 + weight_quant_key: Identifies the weight quantization format. Can be
21 + kMxfp6E2M3Static or kMxfp6E3M2Static.
22 + activation_quant_key: Identifies the activation quantization format,
23 + or `None` when activations must not be quantized.
24 + """
25 +
26 + weight_quant_key: QuantKey
27 + activation_quant_key: QuantKey | None = None
28 +
29 +
30 +class MxFp6LinearKernel(ABC):
31 + """Base class for MXFP6 quantized linear kernels.
32 +
33 + Each subclass implements a specific GEMM backend (CUTLASS, Marlin, etc).
34 + The kernel selection mechanism iterates over registered subclasses in
35 + priority order,calling ``is_supported`` and ``can_implement`` to find the best
36 + match for the current hardware.
37 + """
38 +
39 + def __init__(self, config: MxFp6LinearLayerConfig) -> None:
40 + assert self.can_implement(config)[0]
41 + assert self.is_supported()[0]
42 + self.config = config
43 +
44 + @classmethod
45 + @abstractmethod
46 + def is_supported(
47 + cls, compute_capability: int | None = None
48 + ) -> tuple[bool, str | None]:
49 + """Return whether this kernel can run on the current platform."""
50 + raise NotImplementedError
51 +
52 + @classmethod
53 + @abstractmethod
54 + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]:
55 + """Return whether this kernel can handle *config*."""
56 + raise NotImplementedError
57 +
58 + @abstractmethod
59 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
60 + """Transform weights into the format required by this kernel.
61 +
62 + Called once after checkpoint weights have been loaded onto the
63 + device. Implementations should repack / swizzle / pad weights
64 + and scales in-place on *layer*.
65 + """
66 + raise NotImplementedError
67 +
68 + @abstractmethod
69 + def apply_weights(
70 + self,
71 + layer: torch.nn.Module,
72 + x: torch.Tensor,
73 + bias: torch.Tensor | None = None,
74 + ) -> torch.Tensor:
75 + """Run the quantized GEMM."""
76 + raise NotImplementedError
vllm/model_executor/kernels/linear/mxfp6/emulation.py new
+93
@@ -0,0 +1,93 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Callable
5 +from functools import partial
6 +
7 +import torch
8 +import torch.nn.functional as F
9 +from torch.nn.parameter import Parameter
10 +
11 +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
12 + quant_dequant_mxfp4,
13 +)
14 +from vllm.model_executor.layers.quantization.utils.mxfp6_utils import (
15 + dequant_mxfp6,
16 + quant_dequant_mxfp6,
17 +)
18 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
19 + QuantKey,
20 + kMxfp4Dynamic,
21 + kMxfp6E2M3Dynamic,
22 + kMxfp6E2M3Static,
23 + kMxfp6E3M2Dynamic,
24 + kMxfp6E3M2Static,
25 +)
26 +
27 +from .base import MxFp6LinearKernel, MxFp6LinearLayerConfig
28 +
29 +_WEIGHT_DEQUANT_FUNCS: dict[QuantKey, Callable[..., torch.Tensor]] = {
30 + kMxfp6E3M2Static: partial(dequant_mxfp6, quant_dtype="fp6_e3m2"),
31 + kMxfp6E2M3Static: partial(dequant_mxfp6, quant_dtype="fp6_e2m3"),
32 +}
33 +
34 +_ACTIVATION_QUANT_DEQUANT_FUNCS: dict[
35 + QuantKey, Callable[[torch.Tensor], torch.Tensor]
36 +] = {
37 + kMxfp4Dynamic: quant_dequant_mxfp4,
38 + kMxfp6E3M2Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e3m2"),
39 + kMxfp6E2M3Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e2m3"),
40 +}
41 +
42 +
43 +class EmulationMxfp6LinearKernel(MxFp6LinearKernel):
44 + """Software emulation fallback for OCP MXFP4/MXFP6 (dequant + F.linear)."""
45 +
46 + def __init__(self, config: MxFp6LinearLayerConfig) -> None:
47 + super().__init__(config)
48 + self.dequant_func = _WEIGHT_DEQUANT_FUNCS[config.weight_quant_key]
49 + if config.activation_quant_key is None:
50 + # no input Q/DQ for weight-only
51 + self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = (
52 + lambda x: x
53 + )
54 + else:
55 + self.quant_dequant_func = _ACTIVATION_QUANT_DEQUANT_FUNCS[
56 + config.activation_quant_key
57 + ]
58 +
59 + @classmethod
60 + def is_supported(
61 + cls, compute_capability: int | None = None
62 + ) -> tuple[bool, str | None]:
63 + return True, None
64 +
65 + @classmethod
66 + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]:
67 + if config.weight_quant_key not in (
68 + kMxfp6E2M3Static,
69 + kMxfp6E3M2Static,
70 + ):
71 + return False, "only supports MXFP6 weights"
72 + if config.activation_quant_key not in (
73 + None,
74 + kMxfp4Dynamic,
75 + kMxfp6E3M2Dynamic,
76 + kMxfp6E2M3Dynamic,
77 + ):
78 + return False, "only supports MXFP4 or MXFP6 or unquantized activations"
79 +
80 + return True, None
81 +
82 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
83 + layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
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 + dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype)
92 + qdq_x = self.quant_dequant_func(x)
93 + return F.linear(qdq_x, dq_w, bias)
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py
+6 -1
@@ -9,6 +9,9 @@ from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel
9 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
10 CompressedTensorsScheme,
11 )
12 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
13 + kMxfp4Dynamic,
14 +)
15 from vllm.model_executor.parameter import (
16 GroupQuantScaleParameter,
17 ModelWeightParameter,
@@ -35,7 +38,9 @@ class CompressedTensorsW4A4Mxfp4(CompressedTensorsScheme):
38
39 def __init__(self):
40 self.group_size = 32
38 - self.kernel = init_mxfp4_linear_kernel()
41 + self.kernel = init_mxfp4_linear_kernel(
42 + activation_quant_key=kMxfp4Dynamic,
43 + )
44
45 @classmethod
46 def get_min_capability(cls) -> int:
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+56 -59
@@ -3,25 +3,28 @@
3
4 from collections.abc import Callable
5 from fractions import Fraction
6 -from functools import partial
6 from typing import Any
7
8 import torch
10 -import torch.nn.functional as F
9
10 from vllm.logger import init_logger
13 -from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel
14 -from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
15 - dequant_mxfp4,
16 - quant_dequant_mxfp4,
17 -)
18 -from vllm.model_executor.layers.quantization.utils.mxfp6_utils import (
19 - dequant_mxfp6,
20 - quant_dequant_mxfp6,
11 +from vllm.model_executor.kernels.linear import (
12 + MxFp4LinearKernel,
13 + MxFp6LinearKernel,
14 + init_mxfp4_linear_kernel,
15 + init_mxfp6_linear_kernel,
16 )
17 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
18 OCP_MX_BLOCK_SIZE,
24 - OCP_MX_Scheme,
19 +)
20 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
21 + QuantKey,
22 + kMxfp4Dynamic,
23 + kMxfp4Static,
24 + kMxfp6E2M3Dynamic,
25 + kMxfp6E2M3Static,
26 + kMxfp6E3M2Dynamic,
27 + kMxfp6E3M2Static,
28 )
29 from vllm.model_executor.parameter import (
30 GroupQuantScaleParameter,
@@ -35,8 +38,22 @@ from .quark_scheme import QuarkScheme
38
39 logger = init_logger(__name__)
40
41 +_WEIGHT_QUANT_KEY_MAP: dict[str, QuantKey] = {
42 + "mxfp4": kMxfp4Static,
43 + "mxfp6_e3m2": kMxfp6E3M2Static,
44 + "mxfp6_e2m3": kMxfp6E2M3Static,
45 +}
46 +
47 +_ACTIVATION_QUANT_KEY_MAP: dict[str, QuantKey] = {
48 + "mxfp4": kMxfp4Dynamic,
49 + "mxfp6_e3m2": kMxfp6E3M2Dynamic,
50 + "mxfp6_e2m3": kMxfp6E2M3Dynamic,
51 +}
52 +
53
54 class QuarkOCP_MX(QuarkScheme):
55 + ocp_mx_linear: MxFp6LinearKernel | MxFp4LinearKernel
56 +
57 def __init__(
58 self,
59 weight_quant_spec: dict[str, Any],
@@ -49,35 +66,26 @@ class QuarkOCP_MX(QuarkScheme):
66 self.weight_dtype = weight_quant_spec["dtype"].replace("fp", "mxfp")
67 self.input_dtype: str | None = None
68 if input_quant_spec is not None:
52 - input_quant = input_quant_spec["dtype"]
53 - if input_quant == "fp8_e4m3":
54 - self.input_dtype = "fp8"
55 - else:
56 - self.input_dtype = input_quant.replace("fp", "mxfp")
57 -
58 - self.ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype(
59 - self.input_dtype, self.weight_dtype
69 + self.input_dtype = input_quant_spec["dtype"].replace("fp", "mxfp")
70 +
71 + if self.input_dtype not in [None, *_ACTIVATION_QUANT_KEY_MAP]:
72 + raise ValueError(
73 + f"Unsupported input_dtype={self.input_dtype} for QuarkOCP_MX. "
74 + f"Supported activation dtypes are {_ACTIVATION_QUANT_KEY_MAP.keys()}, "
75 + "or None for weight-only quantization."
76 + )
77 +
78 + self.weight_quant_key = _WEIGHT_QUANT_KEY_MAP[self.weight_dtype]
79 + self.activation_quant_key = (
80 + _ACTIVATION_QUANT_KEY_MAP[self.input_dtype]
81 + if self.input_dtype is not None
82 + else None
83 )
84
85 if self.weight_dtype == "mxfp4":
86 self.packed_factor: int | Fraction = 2
64 - self.dequant_func = dequant_mxfp4
87 else:
88 self.packed_factor = Fraction(numerator=8, denominator=6)
67 - self.dequant_func = partial(
68 - dequant_mxfp6, quant_dtype=self.weight_dtype.replace("mx", "")
69 - )
70 -
71 - if self.input_dtype is None:
72 - self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = (
73 - lambda x: x
74 - ) # no input Q/DQ for weight-only
75 - elif self.input_dtype == "mxfp4":
76 - self.quant_dequant_func = quant_dequant_mxfp4
77 - else:
78 - self.quant_dequant_func = partial(
79 - quant_dequant_mxfp6, quant_dtype=self.input_dtype.replace("mx", "")
80 - )
89
90 if input_quant_spec is None:
91 self.static_input_scales = False
@@ -90,16 +98,6 @@ class QuarkOCP_MX(QuarkScheme):
98 "implemented. Please open an issue."
99 )
100
93 - # TODO: integrate (or test) mixed-precision kernel.
94 - self.emulate = not current_platform.supports_mx() or (
95 - self.input_dtype != "mxfp4" or self.weight_dtype != "mxfp4"
96 - )
97 -
98 - # TODO: Move emulation code path as a kernel, and always
99 - # use init_mxfp4_linear_kernel.
100 - if not self.emulate:
101 - self.ocp_mx_linear = init_mxfp4_linear_kernel()
102 -
101 if not current_platform.supports_mx():
102 logger.warning_once(
103 "The current platform does not support native MXFP4/MXFP6 "
@@ -151,17 +149,10 @@ class QuarkOCP_MX(QuarkScheme):
149 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
150 layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False)
151
154 - if self.emulate:
155 - if self.dynamic_mxfp4_quant:
156 - self.process_dynamic_mxfp4_weights_after_loading(layer)
157 - else:
158 - layer.weight_scale = torch.nn.Parameter(
159 - layer.weight_scale.data, requires_grad=False
160 - )
161 - else:
162 - if self.dynamic_mxfp4_quant:
163 - self.process_dynamic_mxfp4_weights_after_loading(layer)
164 - self.ocp_mx_linear.process_weights_after_loading(layer)
152 + if self.dynamic_mxfp4_quant:
153 + self.process_dynamic_mxfp4_weights_after_loading(layer)
154 +
155 + self.ocp_mx_linear.process_weights_after_loading(layer)
156
157 def create_weights(
158 self,
@@ -218,14 +209,20 @@ class QuarkOCP_MX(QuarkScheme):
209 )
210 layer.register_parameter("weight_scale", weight_scale)
211
212 + if self.weight_quant_key == kMxfp4Static:
213 + self.ocp_mx_linear = init_mxfp4_linear_kernel(
214 + activation_quant_key=self.activation_quant_key,
215 + )
216 + elif self.weight_quant_key in [kMxfp6E2M3Static, kMxfp6E3M2Static]:
217 + self.ocp_mx_linear = init_mxfp6_linear_kernel(
218 + weight_quant_key=self.weight_quant_key,
219 + activation_quant_key=self.activation_quant_key,
220 + )
221 +
222 def apply_weights(
223 self,
224 layer: torch.nn.Module,
225 x: torch.Tensor,
226 bias: torch.Tensor | None = None,
227 ) -> torch.Tensor:
227 - if self.emulate:
228 - dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype)
229 - qdq_x = self.quant_dequant_func(x)
230 - return F.linear(qdq_x, dq_w, bias)
228 return self.ocp_mx_linear.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/utils/quant_utils.py
+32 -2
@@ -90,6 +90,7 @@ class ScaleDesc:
90 GroupShape.PER_CHANNEL: "per_channel",
91 }
92 group_shape = d.get(self.group_shape, str(self.group_shape))
93 +
94 return (
95 f"{fx.graph.dtype_abbrs[self.dtype]},"
96 f"{'static' if self.static else 'dynamic'},{group_shape}"
@@ -106,15 +107,24 @@ class QuantKey:
107 symmetric: symmetric if True, asymmetric if False
108 """
109
109 - dtype: torch.dtype
110 + # TODO: QuantKey.dtype is assumed to be `torch.dtype` in matcher_utils.py,
111 + # but #37990 introduced e.g. `kInt4Static` that uses a `ScalarType` dtype,
112 + # same for kMxfp6 that does not have a native torch representation.
113 + # Logical dtype and storage (torch) dtype should be separated (see #48949).
114 + dtype: torch.dtype | ScalarType
115 scale: ScaleDesc
116 scale2: ScaleDesc | None = None
117 symmetric: bool = True
118
119 def __str__(self):
120 scale2_str = f"scale2({self.scale2})," if self.scale2 else ""
121 + dtype_description = (
122 + fx.graph.dtype_abbrs[self.dtype]
123 + if isinstance(self.dtype, torch.dtype)
124 + else self.dtype
125 + )
126 return (
117 - f"QuantKey({fx.graph.dtype_abbrs[self.dtype]},"
127 + f"QuantKey({dtype_description},"
128 f"scale({self.scale}),{scale2_str}"
129 f"{'a' if not self.symmetric else ''}symmetric)"
130 )
@@ -172,6 +182,26 @@ kMxfp8Dynamic = QuantKey(FP8_DTYPE, scale=kMxfp8DynamicGroupScale, symmetric=Tru
182 kMxfp4StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32))
183 kMxfp4Static = QuantKey(FP4_DTYPE, scale=kMxfp4StaticGroupScale, symmetric=True)
184
185 +kMxfp6E3M2StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32))
186 +kMxfp6E3M2Static = QuantKey(
187 + scalar_types.float6_e3m2f, scale=kMxfp6E3M2StaticGroupScale, symmetric=True
188 +)
189 +
190 +kMxfp6E3M2DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32))
191 +kMxfp6E3M2Dynamic = QuantKey(
192 + scalar_types.float6_e3m2f, scale=kMxfp6E3M2DynamicGroupScale, symmetric=True
193 +)
194 +
195 +kMxfp6E2M3StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32))
196 +kMxfp6E2M3Static = QuantKey(
197 + scalar_types.float6_e2m3f, scale=kMxfp6E2M3StaticGroupScale, symmetric=True
198 +)
199 +
200 +kMxfp6E2M3DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32))
201 +kMxfp6E2M3Dynamic = QuantKey(
202 + scalar_types.float6_e2m3f, scale=kMxfp6E2M3DynamicGroupScale, symmetric=True
203 +)
204 +
205 # TODO: convert this to use SCALAR_TYPE. This is not right.
206 kInt4StaticGroupScale = ScaleDesc(torch.float16, True, GroupShape(1, -1))
207 kInt4Static = QuantKey(INT4_DTYPE, scale=kInt4StaticGroupScale, symmetric=True)