@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3518110f2f

[Online quantization] Add online MXFP4 quantization support (#49347)

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

fxmarty-amd committed Aug 7, 2026 at 21:39 UTC 3518110f2f4532c4605d8b1ea2a9bdc7d61b9f04
18 files changed +1423 -35
docs/features/quantization/online.md
+15
@@ -20,6 +20,16 @@ llm = LLM("meta-llama/Llama-3.1-8B", quantization="fp8_per_block")
20
21 # MXFP8 quantization for weights and activations
22 llm = LLM("meta-llama/Llama-3.1-8B", quantization="mxfp8")
23 +
24 +# MXFP4 weight; activation quantization depends on the `linear_backend` picked
25 +llm = LLM("meta-llama/Llama-3.1-8B", quantization="mxfp4")
26 +
27 +# MXFP4 MOE-only weight and activation quantization
28 +llm = LLM(
29 + "Qwen/Qwen3.5-35B-A3B",
30 + quantization="mxfp4",
31 + quantization_config={"linear": {"activation": None, "weight": None}}
32 +)
33 ```
34
35 Or with the CLI:
@@ -28,6 +38,10 @@ Or with the CLI:
38 vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_tensor
39 vllm serve meta-llama/Llama-3.1-8B --quantization fp8_per_block
40 vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8
41 +vllm serve meta-llama/Llama-3.1-8B --quantization mxfp4
42 +
43 +vllm serve Qwen/Qwen3.5-35B-A3B --quantization mxfp4 \
44 + --quantization-config '{"linear":{"activation":null,"weight":null}}'
45 ```
46
47 ## Supported Schemes
@@ -37,6 +51,7 @@ vllm serve meta-llama/Llama-3.1-8B --quantization mxfp8
51 | `fp8_per_tensor` | fp8_e4m3 data, fp32 per-tensor scale | fp8_e4m3 data, fp32 per-tensor scale | On some GPUs (Ada, Hopper) linear activations use per-token scaling for better performance |
52 | `fp8_per_block` | fp8_e4m3 data, fp32 per-128x128-block scale | fp8_e4m3 data, fp32 per-1x128-block scale | |
53 | `mxfp8` | fp8_e4m3 data, e8m0 per-1x32-block scale | fp8_e4m3 data, e8m0 per-1x32-block scale | Requires SM 100+ (Blackwell or newer) for w8a8, other GPUs use a w8a16 fallback |
54 +| `mxfp4` | fp4_e2m1 data, e8m0 per-1x32-block scale ([OCP MX specs](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf)) | - linear: fp4_e2m1 data, e8m0 per-1x32-block scale in some backends, or BF16. <br> - MOE: fp4_e2m1 data, e8m0 per-1x32-block scale. | Linear MXFP4 backend is auto-selected per platform, not enforcing activation dtype. Some use BF16 activation. Use `--linear-backend` to pin one (e.g. `--linear-backend flashinfer`). |
55
56 ## Advanced Configuration
57
tests/evals/gsm8k/configs/Qwen3-30B-A3B-MXFP4-AITER-TP2-online.yaml new
+15
@@ -0,0 +1,15 @@
1 +model_name: "Qwen/Qwen3-30B-A3B"
2 +accuracy_threshold: 0.89
3 +tolerance: 0.03
4 +num_questions: 1319
5 +num_fewshot: 5
6 +server_args: >-
7 + --max-model-len 4096
8 + --tensor-parallel-size 2
9 + --gpu-memory-utilization 0.35
10 + --moe-backend aiter
11 + --quantization mxfp4
12 + --quantization-config '{"linear":{"activation":null,"weight":null}}'
13 +env:
14 + VLLM_ROCM_USE_AITER: "1"
15 + ENABLE_CK: "0" # Avoid AITER CK-based MHA JIT compilation to save time
tests/evals/gsm8k/configs/Qwen3-30B-A3B-MXFP4-AITER-TP2.yaml new
+14
@@ -0,0 +1,14 @@
1 +# Note: this model was produced from Qwen3-30B-A3B (BF16)
2 +model_name: "amd-quark/Qwen3-30B-A3B-MXFP4"
3 +accuracy_threshold: 0.89
4 +tolerance: 0.03
5 +num_questions: 1319
6 +num_fewshot: 5
7 +server_args: >-
8 + --max-model-len 4096
9 + --tensor-parallel-size 2
10 + --gpu-memory-utilization 0.35
11 + --moe-backend aiter
12 +env:
13 + VLLM_ROCM_USE_AITER: "1"
14 + ENABLE_CK: "0" # Avoid AITER CK-based MHA JIT compilation to save time
tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml
+1
@@ -1,3 +1,4 @@
1 +# NOTE: this model was produced from Qwen/Qwen3.5-35B-A3B-FP8
2 model_name: "amd/Qwen3.5-35B-A3B-MXFP4"
3 accuracy_threshold: 0.89
4 tolerance: 0.03
tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml
+1
@@ -1,3 +1,4 @@
1 +# NOTE: this model was produced from Qwen/Qwen3.5-35B-A3B-FP8
2 model_name: "amd/Qwen3.5-35B-A3B-MXFP4"
3 accuracy_threshold: 0.89
4 tolerance: 0.03
tests/kernels/quantization/test_mxfp4_kernel_selection.py
+14 -2
@@ -114,9 +114,15 @@ def test_weight_only_kernels_reject_non_mxfp4_activation(kernel_cls):
114 "activation_quant_key",
115 [None, kMxfp4Dynamic, kMxfp6E3M2Dynamic, kMxfp6E2M3Dynamic],
116 )
117 -def test_emulation_kernel_accepts_any_config(activation_quant_key):
117 +def test_emulation_kernel_accepts_any_config(activation_quant_key, monkeypatch):
118 """EmulationMxfp4LinearKernel is the universal fallback: it must accept
119 every supported activation format."""
120 + # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`,
121 + # which we are not testing here.
122 + monkeypatch.setattr(
123 + "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark",
124 + lambda: True,
125 + )
126 config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key)
127 with patch(
128 "vllm.model_executor.kernels.linear._get_linear_backend",
@@ -126,9 +132,15 @@ def test_emulation_kernel_accepts_any_config(activation_quant_key):
132 assert can_implement, reason
133
134
129 -def test_emulation_kernel_derives_quant_dequant_func_from_config():
135 +def test_emulation_kernel_derives_quant_dequant_func_from_config(monkeypatch):
136 """quant_dequant_func must be derived purely from the config's activation
137 QuantKey, not set externally."""
138 + # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`,
139 + # which we are not testing here.
140 + monkeypatch.setattr(
141 + "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark",
142 + lambda: True,
143 + )
144 with patch(
145 "vllm.model_executor.kernels.linear._get_linear_backend",
146 return_value="emulation",
tests/quantization/test_online.py
+34 -3
@@ -14,6 +14,7 @@ from tests.quantization.utils import (
14 is_quant_method_supported,
15 )
16 from vllm import _custom_ops as ops
17 +from vllm._aiter_ops import rocm_aiter_ops
18 from vllm._custom_ops import scaled_fp4_quant
19 from vllm.model_executor.layers.linear import UnquantizedLinearMethod
20 from vllm.model_executor.layers.quantization.online.fp8 import (
@@ -27,6 +28,10 @@ from vllm.model_executor.layers.quantization.online.fp8 import (
28 _is_tp_sharded,
29 )
30 from vllm.model_executor.layers.quantization.online.int8 import Int8OnlineMoEMethod
31 +from vllm.model_executor.layers.quantization.online.mxfp4 import (
32 + Mxfp4OnlineLinearMethod,
33 + Mxfp4OnlineMoEMethod,
34 +)
35 from vllm.model_executor.layers.quantization.online.nvfp4 import (
36 Nvfp4OnlineMoEMethod,
37 _quantize_moe_weight_to_nvfp4,
@@ -40,6 +45,17 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
45 from vllm.platforms import current_platform
46 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
47
48 +if current_platform.is_rocm():
49 + from vllm.platforms.rocm import on_gfx942, on_gfx950
50 +else:
51 +
52 + def on_gfx950() -> bool:
53 + return False
54 +
55 + def on_gfx942() -> bool:
56 + return False
57 +
58 +
59 DEVICE = current_platform.device_type
60
61
@@ -84,6 +100,12 @@ DEVICE = current_platform.device_type
100 Fp8PerTensorOnlineLinearMethod,
101 Fp8PerTensorOnlineMoEMethod,
102 ),
103 + (
104 + "mxfp4",
105 + None,
106 + Mxfp4OnlineLinearMethod,
107 + Mxfp4OnlineMoEMethod,
108 + ),
109 ],
110 )
111 @pytest.mark.parametrize(
@@ -106,8 +128,14 @@ def test_online_quantization(
128 Does not test performance, peak memory usage, etc.
129 """
130
109 - if use_rocm_aiter:
110 - monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
131 + # TODO: Relax this condition once there is a native MXFP4_MXFP4
132 + # linear/moe backend supported on cuda.
133 + if quant_scheme == "mxfp4" and not (on_gfx950() or on_gfx942()):
134 + pytest.skip("mxfp4 online quantization is only tested on AMD gfx942, gfx950.")
135 +
136 + if current_platform.is_rocm():
137 + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")
138 + rocm_aiter_ops.refresh_env_variables()
139
140 if current_platform.is_xpu() and quant_scheme == "fp8_per_block":
141 pytest.skip("Skip test for online fp8_per_block on XPU platform.")
@@ -144,7 +172,10 @@ def test_online_quantization(
172 if moe is not None:
173 assert isinstance(moe._quant_method, expected_moe_cls)
174
147 - if current_platform.is_cuda() or current_platform.is_xpu():
175 + if quant_scheme == "mxfp4":
176 + # Packed e2m1 values, two per byte.
177 + assert o_proj.weight.dtype == torch.uint8
178 + elif current_platform.is_cuda() or current_platform.is_xpu():
179 assert o_proj.weight.dtype == torch.float8_e4m3fn
180 elif current_platform.is_rocm():
181 assert o_proj.weight.dtype == current_platform.fp8_dtype()
tests/quantization/test_online_mxfp4.py new
+606
@@ -0,0 +1,606 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Correctness tests for the online MXFP4 weight quantization backends.
4 +
5 +Each backend (Triton, aiter, XPU) quantizes a synthetic bf16/fp16
6 +tensor to MXFP4, the result is dequantized back to bf16/fp16, and compared
7 +against the pure-torch reference in `reference_mxfp4.py`.
8 +"""
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm.config.model import ModelConfig
14 +from vllm.model_executor.kernels.linear import _POSSIBLE_MXFP4_KERNELS
15 +from vllm.model_executor.kernels.linear.mxfp4.aiter import (
16 + AiterMxfp4LinearKernel,
17 +)
18 +from vllm.model_executor.layers.fused_moe import FusedMoEFactory
19 +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
20 +from vllm.model_executor.layers.quantization.online.mxfp4 import (
21 + Mxfp4OnlineLinearMethod,
22 + Mxfp4OnlineMoEMethod,
23 +)
24 +from vllm.model_executor.layers.quantization.quark.quark_moe import (
25 + QuarkOCP_MX_MoEMethod,
26 +)
27 +from vllm.model_executor.layers.quantization.quark.schemes.quark_ocp_mx import (
28 + QuarkOCP_MX,
29 +)
30 +from vllm.model_executor.layers.quantization.quark.utils import (
31 + quark_quantize_weight_to_mxfp4,
32 +)
33 +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
34 + downcast_to_mxfp,
35 + mxfp4_quantize,
36 + quant_dequant_mxfp4,
37 +)
38 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
39 +from vllm.model_executor.utils import replace_parameter
40 +from vllm.platforms import PlatformEnum, current_platform
41 +
42 +from .reference_mxfp4 import dq_mxfp4_torch, qdq_mxfp4_torch
43 +
44 +
45 +def fix_negative_zeros(tensor: torch.Tensor) -> torch.Tensor:
46 + """Normalize FP4 e2m1 negative-zero codewords (0b1000) to positive zero
47 + (0b0000) in a packed uint8 tensor (two e2m1 values per byte, low nibble
48 + first).
49 +
50 + -0.0 and +0.0 dequantize to the same float value, but different MXFP4
51 + quantization backends disagree on which one they emit for values that
52 + round to zero magnitude. This normalizes both packed representations so
53 + tensors from different backends can be compared with `torch.equal`.
54 + """
55 + assert tensor.dtype == torch.uint8, (
56 + f"Expected a torch.uint8 tensor, got {tensor.dtype}"
57 + )
58 +
59 + low_nibble = tensor & 0x0F
60 + high_nibble = tensor & 0xF0
61 +
62 + low_nibble = torch.where(
63 + low_nibble == 0x08, torch.zeros_like(low_nibble), low_nibble
64 + )
65 + high_nibble = torch.where(
66 + high_nibble == 0x80, torch.zeros_like(high_nibble), high_nibble
67 + )
68 +
69 + return low_nibble | high_nibble
70 +
71 +
72 +def assert_quantized_weights_equal(
73 + checkpoint_layer: torch.nn.Module,
74 + online_layer: torch.nn.Module,
75 + keys: tuple[str, ...],
76 + packed_weight_keys: tuple[str, ...],
77 +) -> None:
78 + """Assert the checkpoint and online paths produced byte-identical weights.
79 +
80 + `packed_weight_keys` lists the entries holding packed FP4 codewords, which
81 + need negative-zero normalization before comparison.
82 + """
83 + for key in keys:
84 + checkpoint_tensor = getattr(checkpoint_layer, key).view(torch.uint8)
85 + online_tensor = getattr(online_layer, key).view(torch.uint8)
86 +
87 + # NOTE: AMD Quark checkpoints use exclusively **positive** zeros,
88 + # while other mxfp4_quantize implementations from mxfp4_utils.py may not.
89 + if key in packed_weight_keys:
90 + checkpoint_tensor = fix_negative_zeros(checkpoint_tensor)
91 + online_tensor = fix_negative_zeros(online_tensor)
92 +
93 + assert checkpoint_tensor.shape == online_tensor.shape
94 + assert checkpoint_tensor.dtype == online_tensor.dtype
95 + num_mismatched = (checkpoint_tensor != online_tensor).sum().item()
96 + total = checkpoint_tensor.numel()
97 +
98 + assert torch.equal(checkpoint_tensor, online_tensor), (
99 + f"{key}: {num_mismatched}/{total} "
100 + f"({100 * num_mismatched / total:.4f}%) mismatched bytes"
101 + )
102 +
103 +
104 +def _skip_reason_if_unavailable(backend: str, dtype: torch.dtype) -> str | None:
105 + """Return a skip reason if `backend` cannot run on the current host."""
106 + if backend == "triton":
107 + if not (current_platform.is_cuda() or current_platform.is_rocm()):
108 + return "Triton MXFP4 kernel requires a CUDA or ROCm GPU."
109 + return None
110 + if backend == "aiter":
111 + from vllm._aiter_ops import is_aiter_found_and_supported
112 +
113 + if not is_aiter_found_and_supported():
114 + return "aiter is not available/supported on this platform."
115 + if dtype != torch.bfloat16:
116 + return "aiter's dynamic_mxfp4_quant only supports bfloat16 input."
117 + return None
118 + if backend == "xpu":
119 + if not current_platform.is_xpu():
120 + return "not on XPU platform."
121 + return None
122 + if backend == "quark":
123 + try:
124 + import quark.torch.kernel.mx # noqa: F401
125 + except ImportError:
126 + return "amd-quark is not installed."
127 + return None
128 + raise ValueError(f"Unknown backend {backend}")
129 +
130 +
131 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
132 +@pytest.mark.parametrize("backend", ["triton", "aiter", "xpu", "quark"])
133 +def test_mxfp4_quantization_correctness(backend: str, dtype: torch.dtype):
134 + """
135 + Tests that the different implementations of mxfp4_quantize
136 + in mxfp4_utils.py all match.
137 + """
138 + skip_reason = _skip_reason_if_unavailable(backend, dtype)
139 + if skip_reason is not None:
140 + pytest.skip(skip_reason)
141 +
142 + torch.manual_seed(3)
143 +
144 + num_rows = 64
145 + hidden_size = 32 * 32 # multiple 32-element MXFP4 blocks
146 + device = current_platform.device_type
147 +
148 + x = (torch.rand(num_rows, hidden_size, dtype=dtype, device=device) - 0.5) * 2
149 + # Vary the magnitude block-to-block so several scale exponents are
150 + # exercised, rather than a single one for the whole tensor.
151 + scalings = [2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4]
152 + for i in range(hidden_size // 32):
153 + x[:, i * 32 : (i + 1) * 32] *= scalings[i % len(scalings)]
154 +
155 + if backend == "triton":
156 + x_fp4, x_scale, _ = downcast_to_mxfp(x, axis=-1)
157 + elif backend == "aiter":
158 + x_fp4, x_scale = quark_quantize_weight_to_mxfp4(x)
159 + elif backend == "xpu":
160 + # TODO: enable this test on XPU
161 + pytest.skip("xpu mxfp4 quantization match to reference is untested")
162 + # from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
163 + # xpu_mxfp4_quantize,
164 + # )
165 +
166 + # x_fp4, x_scale = xpu_mxfp4_quantize(x)
167 + elif backend == "quark":
168 + result = quant_dequant_mxfp4(x, scale_calculation_mode="even")
169 + else:
170 + raise ValueError(f"Unknown backend {backend}")
171 +
172 + if backend != "quark":
173 + result = dq_mxfp4_torch(x_fp4, x_scale, x.dtype)
174 + reference = qdq_mxfp4_torch(x, scale_calculation_mode="even")
175 +
176 + assert torch.equal(result, reference)
177 +
178 +
179 +@pytest.mark.skipif(
180 + not current_platform.is_cuda_alike(), reason="Only tested on ROCm/CUDA."
181 +)
182 +@pytest.mark.parametrize("tp_size", [2, 4, 8])
183 +def test_online_mxfp4_tp_weight_quant_matches_unsharded(tp_size: int):
184 + """TP-aligned MXFP4 shards match slices of unsharded quantization."""
185 + torch.manual_seed(3)
186 +
187 + device = current_platform.device_type
188 + hidden_size = 1024
189 + intermediate_size = 512
190 +
191 + dense_weight = torch.randn(
192 + hidden_size, intermediate_size, dtype=torch.bfloat16, device=device
193 + )
194 + w13_weight = torch.randn(
195 + 2,
196 + 2 * intermediate_size,
197 + hidden_size,
198 + dtype=torch.bfloat16,
199 + device=device,
200 + )
201 + w2_weight = torch.randn(
202 + 2,
203 + hidden_size,
204 + intermediate_size,
205 + dtype=torch.bfloat16,
206 + device=device,
207 + )
208 +
209 + dense_quantized, dense_scales = mxfp4_quantize(dense_weight)
210 + w13_quantized, w13_scales = mxfp4_quantize(w13_weight)
211 + w2_quantized, w2_scales = mxfp4_quantize(w2_weight)
212 +
213 + for tp_rank in range(tp_size):
214 + output_per_partition = hidden_size // tp_size
215 + input_per_partition = intermediate_size // tp_size
216 + output_start = tp_rank * output_per_partition
217 + input_start = tp_rank * input_per_partition
218 +
219 + column_quantized, column_scales = mxfp4_quantize(
220 + dense_weight.narrow(0, output_start, output_per_partition)
221 + )
222 + assert torch.equal(
223 + column_quantized,
224 + dense_quantized.narrow(0, output_start, output_per_partition),
225 + )
226 + assert torch.equal(
227 + column_scales,
228 + dense_scales.narrow(0, output_start, output_per_partition),
229 + )
230 +
231 + row_quantized, row_scales = mxfp4_quantize(
232 + dense_weight.narrow(1, input_start, input_per_partition)
233 + )
234 + assert torch.equal(
235 + row_quantized,
236 + dense_quantized.narrow(1, input_start // 2, input_per_partition // 2),
237 + )
238 + assert torch.equal(
239 + row_scales,
240 + dense_scales.narrow(1, input_start // 32, input_per_partition // 32),
241 + )
242 +
243 + w13_quantized_shard, w13_scales_shard = mxfp4_quantize(
244 + w13_weight.narrow(1, 2 * input_start, 2 * input_per_partition)
245 + )
246 + assert torch.equal(
247 + w13_quantized_shard,
248 + w13_quantized.narrow(1, 2 * input_start, 2 * input_per_partition),
249 + )
250 + assert torch.equal(
251 + w13_scales_shard,
252 + w13_scales.narrow(1, 2 * input_start, 2 * input_per_partition),
253 + )
254 +
255 + w2_quantized_shard, w2_scales_shard = mxfp4_quantize(
256 + w2_weight.narrow(2, input_start, input_per_partition)
257 + )
258 + assert torch.equal(
259 + w2_quantized_shard,
260 + w2_quantized.narrow(2, input_start // 2, input_per_partition // 2),
261 + )
262 + assert torch.equal(
263 + w2_scales_shard,
264 + w2_scales.narrow(2, input_start // 32, input_per_partition // 32),
265 + )
266 +
267 +
268 +@pytest.mark.skipif(
269 + not current_platform.is_cuda_alike(), reason="Only tested on ROCm/CUDA."
270 +)
271 +@pytest.mark.parametrize("moe_backend", ["aiter", "emulation"])
272 +@pytest.mark.parametrize(
273 + "unpadded_hidden_size,unpadded_intermediate_size",
274 + [
275 + pytest.param(256, 256, id="no_padding"),
276 + # Not multiples of the MXFP4 block size (32), so that every backend --
277 + # including emulation, whose round-up granularity is 32 -- has to pad.
278 + pytest.param(240, 208, id="padding"),
279 + ],
280 +)
281 +def test_online_mxfp4_moe_matches_quark(
282 + moe_backend: str,
283 + unpadded_hidden_size: int,
284 + unpadded_intermediate_size: int,
285 + default_vllm_config,
286 + dist_init,
287 + monkeypatch,
288 +):
289 + """
290 + Ensures `Mxfp4OnlineMoEMethod` (online quantization)
291 + and `QuarkOCP_MX_MoEMethod` (AMD Quark checkpoints) produce the same weights,
292 + with same MOE backend used.
293 +
294 + The `padding` case additionally covers `maybe_roundup_sizes`: the online
295 + path allocates its bf16 weights with `torch.empty_strided` and the loader
296 + only writes the unpadded slice, so the padding is filled with NaN here to
297 + reproduce that uninitialized state. It must be zeroed before quantization
298 + for the two paths to agree, since Quark allocates its padding with
299 + `torch.zeros`.
300 + """
301 + if moe_backend == "aiter":
302 + from vllm._aiter_ops import rocm_aiter_ops
303 +
304 + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
305 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
306 + monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True)
307 + rocm_aiter_ops.refresh_env_variables()
308 +
309 + if current_platform.is_cuda():
310 + pytest.skip(
311 + "mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4 requires "
312 + "rocm_aiter_ops in weight conversion, not compatible on cuda"
313 + )
314 + elif moe_backend == "emulation":
315 + # `OCP_MXQuantizationEmulationTritonExperts.is_supported_config` gates
316 + # on `has_quark()`, but its weight processing does not require it.
317 + monkeypatch.setattr(
318 + "vllm.model_executor.layers.fused_moe.experts."
319 + "ocp_mx_emulation_moe.has_quark",
320 + lambda: True,
321 + )
322 +
323 + default_vllm_config.model_config = ModelConfig()
324 +
325 + num_experts = 4
326 + device = current_platform.device_type
327 +
328 + def make_layer(prefix: str) -> RoutedExperts:
329 + runner = FusedMoEFactory(
330 + num_experts=num_experts,
331 + top_k=2,
332 + hidden_size=unpadded_hidden_size,
333 + intermediate_size=unpadded_intermediate_size,
334 + prefix=prefix,
335 + )
336 + layer = runner.routed_experts
337 + layer.moe_config.moe_backend = moe_backend
338 + return layer
339 +
340 + # `create_weights` implementations use plain `torch.zeros`/`torch.randn`
341 + # with no explicit device, so without this context they would default
342 + # to CPU and diverge from the (GPU) tensors produced by
343 + # `mxfp4_quantize`/`replace_parameter`.
344 + with torch.device(device):
345 + checkpoint_layer = make_layer("checkpoint_layer")
346 + online_layer = make_layer("online_layer")
347 +
348 + checkpoint_method = QuarkOCP_MX_MoEMethod(
349 + weight_config={"qscheme": "per_group", "dtype": "fp4"},
350 + input_config={"dtype": "fp4", "is_dynamic": True},
351 + moe=checkpoint_layer.moe_config,
352 + )
353 + online_method = Mxfp4OnlineMoEMethod(layer=online_layer)
354 +
355 + # `RoutedExperts.__init__` applies this round-up in production; these
356 + # layers are built without a quant config, so it is applied explicitly.
357 + # Both methods must agree on the padded sizes.
358 + roundup_arguments = dict(
359 + hidden_size=unpadded_hidden_size,
360 + intermediate_size_per_partition=unpadded_intermediate_size,
361 + act_dtype=torch.bfloat16,
362 + moe_parallel_config=online_layer.moe_config.moe_parallel_config,
363 + )
364 + hidden_size, intermediate_size = online_method.maybe_roundup_sizes(
365 + **roundup_arguments
366 + )
367 + assert (
368 + hidden_size,
369 + intermediate_size,
370 + ) == checkpoint_method.maybe_roundup_sizes(**roundup_arguments)
371 +
372 + if unpadded_hidden_size == 240:
373 + # Padded case.
374 + assert hidden_size > unpadded_hidden_size
375 + assert intermediate_size > unpadded_intermediate_size
376 +
377 + for layer in (checkpoint_layer, online_layer):
378 + assert layer.moe_config.hidden_dim_unpadded == unpadded_hidden_size
379 + assert (
380 + layer.moe_config.intermediate_size_per_partition_unpadded
381 + == unpadded_intermediate_size
382 + )
383 + layer.moe_config.hidden_dim = hidden_size
384 + layer.moe_config.intermediate_size_per_partition = intermediate_size
385 +
386 + gate_up_weight = torch.randn(
387 + num_experts,
388 + 2 * unpadded_intermediate_size,
389 + unpadded_hidden_size,
390 + dtype=torch.bfloat16,
391 + )
392 + down_weight = torch.randn(
393 + num_experts,
394 + unpadded_hidden_size,
395 + unpadded_intermediate_size,
396 + dtype=torch.bfloat16,
397 + )
398 +
399 + scalings = [2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4]
400 + for i, start in enumerate(range(0, unpadded_hidden_size, 32)):
401 + gate_up_weight[..., start : start + 32] *= scalings[i % len(scalings)]
402 + for i, start in enumerate(range(0, unpadded_intermediate_size, 32)):
403 + down_weight[..., start : start + 32] *= scalings[i % len(scalings)]
404 +
405 + def scatter_into_padded(
406 + gate_up_weight: torch.Tensor,
407 + down_weight: torch.Tensor,
408 + padding_value: float,
409 + ) -> tuple[torch.Tensor, torch.Tensor]:
410 + """Place the source weights into buffers of the padded size, the
411 + way the weight loader writes only the unpadded slice of a larger
412 + allocation."""
413 + w13 = torch.full(
414 + (num_experts, 2 * intermediate_size, hidden_size),
415 + padding_value,
416 + dtype=torch.bfloat16,
417 + )
418 + w2 = torch.full(
419 + (num_experts, hidden_size, intermediate_size),
420 + padding_value,
421 + dtype=torch.bfloat16,
422 + )
423 + # w13 stacks gate and up, each padded to `intermediate_size`.
424 + w13[:, :unpadded_intermediate_size, :unpadded_hidden_size] = gate_up_weight[
425 + :, :unpadded_intermediate_size, :
426 + ]
427 + w13[
428 + :,
429 + intermediate_size : intermediate_size + unpadded_intermediate_size,
430 + :unpadded_hidden_size,
431 + ] = gate_up_weight[:, unpadded_intermediate_size:, :]
432 + w2[:, :unpadded_hidden_size, :unpadded_intermediate_size] = down_weight
433 + return w13, w2
434 +
435 + # Checkpoint path: pre-quantize the source weights to MXFP4 with the
436 + # ~same RTN recipe a real Quark checkpoint, then feed the
437 + # already-packed tensors into `QuarkOCP_MX_MoEMethod`. Quark's
438 + # `create_weights` allocates with `torch.zeros`, so its padding is zero.
439 + checkpoint_w13, checkpoint_w2 = scatter_into_padded(
440 + gate_up_weight, down_weight, padding_value=0.0
441 + )
442 + checkpoint_w13, checkpoint_w13_scale = mxfp4_quantize(checkpoint_w13)
443 + checkpoint_w2, checkpoint_w2_scale = mxfp4_quantize(checkpoint_w2)
444 +
445 + checkpoint_method.create_weights(
446 + layer=checkpoint_layer,
447 + num_experts=num_experts,
448 + hidden_size=hidden_size,
449 + intermediate_size_per_partition=intermediate_size,
450 + params_dtype=torch.bfloat16,
451 + )
452 + checkpoint_layer.w13_weight.data.copy_(checkpoint_w13)
453 + checkpoint_layer.w2_weight.data.copy_(checkpoint_w2)
454 + checkpoint_layer.w13_weight_scale.data.copy_(checkpoint_w13_scale)
455 + checkpoint_layer.w2_weight_scale.data.copy_(checkpoint_w2_scale)
456 + checkpoint_method.process_weights_after_loading(checkpoint_layer)
457 +
458 + # Online path: feed the *same* raw bf16 source weights and let
459 + # `Mxfp4OnlineMoEMethod` quantize them during
460 + # `process_weights_after_loading`. The padding starts out
461 + # uninitialized, standing in for the materialized meta tensor.
462 + online_w13, online_w2 = scatter_into_padded(
463 + gate_up_weight, down_weight, padding_value=float("nan")
464 + )
465 +
466 + online_method.create_weights(
467 + layer=online_layer,
468 + num_experts=num_experts,
469 + hidden_size=hidden_size,
470 + intermediate_size_per_partition=intermediate_size,
471 + params_dtype=torch.bfloat16,
472 + )
473 + replace_parameter(online_layer, "w13_weight", online_w13)
474 + replace_parameter(online_layer, "w2_weight", online_w2)
475 + online_method.process_weights_after_loading(online_layer)
476 +
477 + assert checkpoint_method.mxfp4_backend == online_method.mxfp4_backend
478 +
479 + assert_quantized_weights_equal(
480 + checkpoint_layer,
481 + online_layer,
482 + keys=("w13_weight", "w13_weight_scale", "w2_weight", "w2_weight_scale"),
483 + packed_weight_keys=("w13_weight", "w2_weight"),
484 + )
485 +
486 +
487 +@pytest.mark.skipif(
488 + not current_platform.is_cuda_alike(), reason="Only tested on ROCm/CUDA."
489 +)
490 +@pytest.mark.parametrize("linear_backend", ["emulation", "aiter", "marlin"])
491 +def test_online_mxfp4_dense_matches_quark(
492 + linear_backend: str, default_vllm_config, dist_init, monkeypatch
493 +):
494 + """
495 + Ensures `Mxfp4OnlineLinearMethod` (online quantization)
496 + and `QuarkOCP_MX` (AMD Quark checkpoints) produce the same weights,
497 + with same linear backend used.
498 + """
499 + if linear_backend == "marlin" and not current_platform.is_cuda():
500 + # `MarlinMxFp4LinearKernel.process_weights_after_loading` is CUDA-only
501 + pytest.skip(
502 + "MarlinMxFp4LinearKernel.process_weights_after_loading is CUDA-only."
503 + )
504 +
505 + if linear_backend == "aiter":
506 + # `AiterMxfp4LinearKernel` is registered only for ROCm in
507 + # `_POSSIBLE_MXFP4_KERNELS`, and `is_supported` gates on
508 + # `current_platform.supports_mx()`. Force it onto the CUDA
509 + # kernel list and bypass the platform gate so this backend can
510 + # be exercised on CUDA hosts too.
511 +
512 + monkeypatch.setattr(
513 + AiterMxfp4LinearKernel,
514 + "is_supported",
515 + classmethod(lambda cls, compute_capability=None: (True, None)),
516 + )
517 + monkeypatch.setitem(
518 + _POSSIBLE_MXFP4_KERNELS,
519 + PlatformEnum.CUDA,
520 + [
521 + AiterMxfp4LinearKernel,
522 + *_POSSIBLE_MXFP4_KERNELS.get(PlatformEnum.CUDA, []),
523 + ],
524 + )
525 + elif linear_backend == "emulation":
526 + # `EmulationMxfp4LinearKernel.can_implement` gates on `has_quark()`,
527 + # EmulationMxfp4LinearKernel.process_weights_after_loading does not require it.
528 + monkeypatch.setattr(
529 + "vllm.model_executor.kernels.linear.mxfp4.emulation.has_quark",
530 + lambda: True,
531 + )
532 +
533 + default_vllm_config.model_config = ModelConfig()
534 + default_vllm_config.kernel_config.linear_backend = linear_backend
535 +
536 + input_size = 256
537 + output_size = 128
538 + device = current_platform.device_type
539 +
540 + # `create_weights` implementations use plain `torch.empty` with no explicit
541 + # device, so without this context they would default to CPU and diverge
542 + # from the (GPU) tensors produced by `mxfp4_quantize`/`replace_parameter`.
543 + with torch.device(device):
544 + checkpoint_layer = torch.nn.Module()
545 + online_layer = torch.nn.Module()
546 + for layer in (checkpoint_layer, online_layer):
547 + # `LinearBase` subclasses normally set these attributes on the
548 + # layer before `create_weights` is called; the marlin kernel's
549 + # `process_weights_after_loading` reads them directly off the
550 + # layer, so a bare `torch.nn.Module` needs them set explicitly.
551 + layer.input_size = input_size
552 + layer.output_size = output_size
553 + layer.input_size_per_partition = input_size
554 + layer.output_size_per_partition = output_size
555 + layer.params_dtype = torch.bfloat16
556 +
557 + weight = torch.randn(output_size, input_size, dtype=torch.bfloat16)
558 +
559 + scalings = [2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4]
560 + for i in range(input_size // 32):
561 + weight[:, i * 32 : (i + 1) * 32] *= scalings[i % len(scalings)]
562 +
563 + # Checkpoint path: pre-quantize the source weight to MXFP4 with the
564 + # ~same RTN recipe as a real Quark checkpoint, then feed the
565 + # already-packed tensors into `QuarkOCP_MX`.
566 + checkpoint_weight, checkpoint_weight_scale = mxfp4_quantize(weight)
567 +
568 + checkpoint_scheme = QuarkOCP_MX(
569 + weight_quant_spec={"qscheme": "per_group", "dtype": "fp4"},
570 + input_quant_spec={"dtype": "fp4", "is_dynamic": True},
571 + )
572 + checkpoint_scheme.create_weights(
573 + layer=checkpoint_layer,
574 + output_partition_sizes=[output_size],
575 + input_size_per_partition=input_size,
576 + params_dtype=torch.bfloat16,
577 + weight_loader=default_weight_loader,
578 + )
579 + checkpoint_layer.weight.data.copy_(checkpoint_weight)
580 + checkpoint_layer.weight_scale.data.copy_(checkpoint_weight_scale)
581 + checkpoint_scheme.process_weights_after_loading(checkpoint_layer)
582 +
583 + # Online path: feed the *same* raw bf16 source weight and let
584 + # `Mxfp4OnlineLinearMethod` quantize it during
585 + # `process_weights_after_loading`.
586 + online_method = Mxfp4OnlineLinearMethod()
587 + online_method.create_weights(
588 + layer=online_layer,
589 + input_size_per_partition=input_size,
590 + output_partition_sizes=[output_size],
591 + input_size=input_size,
592 + output_size=output_size,
593 + params_dtype=torch.bfloat16,
594 + weight_loader=default_weight_loader,
595 + )
596 + replace_parameter(online_layer, "weight", weight.clone())
597 + online_method.process_weights_after_loading(online_layer)
598 +
599 + assert type(checkpoint_scheme.ocp_mx_linear) is type(online_method.kernel)
600 +
601 + assert_quantized_weights_equal(
602 + checkpoint_layer,
603 + online_layer,
604 + keys=("weight", "weight_scale"),
605 + packed_weight_keys=("weight",),
606 + )
vllm/config/quantization.py
+5
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
17 kFp8StaticTensorSym,
18 kInt8StaticChannelSym,
19 kMxfp4Dynamic,
20 + kMxfp4Static,
21 kMxfp8Dynamic,
22 kNvfp4Static,
23 )
@@ -131,6 +132,10 @@ _ONLINE_SHORTHANDS: dict[str, QuantizationConfigArgs] = {
132 linear=QuantSpec(weight=kMxfp8Dynamic),
133 moe=QuantSpec(weight=kMxfp8Dynamic),
134 ),
135 + "mxfp4": QuantizationConfigArgs(
136 + linear=QuantSpec(weight=kMxfp4Static),
137 + moe=QuantSpec(weight=kMxfp4Static),
138 + ),
139 # INT8 weight-only on MoE; linear stays unquantized (no `linear` field).
140 "int8_per_channel_weight_only": QuantizationConfigArgs(
141 moe=QuantSpec(weight=kInt8StaticChannelSym),
vllm/model_executor/kernels/linear/mxfp4/emulation.py
+4
@@ -23,6 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
23 kMxfp6E3M2Dynamic,
24 )
25 from vllm.platforms import current_platform
26 +from vllm.utils.import_utils import has_quark
27
28 from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
29
@@ -68,6 +69,9 @@ class EmulationMxfp4LinearKernel(MxFp4LinearKernel):
69 ):
70 return False, "only supports MXFP4 or MXFP6 or unquantized activations"
71
72 + if not has_quark():
73 + return False, "amd-quark package not available"
74 +
75 if (
76 current_platform.is_rocm()
77 and current_platform.supports_mx()
vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py
+23
@@ -26,7 +26,11 @@ from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mx
26 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
27 OCP_MX_Scheme,
28 )
29 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
30 + QuantKey,
31 +)
32 from vllm.platforms import current_platform
33 +from vllm.utils.import_utils import has_quark
34
35 logger = init_logger(__name__)
36
@@ -92,6 +96,25 @@ class OCP_MXQuantizationEmulationTritonExperts(TritonExperts):
96 ]:
97 self._quant_dtype = current_platform.fp8_dtype()
98
99 + @staticmethod
100 + def is_supported_config(
101 + cls: type["mk.FusedMoEExperts"],
102 + moe_config: FusedMoEConfig,
103 + weight_key: QuantKey | None,
104 + activation_key: QuantKey | None,
105 + activation_format: mk.FusedMoEActivationFormat,
106 + ) -> tuple[bool, str | None]:
107 + if not has_quark():
108 + return False, "kernel requires amd-quark package"
109 +
110 + return TritonExperts.is_supported_config(
111 + cls=cls,
112 + moe_config=moe_config,
113 + weight_key=weight_key,
114 + activation_key=activation_key,
115 + activation_format=activation_format,
116 + )
117 +
118 @property
119 def quant_dtype(self) -> torch.dtype | str | None:
120 return self._quant_dtype
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+15
@@ -1650,6 +1650,21 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1650 w13_bias,
1651 w2_bias,
1652 )
1653 + elif mxfp4_backend in (
1654 + Mxfp4MoeBackend.AITER_MXFP4_MXFP4,
1655 + Mxfp4MoeBackend.AITER_MXFP4_FP8,
1656 + ):
1657 + return convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
1658 + mxfp4_backend=mxfp4_backend,
1659 + layer=layer,
1660 + w13_weight=w13_weight,
1661 + w2_weight=w2_weight,
1662 + w13_weight_scale=w13_weight_scale,
1663 + w2_weight_scale=w2_weight_scale,
1664 + w13_bias=w13_bias,
1665 + w2_bias=w2_bias,
1666 + _cache_permute_indices=_cache_permute_indices,
1667 + )
1668 else:
1669 raise ValueError(
1670 f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. "
vllm/model_executor/layers/quantization/online/base.py
+7
@@ -36,6 +36,10 @@ from vllm.model_executor.layers.quantization.online.fp8 import (
36 from vllm.model_executor.layers.quantization.online.int8 import (
37 Int8OnlineMoEMethod,
38 )
39 +from vllm.model_executor.layers.quantization.online.mxfp4 import (
40 + Mxfp4OnlineLinearMethod,
41 + Mxfp4OnlineMoEMethod,
42 +)
43 from vllm.model_executor.layers.quantization.online.mxfp8 import (
44 Mxfp8OnlineLinearMethod,
45 Mxfp8OnlineMoEMethod,
@@ -49,6 +53,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
53 kFp8StaticChannelSym,
54 kFp8StaticTensorSym,
55 kInt8StaticChannelSym,
56 + kMxfp4Static,
57 kMxfp8Dynamic,
58 kNvfp4Static,
59 )
@@ -64,6 +69,7 @@ _ONLINE_LINEAR_METHODS: dict[QuantKey, type] = {
69 kFp8Static128BlockSym: Fp8PerBlockOnlineLinearMethod,
70 kFp8StaticChannelSym: Fp8PtpcOnlineLinearMethod,
71 kMxfp8Dynamic: Mxfp8OnlineLinearMethod,
72 + kMxfp4Static: Mxfp4OnlineLinearMethod,
73 }
74
75 _ONLINE_MOE_METHODS: dict[QuantKey, type] = {
@@ -71,6 +77,7 @@ _ONLINE_MOE_METHODS: dict[QuantKey, type] = {
77 kFp8Static128BlockSym: Fp8PerBlockOnlineMoEMethod,
78 kFp8StaticChannelSym: Fp8PtpcOnlineMoEMethod,
79 kMxfp8Dynamic: Mxfp8OnlineMoEMethod,
80 + kMxfp4Static: Mxfp4OnlineMoEMethod,
81 kInt8StaticChannelSym: Int8OnlineMoEMethod,
82 kNvfp4Static: Nvfp4OnlineMoEMethod,
83 }
vllm/model_executor/layers/quantization/online/fp8.py
-30
@@ -639,36 +639,6 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
639 round_up(intermediate_size_per_partition, block_size),
640 )
641
642 - def _zero_padding(self, layer: Module) -> None:
643 - hidden_size = layer.moe_config.hidden_dim_unpadded
644 - intermediate_size = layer.moe_config.intermediate_size_per_partition_unpadded
645 -
646 - w13_shard = layer.w13_weight.shape[1] // self.moe.w13_num_shards
647 - if w13_shard > intermediate_size:
648 - for shard in range(self.moe.w13_num_shards):
649 - start = shard * w13_shard + intermediate_size
650 - layer.w13_weight[:, start : (shard + 1) * w13_shard, :] = 0
651 - if layer.w13_weight.shape[2] > hidden_size:
652 - layer.w13_weight[:, :, hidden_size:] = 0
653 -
654 - if layer.w2_weight.shape[1] > hidden_size:
655 - layer.w2_weight[:, hidden_size:, :] = 0
656 - if layer.w2_weight.shape[2] > intermediate_size:
657 - layer.w2_weight[:, :, intermediate_size:] = 0
658 -
659 - if getattr(layer, "w13_bias", None) is not None:
660 - w13_bias_shard = layer.w13_bias.shape[1] // self.moe.w13_num_shards
661 - if w13_bias_shard > intermediate_size:
662 - for shard in range(self.moe.w13_num_shards):
663 - start = shard * w13_bias_shard + intermediate_size
664 - layer.w13_bias[:, start : (shard + 1) * w13_bias_shard] = 0
665 -
666 - if (
667 - getattr(layer, "w2_bias", None) is not None
668 - and layer.w2_bias.shape[1] > hidden_size
669 - ):
670 - layer.w2_bias[:, hidden_size:] = 0
671 -
642 def process_weights_after_loading(self, layer: Module) -> None:
643 if getattr(layer, "_already_called_process_weights_after_loading", False):
644 return
vllm/model_executor/layers/quantization/online/moe_base.py
+30
@@ -95,6 +95,36 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):
95
96 initialize_online_processing(layer)
97
98 + def _zero_padding(self, layer: torch.nn.Module) -> None:
99 + hidden_size = layer.moe_config.hidden_dim_unpadded
100 + intermediate_size = layer.moe_config.intermediate_size_per_partition_unpadded
101 +
102 + w13_shard = layer.w13_weight.shape[1] // self.moe.w13_num_shards
103 + if w13_shard > intermediate_size:
104 + for shard in range(self.moe.w13_num_shards):
105 + start = shard * w13_shard + intermediate_size
106 + layer.w13_weight[:, start : (shard + 1) * w13_shard, :] = 0
107 + if layer.w13_weight.shape[2] > hidden_size:
108 + layer.w13_weight[:, :, hidden_size:] = 0
109 +
110 + if layer.w2_weight.shape[1] > hidden_size:
111 + layer.w2_weight[:, hidden_size:, :] = 0
112 + if layer.w2_weight.shape[2] > intermediate_size:
113 + layer.w2_weight[:, :, intermediate_size:] = 0
114 +
115 + if getattr(layer, "w13_bias", None) is not None:
116 + w13_bias_shard = layer.w13_bias.shape[1] // self.moe.w13_num_shards
117 + if w13_bias_shard > intermediate_size:
118 + for shard in range(self.moe.w13_num_shards):
119 + start = shard * w13_bias_shard + intermediate_size
120 + layer.w13_bias[:, start : (shard + 1) * w13_bias_shard] = 0
121 +
122 + if (
123 + getattr(layer, "w2_bias", None) is not None
124 + and layer.w2_bias.shape[1] > hidden_size
125 + ):
126 + layer.w2_bias[:, hidden_size:] = 0
127 +
128 @abstractmethod
129 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
130 pass
vllm/model_executor/layers/quantization/online/mxfp4.py new
+317
@@ -0,0 +1,317 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Online MXFP4 (microscaling FP4, block-32) quantization methods."""
5 +
6 +from typing import TYPE_CHECKING
7 +
8 +import torch
9 +from torch.nn import Module
10 +
11 +if TYPE_CHECKING:
12 + import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13 + from vllm.model_executor.layers.fused_moe import (
14 + FusedMoEParallelConfig,
15 + FusedMoEQuantConfig,
16 + RoutedExperts,
17 + )
18 +
19 +from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel
20 +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
21 + TRITON_BACKENDS,
22 + Mxfp4MoeBackend,
23 + convert_weight_to_mxfp4_moe_kernel_format,
24 + make_mxfp4_moe_kernel,
25 + make_mxfp4_moe_quant_config,
26 + mxfp4_round_up_hidden_size_and_intermediate_size,
27 + select_mxfp4_moe_backend,
28 +)
29 +from vllm.model_executor.layers.quantization.online.fp8 import (
30 + _Fp8OnlineLinearBase,
31 +)
32 +from vllm.model_executor.layers.quantization.online.moe_base import (
33 + OnlineMoEMethodBase,
34 +)
35 +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
36 + mxfp4_quantize,
37 +)
38 +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic
39 +from vllm.model_executor.utils import replace_parameter
40 +
41 +# MXFP4 block size (elements per shared e8m0 scale), same as the checkpoint
42 +# and MoE oracle paths.
43 +MXFP4_BLOCK_SIZE = 32
44 +
45 +
46 +def _quantize_mxfp4_moe_weight(
47 + weight: torch.Tensor,
48 +) -> tuple[torch.Tensor, torch.Tensor]:
49 + """Batch quantization: bf16/fp16 stacked expert weights -> MXFP4.
50 +
51 + Returns packed FP4 weights (uint8, two values per byte) and per-block
52 + (group-32) e8m0 scales (uint8), one pair per expert.
53 + """
54 + num_experts = weight.size(0)
55 + first_quant, first_scale = mxfp4_quantize(weight[0])
56 + # Pre-allocate the output tensors rather than stacking.
57 + # This is important for consistent memory layout.
58 + w_quant = torch.empty(
59 + (num_experts, *first_quant.shape),
60 + dtype=first_quant.dtype,
61 + device=weight.device,
62 + )
63 + w_scales = torch.empty(
64 + (num_experts, *first_scale.shape),
65 + dtype=first_scale.dtype,
66 + device=weight.device,
67 + )
68 + w_quant[0] = first_quant
69 + w_scales[0] = first_scale
70 + for i in range(1, num_experts):
71 + w_quant[i], w_scales[i] = mxfp4_quantize(weight[i])
72 +
73 + return w_quant, w_scales
74 +
75 +
76 +class Mxfp4OnlineLinearMethod(_Fp8OnlineLinearBase):
77 + """Online MXFP4 linear method.
78 + Loads bf16/fp16 checkpoints and quantizes weights to MXFP4 (microscaling
79 + FP4 with block-32 scales) during weight loading.
80 + """
81 +
82 + activation_quant_key = kMxfp4Dynamic
83 +
84 + def __init__(self):
85 + super().__init__()
86 + self.kernel = init_mxfp4_linear_kernel(
87 + activation_quant_key=self.activation_quant_key
88 + )
89 +
90 + def create_weights(
91 + self,
92 + layer: torch.nn.Module,
93 + input_size_per_partition: int,
94 + output_partition_sizes: list[int],
95 + input_size: int,
96 + output_size: int,
97 + params_dtype: torch.dtype,
98 + **extra_weight_attrs,
99 + ):
100 + if input_size_per_partition % MXFP4_BLOCK_SIZE != 0:
101 + raise ValueError(
102 + f"MXFP4 requires input_size_per_partition "
103 + f"({input_size_per_partition}) to be divisible by "
104 + f"{MXFP4_BLOCK_SIZE}."
105 + )
106 +
107 + super().create_weights(
108 + layer,
109 + input_size_per_partition,
110 + output_partition_sizes,
111 + input_size,
112 + output_size,
113 + params_dtype,
114 + **extra_weight_attrs,
115 + )
116 +
117 + def process_weights_after_loading(self, layer: Module) -> None:
118 + if getattr(layer, "_already_called_process_weights_after_loading", False):
119 + return
120 +
121 + weight_fp4, weight_scale = mxfp4_quantize(layer.weight.contiguous())
122 +
123 + layer.input_scale = None
124 + replace_parameter(layer, "weight", weight_fp4.data)
125 + replace_parameter(layer, "weight_scale", weight_scale.data)
126 +
127 + self.kernel.process_weights_after_loading(layer)
128 +
129 + layer._already_called_process_weights_after_loading = True
130 +
131 + def apply(
132 + self,
133 + layer: torch.nn.Module,
134 + x: torch.Tensor,
135 + bias: torch.Tensor | None = None,
136 + ) -> torch.Tensor:
137 + return self.kernel.apply_weights(layer, x, bias)
138 +
139 +
140 +class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase):
141 + """MoE method for online MXFP4 (block) quantization."""
142 +
143 + mxfp4_backend: Mxfp4MoeBackend
144 + experts_cls: "type[mk.FusedMoEExperts] | None"
145 + activation_quant_key = kMxfp4Dynamic
146 +
147 + def __init__(self, *, layer: torch.nn.Module):
148 + super().__init__(layer.moe_config)
149 + self.weight_block_size: list[int] = [1, MXFP4_BLOCK_SIZE]
150 + self.weight_scale_name = "weight_scale"
151 +
152 + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(
153 + config=self.moe, activation_key=self.activation_quant_key
154 + )
155 +
156 + def maybe_roundup_sizes(
157 + self,
158 + hidden_size: int,
159 + intermediate_size_per_partition: int,
160 + act_dtype: torch.dtype,
161 + moe_parallel_config: "FusedMoEParallelConfig",
162 + ) -> tuple[int, int]:
163 + hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(
164 + hidden_size=hidden_size,
165 + intermediate_size_per_partition=intermediate_size_per_partition,
166 + act_dtype=act_dtype,
167 + moe_parallel_config=moe_parallel_config,
168 + )
169 + return mxfp4_round_up_hidden_size_and_intermediate_size(
170 + self.mxfp4_backend, hidden_size, intermediate_size_per_partition
171 + )
172 +
173 + def create_weights(
174 + self,
175 + layer: Module,
176 + num_experts: int,
177 + hidden_size: int,
178 + intermediate_size_per_partition: int,
179 + params_dtype: torch.dtype,
180 + **extra_weight_attrs,
181 + ):
182 + if (
183 + hidden_size % MXFP4_BLOCK_SIZE != 0
184 + or intermediate_size_per_partition % MXFP4_BLOCK_SIZE != 0
185 + ):
186 + raise ValueError(
187 + "Online MXFP4 MoE requires hidden/intermediate sizes divisible "
188 + f"by {MXFP4_BLOCK_SIZE}."
189 + )
190 +
191 + super().create_weights(
192 + layer=layer,
193 + num_experts=num_experts,
194 + hidden_size=hidden_size,
195 + intermediate_size_per_partition=intermediate_size_per_partition,
196 + params_dtype=params_dtype,
197 + **extra_weight_attrs,
198 + )
199 +
200 + layer.weight_block_size = [1, MXFP4_BLOCK_SIZE]
201 +
202 + def _setup_kernel(
203 + self,
204 + layer: "RoutedExperts",
205 + w13: torch.Tensor,
206 + w2: torch.Tensor,
207 + w13_scale: torch.Tensor,
208 + w2_scale: torch.Tensor,
209 + w13_bias: torch.Tensor | None,
210 + w2_bias: torch.Tensor | None,
211 + ) -> None:
212 + w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (
213 + convert_weight_to_mxfp4_moe_kernel_format(
214 + mxfp4_backend=self.mxfp4_backend,
215 + layer=layer,
216 + w13_weight=w13,
217 + w2_weight=w2,
218 + w13_weight_scale=w13_scale,
219 + w2_weight_scale=w2_scale,
220 + w13_bias=w13_bias,
221 + w2_bias=w2_bias,
222 + )
223 + )
224 +
225 + # Handle weight/scale assignment based on backend type
226 + # Copied from `quark_moe.py`.
227 + # TODO: This should not be here, replace_parameter should handle
228 + # triton_kernels.tensor.Tensor.
229 + if self.mxfp4_backend in TRITON_BACKENDS or self.mxfp4_backend in (
230 + Mxfp4MoeBackend.AITER_MXFP4_FP8,
231 + ):
232 + # Triton-based backends: w13/w2 are triton_kernels.tensor.Tensor
233 + # Store on layer for apply(), scales are PrecisionConfig
234 + layer.w13_weight = w13
235 + layer.w2_weight = w2
236 + self.w13_precision_config = w13_scale
237 + self.w2_precision_config = w2_scale
238 + else:
239 + # Standard backends: replace parameters
240 + replace_parameter(layer, "w13_weight", w13)
241 + replace_parameter(layer, "w2_weight", w2)
242 + replace_parameter(layer, "w13_weight_scale", w13_scale)
243 + replace_parameter(layer, "w2_weight_scale", w2_scale)
244 +
245 + if w13_bias is not None and w2_bias is not None:
246 + replace_parameter(layer, "w13_bias", w13_bias)
247 + replace_parameter(layer, "w2_bias", w2_bias)
248 +
249 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
250 + if self.moe_quant_config is not None and self.experts_cls is not None:
251 + self.moe_kernel = make_mxfp4_moe_kernel(
252 + moe_quant_config=self.moe_quant_config,
253 + moe_config=self.moe,
254 + mxfp4_backend=self.mxfp4_backend,
255 + experts_cls=self.experts_cls,
256 + routing_tables=layer._expert_routing_tables(),
257 + layer=layer,
258 + )
259 +
260 + def get_fused_moe_quant_config(
261 + self, layer: torch.nn.Module
262 + ) -> "FusedMoEQuantConfig | None":
263 + # NOTE: unlike `QuarkOCP_MX_MoEMethod.get_fused_moe_quant_config`, this
264 + # does not branch on `self.mxfp4_backend in TRITON_BACKENDS or
265 + # AITER_MXFP4_FP8` to read `self.w13_precision_config`/
266 + # `self.w2_precision_config` instead of `layer.w13_weight_scale`/
267 + # `layer.w2_weight_scale`. For those backends,
268 + # `convert_weight_to_mxfp4_moe_kernel_format` deletes the latter, so
269 + # `getattr` below would raise `AttributeError`. This is currently
270 + # unreachable because `activation_quant_key` is hardcoded to
271 + # `kMxfp4Dynamic`
272 + # TODO: When supporting online activation quant key override,
273 + # fix this convert_weight_to_mxfp4_moe_kernel_format issue.
274 + w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")
275 + w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
276 +
277 + return make_mxfp4_moe_quant_config(
278 + mxfp4_backend=self.mxfp4_backend,
279 + w1_scale=w1_scale,
280 + w2_scale=w2_scale,
281 + w1_bias=getattr(layer, "w13_bias", None),
282 + w2_bias=getattr(layer, "w2_bias", None),
283 + a1_scale=layer.w13_input_scale,
284 + a2_scale=layer.w2_input_scale,
285 + swiglu_limit=getattr(layer, "swiglu_limit", None),
286 + gemm1_alpha=getattr(layer, "swiglu_alpha", None),
287 + gemm1_beta=getattr(layer, "swiglu_beta", None),
288 + layer=layer,
289 + )
290 +
291 + def process_weights_after_loading(self, layer: Module) -> None:
292 + if getattr(layer, "_already_called_process_weights_after_loading", False):
293 + return
294 +
295 + self._zero_padding(layer)
296 +
297 + if self.mxfp4_backend == Mxfp4MoeBackend.NONE:
298 + layer._already_called_process_weights_after_loading = True
299 + return
300 +
301 + layer.w13_input_scale = None
302 + layer.w2_input_scale = None
303 +
304 + w13, w13_scale = _quantize_mxfp4_moe_weight(layer.w13_weight)
305 + w2, w2_scale = _quantize_mxfp4_moe_weight(layer.w2_weight)
306 +
307 + self._setup_kernel(
308 + layer,
309 + w13,
310 + w2,
311 + w13_scale,
312 + w2_scale,
313 + getattr(layer, "w13_bias", None),
314 + getattr(layer, "w2_bias", None),
315 + )
316 +
317 + layer._already_called_process_weights_after_loading = True
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+317
@@ -6,6 +6,7 @@ import torch
6
7 from vllm.logger import init_logger
8 from vllm.platforms import current_platform
9 +from vllm.triton_utils import tl, triton
10 from vllm.utils.import_utils import has_triton_kernels
11 from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
12
@@ -184,3 +185,319 @@ except AttributeError as error:
185
186 def xpu_mxfp4_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
187 return torch.ops.vllm.xpu_mxfp4_quantize(x)
188 +
189 +
190 +def permute_shape(shape: tuple[int, ...], permute_order: list[int]) -> tuple[int, ...]:
191 + return tuple(shape[i] for i in permute_order)
192 +
193 +
194 +@triton.jit
195 +def _compute_quant_and_scale(src_tensor, valid_src_mask):
196 + BLOCK_SIZE_OUT_DIM: tl.constexpr = src_tensor.shape[0]
197 + BLOCK_SIZE_QUANT_DIM: tl.constexpr = src_tensor.shape[1]
198 + BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = src_tensor.shape[1] // 32
199 +
200 + # Explicit cast to fp32 since most ops are not supported on bfloat16. We
201 + # avoid needless conversions to and from bf16
202 + f32_tensor = src_tensor.to(tl.float32)
203 + abs_tensor = tl.abs(f32_tensor)
204 + abs_tensor = tl.where(
205 + valid_src_mask, abs_tensor, -1.0
206 + ) # Don't consider padding tensors in scale computation
207 + abs_tensor = tl.reshape(
208 + abs_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 32]
209 + )
210 + max_val = tl.max(abs_tensor, axis=2, keep_dims=True)
211 + # Round-to-nearest-even block scale (OCP MX spec), matching hardware/aiter
212 + # MXFP4 quantizers: compute 2 ** (floor(log2(rounding(max_abs(v)))) - max_exp).
213 + max_val = max_val.to(tl.int32, bitcast=True)
214 + max_val = (max_val + 0x200000).to(tl.uint32, bitcast=True) & 0x7F800000
215 + max_val = max_val.to(tl.float32, bitcast=True)
216 + # Add epsilon to prevent log2(0) = -inf when max_val is zero (all-zero input block)
217 + eps = tl.where(max_val == 0.0, 2 ** (-126), 0.0)
218 + scale_e8m0_unbiased = tl.log2(max_val + eps).floor() - 2
219 + scale_e8m0_unbiased = tl.clamp(scale_e8m0_unbiased, min=-127, max=127)
220 + dequant_scale_rounded = tl.exp2(scale_e8m0_unbiased)
221 + dequant_scale_exponent = dequant_scale_rounded.to(tl.uint32, bitcast=True)
222 +
223 + dequant_scale_rounded = dequant_scale_exponent.to(tl.float32, bitcast=True)
224 + quant_scale = tl.where(dequant_scale_rounded == 0, 0, 1.0 / dequant_scale_rounded)
225 +
226 + f32_tensor = tl.reshape(
227 + f32_tensor, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE, 32]
228 + )
229 + quant_tensor = f32_tensor * quant_scale
230 +
231 + # Reshape the tensors after scaling
232 + quant_tensor = quant_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM])
233 + # Set the invalid portions of the tensor to 0. This will ensure that any
234 + # padding tensors are 0 in the mx format.
235 + quant_tensor = tl.where(valid_src_mask, quant_tensor, 0)
236 + dequant_scale_exponent = dequant_scale_exponent.reshape(
237 + [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE]
238 + )
239 +
240 + # First, we simply extract the exponent part of the scales and store the result
241 + dequant_scale_exponent = (dequant_scale_exponent >> 23).to(tl.uint8)
242 + # Now we must convert the tensor to packed e2m1 (mxfp4).
243 + quant_tensor = quant_tensor.to(tl.uint32, bitcast=True)
244 + signs = quant_tensor & 0x80000000
245 + exponents = (quant_tensor >> 23) & 0xFF
246 + mantissas = quant_tensor & 0x7FFFFF
247 +
248 + # 0.25 <= x < 0.75 maps to 0.5, a denormal number
249 + E8_BIAS = 127
250 + E2_BIAS = 1
251 + # Move implicit bit 1 at the beginning to mantissa for denormals. Bits
252 + # shifted out are OR-reduced into the LSB as a sticky bit, so the
253 + # round-to-even logic below can tell an exact tie from a value that is
254 + # merely close to one.
255 + adjusted_exponents = tl.core.sub(E8_BIAS, exponents + 1, sanitize_overflow=False)
256 + folded_mantissas = 0x400000 | (mantissas >> 1)
257 + denormal_mantissas = folded_mantissas >> adjusted_exponents
258 + denormal_mantissas |= (
259 + (denormal_mantissas << adjusted_exponents) != folded_mantissas
260 + ).to(tl.uint32)
261 + mantissas = tl.where(exponents < E8_BIAS, denormal_mantissas, mantissas)
262 +
263 + # For normal numbers, we change the bias from 127 to 1, and for subnormals,
264 + # we keep exponent as 0.
265 + exponents = tl.maximum(exponents, E8_BIAS - E2_BIAS) - (E8_BIAS - E2_BIAS)
266 +
267 + # Combine sign, exponent, and mantissa, rounding to nearest with ties
268 + # broken to even (OCP MX spec, matching hardware/aiter MXFP4
269 + # quantizers) via the standard "add a rounding bias, then truncate"
270 + # trick: the bias's LSB is the kept mantissa bit, so an exact tie
271 + # rounds to whichever candidate is even, while any lower (sticky) bit
272 + # being set always rounds up.
273 + rounding_bias = ((mantissas >> 22) & 1) + 0x1FFFFF
274 + rounded_mantissa = (mantissas + rounding_bias) >> 22 # in {0, 1, 2}
275 + e2m1_tmp = tl.minimum((exponents << 1) + rounded_mantissa, 0x7)
276 + e2m1_value = ((signs >> 28) | e2m1_tmp).to(tl.uint8)
277 +
278 + e2m1_value = tl.reshape(
279 + e2m1_value, [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_DIM // 2, 2]
280 + )
281 + evens, odds = tl.split(e2m1_value)
282 + out_tensor = evens | (odds << 4)
283 +
284 + return out_tensor, dequant_scale_exponent
285 +
286 +
287 +@triton.jit
288 +def _downcast_to_mxfp(
289 + mx_tensor_ptr,
290 + stride_mxt_outer,
291 + stride_mxt_quant: tl.constexpr,
292 + mx_scale_ptr,
293 + stride_mx_scale_outer,
294 + stride_mx_scale_quant,
295 + src_ptr,
296 + stride_src_outer,
297 + stride_src_quant,
298 + outer_dim,
299 + quant_dim,
300 + BLOCK_SIZE_OUT_DIM: tl.constexpr,
301 + BLOCK_SIZE_QUANT_DIM: tl.constexpr,
302 +):
303 + tl.static_assert(
304 + stride_mxt_quant == 1, f"Output stride, {stride_mxt_quant=} must be 1."
305 + )
306 + tl.static_assert(
307 + BLOCK_SIZE_QUANT_DIM % 32 == 0,
308 + f"{BLOCK_SIZE_QUANT_DIM=} must be a multiple of 32",
309 + )
310 +
311 + # uint8 signifies two fp4 e2m1 values packed into a single byte
312 + tl.static_assert(
313 + mx_tensor_ptr.dtype.element_ty == tl.uint8,
314 + f"Invalid {mx_tensor_ptr.dtype.element_ty=}. Must be uint8.",
315 + )
316 +
317 + src_dtype: tl.constexpr = src_ptr.dtype.element_ty
318 + tl.static_assert(
319 + mx_scale_ptr.dtype.element_ty == tl.uint8,
320 + f"{mx_scale_ptr.dtype.element_ty=} must be uint8",
321 + )
322 + tl.static_assert(
323 + (src_dtype == tl.bfloat16) or (src_dtype == tl.float16),
324 + f"{src_dtype=} must be bfloat16 or float16",
325 + )
326 +
327 + outer_block = tl.program_id(0).to(tl.int64)
328 + quant_block = tl.program_id(1).to(tl.int64)
329 +
330 + BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = BLOCK_SIZE_QUANT_DIM // 32
331 + BLOCK_SIZE_QUANT_MX_TENSOR: tl.constexpr = BLOCK_SIZE_QUANT_DIM // 2
332 +
333 + start_src_quant = quant_block * BLOCK_SIZE_QUANT_DIM
334 + start_mx_scale_quant = quant_block * BLOCK_SIZE_QUANT_MX_SCALE
335 + start_mx_quant = quant_block * BLOCK_SIZE_QUANT_MX_TENSOR
336 + start_out = outer_block * BLOCK_SIZE_OUT_DIM
337 +
338 + src_ptr += start_src_quant * stride_src_quant + start_out * stride_src_outer
339 + mx_scale_ptr += (
340 + start_mx_scale_quant * stride_mx_scale_quant + start_out * stride_mx_scale_outer
341 + )
342 + mx_tensor_ptr += start_mx_quant * stride_mxt_quant + start_out * stride_mxt_outer
343 +
344 + offs_src_quant = tl.arange(0, BLOCK_SIZE_QUANT_DIM)[None, :].to(tl.int64)
345 + offs_mxt_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_TENSOR)[None, :].to(tl.int64)
346 + offs_scale_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_SCALE)[None, :].to(tl.int64)
347 + offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None].to(tl.int64)
348 +
349 + mask_src_quant = start_src_quant + offs_src_quant < quant_dim
350 + mask_n = start_out + offs_outer < outer_dim
351 + full_mask_src = mask_src_quant and mask_n
352 +
353 + mask_mxt_quant = start_mx_quant + offs_mxt_quant < tl.cdiv(quant_dim, 2)
354 + full_mask_mxt = mask_mxt_quant and mask_n
355 +
356 + scale_mask_k = start_mx_scale_quant + offs_scale_quant < tl.cdiv(quant_dim, 32)
357 + full_scale_mask = scale_mask_k and mask_n
358 +
359 + src_tensor_offsets = (
360 + offs_src_quant * stride_src_quant + offs_outer * stride_src_outer
361 + )
362 + mx_scale_offsets = (
363 + offs_scale_quant * stride_mx_scale_quant + offs_outer * stride_mx_scale_outer
364 + )
365 + mx_tensor_offsets = (
366 + offs_mxt_quant * stride_mxt_quant + offs_outer * stride_mxt_outer
367 + )
368 + src_tensor = tl.load(src_ptr + src_tensor_offsets, mask=full_mask_src)
369 +
370 + out_tensor, scale_tensor = _compute_quant_and_scale(src_tensor, full_mask_src)
371 +
372 + tl.store(mx_scale_ptr + mx_scale_offsets, scale_tensor, mask=full_scale_mask)
373 + tl.store(mx_tensor_ptr + mx_tensor_offsets, out_tensor, mask=full_mask_mxt)
374 +
375 +
376 +def downcast_to_mxfp(
377 + src_tensor: torch.Tensor,
378 + axis: int,
379 + out_quant_tensor: torch.Tensor | None = None,
380 + out_scale: torch.Tensor | None = None,
381 + BLOCK_OUT_DIM: int = 128,
382 + BLOCK_QUANT_DIM: int = 32,
383 +):
384 + """
385 + Convert the src weights to MXFP4. The src weight is quantized along the
386 + axis dimension into packed e2m1 values (torch.uint8, two values per
387 + byte), so the size of that dimension in the output is half of the
388 + logical (unpacked) size.
389 + """
390 + out_quant_type = torch.uint8
391 + ndim = src_tensor.ndim
392 + assert -ndim <= axis < ndim, f"Invalid axis {axis=}"
393 + axis = axis if axis >= 0 else axis + ndim
394 +
395 + L = src_tensor.shape[axis]
396 + # We make this assertion since we can't track if the "real" shape was odd,
397 + # and we padded it to be even.
398 + # We want to maintain the property dequant(quant(x)).shape == x.shape
399 + assert L % 2 == 0, f"axis dim must be divisible by 2 for e2m1. Got {L}"
400 +
401 + device = src_tensor.device
402 +
403 + packed_quant_dim = triton.cdiv(L, 2)
404 + out_scale_dim = triton.cdiv(L, 32)
405 +
406 + # Move the quantization axis to the end for the kernel, then permute back.
407 + permute_order = list(range(ndim))
408 + permute_order[axis], permute_order[-1] = permute_order[-1], permute_order[axis]
409 +
410 + prmted_quant_tensor_shape = permute_shape(src_tensor.shape, permute_order)[:-1] + (
411 + packed_quant_dim,
412 + )
413 + prmted_scale_shape = permute_shape(src_tensor.shape, permute_order)[:-1] + (
414 + out_scale_dim,
415 + )
416 + prmted_src_tensor = src_tensor.permute(permute_order)
417 +
418 + if out_quant_tensor is None:
419 + out_quant_tensor = torch.empty(
420 + prmted_quant_tensor_shape, dtype=out_quant_type, device=device
421 + )
422 + else:
423 + expected_shape = (
424 + src_tensor.shape[:axis] + (packed_quant_dim,) + src_tensor.shape[axis + 1 :]
425 + )
426 + assert out_quant_tensor.shape == expected_shape, (
427 + f"{out_quant_tensor.shape=} != {expected_shape=}"
428 + )
429 + assert out_quant_tensor.dtype == out_quant_type, (
430 + f"{out_quant_tensor.dtype=} != {out_quant_type=}"
431 + )
432 + assert out_quant_tensor.stride(axis) == 1, (
433 + f"{out_quant_tensor.stride(axis)=} != 1"
434 + )
435 + # We expect the axis dimension to be last, so permute the tensor
436 + out_quant_tensor = out_quant_tensor.permute(permute_order)
437 +
438 + if out_scale is None:
439 + out_scale = torch.empty(prmted_scale_shape, dtype=torch.uint8, device=device)
440 + else:
441 + expected_scale_shape = permute_shape(prmted_scale_shape, permute_order)
442 + assert out_scale.shape == expected_scale_shape, (
443 + f"{out_scale.shape=} {expected_scale_shape=}"
444 + )
445 + assert out_scale.dtype == torch.uint8, f"{out_scale.dtype=} != torch.uint8"
446 + out_scale = out_scale.permute(permute_order)
447 +
448 + # Flatten input tensor for kernel. This will typically make a copy
449 + reshaped_src_tensor = prmted_src_tensor.reshape(-1, L)
450 + blocks_quant_dim = triton.cdiv(reshaped_src_tensor.shape[-1], BLOCK_QUANT_DIM)
451 + blocks_out_dim = triton.cdiv(reshaped_src_tensor.shape[0], BLOCK_OUT_DIM)
452 +
453 + # Flatten the output tensors for the kernel, this should be a view always
454 + kernel_quant_tensor = out_quant_tensor.reshape(-1, packed_quant_dim)
455 + kernel_scale = out_scale.reshape(-1, out_scale_dim)
456 + assert kernel_quant_tensor.data_ptr() == out_quant_tensor.data_ptr()
457 + assert kernel_scale.data_ptr() == out_scale.data_ptr()
458 +
459 + _downcast_to_mxfp[(blocks_out_dim, blocks_quant_dim)](
460 + kernel_quant_tensor,
461 + *kernel_quant_tensor.stride(),
462 + kernel_scale,
463 + *kernel_scale.stride(),
464 + reshaped_src_tensor,
465 + *reshaped_src_tensor.stride(),
466 + *reshaped_src_tensor.shape,
467 + BLOCK_OUT_DIM,
468 + BLOCK_QUANT_DIM,
469 + num_warps=8,
470 + )
471 +
472 + out_quant_tensor = out_quant_tensor.permute(permute_order)
473 + out_scale = out_scale.permute(permute_order).contiguous()
474 + return out_quant_tensor, out_scale, permute_shape(prmted_scale_shape, permute_order)
475 +
476 +
477 +def mxfp4_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
478 + """Quantize a bf16/fp16 tensor to MXFP4 along its last dimension.
479 +
480 + Dispatches to the fastest backend available on the current platform:
481 +
482 + - the native XPU custom op on XPU,
483 + - aiter on ROCm when `aiter` is installed,
484 + - and the portable Triton kernel (``downcast_to_mxfp``) otherwise.
485 +
486 + Returns packed FP4 values (uint8, two values per byte) and per-block
487 + (group-32) e8m0 scales (uint8).
488 + """
489 + if current_platform.is_xpu():
490 + return xpu_mxfp4_quantize(x)
491 +
492 + from vllm._aiter_ops import is_aiter_found_and_supported
493 +
494 + if is_aiter_found_and_supported() and x.dtype == torch.bfloat16:
495 + from vllm.model_executor.layers.quantization.quark.utils import (
496 + quark_quantize_weight_to_mxfp4,
497 + )
498 +
499 + return quark_quantize_weight_to_mxfp4(x)
500 +
501 + quant_tensor, scale, _ = downcast_to_mxfp(x, axis=-1)
502 +
503 + return quant_tensor, scale
vllm/utils/import_utils.py
+5
@@ -565,6 +565,11 @@ def has_humming() -> bool:
565 return _has_module("humming")
566
567
568 +def has_quark():
569 + """Whether the optional `quark` package is available."""
570 + return _has_module("quark")
571 +
572 +
573 def check_torchcodec_available():
574 """Whether the optional `torchcodec` package is available."""
575 try: