@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 41f1cf38f2
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Feature][OCP MX] Support mxfp6 and mixed mxfp6-mxfp4 (#21166)

fxmarty-amd committed Oct 7, 2025 at 15:35 UTC 41f1cf38f2e14b44842ea76011282efec2e8fd56
18 files changed +658 -182
docs/features/quantization/quark.md
+8 -4
@@ -231,9 +231,9 @@ python3 quantize_quark.py --model_dir meta-llama/Llama-2-70b-chat-hf \
231 --tasks gsm8k
232 ```
233
234 -## Using MXFP4 models
234 +## Using OCP MX (MXFP4, MXFP6) models
235
236 -vLLM supports loading MXFP4 models quantized offline through AMD Quark, compliant with [Open Compute Project (OCP) specification](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf).
236 +vLLM supports loading MXFP4 and MXFP6 models quantized offline through AMD Quark, compliant with [Open Compute Project (OCP) specification](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf).
237
238 The scheme currently only supports dynamic quantization for activations.
239
@@ -241,17 +241,21 @@ Example usage, after installing the latest AMD Quark release:
241
242 ```bash
243 vllm serve fxmarty/qwen_1.5-moe-a2.7b-mxfp4 --tensor-parallel-size 1
244 +# or, for a model using fp6 activations and fp4 weights:
245 +vllm serve fxmarty/qwen1.5_moe_a2.7b_chat_w_fp4_a_fp6_e2m3 --tensor-parallel-size 1
246 ```
247
246 -A simulation of the matrix multiplication execution in MXFP4 can be run on devices that do not support MXFP4 operations natively (e.g. AMD Instinct MI325, MI300 and MI250), dequantizing weights from MXFP4 to half precision on the fly, using a fused kernel. This is useful e.g. to evaluate MXFP4 models using vLLM, or alternatively to benefit from the ~4x memory savings (compared to float16 and bfloat16).
248 +A simulation of the matrix multiplication execution in MXFP4/MXFP6 can be run on devices that do not support OCP MX operations natively (e.g. AMD Instinct MI325, MI300 and MI250), dequantizing weights from FP4/FP6 to half precision on the fly, using a fused kernel. This is useful e.g. to evaluate FP4/FP6 models using vLLM, or alternatively to benefit from the ~2.5-4x memory savings (compared to float16 and bfloat16).
249
250 To generate offline models quantized using MXFP4 data type, the easiest approach is to use AMD Quark's [quantization script](https://quark.docs.amd.com/latest/pytorch/example_quark_torch_llm_ptq.html), as an example:
251
252 ```bash
253 python quantize_quark.py --model_dir Qwen/Qwen1.5-MoE-A2.7B-Chat \
252 - --quant_scheme w_mxfp4_a_mxfp4_sym \
254 + --quant_scheme w_mxfp4_a_mxfp4 \
255 --output_dir qwen_1.5-moe-a2.7b-mxfp4 \
256 --skip_evaluation \
257 --model_export hf_format \
258 --group_size 32
259 ```
260 +
261 +The current integration supports [all combination of FP4, FP6_E3M2, FP6_E2M3](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py) used for either weights or activations. Eventually, some target hardware support mixed precision GEMM, as AMD Instinct MI350/MI355, for example using FP6 for activations and FP4 for weights.
tests/kernels/moe/test_ocp_mx_moe.py renamed
+23 -17
@@ -10,13 +10,6 @@ import pytest
10 import torch
11 from packaging import version
12
13 -from vllm.model_executor.layers.quantization.quark.quark import (
14 - QuarkLinearMethod,
15 - QuarkW4A4MXFP4,
16 -)
17 -from vllm.model_executor.layers.quantization.quark.quark_moe import (
18 - QuarkW4A4MXFp4MoEMethod,
19 -)
13 from vllm.platforms import current_platform
14 from vllm.utils.flashinfer import has_flashinfer
15
@@ -63,9 +56,11 @@ def enable_pickle(monkeypatch):
56 @pytest.mark.parametrize(
57 "model_case",
58 [
66 - ModelCase("fxmarty/qwen_1.5-moe-a2.7b-mxfp4", tp=1),
59 + ModelCase("fxmarty/qwen_1.5-moe-a2.7b-mxfp4", tp=2),
60 ModelCase("fxmarty/deepseek_r1_3_layers_mxfp4", tp=8),
61 ModelCase("fxmarty/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1),
62 + ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=1),
63 + ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=4),
64 ],
65 )
66 @pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
@@ -76,22 +71,33 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
71 f"{torch.cuda.device_count()}"
72 )
73
74 + # `cuda_graph_sizes=[16]` to reduce load time.
75 with vllm_runner(
80 - model_case.model_id, tensor_parallel_size=model_case.tp, load_format="dummy"
76 + model_case.model_id,
77 + tensor_parallel_size=model_case.tp,
78 + load_format="dummy",
79 + cuda_graph_sizes=[16],
80 ) as llm:
81 + # Disabled as check_model is broken: https://github.com/vllm-project/vllm/pull/18465#issuecomment-3329880562
82 + # def check_model(model):
83 + # from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501
84 + # QuarkLinearMethod)
85 + # from vllm.model_executor.layers.quantization.quark.schemes.quark_ocp_mx import QuarkOCP_MX # noqa: E501
86 + # from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E501
87 + # QuarkOCP_MX_MoEMethod)
88
83 - def check_model(model):
84 - layer = model.model.layers[0]
89 + # layer = model.model.layers[0]
90
86 - qkv_proj = layer.self_attn.qkv_proj
91 + # qkv_proj = layer.self_attn.qkv_proj
92
88 - assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
89 - assert isinstance(qkv_proj.scheme, QuarkW4A4MXFP4)
93 + # assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
94 + # assert isinstance(qkv_proj.scheme, QuarkOCP_MX)
95
91 - assert isinstance(layer.mlp.experts.quant_method, QuarkW4A4MXFp4MoEMethod)
96 + # assert isinstance(layer.mlp.experts.quant_method,
97 + # QuarkOCP_MX_MoEMethod)
98
93 - if model_case.model_id == "fxmarty/qwen_1.5-moe-a2.7b-mxfp4":
94 - llm.apply_model(check_model)
99 + # if model_case.model_id == "fxmarty/qwen_1.5-moe-a2.7b-mxfp4":
100 + # llm.apply_model(check_model)
101
102 output = llm.generate_greedy("Today I am in the French Alps and", max_tokens=20)
103 assert output
tests/quantization/test_quark.py
+74 -19
@@ -11,6 +11,7 @@ import importlib.metadata
11 import os
12 from dataclasses import dataclass
13 from importlib.util import find_spec
14 +from typing import Optional
15
16 import huggingface_hub
17 import lm_eval
@@ -148,39 +149,93 @@ def test_quark_fp8_parity(vllm_runner):
149
150
151 @dataclass
151 -class ModelCase:
152 - model_id: str
153 - tp: int
154 -
155 -
156 -@dataclass
157 -class GSM8KAccuracyTestConfig:
152 +class AccuracyTestConfig:
153 model_name: str
154 excepted_value: float
155
161 - def get_model_args(self) -> str:
162 - return (
163 - f"pretrained={self.model_name},"
164 - "dtype=auto,add_bos_token=True,tensor_parallel_size=8,gpu_memory_utilization=0.7,max_model_len=38768"
165 - )
166 -
167 -
168 -ACCURACY_CONFIGS = [
156 + def get_model_args(
157 + self,
158 + tp_size: int,
159 + model_max_len: Optional[int] = None,
160 + kwargs: Optional[dict] = None,
161 + ) -> dict:
162 + if kwargs is None:
163 + kwargs = {}
164 +
165 + model_args = {
166 + "pretrained": self.model_name,
167 + "dtype": "auto",
168 + "add_bos_token": True,
169 + "tensor_parallel_size": tp_size,
170 + "gpu_memory_utilization": 0.7,
171 + **kwargs,
172 + }
173 + if model_max_len is not None:
174 + model_args["max_model_len"] = model_max_len
175 +
176 + return model_args
177 +
178 +
179 +GSM8K_ACCURACY_CONFIGS = [
180 # Private model.
170 - GSM8KAccuracyTestConfig(
181 + AccuracyTestConfig(
182 model_name="amd/DeepSeek-R1-WMXFP4-AMXFP4-Scale-UINT8-MoE-Quant",
183 excepted_value=0.96,
184 ),
185 ]
186
187 +WIKITEXT_ACCURACY_CONFIGS = [
188 + AccuracyTestConfig(
189 + model_name="fxmarty/qwen1.5_moe_a2.7b_chat_w_fp4_a_fp6_e2m3",
190 + excepted_value=11.3,
191 + ),
192 + AccuracyTestConfig(
193 + model_name="fxmarty/qwen1.5_moe_a2.7b_chat_w_fp6_e3m2_a_fp6_e3m2",
194 + excepted_value=10.6,
195 + ),
196 + AccuracyTestConfig(
197 + model_name="fxmarty/qwen_1.5-moe-a2.7b-mxfp4", excepted_value=12.4
198 + ),
199 +]
200 +
201 +
202 +@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
203 +@pytest.mark.parametrize("config", WIKITEXT_ACCURACY_CONFIGS)
204 +@pytest.mark.parametrize("tp_size", [1, 2])
205 +def test_ocp_mx_wikitext_correctness(config: AccuracyTestConfig, tp_size: int):
206 + if torch.cuda.device_count() < tp_size:
207 + pytest.skip(
208 + f"This test requires >={tp_size} gpus, got only {torch.cuda.device_count()}"
209 + )
210 +
211 + task = "wikitext"
212 + rtol = 0.1
213 +
214 + # Smaller cuda_graph_sizes to speed up the test.
215 + results = lm_eval.simple_evaluate(
216 + model="vllm",
217 + model_args=config.get_model_args(
218 + tp_size=tp_size, kwargs={"cuda_graph_sizes": [16]}
219 + ),
220 + tasks=task,
221 + batch_size=64,
222 + )
223 +
224 + EXPECTED_VALUE = config.excepted_value
225 + measured_value = results["results"][task]["word_perplexity,none"]
226 + assert (
227 + measured_value < EXPECTED_VALUE + rtol
228 + and measured_value > EXPECTED_VALUE - rtol
229 + ), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
230 +
231
177 -@pytest.mark.parametrize("config", ACCURACY_CONFIGS)
232 +@pytest.mark.parametrize("config", GSM8K_ACCURACY_CONFIGS)
233 @pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
234 @pytest.mark.skipif(
235 not HF_HUB_AMD_ORG_ACCESS,
236 reason="Read access to huggingface.co/amd is required for this test.",
237 )
183 -def test_mxfp4_gsm8k_correctness(config: GSM8KAccuracyTestConfig):
238 +def test_mxfp4_gsm8k_correctness(config: AccuracyTestConfig):
239 if torch.cuda.device_count() < 8:
240 pytest.skip(
241 f"This test requires >=8 gpus, got only {torch.cuda.device_count()}"
@@ -193,7 +248,7 @@ def test_mxfp4_gsm8k_correctness(config: GSM8KAccuracyTestConfig):
248
249 results = lm_eval.simple_evaluate(
250 model="vllm",
196 - model_args=config.get_model_args(),
251 + model_args=config.get_model_args(tp_size=8, model_max_len=38768),
252 tasks=task,
253 batch_size=64,
254 num_fewshot=8,
vllm/model_executor/layers/fused_moe/config.py
+56 -17
@@ -9,6 +9,10 @@ import vllm.envs as envs
9 from vllm.config import ParallelConfig
10 from vllm.distributed import get_dp_group, get_tensor_model_parallel_rank
11 from vllm.logger import init_logger
12 +from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
13 + OCP_MX_DTYPES,
14 + OCP_MX_Scheme,
15 +)
16 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
17 from vllm.utils import cdiv, has_triton_kernels
18 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
@@ -30,7 +34,7 @@ def _get_config_dtype_str(
34 use_fp8_w8a8: bool = False,
35 use_int8_w8a16: bool = False,
36 use_int4_w4a16: bool = False,
33 - use_mxfp4_w4a4: bool = False,
37 + ocp_mx_scheme: Optional[str] = None,
38 ) -> Optional[str]:
39 """
40 Return a string used to construct the filename that contains the
@@ -43,8 +47,11 @@ def _get_config_dtype_str(
47 return "int8_w8a16"
48 elif use_int4_w4a16:
49 return "int4_w4a16"
46 - elif use_mxfp4_w4a4:
47 - return "mxfp4_w4a4"
50 + elif ocp_mx_scheme is not None:
51 + # The output of this function is passed to `try_get_optimal_moe_config`,
52 + # and as we only simulate OCP MX execution in fused_moe for now,
53 + # we will NOT look for `*,dtype=w_mxfp4_a_mxfp4.json` for now.
54 + return None
55 elif dtype == torch.float:
56 # avoiding cases where kernel fails when float32 MoE
57 # use fp16/bfloat16 configs
@@ -289,8 +296,23 @@ class FusedMoEQuantConfig:
296 return self._a1.dtype is None and self._w1.dtype == "int4"
297
298 @property
292 - def use_mxfp4_w4a4(self) -> bool:
293 - return self._a1.dtype == "mxfp4" and self._w1.dtype == "mxfp4"
299 + def ocp_mx_scheme(self) -> Union[str, None]:
300 + if not hasattr(self, "_ocp_mx_scheme"):
301 + if (self._a1.dtype is not None and not isinstance(self._a1.dtype, str)) or (
302 + self._w1.dtype is not None and not isinstance(self._w1.dtype, str)
303 + ):
304 + self._ocp_mx_scheme = None
305 + else:
306 + ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype(
307 + self._a1.dtype, self._w1.dtype
308 + )
309 +
310 + if ocp_mx_scheme is not None:
311 + ocp_mx_scheme = ocp_mx_scheme.value
312 +
313 + self._ocp_mx_scheme = ocp_mx_scheme
314 +
315 + return self._ocp_mx_scheme
316
317 @property
318 def use_mxfp4_w4a16(self) -> bool:
@@ -310,7 +332,7 @@ class FusedMoEQuantConfig:
332 use_fp8_w8a8=self.use_fp8_w8a8,
333 use_int8_w8a16=self.use_int8_w8a16,
334 use_int4_w4a16=self.use_int4_w4a16,
313 - use_mxfp4_w4a4=self.use_mxfp4_w4a4,
335 + ocp_mx_scheme=self.ocp_mx_scheme,
336 dtype=dtype,
337 )
338
@@ -371,12 +393,14 @@ class FusedMoEQuantConfig:
393 w2_bias: Optional[torch.Tensor] = None,
394 w1_zp: Optional[torch.Tensor] = None,
395 w2_zp: Optional[torch.Tensor] = None,
396 + weight_dtype: Union[torch.dtype, str, None] = None,
397 ) -> "FusedMoEQuantConfig":
398 """
399 General builder function for a FusedMoEQuantConfig.
400 - quant_dtype: Optional quantization type. None if activations are
378 - unquantized or quantized prior to calling. Note: "nvfp4" and
379 - "mxfp4" are the only valid string values for quant_dtype.
401 + unquantized or quantized prior to calling. Note: "nvfp4", "mxfp4",
402 + "mxfp6_e3m2", "mxfp6_e2m3" are the only valid string values
403 + for quant_dtype.
404 - per_act_token_quant: Activations have per token quantization.
405 - per_out_ch_quant: Outputs have per channel quantization. (only
406 for cutlass).
@@ -395,11 +419,22 @@ class FusedMoEQuantConfig:
419 - w1_zp: Optional w1 zero points for int4/int8 quantization.
420 - w2_zp: Optional w2 zero points for int4/int8 quantization.
421 """
398 - assert (
399 - not isinstance(quant_dtype, str)
400 - or quant_dtype == "nvfp4"
401 - or quant_dtype == "mxfp4"
402 - )
422 + assert not isinstance(quant_dtype, str) or quant_dtype in {
423 + "nvfp4",
424 + "mxfp4",
425 + "mxfp6_e3m2",
426 + "mxfp6_e2m3",
427 + }
428 + assert not isinstance(weight_dtype, str) or weight_dtype in {
429 + "nvfp4",
430 + "mxfp4",
431 + "mxfp6_e3m2",
432 + "mxfp6_e2m3",
433 + }
434 +
435 + if weight_dtype is None:
436 + weight_dtype = quant_dtype
437 +
438 a_shape, w_shape = _quant_flags_to_group_shape(
439 quant_dtype, per_act_token_quant, per_out_ch_quant, block_shape
440 )
@@ -407,10 +442,10 @@ class FusedMoEQuantConfig:
442 _a1=FusedMoEQuantDesc(quant_dtype, a_shape, a1_scale, a1_gscale),
443 _a2=FusedMoEQuantDesc(quant_dtype, a_shape, a2_scale, a2_gscale),
444 _w1=FusedMoEQuantDesc(
410 - quant_dtype, w_shape, w1_scale, g1_alphas, w1_zp, w1_bias
445 + weight_dtype, w_shape, w1_scale, g1_alphas, w1_zp, w1_bias
446 ),
447 _w2=FusedMoEQuantDesc(
413 - quant_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias
448 + weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias
449 ),
450 )
451 assert quant_config.per_act_token_quant == per_act_token_quant
@@ -482,9 +517,11 @@ def mxfp4_w4a16_moe_quant_config(
517 )
518
519
485 -def mxfp4_w4a4_moe_quant_config(
520 +def ocp_mx_moe_quant_config(
521 + quant_dtype: str,
522 w1_scale: Union[torch.Tensor, "PrecisionConfig"],
523 w2_scale: Union[torch.Tensor, "PrecisionConfig"],
524 + weight_dtype: Optional[str] = None,
525 a1_scale: Optional[torch.Tensor] = None,
526 a2_scale: Optional[torch.Tensor] = None,
527 w1_bias: Optional[torch.Tensor] = None,
@@ -494,8 +531,10 @@ def mxfp4_w4a4_moe_quant_config(
531 """
532 Construct a quant config for mxfp4 activations and mxfp4 weights.
533 """
534 + assert quant_dtype in OCP_MX_DTYPES
535 return FusedMoEQuantConfig.make(
498 - "mxfp4",
536 + quant_dtype=quant_dtype,
537 + weight_dtype=weight_dtype,
538 w1_scale=w1_scale,
539 w2_scale=w2_scale,
540 a1_scale=a1_scale,
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+2 -2
@@ -640,7 +640,7 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
640 assert not self.quant_config.use_int8_w8a8, "NYI"
641 assert not self.quant_config.use_int8_w8a16, "NYI"
642 assert not self.quant_config.use_int4_w4a16, "NYI"
643 - assert not self.quant_config.use_mxfp4_w4a4, "NYI"
643 + assert self.quant_config.ocp_mx_scheme is None, "NYI"
644 self.max_num_tokens = max_num_tokens
645 self.num_dispatchers = num_dispatchers
646
@@ -835,7 +835,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
835 assert not self.quant_config.use_int8_w8a8, "NYI"
836 assert not self.quant_config.use_int8_w8a16, "NYI"
837 assert not self.quant_config.use_int4_w4a16, "NYI"
838 - assert not self.quant_config.use_mxfp4_w4a4, "NYI"
838 + assert self.quant_config.ocp_mx_scheme is None, "NYI"
839 assert max_num_tokens > 0
840 assert num_dispatchers > 0
841 self.max_num_tokens = max_num_tokens
vllm/model_executor/layers/fused_moe/fused_moe.py
+69 -21
@@ -42,6 +42,8 @@ from vllm.model_executor.layers.fused_moe.utils import (
42 moe_kernel_quantize_input,
43 )
44 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4
45 +from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6
46 +from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_Scheme
47 from vllm.platforms import current_platform
48 from vllm.triton_utils import tl, triton
49 from vllm.utils import direct_register_custom_op, is_torch_equal_or_newer
@@ -1323,7 +1325,7 @@ def inplace_fused_experts(
1325 use_int8_w8a8: bool = False,
1326 use_int8_w8a16: bool = False,
1327 use_int4_w4a16: bool = False,
1326 - use_mxfp4_w4a4: bool = False,
1328 + ocp_mx_scheme: Optional[str] = None,
1329 per_channel_quant: bool = False,
1330 global_num_experts: int = -1,
1331 expert_map: Optional[torch.Tensor] = None,
@@ -1350,7 +1352,7 @@ def inplace_fused_experts(
1352 use_int8_w8a8,
1353 use_int8_w8a16,
1354 use_int4_w4a16,
1353 - use_mxfp4_w4a4,
1355 + ocp_mx_scheme,
1356 per_channel_quant,
1357 global_num_experts,
1358 expert_map,
@@ -1378,7 +1380,7 @@ def inplace_fused_experts_fake(
1380 use_int8_w8a8: bool = False,
1381 use_int8_w8a16: bool = False,
1382 use_int4_w4a16: bool = False,
1381 - use_mxfp4_w4a4: bool = False,
1383 + ocp_mx_scheme: Optional[str] = None,
1384 per_channel_quant: bool = False,
1385 global_num_experts: int = -1,
1386 expert_map: Optional[torch.Tensor] = None,
@@ -1420,7 +1422,7 @@ def outplace_fused_experts(
1422 use_int8_w8a8: bool = False,
1423 use_int8_w8a16: bool = False,
1424 use_int4_w4a16: bool = False,
1423 - use_mxfp4_w4a4: bool = False,
1425 + ocp_mx_scheme: Optional[str] = None,
1426 per_channel_quant: bool = False,
1427 global_num_experts: int = -1,
1428 expert_map: Optional[torch.Tensor] = None,
@@ -1447,7 +1449,7 @@ def outplace_fused_experts(
1449 use_int8_w8a8,
1450 use_int8_w8a16,
1451 use_int4_w4a16,
1450 - use_mxfp4_w4a4,
1452 + ocp_mx_scheme,
1453 per_channel_quant,
1454 global_num_experts,
1455 expert_map,
@@ -1474,7 +1476,7 @@ def outplace_fused_experts_fake(
1476 use_int8_w8a8: bool = False,
1477 use_int8_w8a16: bool = False,
1478 use_int4_w4a16: bool = False,
1477 - use_mxfp4_w4a4: bool = False,
1479 + ocp_mx_scheme: Optional[str] = None,
1480 per_channel_quant: bool = False,
1481 global_num_experts: int = -1,
1482 expert_map: Optional[torch.Tensor] = None,
@@ -1599,7 +1601,7 @@ def fused_experts(
1601 use_int8_w8a8=quant_config.use_int8_w8a8,
1602 use_int8_w8a16=quant_config.use_int8_w8a16,
1603 use_int4_w4a16=quant_config.use_int4_w4a16,
1602 - use_mxfp4_w4a4=quant_config.use_mxfp4_w4a4,
1604 + ocp_mx_scheme=quant_config.ocp_mx_scheme,
1605 per_channel_quant=quant_config.per_act_token_quant,
1606 global_num_experts=global_num_experts,
1607 expert_map=expert_map,
@@ -1622,7 +1624,7 @@ GELU_NO_MUL: str = activation_without_mul("gelu")
1624 def _get_config_quant_dtype(
1625 use_fp8_w8a8: bool,
1626 use_int8_w8a8: bool,
1625 - use_mxfp4_w4a4: bool,
1627 + ocp_mx_scheme: Optional[str],
1628 ) -> Union[None, torch.dtype, str]:
1629 """
1630 Get the quantization type based on the quantization strategy flags.
@@ -1635,8 +1637,12 @@ def _get_config_quant_dtype(
1637 return torch.float8_e4m3fn
1638 elif use_int8_w8a8:
1639 return torch.int8
1638 - elif use_mxfp4_w4a4:
1640 + elif ocp_mx_scheme == "w_mxfp4_a_mxfp4":
1641 return "mxfp4"
1642 + elif ocp_mx_scheme in {"w_mxfp4_a_mxfp6_e3m2", "w_mxfp6_e3m2_a_mxfp6_e3m2"}:
1643 + return "mxfp6_e3m2"
1644 + elif ocp_mx_scheme in {"w_mxfp4_a_mxfp6_e2m3", "w_mxfp6_e2m3_a_mxfp6_e2m3"}:
1645 + return "mxfp6_e2m3"
1646 return None
1647
1648
@@ -1653,7 +1659,7 @@ def fused_experts_impl(
1659 use_int8_w8a8: bool = False,
1660 use_int8_w8a16: bool = False,
1661 use_int4_w4a16: bool = False,
1656 - use_mxfp4_w4a4: bool = False,
1662 + ocp_mx_scheme: Optional[str] = None,
1663 per_channel_quant: bool = False,
1664 global_num_experts: int = -1,
1665 expert_map: Optional[torch.Tensor] = None,
@@ -1670,9 +1676,23 @@ def fused_experts_impl(
1676 # Check constraints.
1677 if use_int4_w4a16:
1678 assert hidden_states.size(1) // 2 == w1.size(2), "Hidden size mismatch"
1673 - elif use_mxfp4_w4a4:
1674 - # 16bit activation and fp4x2 packed weight
1675 - assert hidden_states.size(1) // 2 == w1.size(2), "hidden size mismatch"
1679 + elif ocp_mx_scheme is not None:
1680 + if ocp_mx_scheme in {
1681 + "w_mxfp4_a_mxfp4",
1682 + "w_mxfp4_a_mxfp6_e3m2",
1683 + "w_mxfp4_a_mxfp6_e2m3",
1684 + }:
1685 + # 16bit activation and fp4x2 packed weight
1686 + assert hidden_states.size(1) == w1.size(2) * 2, "hidden size mismatch"
1687 + elif ocp_mx_scheme in {
1688 + "w_mxfp6_e3m2_a_mxfp6_e3m2",
1689 + "w_mxfp6_e2m3_a_mxfp6_e2m3",
1690 + }:
1691 + assert hidden_states.size(1) == (w1.size(2) * 4) // 3, (
1692 + "hidden size mismatch"
1693 + )
1694 + else:
1695 + raise NotImplementedError(f"Unsupported ocp_mx_scheme={ocp_mx_scheme}")
1696 else:
1697 assert hidden_states.size(1) == w1.size(2), (
1698 f"Hidden size mismatch {hidden_states.size(1)} != {w1.size(2)}"
@@ -1699,7 +1719,7 @@ def fused_experts_impl(
1719 use_fp8_w8a8=use_fp8_w8a8,
1720 use_int8_w8a16=use_int8_w8a16,
1721 use_int4_w4a16=use_int4_w4a16,
1702 - use_mxfp4_w4a4=use_mxfp4_w4a4,
1722 + ocp_mx_scheme=ocp_mx_scheme,
1723 dtype=hidden_states.dtype,
1724 )
1725
@@ -1708,7 +1728,7 @@ def fused_experts_impl(
1728 quant_dtype = _get_config_quant_dtype(
1729 use_fp8_w8a8=use_fp8_w8a8,
1730 use_int8_w8a8=use_int8_w8a8,
1711 - use_mxfp4_w4a4=use_mxfp4_w4a4,
1731 + ocp_mx_scheme=ocp_mx_scheme,
1732 )
1733
1734 get_config_func = functools.partial(
@@ -1748,12 +1768,40 @@ def fused_experts_impl(
1768
1769 out_hidden_states = hidden_states if inplace else torch.empty_like(hidden_states)
1770
1751 - if use_mxfp4_w4a4:
1752 - # Weight has to be dequantized for mxfp4 emulation.
1753 - w1 = dequant_mxfp4(w1, w1_scale, hidden_states.dtype)
1754 - w1_scale = None
1755 - w2 = dequant_mxfp4(w2, w2_scale, hidden_states.dtype)
1756 - w2_scale = None
1771 + if ocp_mx_scheme is not None:
1772 + # TODO: On platforms for which `current_platform.supports_mx()` is True
1773 + # and for which we have a native OCP mx fused MOE kernel,
1774 + # this dequantization step should not be done.
1775 + if ocp_mx_scheme in {
1776 + OCP_MX_Scheme.w_mxfp4_a_mxfp4,
1777 + OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2,
1778 + OCP_MX_Scheme.w_mxfp4_a_mxfp6_e2m3,
1779 + }:
1780 + # Weight has to be dequantized for mxfp4 emulation.
1781 + w1 = dequant_mxfp4(w1, w1_scale, hidden_states.dtype)
1782 + w1_scale = None
1783 + w2 = dequant_mxfp4(w2, w2_scale, hidden_states.dtype)
1784 + w2_scale = None
1785 + elif ocp_mx_scheme == OCP_MX_Scheme.w_mxfp6_e3m2_a_mxfp6_e3m2:
1786 + w1 = dequant_mxfp6(
1787 + w1, w1_scale, quant_dtype="fp6_e3m2", float_dtype=hidden_states.dtype
1788 + )
1789 + w1_scale = None
1790 + w2 = dequant_mxfp6(
1791 + w2, w2_scale, quant_dtype="fp6_e3m2", float_dtype=hidden_states.dtype
1792 + )
1793 + w2_scale = None
1794 + elif ocp_mx_scheme == OCP_MX_Scheme.w_mxfp6_e2m3_a_mxfp6_e2m3:
1795 + w1 = dequant_mxfp6(
1796 + w1, w1_scale, quant_dtype="fp6_e2m3", float_dtype=hidden_states.dtype
1797 + )
1798 + w1_scale = None
1799 + w2 = dequant_mxfp6(
1800 + w2, w2_scale, quant_dtype="fp6_e2m3", float_dtype=hidden_states.dtype
1801 + )
1802 + w2_scale = None
1803 + else:
1804 + raise NotImplementedError(f"Unsupported ocp_mx_scheme={ocp_mx_scheme}")
1805
1806 for chunk in range((num_tokens // CHUNK_SIZE) + 1):
1807 begin_chunk_idx, end_chunk_idx = (
vllm/model_executor/layers/fused_moe/utils.py
+60 -13
@@ -16,11 +16,15 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
16 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
17 quant_dequant_mxfp4,
18 )
19 -from vllm.model_executor.layers.quantization.utils.mxfp8_utils import mxfp8_quantize
20 -from vllm.platforms import current_platform
19 +from vllm.model_executor.layers.quantization.utils.mxfp6_utils import (
20 + quant_dequant_mxfp6,
21 +)
22 +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
23 + mxfp8_e4m3_quantize,
24 +)
25 from vllm.triton_utils import tl, triton
26 from vllm.utils import cdiv
23 -from vllm.utils.flashinfer import fp4_quantize
27 +from vllm.utils.flashinfer import flashinfer_fp4_quantize
28
29
30 @triton.jit
@@ -106,12 +110,14 @@ def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor:
110 return x.flatten()[: prod(v)].view(*v)
111
112
109 -def _fp4_quantize(
113 +def _nvfp4_quantize(
114 A: torch.Tensor,
115 A_scale: Optional[torch.Tensor],
116 is_sf_swizzled_layout: bool,
117 ) -> tuple[torch.Tensor, torch.Tensor]:
114 - return fp4_quantize(A, A_scale, is_sf_swizzled_layout=is_sf_swizzled_layout)
118 + return flashinfer_fp4_quantize(
119 + A, A_scale, is_sf_swizzled_layout=is_sf_swizzled_layout
120 + )
121
122
123 def _fp8_quantize(
@@ -174,15 +180,16 @@ def _mxfp4_quantize(
180 block_shape: Optional[list[int]] = None,
181 ) -> tuple[torch.Tensor, None]:
182 assert block_shape is None
177 - if not current_platform.supports_mx():
178 - A = quant_dequant_mxfp4(A)
179 - else:
180 - raise NotImplementedError()
183 + # TODO: native mxfp4 is currently not integrated in vllm,
184 + # so simulating even on devices supporting this data type natively.
185 + # Once integrated, `current_platform.supports_mx()` should be used to
186 + # control quantize+dequantize, or simply quantize here down to mxfp4.
187 + A = quant_dequant_mxfp4(A)
188
189 return A, None
190
191
185 -def _mxfp8_quantize(
192 +def _mxfp8_e4m3_quantize(
193 A: torch.Tensor,
194 A_scale: Optional[torch.Tensor],
195 per_act_token_quant: bool,
@@ -191,7 +198,41 @@ def _mxfp8_quantize(
198 assert A_scale is None
199 assert not per_act_token_quant
200 assert block_shape is None
194 - return mxfp8_quantize(A)
201 + return mxfp8_e4m3_quantize(A)
202 +
203 +
204 +def _mxfp6_e3m2_quantize(
205 + A: torch.Tensor,
206 + A_scale: Optional[torch.Tensor],
207 + per_act_token_quant: bool,
208 + block_shape: Optional[list[int]] = None,
209 +) -> tuple[torch.Tensor, None]:
210 + assert block_shape is None
211 +
212 + # TODO: native mxfp6 is currently not integrated in vllm,
213 + # so simulating even on devices supporting this data type natively.
214 + # Eventually, there should be a check based on
215 + # `current_platform.supports_mx()` here.
216 + A = quant_dequant_mxfp6(A, quant_dtype="fp6_e3m2")
217 +
218 + return A, None
219 +
220 +
221 +def _mxfp6_e2m3_quantize(
222 + A: torch.Tensor,
223 + A_scale: Optional[torch.Tensor],
224 + per_act_token_quant: bool,
225 + block_shape: Optional[list[int]] = None,
226 +) -> tuple[torch.Tensor, None]:
227 + assert block_shape is None
228 +
229 + # TODO: native mxfp6 is currently not integrated in vllm,
230 + # so simulating even on devices supporting this data type natively.
231 + # Eventually, there should be a check based on
232 + # `current_platform.supports_mx()` here.
233 + A = quant_dequant_mxfp6(A, quant_dtype="fp6_e2m3")
234 +
235 + return A, None
236
237
238 def moe_kernel_quantize_input(
@@ -207,11 +248,17 @@ def moe_kernel_quantize_input(
248 elif quant_dtype == torch.int8:
249 return _int8_quantize(A, A_scale, per_act_token_quant, block_shape)
250 elif quant_dtype == "nvfp4":
210 - return _fp4_quantize(A, A_scale, is_sf_swizzled_layout=is_fp4_scale_swizzled)
251 + return _nvfp4_quantize(A, A_scale, is_sf_swizzled_layout=is_fp4_scale_swizzled)
252 elif quant_dtype == "mxfp4":
253 return _mxfp4_quantize(A, A_scale, per_act_token_quant, block_shape)
254 elif quant_dtype == "mxfp8":
214 - return _mxfp8_quantize(A, A_scale, per_act_token_quant, block_shape)
255 + # TODO: `quant_dtype == "mxfp8"` is ambiguous,
256 + # should be fp8_e4m3. OCP MX also defines `fp8_e5m2`.
257 + return _mxfp8_e4m3_quantize(A, A_scale, per_act_token_quant, block_shape)
258 + elif quant_dtype == "mxfp6_e3m2":
259 + return _mxfp6_e3m2_quantize(A, A_scale, per_act_token_quant, block_shape)
260 + elif quant_dtype == "mxfp6_e2m3":
261 + return _mxfp6_e2m3_quantize(A, A_scale, per_act_token_quant, block_shape)
262 else:
263 return A, A_scale
264
vllm/model_executor/layers/quantization/mxfp4.py
+3 -2
@@ -17,8 +17,8 @@ from vllm.model_executor.layers.fused_moe import (
17 from vllm.model_executor.layers.fused_moe import modular_kernel as mk
18 from vllm.model_executor.layers.fused_moe.config import (
19 FusedMoEQuantConfig,
20 - mxfp4_w4a4_moe_quant_config,
20 mxfp4_w4a16_moe_quant_config,
21 + ocp_mx_moe_quant_config,
22 )
23 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import MarlinExperts
24 from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
@@ -776,7 +776,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
776 else:
777 w1_scale = layer.w13_weight_scale
778 w2_scale = layer.w2_weight_scale
779 - return mxfp4_w4a4_moe_quant_config(
779 + return ocp_mx_moe_quant_config(
780 + quant_dtype="mxfp4",
781 w1_bias=layer.w13_bias,
782 w2_bias=layer.w2_bias,
783 w1_scale=w1_scale,
vllm/model_executor/layers/quantization/quark/quark.py
+20 -18
@@ -23,8 +23,8 @@ from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E
23 QuarkMoEMethod,
24 )
25 from vllm.model_executor.layers.quantization.quark.schemes import (
26 + QuarkOCP_MX,
27 QuarkScheme,
27 - QuarkW4A4MXFP4,
28 QuarkW8A8Fp8,
29 QuarkW8A8Int8,
30 )
@@ -235,7 +235,7 @@ class QuarkConfig(QuantizationConfig):
235 # Only symmetric weight quantization supported.
236 return is_int8_dtype and is_tensor and is_weight_symmetric and is_static
237
238 - def _is_mx_fp4(
238 + def _is_ocp_mx(
239 self,
240 weight_quant: Optional[dict[str, Any]],
241 input_quant: Optional[dict[str, Any]],
@@ -243,32 +243,22 @@ class QuarkConfig(QuantizationConfig):
243 # Confirm weights and input quantized.
244 if weight_quant is None or input_quant is None:
245 logger.debug(
246 - "Quark model is not in MX-FP4 format: "
246 + "Quark model is not in OCP MX format: "
247 "weight_quant or input_quant not set"
248 )
249 return False
250
251 - # Input and weight dtype needs to be fp4.
252 - if weight_quant.get("dtype") != "fp4" or input_quant.get("dtype") != "fp4":
253 - logger.debug("Quark model is not in MX-FP4 format: dtype not fp4")
254 - return False
255 -
251 # Input and weight qscheme needs to be per group.
252 if (
253 weight_quant.get("qscheme") != "per_group"
254 or input_quant.get("qscheme") != "per_group"
255 ):
261 - logger.debug("Quark model is not in MX-FP4 format: not per_group")
256 + logger.debug("Quark model is not in OCP MX format: not per_group")
257 return False
258
259 # Input and weight group size needs to be 32.
260 if weight_quant.get("group_size") != 32 or input_quant.get("group_size") != 32:
266 - logger.debug("Quark model is not in MX-FP4 format: not group_size=32")
267 - return False
268 -
269 - # Activations need to use dynamic quantization.
270 - if input_quant.get("is_dynamic") is False:
271 - logger.debug("Quark model is not in MX-FP4 format: not activation dynamic")
261 + logger.debug("Quark model is not in OCP MX format: not group_size=32")
262 return False
263
264 # Activations and weight scales need to be in e8m0 format.
@@ -276,7 +266,19 @@ class QuarkConfig(QuantizationConfig):
266 weight_quant.get("scale_format") != "e8m0"
267 or input_quant.get("scale_format") != "e8m0"
268 ):
279 - logger.debug("Quark model is not in MX-FP4 format: not scale_format e8m0")
269 + logger.debug("Quark model is not in OCP MX format: not scale_format e8m0")
270 + return False
271 +
272 + # Input and weight dtypes need to be any of fp4,
273 + # fp6_e3m2 or fp6_e3m2, possibly mixed.
274 + if weight_quant.get("dtype") not in {
275 + "fp4",
276 + "fp6_e3m2",
277 + "fp6_e2m3",
278 + } or input_quant.get("dtype") not in {"fp4", "fp6_e3m2", "fp6_e2m3"}:
279 + logger.debug(
280 + "Quark model is not in OCP MX format: dtype not fp4, fp6_e3m2, fp6_e2m3"
281 + )
282 return False
283
284 return True
@@ -348,8 +350,8 @@ class QuarkConfig(QuantizationConfig):
350 is_static_input_scheme=True,
351 input_symmetric=input_config.get("symmetric"),
352 )
351 - elif self._is_mx_fp4(weight_config, input_config):
352 - return QuarkW4A4MXFP4(weight_config, input_config)
353 + elif self._is_ocp_mx(weight_config, input_config):
354 + return QuarkOCP_MX(weight_config, input_config)
355
356 raise NotImplementedError(
357 "No quark compatible scheme was found. "
vllm/model_executor/layers/quantization/quark/quark_moe.py
+34 -15
@@ -17,7 +17,7 @@ from vllm.model_executor.layers.fused_moe import (
17 from vllm.model_executor.layers.fused_moe.config import (
18 FusedMoEQuantConfig,
19 fp8_w8a8_moe_quant_config,
20 - mxfp4_w4a4_moe_quant_config,
20 + ocp_mx_moe_quant_config,
21 )
22 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
23 is_rocm_aiter_moe_enabled,
@@ -25,7 +25,10 @@ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
25 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
26 prepare_moe_fp8_layer_for_marlin,
27 )
28 -from vllm.model_executor.layers.quantization.utils.mxfp4_utils import OCP_MX_BLOCK_SIZE
28 +from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
29 + OCP_MX_BLOCK_SIZE,
30 + OCP_MX_Scheme,
31 +)
32 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
33 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
34 all_close_1d,
@@ -38,7 +41,7 @@ from vllm.scalar_type import scalar_types
41
42 logger = init_logger(__name__)
43
41 -__all__ = ["QuarkMoEMethod", "QuarkW8A8Fp8MoEMethod", "QuarkW4A4MXFp4MoEMethod"]
44 +__all__ = ["QuarkMoEMethod", "QuarkW8A8Fp8MoEMethod", "QuarkOCP_MX_MoEMethod"]
45
46
47 class QuarkMoEMethod(FusedMoEMethodBase):
@@ -64,10 +67,8 @@ class QuarkMoEMethod(FusedMoEMethodBase):
67
68 if quant_config._is_fp8_w8a8(weight_config, input_config):
69 return QuarkW8A8Fp8MoEMethod(weight_config, input_config, module.moe_config)
67 - elif quant_config._is_mx_fp4(weight_config, input_config):
68 - return QuarkW4A4MXFp4MoEMethod(
69 - weight_config, input_config, module.moe_config
70 - )
70 + elif quant_config._is_ocp_mx(weight_config, input_config):
71 + return QuarkOCP_MX_MoEMethod(weight_config, input_config, module.moe_config)
72 else:
73 raise RuntimeError("Unsupported FusedMoe scheme")
74
@@ -434,7 +435,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
435 )
436
437
437 -class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
438 +class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
439 def __init__(
440 self,
441 weight_config: dict[str, Any],
@@ -456,16 +457,23 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
457
458 self.static_input_scales = not self.input_quant.get("is_dynamic")
459
460 + self.weight_dtype = self.weight_quant["dtype"].replace("fp", "mxfp")
461 + self.input_dtype = self.input_quant["dtype"].replace("fp", "mxfp")
462 +
463 + self.ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype(
464 + self.input_dtype, self.weight_dtype
465 + )
466 +
467 if self.static_input_scales:
468 raise NotImplementedError(
461 - "QuarkW4A4MXFp4MoEMethod with static input scales is currently "
469 + "QuarkOCP_MX_MoEMethod with static input scales is currently "
470 "not implemented. Please open an issue."
471 )
472
473 if not current_platform.supports_mx():
474 self.emulate = True
475 logger.warning_once(
468 - "The current platform does not support native MXFP4 "
476 + "The current platform does not support native MXFP4/MXFP6 "
477 "computation. Simulated weight dequantization and activation "
478 "QDQ (quantize and dequantize) will be used, with the linear "
479 "layers computed in high precision."
@@ -473,13 +481,22 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
481 else:
482 self.emulate = True
483 logger.warning_once(
476 - "The current platform supports native MXFP4 "
484 + "The current platform supports native MXFP4/MXFP6 "
485 "computation, but kernels are not yet integrated in vLLM. "
486 "Simulated weight dequantization and activation "
487 "QDQ (quantize and dequantize) will be used, with the linear "
488 "layers computed in high precision."
489 )
490
491 + def get_packed_dim(self, dim: int, quant_dtype: str):
492 + if quant_dtype == "mxfp4":
493 + assert dim % 2 == 0
494 + return dim // 2
495 + else:
496 + # FP6 packs 4 * 6 = 24 bits on 3 bytes.
497 + assert (dim * 3) % 4 == 0
498 + return (dim * 3) // 4
499 +
500 def create_weights(
501 self,
502 layer: torch.nn.Module,
@@ -502,7 +519,7 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
519 torch.empty(
520 num_experts,
521 2 * intermediate_size_per_partition,
505 - hidden_size // 2,
522 + self.get_packed_dim(hidden_size, self.weight_dtype),
523 dtype=params_dtype,
524 ),
525 requires_grad=False,
@@ -515,7 +532,7 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
532 torch.empty(
533 num_experts,
534 hidden_size,
518 - intermediate_size_per_partition // 2,
535 + self.get_packed_dim(intermediate_size_per_partition, self.weight_dtype),
536 dtype=params_dtype,
537 ),
538 requires_grad=False,
@@ -552,7 +569,9 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
569 def get_fused_moe_quant_config(
570 self, layer: torch.nn.Module
571 ) -> Optional[FusedMoEQuantConfig]:
555 - return mxfp4_w4a4_moe_quant_config(
572 + return ocp_mx_moe_quant_config(
573 + quant_dtype=self.input_dtype,
574 + weight_dtype=self.weight_dtype,
575 w1_scale=layer.w13_weight_scale,
576 w2_scale=layer.w2_weight_scale,
577 a1_scale=None,
@@ -587,7 +606,7 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
606
607 if enable_eplb:
608 raise NotImplementedError(
590 - "EPLB not supported for `QuarkW4A4MXFp4MoEMethod` yet."
609 + "EPLB not supported for `QuarkOCP_MX_MoEMethod` yet."
610 )
611
612 from vllm.model_executor.layers.fused_moe import fused_experts
vllm/model_executor/layers/quantization/quark/schemes/__init__.py
+2 -2
@@ -1,9 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from .quark_ocp_mx import QuarkOCP_MX
5 from .quark_scheme import QuarkScheme
5 -from .quark_w4a4_mxfp4 import QuarkW4A4MXFP4
6 from .quark_w8a8_fp8 import QuarkW8A8Fp8
7 from .quark_w8a8_int8 import QuarkW8A8Int8
8
9 -__all__ = ["QuarkScheme", "QuarkW8A8Fp8", "QuarkW8A8Int8", "QuarkW4A4MXFP4"]
9 +__all__ = ["QuarkScheme", "QuarkW8A8Fp8", "QuarkW8A8Int8", "QuarkOCP_MX"]
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py renamed
+95 -45
@@ -1,22 +1,34 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from functools import cache
5 -from typing import Any, Callable, Optional
4 +from fractions import Fraction
5 +from functools import cache, partial
6 +from typing import Any, Callable, Optional, Union
7
8 import torch
9 import torch.nn.functional as F
10
11 from vllm import envs
11 -from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
12 +from vllm.logger import init_logger
13 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
13 - OCP_MX_BLOCK_SIZE,
14 dequant_mxfp4,
15 quant_dequant_mxfp4,
16 )
17 +from vllm.model_executor.layers.quantization.utils.mxfp6_utils import (
18 + dequant_mxfp6,
19 + quant_dequant_mxfp6,
20 +)
21 +from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
22 + OCP_MX_BLOCK_SIZE,
23 + OCP_MX_Scheme,
24 +)
25 from vllm.model_executor.parameter import GroupQuantScaleParameter, PackedvLLMParameter
26 from vllm.platforms import current_platform
27
28 +from .quark_scheme import QuarkScheme
29 +
30 +logger = init_logger(__name__)
31 +
32
33 @cache
34 def is_rocm_aiter_fp4_asm_gemm_enabled() -> bool:
@@ -96,14 +108,11 @@ try:
108 fake_impl=gemm_with_dynamic_quant_fake,
109 dispatch_key=current_platform.dispatch_key,
110 )
99 -
111 except (ImportError, AttributeError):
112 dynamic_mxfp4_quant = gemm_afp4wfp4 = None
113
103 -__all__ = ["QuarkW4A4MXFP4"]
104 -
114
106 -class QuarkW4A4MXFP4(QuarkScheme):
115 +class QuarkOCP_MX(QuarkScheme):
116 def __init__(
117 self, weight_quant_spec: dict[str, Any], input_quant_spec: dict[str, Any]
118 ):
@@ -111,8 +120,45 @@ class QuarkW4A4MXFP4(QuarkScheme):
120 self.qscheme = "per_group"
121 self.weight_quant_spec = weight_quant_spec
122 self.input_quant_spec = input_quant_spec
114 - self.emulate = not current_platform.supports_mx()
123 +
124 + self.weight_dtype = weight_quant_spec["dtype"].replace("fp", "mxfp")
125 + self.input_dtype = input_quant_spec["dtype"].replace("fp", "mxfp")
126 +
127 + self.ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype(
128 + self.input_dtype, self.weight_dtype
129 + )
130 +
131 + if self.weight_dtype == "mxfp4":
132 + self.packed_factor: Union[int, Fraction] = 2
133 + self.dequant_func = dequant_mxfp4
134 + else:
135 + self.packed_factor = Fraction(numerator=8, denominator=6)
136 + self.dequant_func = partial(
137 + dequant_mxfp6, quant_dtype=self.weight_dtype.replace("mx", "")
138 + )
139 +
140 + if self.input_dtype == "mxfp4":
141 + self.quant_dequant_func = quant_dequant_mxfp4
142 + else:
143 + self.quant_dequant_func = partial(
144 + quant_dequant_mxfp6, quant_dtype=self.input_dtype.replace("mx", "")
145 + )
146 +
147 + self.static_input_scales = not input_quant_spec.get("is_dynamic")
148 +
149 + if self.static_input_scales:
150 + raise NotImplementedError(
151 + "QuarkOCP_MX with static input scales is currently not "
152 + "implemented. Please open an issue."
153 + )
154 +
155 + # TODO: integrate (or test) mixed-precision kernel.
156 + self.emulate = not current_platform.supports_mx() or (
157 + self.input_dtype != "mxfp4" or self.weight_dtype != "mxfp4"
158 + )
159 +
160 self.rocm_use_aiter_fp4_asm_gemm = is_rocm_aiter_fp4_asm_gemm_enabled()
161 +
162 if not self.emulate and (dynamic_mxfp4_quant is None or gemm_afp4wfp4 is None):
163 # Currently need these kernels if not emulating
164 raise NotImplementedError(
@@ -121,6 +167,41 @@ class QuarkW4A4MXFP4(QuarkScheme):
167 "https://github.com/ROCm/aiter for installation details."
168 )
169
170 + if not current_platform.supports_mx():
171 + logger.warning_once(
172 + "The current platform does not support native MXFP4/MXFP6 "
173 + "computation. Simulated weight dequantization and activation "
174 + "QDQ (quantize and dequantize) will be used, with the linear "
175 + "layers computed in high precision."
176 + )
177 +
178 + if current_platform.supports_mx() and (
179 + self.input_dtype != "mxfp4" or self.weight_dtype != "mxfp4"
180 + ):
181 + logger.warning_once(
182 + "The current platform supports native MXFP4/MXFP6 "
183 + f"computation, but kernels for input_dtype={self.input_dtype} "
184 + f"and weight_dtype={self.weight_dtype} are not yet integrated "
185 + "in vLLM. Simulated weight dequantization and activation "
186 + "QDQ (quantize and dequantize) will be used, with the linear "
187 + "layers computed in high precision."
188 + )
189 +
190 + def get_packed_dim(self, dim: int, quant_dtype: str):
191 + if quant_dtype == "mxfp4":
192 + assert dim % 2 == 0
193 + return dim // 2
194 + elif quant_dtype in {"mxfp6_e3m2", "mxfp6_e2m3"}:
195 + # FP6 packs 4 * 6 = 24 bits on 3 bytes.
196 + assert (dim * 3) % 4 == 0
197 + return (dim * 3) // 4
198 + else:
199 + raise NotImplementedError(
200 + "Unsupported quant_dtype in QuarkOCP_MX.get_packed_dim, "
201 + f"got quant_dtype={quant_dtype}. Something is wrong, please "
202 + "open an issue."
203 + )
204 +
205 @classmethod
206 def get_min_capability(cls) -> int:
207 return 70
@@ -132,37 +213,6 @@ class QuarkW4A4MXFP4(QuarkScheme):
213 layer.weight_scale = torch.nn.Parameter(
214 layer.weight_scale.data, requires_grad=False
215 )
135 - try:
136 - from quark.torch.export.nn.modules import realquantizer
137 - from quark.torch.quantization.config.config import QuantizationSpec
138 - except ImportError as err:
139 - raise ImportError(
140 - "The package `amd-quark` is required to use AMD Quark "
141 - "MX-FP4 models. Please install it with `pip install "
142 - "amd-quark`."
143 - ) from err
144 -
145 - weight_quant_spec = QuantizationSpec.from_dict(self.weight_quant_spec)
146 -
147 - weight_quantizer = realquantizer.get_real_quantizer(
148 - qspec=weight_quant_spec,
149 - quantizer=None,
150 - real_quantized=True,
151 - reorder=False,
152 - float_dtype=self.out_dtype,
153 - scale_shape=layer.weight_scale.shape,
154 - zero_point_shape=None,
155 - )
156 - weight_quantizer.scale.data = layer.weight_scale.data
157 -
158 - layer.weight = torch.nn.Parameter(
159 - weight_quantizer(layer.weight.data).to(self.out_dtype),
160 - requires_grad=False,
161 - )
162 - layer.weight_scale = None
163 -
164 - # This call is necessary to release the scales memory.
165 - torch.cuda.empty_cache()
216 else:
217 if self.rocm_use_aiter_fp4_asm_gemm:
218 # shuffle weight scale
@@ -204,13 +254,13 @@ class QuarkW4A4MXFP4(QuarkScheme):
254 weight = PackedvLLMParameter(
255 data=torch.empty(
256 output_size_per_partition,
207 - input_size_per_partition // 2,
257 + self.get_packed_dim(input_size_per_partition, self.weight_dtype),
258 dtype=torch.uint8,
259 ),
260 input_dim=1,
261 output_dim=0,
262 packed_dim=1,
213 - packed_factor=2,
263 + packed_factor=self.packed_factor,
264 weight_loader=weight_loader,
265 )
266 layer.register_parameter("weight", weight)
@@ -235,9 +285,9 @@ class QuarkW4A4MXFP4(QuarkScheme):
285 bias: Optional[torch.Tensor] = None,
286 ) -> torch.Tensor:
287 if self.emulate:
238 - dq_w = dequant_mxfp4(layer.weight, layer.weight_scale, x.dtype)
239 - x = quant_dequant_mxfp4(x)
240 - return F.linear(x, dq_w, bias)
288 + dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype)
289 + qdq_x = self.quant_dequant_func(x)
290 + return F.linear(qdq_x, dq_w, bias)
291 else:
292 return torch.ops.vllm.gemm_with_dynamic_quant(
293 x,
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+8 -2
@@ -10,8 +10,6 @@ from vllm.utils import direct_register_custom_op, is_torch_equal_or_newer
10
11 logger = init_logger(__name__)
12
13 -OCP_MX_BLOCK_SIZE = 32
14 -
13
14 def _swizzle_mxfp4(quant_tensor, scale, num_warps):
15 """weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel"""
@@ -144,6 +142,14 @@ def _quant_dequant_mxfp4_fake(
142 return torch.empty_like(x)
143
144
145 +# Protect these operations into a torch custom op to avoid errors as
146 +# torch._dynamo.exc.Unsupported: Attempted to call function marked as skipped
147 +# Explanation: Dynamo does not know how to trace the builtin
148 +# `kernel_ext.PyCapsule.dq_uint8_mxfp4_to_half.` This function is either a
149 +# Python builtin (e.g. _warnings.warn) or a third-party C/C++ Python
150 +# extension (perhaps created with pybind).
151 +# TODO: Make sure there is no way to avoid having these functions
152 +# marked as skipped by dynamo.
153 try:
154 direct_register_custom_op(
155 op_name="dequant_mxfp4",
vllm/model_executor/layers/quantization/utils/mxfp6_utils.py new
+142
@@ -0,0 +1,142 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_BLOCK_SIZE
6 +from vllm.utils import direct_register_custom_op
7 +
8 +
9 +def _quant_dequant_mxfp6(
10 + x: torch.Tensor,
11 + quant_dtype: str,
12 + scale_calculation_mode: str = "even",
13 +) -> torch.Tensor:
14 + try:
15 + from quark.torch.kernel.hw_emulation.hw_emulation_interface import (
16 + fake_quantize_fp4_fp6_per_group_with_scale,
17 + )
18 + from quark.torch.quantization.utils import even_round, reshape_to_blocks
19 + except ImportError as err:
20 + raise ImportError(
21 + "The package `amd-quark` is required to use "
22 + "MX-FP6 models. Please install it with `pip install "
23 + "amd-quark`."
24 + ) from err
25 +
26 + axis = -1
27 + block_x = reshape_to_blocks(x, OCP_MX_BLOCK_SIZE, axis)
28 + amax, _ = torch.max(torch.abs(block_x), dim=-1, keepdim=True)
29 + amax = amax.squeeze(-1)
30 +
31 + # TODO: there are other rounding strategies supported in quark and in the
32 + # config.json that we do not check for here!
33 + if scale_calculation_mode != "even":
34 + raise NotImplementedError(
35 + f"Scale calculation mode {scale_calculation_mode} is not yet "
36 + "supported in MX-FP6 quantization"
37 + )
38 + scale = even_round(amax, quant_dtype)
39 +
40 + # Apply dequantize(quantize(x)).
41 + x = fake_quantize_fp4_fp6_per_group_with_scale(
42 + x,
43 + scale.to(x.device),
44 + axis=axis,
45 + group_size=OCP_MX_BLOCK_SIZE,
46 + quant_dtype=quant_dtype,
47 + )
48 +
49 + return x
50 +
51 +
52 +def _quant_dequant_mxfp6_fake(
53 + x: torch.Tensor,
54 + quant_dtype: str,
55 + scale_calculation_mode: str = "even",
56 +) -> torch.Tensor:
57 + return torch.empty_like(x)
58 +
59 +
60 +def _dequant_mxfp6(
61 + x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype, quant_dtype: str
62 +) -> torch.Tensor:
63 + try:
64 + from quark.torch.kernel.hw_emulation.hw_emulation_interface import (
65 + dequantize_fp4_fp6_per_group,
66 + )
67 + from quark.torch.utils.pack import create_pack_method
68 + except ImportError as e:
69 + raise ImportError(
70 + "The package `amd-quark` is required to use "
71 + "MX-FP6 models. Please install it with `pip install "
72 + "amd-quark`."
73 + ) from e
74 +
75 + pack_method = create_pack_method(None, dtype=quant_dtype)
76 + unpacked_x = pack_method.unpack(x, reorder=False)
77 +
78 + scale = 2 ** (scale.view(torch.uint8).to(torch.int16) - 127).to(float_dtype)
79 +
80 + # TODO: `dequantize_fp4_fp6_per_group` and `prepare_inputs_per_group`
81 + # always return fp32.
82 + return dequantize_fp4_fp6_per_group(
83 + unpacked_x,
84 + scale,
85 + axis=-1,
86 + group_size=OCP_MX_BLOCK_SIZE,
87 + quant_dtype=quant_dtype,
88 + ).to(float_dtype)
89 +
90 +
91 +def _dequant_mxfp6_fake(
92 + x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype, quant_dtype: str
93 +) -> torch.Tensor:
94 + assert (x.shape[-1] * 4) % 3 == 0
95 + return torch.empty(
96 + (*x.shape[:-1], (x.shape[-1] * 4) // 3), dtype=float_dtype, device=x.device
97 + )
98 +
99 +
100 +# Protect these operations into a torch custom op to avoid errors as
101 +# torch._dynamo.exc.Unsupported: Attempted to call function marked as skipped
102 +# Explanation: Dynamo does not know how to trace the builtin
103 +# `kernel_ext.PyCapsule.dq_uint8_mxfp4_to_half.` This function is either a
104 +# Python builtin (e.g. _warnings.warn) or a third-party C/C++ Python
105 +# extension (perhaps created with pybind).
106 +# TODO: Make sure there is no way to avoid having these functions
107 +# marked as skipped by dynamo.
108 +try:
109 + direct_register_custom_op(
110 + op_name="quant_dequant_mxfp6",
111 + op_func=_quant_dequant_mxfp6,
112 + mutates_args=[],
113 + fake_impl=_quant_dequant_mxfp6_fake,
114 + )
115 +except AttributeError as error:
116 + raise error
117 +
118 +
119 +# Expose keyword arguments.
120 +def quant_dequant_mxfp6(
121 + x: torch.Tensor,
122 + quant_dtype: str,
123 + scale_calculation_mode: str = "even",
124 +) -> torch.Tensor:
125 + return torch.ops.vllm.quant_dequant_mxfp6(x, quant_dtype, scale_calculation_mode)
126 +
127 +
128 +try:
129 + direct_register_custom_op(
130 + op_name="dequant_mxfp6",
131 + op_func=_dequant_mxfp6,
132 + mutates_args=[],
133 + fake_impl=_dequant_mxfp6_fake,
134 + )
135 +except AttributeError as error:
136 + raise error
137 +
138 +
139 +def dequant_mxfp6(
140 + x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype, quant_dtype: str
141 +) -> torch.Tensor:
142 + return torch.ops.vllm.dequant_mxfp6(x, scale, float_dtype, quant_dtype)
vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
+3 -3
@@ -8,9 +8,9 @@ from vllm.logger import init_logger
8 logger = init_logger(__name__)
9
10
11 -def mxfp8_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
11 +def mxfp8_e4m3_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
12 try:
13 - from flashinfer import mxfp8_quantize
13 + from flashinfer import mxfp8_quantize as mxfp8_e4m3_quantize
14 except ImportError as err:
15 raise ImportError(
16 "The package `flashinfer` is required to do "
@@ -18,4 +18,4 @@ def mxfp8_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
18 "`pip install flashinfer`"
19 ) from err
20
21 - return mxfp8_quantize(x, is_sf_swizzled_layout=False)
21 + return mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=False)
vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py new
+54
@@ -0,0 +1,54 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from enum import Enum
4 +from typing import Union
5 +
6 +from vllm.logger import init_logger
7 +
8 +logger = init_logger(__name__)
9 +
10 +OCP_MX_BLOCK_SIZE = 32
11 +
12 +OCP_MX_DTYPES = {
13 + "mxfp4",
14 + "mxfp6_e3m2",
15 + "mxfp6_e2m3",
16 + "mxfp8_e4m3",
17 + "mxfp8_e5m2",
18 + "mxint8",
19 +}
20 +SUPPORTED_OCP_MX_DTYPES = {"mxfp4", "mxfp6_e3m2", "mxfp6_e2m3"}
21 +
22 +
23 +class OCP_MX_Scheme(str, Enum):
24 + w_mxfp4_a_mxfp4 = "w_mxfp4_a_mxfp4"
25 + w_mxfp4_a_mxfp6_e3m2 = "w_mxfp4_a_mxfp6_e3m2"
26 + w_mxfp4_a_mxfp6_e2m3 = "w_mxfp4_a_mxfp6_e2m3"
27 + w_mxfp6_e3m2_a_mxfp6_e3m2 = "w_mxfp6_e3m2_a_mxfp6_e3m2"
28 + w_mxfp6_e2m3_a_mxfp6_e2m3 = "w_mxfp6_e2m3_a_mxfp6_e2m3"
29 +
30 + @classmethod
31 + def from_quant_dtype(
32 + cls, input_dtype: Union[str, None], weight_dtype: Union[str, None]
33 + ):
34 + if input_dtype not in OCP_MX_DTYPES or weight_dtype not in OCP_MX_DTYPES:
35 + return None
36 + elif input_dtype == "mxfp4" and weight_dtype == "mxfp4":
37 + return cls.w_mxfp4_a_mxfp4
38 + elif input_dtype == "mxfp6_e3m2" and weight_dtype == "mxfp4":
39 + return cls.w_mxfp4_a_mxfp6_e3m2
40 + elif input_dtype == "mxfp6_e2m3" and weight_dtype == "mxfp4":
41 + return cls.w_mxfp4_a_mxfp6_e2m3
42 + elif input_dtype == "mxfp6_e3m2" and weight_dtype == "mxfp6_e3m2":
43 + return cls.w_mxfp6_e3m2_a_mxfp6_e3m2
44 + elif input_dtype == "mxfp6_e2m3" and weight_dtype == "mxfp6_e2m3":
45 + return cls.w_mxfp6_e2m3_a_mxfp6_e2m3
46 + else:
47 + logger.warning(
48 + "input_dtype='%s' and"
49 + " weight_dtype='%s' is not supported "
50 + "in OCP_MX_Scheme at the moment.",
51 + input_dtype,
52 + weight_dtype,
53 + )
54 + return None
vllm/scalar_type.py
+3
@@ -337,8 +337,11 @@ class scalar_types:
337 float16_e5m10 = ScalarType.float_IEEE754(5, 10)
338
339 # fp6, https://github.com/usyd-fsalab/fp6_llm/tree/main
340 + # and https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
341 float6_e3m2f = ScalarType.float_(3, 2, True, NanRepr.NONE)
342
343 + float6_e2m3f = ScalarType.float_(2, 3, True, NanRepr.NONE)
344 +
345 # fp4, https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf
346 float4_e2m1f = ScalarType.float_(2, 1, True, NanRepr.NONE)
347
vllm/utils/flashinfer.py
+2 -2
@@ -89,7 +89,7 @@ flashinfer_trtllm_fp8_per_tensor_scale_moe = _lazy_import_wrapper(
89 flashinfer_cutlass_fused_moe = _lazy_import_wrapper(
90 "flashinfer.fused_moe", "cutlass_fused_moe"
91 )
92 -fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize")
92 +flashinfer_fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize")
93 nvfp4_block_scale_interleave = _lazy_import_wrapper(
94 "flashinfer", "nvfp4_block_scale_interleave"
95 )
@@ -442,7 +442,7 @@ __all__ = [
442 "has_flashinfer",
443 "flashinfer_trtllm_fp8_block_scale_moe",
444 "flashinfer_cutlass_fused_moe",
445 - "fp4_quantize",
445 + "flashinfer_fp4_quantize",
446 "nvfp4_block_scale_interleave",
447 "trtllm_fp4_block_scale_moe",
448 "autotune",