| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for MXFP4 linear kernel selection logic (CPU-only). |
| 4 | |
| 5 | Run `pytest tests/kernels/quantization/test_mxfp4_kernel_selection.py`. |
| 6 | """ |
| 7 | |
| 8 | from unittest.mock import MagicMock, patch |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | |
| 13 | import vllm.envs as envs |
| 14 | from vllm.model_executor.kernels.linear import ( |
| 15 | AiterMxfp4LinearKernel, |
| 16 | EmulationMxfp4LinearKernel, |
| 17 | FlashInferMxFp4LinearKernel, |
| 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] |
| 48 | |
| 49 | |
| 50 | def _make_emulation_layer(): |
| 51 | layer = torch.nn.Module() |
| 52 | layer.register_parameter( |
| 53 | "weight", |
| 54 | torch.nn.Parameter(torch.ones((2, 2), dtype=torch.uint8), False), |
| 55 | ) |
| 56 | layer.register_parameter( |
| 57 | "weight_scale", |
| 58 | torch.nn.Parameter(torch.ones((2, 1), dtype=torch.uint8), False), |
| 59 | ) |
| 60 | return layer |
| 61 | |
| 62 | |
| 63 | def test_can_implement_is_abstract(): |
| 64 | """Test that can_implement()/is_supported() are properly defined.""" |
| 65 | assert hasattr(MxFp4LinearKernel, "can_implement") |
| 66 | assert hasattr(MxFp4LinearKernel, "is_supported") |
| 67 | |
| 68 | |
| 69 | @pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS) |
| 70 | def test_true_w4a4_kernels_accept_dynamic_mxfp4_activation(kernel_cls): |
| 71 | config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic) |
| 72 | can_implement, reason = kernel_cls.can_implement(config) |
| 73 | assert can_implement, reason |
| 74 | |
| 75 | |
| 76 | @pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS) |
| 77 | def test_true_w4a4_kernels_reject_unset_activation(kernel_cls): |
| 78 | """None means weight-only/unquantized activations, not dynamic MXFP4.""" |
| 79 | config = MxFp4LinearLayerConfig() |
| 80 | can_implement, reason = kernel_cls.can_implement(config) |
| 81 | assert not can_implement |
| 82 | assert reason |
| 83 | |
| 84 | |
| 85 | @pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS) |
| 86 | def test_true_w4a4_kernels_reject_explicit_non_mxfp4_activation(kernel_cls): |
| 87 | """FlashInfer/XPU/Aiter quantize activations to MXFP4 internally, so an |
| 88 | explicit request for a different activation format must be rejected.""" |
| 89 | config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp6E3M2Dynamic) |
| 90 | can_implement, reason = kernel_cls.can_implement(config) |
| 91 | assert not can_implement |
| 92 | assert reason |
| 93 | |
| 94 | |
| 95 | @pytest.mark.parametrize("kernel_cls", _WEIGHT_ONLY_KERNELS) |
| 96 | @pytest.mark.parametrize("activation_quant_key", [None, kMxfp4Dynamic]) |
| 97 | def test_weight_only_kernels_accept_unquantized_or_mxfp4_activation( |
| 98 | kernel_cls, activation_quant_key |
| 99 | ): |
| 100 | """Marlin/Humming never quantize activations, so an unset activation key, |
| 101 | or one that already describes MXFP4-shaped data, is tolerated. When an |
| 102 | activation key is explicitly set, a warning must be logged noting that it |
| 103 | is ignored, since these kernels are weight-only (A16).""" |
| 104 | config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key) |
| 105 | with patch(f"{kernel_cls.__module__}.logger.warning_once") as warning_once: |
| 106 | can_implement, reason = kernel_cls.can_implement(config) |
| 107 | assert can_implement, reason |
| 108 | |
| 109 | if activation_quant_key is None: |
| 110 | warning_once.assert_not_called() |
| 111 | else: |
| 112 | warning_once.assert_called_once() |
| 113 | message = warning_once.call_args.args[0] |
| 114 | assert "the requested activation quantization" in message |
| 115 | assert "is ignored" in message |
| 116 | |
| 117 | |
| 118 | @pytest.mark.parametrize("kernel_cls", _WEIGHT_ONLY_KERNELS) |
| 119 | def test_weight_only_kernels_reject_non_mxfp4_activation(kernel_cls): |
| 120 | config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp6E3M2Dynamic) |
| 121 | can_implement, reason = kernel_cls.can_implement(config) |
| 122 | assert not can_implement |
| 123 | assert reason |
| 124 | |
| 125 | |
| 126 | @pytest.mark.parametrize( |
| 127 | "activation_quant_key", |
| 128 | [None, kMxfp4Dynamic, kMxfp6E3M2Dynamic, kMxfp6E2M3Dynamic], |
| 129 | ) |
| 130 | def test_emulation_kernel_accepts_any_config(activation_quant_key, monkeypatch): |
| 131 | """EmulationMxfp4LinearKernel is the universal fallback: it must accept |
| 132 | every supported activation format.""" |
| 133 | # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`, |
| 134 | # which we are not testing here. |
| 135 | monkeypatch.setattr( |
| 136 | "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark", |
| 137 | lambda: True, |
| 138 | ) |
| 139 | config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key) |
| 140 | with patch( |
| 141 | "vllm.model_executor.kernels.linear._get_linear_backend", |
| 142 | return_value="emulation", |
| 143 | ): |
| 144 | can_implement, reason = EmulationMxfp4LinearKernel.can_implement(config) |
| 145 | assert can_implement, reason |
| 146 | |
| 147 | |
| 148 | def test_emulation_kernel_derives_quant_dequant_func_from_config(monkeypatch): |
| 149 | """quant_dequant_func must be derived purely from the config's activation |
| 150 | QuantKey, not set externally.""" |
| 151 | # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`, |
| 152 | # which we are not testing here. |
| 153 | monkeypatch.setattr( |
| 154 | "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark", |
| 155 | lambda: True, |
| 156 | ) |
| 157 | with patch( |
| 158 | "vllm.model_executor.kernels.linear._get_linear_backend", |
| 159 | return_value="emulation", |
| 160 | ): |
| 161 | weight_only_config = MxFp4LinearLayerConfig() |
| 162 | kernel = EmulationMxfp4LinearKernel(weight_only_config) |
| 163 | x = torch.randn(4) |
| 164 | # identity for weight-only |
| 165 | assert torch.equal(kernel.quant_dequant_func(x), x) |
| 166 | |
| 167 | w4a4_config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic) |
| 168 | kernel = EmulationMxfp4LinearKernel(w4a4_config) |
| 169 | assert kernel.quant_dequant_func is quant_dequant_mxfp4 |
| 170 | |
| 171 | |
| 172 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 173 | def test_emulation_kernel_dequantizes_at_load_and_keeps_activation_qdq( |
| 174 | dtype, monkeypatch |
| 175 | ): |
| 176 | monkeypatch.setattr( |
| 177 | "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark", |
| 178 | lambda: True, |
| 179 | ) |
| 180 | kernel = EmulationMxfp4LinearKernel( |
| 181 | MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic) |
| 182 | ) |
| 183 | kernel.quant_dequant_func = MagicMock(side_effect=lambda x: x * 0.5) |
| 184 | layer = _make_emulation_layer() |
| 185 | dequantized = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.bfloat16) |
| 186 | x = torch.tensor([[2.0, 4.0]], dtype=dtype) |
| 187 | |
| 188 | with ( |
| 189 | patch.object(envs, "VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD", True), |
| 190 | patch( |
| 191 | "vllm.model_executor.kernels.linear.mxfp4.emulation.dequant_mxfp4", |
| 192 | return_value=dequantized, |
| 193 | ) as dequant, |
| 194 | ): |
| 195 | kernel.process_weights_after_loading(layer) |
| 196 | actual = kernel.apply_weights(layer, x) |
| 197 | |
| 198 | expected = torch.nn.functional.linear(x * 0.5, dequantized.to(dtype)) |
| 199 | torch.testing.assert_close(actual, expected) |
| 200 | dequant.assert_called_once() |
| 201 | kernel.quant_dequant_func.assert_called_once_with(x) |
| 202 | assert layer.weight.dtype == torch.bfloat16 |
| 203 | assert layer.weight.device == dequantized.device |
| 204 | assert not layer.weight.requires_grad |
| 205 | assert layer.weight_scale is not None |
| 206 | assert not layer.weight_scale.requires_grad |
| 207 | |
| 208 | |
| 209 | def test_emulation_kernel_opt_out_dequantizes_per_invocation(monkeypatch): |
| 210 | monkeypatch.setattr( |
| 211 | "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark", |
| 212 | lambda: True, |
| 213 | ) |
| 214 | kernel = EmulationMxfp4LinearKernel( |
| 215 | MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic) |
| 216 | ) |
| 217 | kernel.quant_dequant_func = MagicMock(side_effect=lambda x: x * 0.5) |
| 218 | layer = _make_emulation_layer() |
| 219 | dequantized = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.bfloat16) |
| 220 | x = torch.tensor([[2.0, 4.0]], dtype=torch.bfloat16) |
| 221 | |
| 222 | with ( |
| 223 | patch.object(envs, "VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD", False), |
| 224 | patch( |
| 225 | "vllm.model_executor.kernels.linear.mxfp4.emulation.dequant_mxfp4", |
| 226 | return_value=dequantized, |
| 227 | ) as dequant, |
| 228 | ): |
| 229 | kernel.process_weights_after_loading(layer) |
| 230 | dequant.assert_not_called() |
| 231 | first = kernel.apply_weights(layer, x) |
| 232 | second = kernel.apply_weights(layer, x) |
| 233 | |
| 234 | expected = torch.nn.functional.linear(x * 0.5, dequantized) |
| 235 | torch.testing.assert_close(first, expected) |
| 236 | torch.testing.assert_close(second, expected) |
| 237 | assert dequant.call_count == 2 |
| 238 | assert kernel.quant_dequant_func.call_count == 2 |
| 239 | assert layer.weight.dtype == torch.uint8 |
| 240 | assert layer.weight_scale is not None |
| 241 | |
| 242 | |
| 243 | def test_aiter_kernel_is_supported_requires_native_mx_support(): |
| 244 | """AiterMxfp4LinearKernel must not be selected on platforms without |
| 245 | native MX compute, even if AITER itself is importable.""" |
| 246 | with patch( |
| 247 | "vllm.model_executor.kernels.linear.mxfp4.aiter.current_platform.supports_mx", |
| 248 | return_value=False, |
| 249 | ): |
| 250 | is_supported, reason = AiterMxfp4LinearKernel.is_supported() |
| 251 | assert not is_supported |
| 252 | assert reason |
| 253 | |
| 254 | |
| 255 | class OOTMxFp4LinearKernel(MxFp4LinearKernel): |
| 256 | @classmethod |
| 257 | def is_supported( |
| 258 | cls, compute_capability: int | None = None |
| 259 | ) -> tuple[bool, str | None]: |
| 260 | return True, None |
| 261 | |
| 262 | @classmethod |
| 263 | def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: |
| 264 | return True, None |
| 265 | |
| 266 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 267 | pass |
| 268 | |
| 269 | def apply_weights( |
| 270 | self, |
| 271 | layer: torch.nn.Module, |
| 272 | x: torch.Tensor, |
| 273 | bias: torch.Tensor | None = None, |
| 274 | ) -> torch.Tensor: |
| 275 | pass |
| 276 | |
| 277 | |
| 278 | @patch("vllm.model_executor.kernels.linear.current_platform") |
| 279 | def test_init_mxfp4_linear_kernel_dispatches_to_registered_kernel(platform_mock): |
| 280 | """init_mxfp4_linear_kernel should select a registered kernel that |
| 281 | reports itself as supported, and construct it with a fresh config.""" |
| 282 | platform_mock._enum = PlatformEnum.OOT |
| 283 | register_linear_kernel(OOTMxFp4LinearKernel, PlatformEnum.OOT, "mxfp4") |
| 284 | |
| 285 | kernel = init_mxfp4_linear_kernel(activation_quant_key=kMxfp4Dynamic) |
| 286 | |
| 287 | assert isinstance(kernel, OOTMxFp4LinearKernel) |
| 288 | assert kernel.config == MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic) |
| 289 | |
| 290 | |
| 291 | class UnsupportedMxFp4LinearKernel(MxFp4LinearKernel): |
| 292 | @classmethod |
| 293 | def is_supported( |
| 294 | cls, compute_capability: int | None = None |
| 295 | ) -> tuple[bool, str | None]: |
| 296 | return False, "never supported" |
| 297 | |
| 298 | @classmethod |
| 299 | def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: |
| 300 | return True, None |
| 301 | |
| 302 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 303 | pass |
| 304 | |
| 305 | def apply_weights( |
| 306 | self, |
| 307 | layer: torch.nn.Module, |
| 308 | x: torch.Tensor, |
| 309 | bias: torch.Tensor | None = None, |
| 310 | ) -> torch.Tensor: |
| 311 | pass |
| 312 | |
| 313 | |
| 314 | @patch("vllm.model_executor.kernels.linear.current_platform") |
| 315 | def test_init_mxfp4_linear_kernel_raises_when_no_kernel_matches(platform_mock): |
| 316 | platform_mock._enum = PlatformEnum.UNSPECIFIED |
| 317 | register_linear_kernel( |
| 318 | UnsupportedMxFp4LinearKernel, PlatformEnum.UNSPECIFIED, "mxfp4" |
| 319 | ) |
| 320 | |
| 321 | with pytest.raises(ValueError, match="Failed to find a kernel"): |
| 322 | init_mxfp4_linear_kernel() |