[Kernel] Add opt-in load-time MXFP4 dequantization (#50814)
Signed-off-by: Liuyinfeng01 <yinfeliu@amd.com> Co-authored-by: Shanshan Shen <467638484@qq.com>
yinfengLiu committed
Sep 22, 2026 at 14:51 UTC
2af05512cf94eda87b4eb51e9f1345a12337a36d
3 files changed
+112
-3
tests/kernels/quantization/test_mxfp4_kernel_selection.py
+86
-1
@@ -5,11 +5,12 @@
5
Run `pytest tests/kernels/quantization/test_mxfp4_kernel_selection.py`.
6
"""
7
8
-from unittest.mock import patch
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,
@@ -46,6 +47,19 @@ _TRUE_W4A4_KERNELS = [
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")
@@ -155,6 +169,77 @@ def test_emulation_kernel_derives_quant_dequant_func_from_config(monkeypatch):
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."""
vllm/envs.py
+8
@@ -131,6 +131,7 @@ if TYPE_CHECKING:
131
VLLM_GDN_DECODE_KERNEL: Literal["cuda", "triton"] = "cuda"
132
VLLM_DISABLE_PYNCCL: bool = False
133
VLLM_USE_OINK_OPS: bool = False
134
+ VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD: bool = False
135
VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True
136
VLLM_ROCM_USE_AITER: bool = False
137
VLLM_ROCM_USE_AITER_CUSTOM_AR: bool = True
@@ -1229,6 +1230,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
1230
),
1231
# Disable aiter ops unless specifically enabled.
1232
# Acts as a parent switch to enable the rest of the other operations.
1233
+ # Set to 1 to dequantize MXFP4 weights to BF16 once at load time and run as
1234
+ # a BF16 checkpoint (no per-step weight dequant). This improves emulation
1235
+ # latency at the cost of additional device memory. Default off.
1236
+ "VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD": lambda: (
1237
+ os.getenv("VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD", "False").lower()
1238
+ in ("true", "1")
1239
+ ),
1240
# On hardware without a native MXFP8 kernel (e.g. ROCm gfx942 / MI300), the
1241
# MXFP8 emulation path dequantizes weights MXFP8->BF16 once at load time and
1242
# runs as a BF16 checkpoint (no per-step dequant). Set to 0 to fall back to
vllm/model_executor/kernels/linear/mxfp4/emulation.py
+18
-2
@@ -8,6 +8,7 @@ import torch
8
import torch.nn.functional as F
9
from torch.nn.parameter import Parameter
10
11
+import vllm.envs as envs
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
14
dequant_mxfp4,
@@ -97,7 +98,16 @@ class EmulationMxfp4LinearKernel(MxFp4LinearKernel):
98
return True, None
99
100
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
100
- layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
101
+ weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
102
+
103
+ # Match the MXFP8 emulation path when explicitly enabled: dequantize
104
+ # once at load time so inference uses a plain high-precision linear.
105
+ # Otherwise retain packed weights and dequantize per invocation.
106
+ if envs.VLLM_MXFP4_EMULATION_DEQUANT_AT_LOAD:
107
+ weight = dequant_mxfp4(layer.weight, weight_scale, torch.bfloat16)
108
+ layer.weight = Parameter(weight.contiguous(), requires_grad=False)
109
+
110
+ layer.weight_scale = weight_scale
111
112
def apply_weights(
113
self,
@@ -105,6 +115,12 @@ class EmulationMxfp4LinearKernel(MxFp4LinearKernel):
115
x: torch.Tensor,
116
bias: torch.Tensor | None = None,
117
) -> torch.Tensor:
108
- dq_w = dequant_mxfp4(layer.weight, layer.weight_scale, x.dtype)
118
qdq_x = self.quant_dequant_func(x)
119
+
120
+ # Packed MXFP4 weights use one byte for two values. A weight with
121
+ # elements wider than one byte was already dequantized at load time.
122
+ if layer.weight.element_size() >= 2:
123
+ return F.linear(qdq_x, layer.weight.to(x.dtype), bias)
124
+
125
+ dq_w = dequant_mxfp4(layer.weight, layer.weight_scale, x.dtype)
126
return F.linear(qdq_x, dq_w, bias)