@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 332d4cb17b
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 15, 02:15 UTC

[Feature][Quantization] MXFP4 support for MOE models (#17888)

Signed-off-by: Felix Marty <felmarty@amd.com> Signed-off-by: Bowen Bao <bowenbao@amd.com> Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Bowen Bao <bowenbao@amd.com>

fxmarty-amd committed Jul 9, 2025 at 22:19 UTC 332d4cb17b7c78ae69a5c96595b6c765ae295abc
15 files changed +875 -106
docs/features/quantization/quark.md
+25
@@ -229,3 +229,28 @@ python3 quantize_quark.py --model_dir meta-llama/Llama-2-70b-chat-hf \
229 --model_export hf_format \
230 --tasks gsm8k
231 ```
232 +
233 +## Using MXFP4 models
234 +
235 +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 +
237 +The scheme currently only supports dynamic quantization for activations.
238 +
239 +Example usage, after installing the latest AMD Quark release:
240 +
241 +```bash
242 +vllm serve fxmarty/qwen_1.5-moe-a2.7b-mxfp4 --tensor-parallel-size 1
243 +```
244 +
245 +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).
246 +
247 +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:
248 +
249 +```bash
250 +python quantize_quark.py --model_dir Qwen/Qwen1.5-MoE-A2.7B-Chat \
251 + --quant_scheme w_mxfp4_a_mxfp4_sym \
252 + --output_dir qwen_1.5-moe-a2.7b-mxfp4 \
253 + --skip_evaluation \
254 + --model_export hf_format \
255 + --group_size 32
256 +```
tests/kernels/moe/test_moe.py
+1
@@ -174,6 +174,7 @@ def test_fused_moe(
174 use_int8_w8a8=False,
175 use_int8_w8a16=False,
176 use_int4_w4a16=False,
177 + use_mxfp4_w4a4=False,
178 per_act_token_quant=False,
179 block_shape=None)
180
tests/kernels/moe/test_mxfp4_moe.py new
+57
@@ -0,0 +1,57 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import importlib
5 +import importlib.metadata
6 +from dataclasses import dataclass
7 +
8 +import pytest
9 +import torch
10 +from packaging import version
11 +
12 +QUARK_MXFP4_AVAILABLE = importlib.util.find_spec(
13 + "quark") is not None and version.parse(
14 + importlib.metadata.version("amd-quark")) >= version.parse('0.8.99')
15 +
16 +
17 +@dataclass
18 +class ModelCase:
19 + model_id: str
20 + tp: int
21 +
22 +
23 +@pytest.mark.parametrize('model_case', [
24 + ModelCase("fxmarty/qwen_1.5-moe-a2.7b-mxfp4", tp=1),
25 + ModelCase("fxmarty/deepseek_r1_3_layers_mxfp4", tp=8),
26 + ModelCase("fxmarty/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1)
27 +])
28 +@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE,
29 + reason="amd-quark>=0.9 is not available")
30 +def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
31 + if torch.cuda.device_count() < model_case.tp:
32 + pytest.skip(f"This test requires >={model_case.tp} gpus, got only "
33 + f"{torch.cuda.device_count()}")
34 +
35 + with vllm_runner(model_case.model_id,
36 + tensor_parallel_size=model_case.tp,
37 + load_format="dummy") as llm:
38 +
39 + # TODO: llm.apply_model(check_model) currently relies on V0 internals.
40 + # Re-enable this later.
41 + # def check_model(model):
42 + # layer = model.model.layers[0]
43 +
44 + # qkv_proj = layer.self_attn.qkv_proj
45 +
46 + # assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
47 + # assert isinstance(qkv_proj.scheme, QuarkW4A4MXFP4)
48 +
49 + # assert isinstance(layer.mlp.experts.quant_method,
50 + # QuarkW4A4MXFp4MoEMethod)
51 +
52 + # if model_case.model_id == "fxmarty/qwen_1.5-moe-a2.7b-mxfp4":
53 + # llm.apply_model(check_model)
54 +
55 + output = llm.generate_greedy("Today I am in the French Alps and",
56 + max_tokens=20)
57 + assert output
\ No newline at end of file
tests/quantization/reference_mxfp4.py new
+287
@@ -0,0 +1,287 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +BFLOAT16_EXP_BIAS = 127
6 +BFLOAT16_MANTISSA_BITS = 7
7 +BFLOAT16_EXP_BITS = 8
8 +
9 +FLOAT16_EXP_BIAS = 15
10 +FLOAT16_MANTISSA_BITS = 10
11 +FLOAT16_EXP_BITS = 5
12 +
13 +FLOAT8_E8M0_MAX_EXP = 127
14 +FLOAT4_EXP_BIAS = 1
15 +FLOAT4_MANTISSA_BITS = 1
16 +
17 +FLOAT16_VAL_TO_ADD = (1 << (FLOAT16_MANTISSA_BITS - FLOAT4_MANTISSA_BITS - 1))
18 +FLOAT16_SIGN_EXPONENT_MASK = ((
19 + (1 << (FLOAT16_EXP_BITS + 1)) - 1) << FLOAT16_MANTISSA_BITS)
20 +
21 +BFLOAT16_VAL_TO_ADD = (1 <<
22 + (BFLOAT16_MANTISSA_BITS - FLOAT4_MANTISSA_BITS - 1))
23 +BFLOAT16_SIGN_EXPONENT_MASK = ((
24 + (1 << (BFLOAT16_EXP_BITS + 1)) - 1) << BFLOAT16_MANTISSA_BITS)
25 +
26 +
27 +def e8m0_to_half(scale, half_dtype: torch.dtype):
28 + assert scale.dtype == torch.uint8
29 +
30 + scale_exp = scale.to(torch.int16) - 127
31 +
32 + # This can be implemented with bitwise operations in a proper kernel.
33 + scale_half = 2.0**(scale_exp.to(torch.float))
34 +
35 + return scale_half.to(half_dtype)
36 +
37 +
38 +def upcast_fp4_to_fp16_or_bf16(val, float_dtype: torch.dtype,
39 + half_exp_bias: int, half_mantissa_bits: int):
40 + assert val.dtype == torch.uint8
41 +
42 + unpacked = torch.zeros(*val.shape[:-1],
43 + val.shape[-1] * 2,
44 + dtype=torch.uint8,
45 + device=val.device)
46 + unpacked[..., 1::2] = (val >> 4) & 0x0F # Extract high 4 bits.
47 + unpacked[..., ::2] = val & 0x0F # Extract low 4 bits.
48 +
49 + # Takes one float4 values represented as b0000xxxx,
50 + # and converts it to the corresponding float16 value.
51 +
52 + sign = unpacked >> 3
53 +
54 + exp = (unpacked >> 1) & 3
55 + new_mantissa = unpacked & 1
56 +
57 + # if exp == 0 and new_mantissa == 0:
58 + # new_exp = 0
59 + # else:
60 + # new_exp = exp - FLOAT4_EXP_BIAS + FLOAT16_EXP_BIAS
61 +
62 + # int8_t works with float16, but may overflow with bfloat16.
63 + new_exp = exp - FLOAT4_EXP_BIAS + half_exp_bias
64 +
65 + # Cast b0000 to 0. in fp16/bf16.
66 + new_exp = new_exp * torch.logical_or(exp > 0, new_mantissa > 0)
67 +
68 + # Cast b0001 to 0.5 in fp16/bf16.
69 + new_mantissa = torch.logical_and(new_mantissa, exp > 0)
70 +
71 + new_mantissa = new_mantissa.to(torch.int32)
72 + new_exp = new_exp.to(torch.int32)
73 + sign = sign.to(torch.int32)
74 +
75 + qdq_val = (sign << 15) + (new_exp << half_mantissa_bits) + (
76 + new_mantissa << (half_mantissa_bits - 1))
77 +
78 + assert qdq_val.max() <= 65535
79 + assert qdq_val.min() >= 0
80 + qdq_val = qdq_val.to(torch.uint16)
81 +
82 + result = qdq_val.view(float_dtype)
83 +
84 + return result
85 +
86 +
87 +def dq_mxfp4_torch(x: torch.Tensor, scale: torch.Tensor,
88 + float_dtype: torch.dtype) -> torch.Tensor:
89 + assert x.dtype == torch.uint8
90 + assert scale.dtype == torch.uint8
91 +
92 + if float_dtype == torch.float16:
93 + half_exp_bias = FLOAT16_EXP_BIAS
94 + half_mantissa_bits = FLOAT16_MANTISSA_BITS
95 + elif float_dtype == torch.bfloat16:
96 + half_exp_bias = BFLOAT16_EXP_BIAS
97 + half_mantissa_bits = BFLOAT16_MANTISSA_BITS
98 +
99 + scale_half = e8m0_to_half(scale, half_dtype=float_dtype)
100 +
101 + x_half = upcast_fp4_to_fp16_or_bf16(x,
102 + float_dtype=float_dtype,
103 + half_exp_bias=half_exp_bias,
104 + half_mantissa_bits=half_mantissa_bits)
105 +
106 + x_half = x_half.reshape(*x_half.shape[:-1], -1, 32)
107 + x_half = x_half * scale_half[..., None]
108 + x_half = x_half.reshape(*x_half.shape[:-2], -1)
109 +
110 + return x_half
111 +
112 +
113 +def fp16_to_fp4_simulate(val, half_mantissa_bits: int, half_exp_bits: int,
114 + half_exp_bias: int):
115 + # Casts an fp16/bf16 input to the restricted values of float4_e2m1,
116 + # that is to say [0., 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0,
117 + # -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0].
118 +
119 + float_type = val.dtype
120 +
121 + # "rshift_cuda" not implemented for 'UInt16'
122 + val_view = val.view(torch.int16) #.to(torch.int32)
123 +
124 + exp = val_view >> half_mantissa_bits
125 + exp = exp & ((1 << half_exp_bits) - 1)
126 +
127 + exp = exp.view(torch.uint16).to(torch.int32)
128 +
129 + sign = (val_view >> (half_mantissa_bits + half_exp_bits)) & 1
130 +
131 + mantissa_last = (val_view >> (half_mantissa_bits - 1)) & 1
132 +
133 + exp_unbias = exp - half_exp_bias
134 + new_exp = exp_unbias + FLOAT4_EXP_BIAS
135 +
136 + exp_shift = (new_exp <= 0) * (1 - new_exp)
137 +
138 + # Typically 9.
139 + # Take the min to prevent overflow on `uint16_t half`. This is the case for
140 + # very small values, correctly mapped to `round_close`.
141 + tail_bits = half_mantissa_bits - FLOAT4_MANTISSA_BITS + exp_shift
142 + tail_bits[tail_bits >= 16] = 16
143 +
144 + mantissa_plus_one = val_view & ((1 << (half_mantissa_bits + 1)) - 1)
145 +
146 + half = 1 << (tail_bits - 1)
147 +
148 + tail = mantissa_plus_one & ((1 << tail_bits) - 1)
149 +
150 + round_close = (tail < half) # round towards 0
151 + round_away = (tail > half) # round away from 0
152 + tie = tail == half
153 +
154 + new_mantissa_close = torch.zeros(val.shape,
155 + device=val.device,
156 + dtype=torch.bool)
157 + new_exp_close = torch.zeros(val.shape,
158 + device=val.device,
159 + dtype=torch.uint16)
160 +
161 + new_mantissa_away = torch.zeros(val.shape,
162 + device=val.device,
163 + dtype=torch.bool)
164 + new_exp_away = torch.zeros(val.shape,
165 + device=val.device,
166 + dtype=torch.uint16)
167 +
168 + new_exp_tie = torch.zeros(val.shape, device=val.device, dtype=torch.uint16)
169 +
170 + # 1. round down
171 + # if new_exp == 0: # case [0.5, 0.749999]
172 + # new_mantissa = 0
173 + # elif new_exp < 0: # case [0, 0.24999]
174 + # new_mantissa = 0
175 + # else:
176 + # new_mantissa = mantissa_last
177 +
178 + new_mantissa_close = (new_exp > 0) * mantissa_last
179 + new_exp_close = exp
180 +
181 + # # 2. round up
182 + # if new_exp <= 0: # case [0.250001, 0.499999] and [0.75001, 0.99999]
183 + # new_mantissa = 0
184 + # new_exp += 1
185 + # elif mantissa_last == 0:
186 + # new_mantissa = 1
187 + # else:
188 + # new_mantissa = 0
189 + # new_exp += 1
190 +
191 + new_mantissa_away = torch.logical_and(new_exp > 0, mantissa_last == 0)
192 + new_exp_away = exp + torch.logical_or(new_exp <= 0, mantissa_last == 1)
193 +
194 + # # 3. tie
195 + # 0.25 -> 0. (handled by `exp > (half_exp_bias - 2)`)
196 + # 0.75 -> 1.
197 + # 1.25 -> 1.
198 + # 1.75 -> 2.
199 + # 2.5 -> 2.
200 + # 3.5 -> 4.
201 + # 5. -> 4.
202 + new_exp_tie = (exp > (half_exp_bias - 2)) * (exp + (mantissa_last == 1))
203 +
204 + # Gather round up, round down and tie.
205 + new_exp = round_away * new_exp_away \
206 + + round_close * new_exp_close \
207 + + tie * new_exp_tie
208 +
209 + new_mantissa = round_away * new_mantissa_away \
210 + + round_close * new_mantissa_close
211 +
212 + # if new_exp > 3:
213 + # new_mantissa = 1
214 + new_mantissa = new_mantissa + (new_exp >
215 + (2 + half_exp_bias)) * (new_mantissa == 0)
216 +
217 + # Clamp the exponent to acceptable values.
218 + new_exp = (new_exp >= (half_exp_bias - 2)) * torch.clamp(
219 + new_exp, half_exp_bias - 2, half_exp_bias + 2)
220 +
221 + sign = sign.to(torch.int32)
222 + new_mantissa = new_mantissa.to(torch.int32)
223 +
224 + qdq_val = (sign << 15) + (new_exp << half_mantissa_bits) + (
225 + new_mantissa << (half_mantissa_bits - 1))
226 +
227 + assert qdq_val.max() <= 65535
228 + assert qdq_val.min() >= 0
229 + assert qdq_val.dtype == torch.int32
230 + qdq_val = qdq_val.to(torch.uint16)
231 +
232 + result = qdq_val.view(float_type)
233 + return result
234 +
235 +
236 +def qdq_mxfp4_torch(x: torch.Tensor,
237 + scale_calculation_mode: str = "even") -> torch.Tensor:
238 + half_dtype = x.dtype
239 +
240 + if half_dtype == torch.float16:
241 + half_mantissa_bits = FLOAT16_MANTISSA_BITS
242 + half_exp_bits = FLOAT16_EXP_BITS
243 + half_exp_bias = FLOAT16_EXP_BIAS
244 + val_to_add = FLOAT16_VAL_TO_ADD
245 + sign_exponent_mask = FLOAT16_SIGN_EXPONENT_MASK
246 + elif half_dtype == torch.bfloat16:
247 + half_mantissa_bits = BFLOAT16_MANTISSA_BITS
248 + half_exp_bits = BFLOAT16_EXP_BITS
249 + half_exp_bias = BFLOAT16_EXP_BIAS
250 + val_to_add = BFLOAT16_VAL_TO_ADD
251 + sign_exponent_mask = BFLOAT16_SIGN_EXPONENT_MASK
252 + else:
253 + raise ValueError("not implemented")
254 +
255 + x = x.reshape(*x.shape[:-1], -1, 32)
256 +
257 + block_max = torch.max(torch.abs(x), dim=-1).values
258 +
259 + block_max = block_max.view(torch.uint16).to(torch.int32)
260 +
261 + block_max_uint = torch.bitwise_and(block_max + val_to_add,
262 + sign_exponent_mask)
263 +
264 + assert block_max_uint.max() <= 65535
265 + assert block_max_uint.min() >= 0
266 + assert block_max_uint.dtype == torch.int32
267 + block_max_uint = block_max_uint.to(torch.uint16)
268 +
269 + block_max = block_max_uint.view(half_dtype)
270 +
271 + scale_exp = FLOAT8_E8M0_MAX_EXP + torch.floor(torch.log2(block_max)).to(
272 + torch.int32) - 2
273 +
274 + scale_exp = torch.clamp(scale_exp, 0, 2 * FLOAT8_E8M0_MAX_EXP)
275 +
276 + scale = 2.0**(scale_exp - FLOAT8_E8M0_MAX_EXP)
277 + scale = scale.to(half_dtype)
278 +
279 + x = x / scale[..., None]
280 +
281 + x_fp4 = fp16_to_fp4_simulate(x,
282 + half_exp_bits=half_exp_bits,
283 + half_mantissa_bits=half_mantissa_bits,
284 + half_exp_bias=half_exp_bias)
285 +
286 + x_fp4 = x_fp4 * scale[..., None]
287 + return x_fp4.reshape(*x_fp4.shape[:-2], -1)
tests/quantization/test_quark.py
+171
@@ -3,15 +3,44 @@
3 """Test model set-up and weight loading for quark-quantized models.
4
5 Run `pytest tests/quantization/test_quark.py`.
6 +
7 +See also `tests/kernels/moe/test_mxfp4_moe.py`.
8 """
9
10 +import importlib
11 +import importlib.metadata
12 +import os
13 +from dataclasses import dataclass
14 +
15 +import huggingface_hub
16 +import lm_eval
17 import pytest
18 import torch
19 +from packaging import version
20
21 from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501
22 QuarkLinearMethod, QuarkW8A8Fp8, QuarkW8A8Int8)
23 from vllm.platforms import current_platform
24
25 +from .reference_mxfp4 import dq_mxfp4_torch, qdq_mxfp4_torch
26 +
27 +QUARK_MXFP4_AVAILABLE = importlib.util.find_spec(
28 + "quark") is not None and version.parse(
29 + importlib.metadata.version("amd-quark")) >= version.parse('0.8.99')
30 +
31 +if QUARK_MXFP4_AVAILABLE:
32 + from quark.torch.export.nn.modules.realquantizer import (
33 + StaticScaledRealQuantizer)
34 + from quark.torch.kernel import mx as mx_kernel
35 + from quark.torch.quantization.config.config import FP4PerGroupSpec
36 +
37 +try:
38 + huggingface_hub.list_repo_refs(
39 + "amd/Llama-3.3-70B-Instruct-WMXFP4-AMXFP4-KVFP8-Scale-UINT8-SQ")
40 + HF_HUB_AMD_ORG_ACCESS = True
41 +except huggingface_hub.errors.RepositoryNotFoundError:
42 + HF_HUB_AMD_ORG_ACCESS = False
43 +
44
45 @pytest.fixture(scope="function", autouse=True)
46 def use_v0_only(monkeypatch):
@@ -90,3 +119,145 @@ def test_quark_fp8_parity(vllm_runner):
119
120 for key in fp8_state_dict:
121 assert torch.equal(fp8_state_dict[key], quark_state_dict[key])
122 +
123 +
124 +@dataclass
125 +class ModelCase:
126 + model_id: str
127 + tp: int
128 +
129 +
130 +@dataclass
131 +class GSM8KAccuracyTestConfig:
132 + model_name: str
133 + excepted_value: float
134 +
135 + def get_model_args(self) -> str:
136 + return (
137 + f"pretrained={self.model_name},"
138 + "dtype=auto,add_bos_token=True,tensor_parallel_size=8,gpu_memory_utilization=0.7,max_model_len=38768"
139 + )
140 +
141 +
142 +ACCURACY_CONFIGS = [
143 + # Private model.
144 + GSM8KAccuracyTestConfig(
145 + model_name="amd/DeepSeek-R1-WMXFP4-AMXFP4-Scale-UINT8-MoE-Quant",
146 + excepted_value=0.96),
147 +]
148 +
149 +
150 +@pytest.mark.parametrize("config", ACCURACY_CONFIGS)
151 +@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE,
152 + reason="amd-quark>=0.9 is not available")
153 +@pytest.mark.skipif(
154 + not HF_HUB_AMD_ORG_ACCESS,
155 + reason="Read access to huggingface.co/amd is required for this test.")
156 +def test_mxfp4_gsm8k_correctness(config: GSM8KAccuracyTestConfig):
157 + if torch.cuda.device_count() < 8:
158 + pytest.skip(
159 + f"This test requires >=8 gpus, got only {torch.cuda.device_count()}"
160 + )
161 +
162 + task = "gsm8k"
163 + rtol = 0.03
164 +
165 + os.environ["VLLM_USE_TRITON_FLASH_ATTN"] = "0"
166 +
167 + results = lm_eval.simple_evaluate(
168 + model="vllm",
169 + model_args=config.get_model_args(),
170 + tasks=task,
171 + batch_size=64,
172 + num_fewshot=8,
173 + )
174 +
175 + EXPECTED_VALUE = config.excepted_value
176 + measured_value = results["results"][task]["exact_match,strict-match"]
177 + assert (measured_value - rtol < EXPECTED_VALUE
178 + and measured_value + rtol > EXPECTED_VALUE
179 + ), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
180 +
181 + del os.environ["VLLM_USE_TRITON_FLASH_ATTN"]
182 +
183 +
184 +@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE,
185 + reason="amd-quark>=0.9 is not available")
186 +@pytest.mark.parametrize("float_dtype", [torch.bfloat16, torch.float16])
187 +@pytest.mark.parametrize("scalings",
188 + [[2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4]])
189 +def test_mxfp4_fused_qdq_match_quark(float_dtype: torch.dtype,
190 + scalings: list[int]):
191 + torch.manual_seed(0)
192 +
193 + hidden_size = 64 * 32
194 + inp = (torch.rand(1, hidden_size, dtype=float_dtype, device="cuda") -
195 + 0.5) * 2
196 + for i in range(hidden_size // 32):
197 + inp[:, i * 32:(i + 1) *
198 + 32] = inp[:, i * 32:(i + 1) * 32] * scalings[i % len(scalings)]
199 +
200 + inp_kernel = inp.clone()
201 + inp_kernel_clone = inp_kernel.clone()
202 +
203 + res_hip = mx_kernel.qdq_mxfp4_hip(inp_kernel_clone, "even")
204 + res_torch = qdq_mxfp4_torch(inp_kernel, "even")
205 +
206 + for i in range(hidden_size // 32):
207 + assert torch.all(torch.isfinite(res_hip[:, i * 32:(i + 1) * 32]))
208 + assert torch.all(torch.isfinite(res_torch[:, i * 32:(i + 1) * 32]))
209 +
210 + torch.testing.assert_close(res_hip[:, i * 32:(i + 1) * 32],
211 + res_torch[:, i * 32:(i + 1) * 32])
212 +
213 +
214 +@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE,
215 + reason="amd-quark>=0.9 is not available")
216 +@pytest.mark.parametrize("float_dtype", [torch.bfloat16, torch.float16])
217 +@pytest.mark.parametrize("scalings",
218 + [[2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4]])
219 +def test_mxfp4_dequant_kernel_match_quark(float_dtype: torch.dtype,
220 + scalings: list[int]):
221 + qspec = FP4PerGroupSpec(
222 + ch_axis=-1,
223 + group_size=32,
224 + scale_format="e8m0",
225 + scale_calculation_mode="even",
226 + is_dynamic=False,
227 + ).to_quantization_spec()
228 +
229 + weight_quantizer = StaticScaledRealQuantizer(
230 + qspec=qspec,
231 + quantizer=None,
232 + reorder=False,
233 + real_quantized=True,
234 + float_dtype=float_dtype,
235 + device="cuda",
236 + )
237 +
238 + observer = qspec.observer_cls(qspec, device="cuda")
239 +
240 + hidden_size = 512
241 + shape = (11008, hidden_size)
242 +
243 + w = (torch.rand(shape, device="cuda", dtype=float_dtype) - 0.5) * 2
244 +
245 + # Make it so that different groups have different scales.
246 + for i in range(hidden_size // 32):
247 + w[:, i * 32:(i + 1) *
248 + 32] = w[:, i * 32:(i + 1) * 32] * scalings[i % len(scalings)]
249 +
250 + observer(w)
251 + scale, _ = observer._calculate_qparams()
252 + weight_quantizer.scale = scale
253 +
254 + w_mxfp4 = weight_quantizer.to_real_quantize_params(w).to("cuda")
255 + weight_quantizer.maybe_convert_and_transpose_scale()
256 +
257 + scale = weight_quantizer.scale
258 +
259 + out_hip = mx_kernel.dq_mxfp4_hip(w_mxfp4, scale, float_dtype)
260 +
261 + out_torch = dq_mxfp4_torch(w_mxfp4, scale, float_dtype)
262 +
263 + assert torch.equal(out_hip, out_torch)
vllm/envs.py
-9
@@ -94,7 +94,6 @@ if TYPE_CHECKING:
94 VLLM_ROCM_FP8_PADDING: bool = True
95 VLLM_ROCM_MOE_PADDING: bool = True
96 VLLM_ROCM_CUSTOM_PAGED_ATTN: bool = True
97 - VLLM_QUARK_EMU_MEM_OPT: bool = False
97 VLLM_ENABLE_V1_MULTIPROCESSING: bool = True
98 VLLM_LOG_BATCHSIZE_INTERVAL: float = -1
99 VLLM_DISABLE_COMPILE_CACHE: bool = False
@@ -723,14 +722,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
722 lambda: maybe_convert_int(
723 os.environ.get("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", None)),
724
726 - # If set, when running in Quark emulation mode, do not dequantize the
727 - # weights at load time. Instead, dequantize weights on-the-fly during
728 - # kernel execution.
729 - # This allows running larger models at the cost of slower inference.
730 - # This flag has no effect when not running in Quark emulation mode.
731 - "VLLM_QUARK_EMU_MEM_OPT":
732 - lambda: bool(int(os.getenv("VLLM_QUARK_EMU_MEM_OPT", "0"))),
733 -
725 # Divisor for dynamic query scale factor calculation for FP8 KV Cache
726 "Q_SCALE_CONSTANT":
727 lambda: int(os.getenv("Q_SCALE_CONSTANT", "200")),
vllm/model_executor/layers/fused_moe/config.py
+6 -1
@@ -50,11 +50,14 @@ def get_config_quant_dtype(
50 use_int8_w8a8: bool,
51 use_int8_w8a16: bool,
52 use_int4_w4a16: bool,
53 -) -> Optional[torch.dtype]:
53 + use_mxfp4_w4a4: bool,
54 +) -> Union[None, torch.dtype, str]:
55 if use_fp8_w8a8:
56 return torch.float8_e4m3fn
57 elif use_int8_w8a8:
58 return torch.int8
59 + elif use_mxfp4_w4a4:
60 + return "mxfp4"
61 return None
62
63
@@ -126,6 +129,7 @@ class FusedMoEQuantConfig:
129 use_int8_w8a8: bool = False,
130 use_int8_w8a16: bool = False,
131 use_int4_w4a16: bool = False,
132 + use_mxfp4_w4a4: bool = False,
133 per_act_token_quant: bool = False,
134 per_out_ch_quant: bool = False,
135 block_shape: Optional[list[int]] = None,
@@ -144,6 +148,7 @@ class FusedMoEQuantConfig:
148 use_int8_w8a8=use_int8_w8a8,
149 use_int8_w8a16=use_int8_w8a16,
150 use_int4_w4a16=use_int4_w4a16,
151 + use_mxfp4_w4a4=use_mxfp4_w4a4,
152 )
153 return FusedMoEQuantConfig(
154 quant_dtype,
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+8
@@ -632,6 +632,7 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
632 use_int8_w8a8: bool = False,
633 use_int8_w8a16: bool = False,
634 use_int4_w4a16: bool = False,
635 + use_mxfp4_w4a4: bool = False,
636 block_shape: Optional[list[int]] = None,
637 per_act_token_quant: bool = False,
638 ):
@@ -641,12 +642,14 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
642 use_int8_w8a8=use_int8_w8a8,
643 use_int8_w8a16=use_int8_w8a16,
644 use_int4_w4a16=use_int4_w4a16,
645 + use_mxfp4_w4a4=use_mxfp4_w4a4,
646 per_act_token_quant=per_act_token_quant,
647 block_shape=block_shape,
648 ))
649 assert not use_int8_w8a8, "NYI"
650 assert not use_int8_w8a16, "NYI"
651 assert not use_int4_w4a16, "NYI"
652 + assert not use_mxfp4_w4a4, "NYI"
653 self.max_num_tokens = max_num_tokens
654 self.num_dispatchers = num_dispatchers
655
@@ -838,6 +841,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
841 use_int8_w8a8: bool = False,
842 use_int8_w8a16: bool = False,
843 use_int4_w4a16: bool = False,
844 + use_mxfp4_w4a4: bool = False,
845 per_act_token_quant: bool = False,
846 block_shape: Optional[list[int]] = None,
847 ):
@@ -847,18 +851,21 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
851 use_int8_w8a8=use_int8_w8a8,
852 use_int8_w8a16=use_int8_w8a16,
853 use_int4_w4a16=use_int4_w4a16,
854 + use_mxfp4_w4a4=use_mxfp4_w4a4,
855 per_act_token_quant=per_act_token_quant,
856 block_shape=block_shape,
857 ))
858 assert not use_int8_w8a8, "NYI"
859 assert not use_int8_w8a16, "NYI"
860 assert not use_int4_w4a16, "NYI"
861 + assert not use_mxfp4_w4a4, "NYI"
862 assert max_num_tokens > 0
863 assert num_dispatchers > 0
864 self.use_fp8_w8a8 = use_fp8_w8a8
865 self.use_int8_w8a8 = use_int8_w8a8
866 self.use_int4_w4a16 = use_int4_w4a16
867 self.use_int8_w8a16 = use_int8_w8a16
868 + self.use_mxfp4_w4a4 = use_mxfp4_w4a4
869 self.max_num_tokens = max_num_tokens
870 self.num_dispatchers = num_dispatchers
871
@@ -941,6 +948,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
948 config_dtype = get_config_dtype_str(use_fp8_w8a8=self.use_fp8_w8a8,
949 use_int8_w8a16=self.use_int8_w8a16,
950 use_int4_w4a16=self.use_int4_w4a16,
951 + use_mxfp4_w4a4=self.use_mxfp4_w4a4,
952 dtype=hidden_states.dtype)
953
954 config = try_get_optimal_moe_config(
vllm/model_executor/layers/fused_moe/fused_moe.py
+44 -9
@@ -27,6 +27,8 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
27 MoEPrepareAndFinalizeNoEP)
28 from vllm.model_executor.layers.fused_moe.utils import (
29 _resize_cache, moe_kernel_quantize_input)
30 +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
31 + dequant_mxfp4)
32 from vllm.platforms import current_platform
33 from vllm.triton_utils import tl, triton
34 from vllm.utils import direct_register_custom_op
@@ -973,13 +975,16 @@ def get_config_dtype_str(
975 dtype: torch.dtype,
976 use_int4_w4a16: Optional[bool] = False,
977 use_int8_w8a16: Optional[bool] = False,
976 - use_fp8_w8a8: Optional[bool] = False) -> Optional[str]:
978 + use_fp8_w8a8: Optional[bool] = False,
979 + use_mxfp4_w4a4: Optional[bool] = False) -> Optional[str]:
980 if use_fp8_w8a8:
981 return "fp8_w8a8"
982 elif use_int8_w8a16:
983 return "int8_w8a16"
984 elif use_int4_w4a16:
985 return "int4_w4a16"
986 + elif use_mxfp4_w4a4:
987 + return "mxfp4_w4a4"
988 elif dtype == torch.float:
989 # avoiding cases where kernel fails when float32 MoE
990 # use fp16/bfloat16 configs
@@ -998,6 +1003,7 @@ def inplace_fused_experts(hidden_states: torch.Tensor,
1003 use_int8_w8a8: bool = False,
1004 use_int8_w8a16: bool = False,
1005 use_int4_w4a16: bool = False,
1006 + use_mxfp4_w4a4: bool = False,
1007 per_channel_quant: bool = False,
1008 global_num_experts: int = -1,
1009 expert_map: Optional[torch.Tensor] = None,
@@ -1011,9 +1017,9 @@ def inplace_fused_experts(hidden_states: torch.Tensor,
1017 fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,
1018 activation, apply_router_weight_on_input, use_fp8_w8a8,
1019 use_int8_w8a8, use_int8_w8a16, use_int4_w4a16,
1014 - per_channel_quant, global_num_experts, expert_map,
1015 - w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
1016 - block_shape)
1020 + use_mxfp4_w4a4, per_channel_quant, global_num_experts,
1021 + expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,
1022 + a2_scale, block_shape)
1023
1024
1025 def inplace_fused_experts_fake(
@@ -1028,6 +1034,7 @@ def inplace_fused_experts_fake(
1034 use_int8_w8a8: bool = False,
1035 use_int8_w8a16: bool = False,
1036 use_int4_w4a16: bool = False,
1037 + use_mxfp4_w4a4: bool = False,
1038 per_channel_quant: bool = False,
1039 global_num_experts: int = -1,
1040 expert_map: Optional[torch.Tensor] = None,
@@ -1062,6 +1069,7 @@ def outplace_fused_experts(
1069 use_int8_w8a8: bool = False,
1070 use_int8_w8a16: bool = False,
1071 use_int4_w4a16: bool = False,
1072 + use_mxfp4_w4a4: bool = False,
1073 per_channel_quant: bool = False,
1074 global_num_experts: int = -1,
1075 expert_map: Optional[torch.Tensor] = None,
@@ -1075,10 +1083,10 @@ def outplace_fused_experts(
1083 return fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids,
1084 False, activation, apply_router_weight_on_input,
1085 use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16,
1078 - use_int4_w4a16, per_channel_quant,
1079 - global_num_experts, expert_map, w1_scale,
1080 - w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
1081 - block_shape)
1086 + use_int4_w4a16, use_mxfp4_w4a4,
1087 + per_channel_quant, global_num_experts,
1088 + expert_map, w1_scale, w2_scale, w1_zp, w2_zp,
1089 + a1_scale, a2_scale, block_shape)
1090
1091
1092 def outplace_fused_experts_fake(
@@ -1092,6 +1100,7 @@ def outplace_fused_experts_fake(
1100 use_int8_w8a8: bool = False,
1101 use_int8_w8a16: bool = False,
1102 use_int4_w4a16: bool = False,
1103 + use_mxfp4_w4a4: bool = False,
1104 per_channel_quant: bool = False,
1105 global_num_experts: int = -1,
1106 expert_map: Optional[torch.Tensor] = None,
@@ -1145,6 +1154,7 @@ def fused_experts(
1154 use_int8_w8a8: bool = False,
1155 use_int8_w8a16: bool = False,
1156 use_int4_w4a16: bool = False,
1157 + use_mxfp4_w4a4: bool = False,
1158 per_channel_quant: bool = False,
1159 global_num_experts: int = -1,
1160 expert_map: Optional[torch.Tensor] = None,
@@ -1203,6 +1213,7 @@ def fused_experts(
1213 use_int8_w8a8=use_int8_w8a8,
1214 use_int8_w8a16=use_int8_w8a16,
1215 use_int4_w4a16=use_int4_w4a16,
1216 + use_mxfp4_w4a4=use_mxfp4_w4a4,
1217 per_channel_quant=per_channel_quant,
1218 global_num_experts=global_num_experts,
1219 expert_map=expert_map,
@@ -1228,6 +1239,7 @@ def fused_experts_impl(
1239 use_int8_w8a8: bool = False,
1240 use_int8_w8a16: bool = False,
1241 use_int4_w4a16: bool = False,
1242 + use_mxfp4_w4a4: bool = False,
1243 per_channel_quant: bool = False,
1244 global_num_experts: int = -1,
1245 expert_map: Optional[torch.Tensor] = None,
@@ -1243,6 +1255,9 @@ def fused_experts_impl(
1255 if use_int4_w4a16:
1256 assert hidden_states.size(1) // 2 == w1.size(2), (
1257 "Hidden size mismatch")
1258 + elif use_mxfp4_w4a4:
1259 + # 16bit activation and fp4x2 packed weight
1260 + assert hidden_states.size(1) // 2 == w1.size(2), "hidden size mismatch"
1261 else:
1262 assert hidden_states.size(1) == w1.size(2), (
1263 f"Hidden size mismatch {hidden_states.size(1)} != {w1.size(2)}")
@@ -1268,12 +1283,14 @@ def fused_experts_impl(
1283 config_dtype = get_config_dtype_str(use_fp8_w8a8=use_fp8_w8a8,
1284 use_int8_w8a16=use_int8_w8a16,
1285 use_int4_w4a16=use_int4_w4a16,
1286 + use_mxfp4_w4a4=use_mxfp4_w4a4,
1287 dtype=hidden_states.dtype)
1288
1289 qtype = get_config_quant_dtype(use_fp8_w8a8=use_fp8_w8a8,
1290 use_int8_w8a8=use_int8_w8a8,
1291 use_int8_w8a16=use_int8_w8a16,
1276 - use_int4_w4a16=use_int4_w4a16)
1292 + use_int4_w4a16=use_int4_w4a16,
1293 + use_mxfp4_w4a4=use_mxfp4_w4a4)
1294
1295 get_config_func = functools.partial(
1296 try_get_optimal_moe_config,
@@ -1313,6 +1330,13 @@ def fused_experts_impl(
1330 else:
1331 out_hidden_states = torch.empty_like(hidden_states)
1332
1333 + if use_mxfp4_w4a4:
1334 + # Weight has to be dequantized for mxfp4 emulation.
1335 + w1 = dequant_mxfp4(w1, w1_scale, hidden_states.dtype)
1336 + w1_scale = None
1337 + w2 = dequant_mxfp4(w2, w2_scale, hidden_states.dtype)
1338 + w2_scale = None
1339 +
1340 for chunk in range((num_tokens // CHUNK_SIZE) + 1):
1341 begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
1342 min((chunk + 1) * CHUNK_SIZE,
@@ -1429,6 +1453,7 @@ def fused_moe(
1453 use_int8_w8a8: bool = False,
1454 use_int8_w8a16: bool = False,
1455 use_int4_w4a16: bool = False,
1456 + use_mxfp4_w4a4: bool = False,
1457 per_channel_quant: bool = False,
1458 global_num_experts: int = -1,
1459 expert_map: Optional[torch.Tensor] = None,
@@ -1470,6 +1495,9 @@ def fused_moe(
1495 - use_int4_w4a16 (bool): If True, use matmul of int4 weight and bf16/fp16
1496 activation to compute the inner products for w1 and w2.
1497 Defaults to False.
1498 + - use_mxfp4_w4a4 (bool): If True, use matmul of OCP MXFP4 weight and
1499 + OCP MXFP4 activation to compute the inner products for w1 and w2.
1500 + Defaults to False.
1501 - global_num_experts (int): The total number of experts in the global
1502 expert space.
1503 - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
@@ -1513,6 +1541,7 @@ def fused_moe(
1541 use_int8_w8a8=use_int8_w8a8,
1542 use_int8_w8a16=use_int8_w8a16,
1543 use_int4_w4a16=use_int4_w4a16,
1544 + use_mxfp4_w4a4=use_mxfp4_w4a4,
1545 per_channel_quant=per_channel_quant,
1546 global_num_experts=global_num_experts,
1547 expert_map=expert_map,
@@ -1533,6 +1562,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1562 use_int8_w8a8: bool = False,
1563 use_int8_w8a16: bool = False,
1564 use_int4_w4a16: bool = False,
1565 + use_mxfp4_w4a4: bool = False,
1566 per_act_token_quant: bool = False,
1567 block_shape: Optional[list[int]] = None,
1568 ):
@@ -1542,6 +1572,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1572 use_int8_w8a8=use_int8_w8a8,
1573 use_int8_w8a16=use_int8_w8a16,
1574 use_int4_w4a16=use_int4_w4a16,
1575 + use_mxfp4_w4a4=use_mxfp4_w4a4,
1576 per_act_token_quant=per_act_token_quant,
1577 block_shape=block_shape,
1578 ))
@@ -1550,6 +1581,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1581 self.use_int4_w4a16 = use_int4_w4a16
1582 self.use_int8_w8a8 = use_int8_w8a8
1583 self.use_int8_w8a16 = use_int8_w8a16
1584 + self.use_mxfp4_w4a4 = use_mxfp4_w4a4
1585
1586 @property
1587 def activation_formats(
@@ -1627,6 +1659,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1659 config_dtype = get_config_dtype_str(use_fp8_w8a8=self.use_fp8_w8a8,
1660 use_int8_w8a16=self.use_int8_w8a16,
1661 use_int4_w4a16=self.use_int4_w4a16,
1662 + use_mxfp4_w4a4=self.use_mxfp4_w4a4,
1663 dtype=hidden_states.dtype)
1664
1665 config = try_get_optimal_moe_config(
@@ -1718,6 +1751,7 @@ def modular_triton_fused_moe(
1751 use_int8_w8a8: bool,
1752 use_int8_w8a16: bool,
1753 use_int4_w4a16: bool,
1754 + use_mxfp4_w4a4: bool,
1755 per_act_token_quant: bool,
1756 block_shape: Optional[list[int]] = None,
1757 ) -> mk.FusedMoEModularKernel:
@@ -1728,6 +1762,7 @@ def modular_triton_fused_moe(
1762 use_int8_w8a8=use_int8_w8a8,
1763 use_int8_w8a16=use_int8_w8a16,
1764 use_int4_w4a16=use_int4_w4a16,
1765 + use_mxfp4_w4a4=use_mxfp4_w4a4,
1766 per_act_token_quant=per_act_token_quant,
1767 block_shape=block_shape,
1768 ),
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+3
@@ -19,6 +19,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
19 use_int8_w8a8: bool = False,
20 use_int8_w8a16: bool = False,
21 use_int4_w4a16: bool = False,
22 + use_mxfp4_w4a4: bool = False,
23 per_act_token_quant: bool = False,
24 block_shape: Optional[list[int]] = None,
25 allow_deep_gemm: bool = False,
@@ -29,6 +30,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
30 use_int8_w8a8=use_int8_w8a8,
31 use_int8_w8a16=use_int8_w8a16,
32 use_int4_w4a16=use_int4_w4a16,
33 + use_mxfp4_w4a4=use_mxfp4_w4a4,
34 per_act_token_quant=per_act_token_quant,
35 block_shape=block_shape,
36 ))
@@ -37,6 +39,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
39 use_int8_w8a8=use_int8_w8a8,
40 use_int4_w4a16=use_int4_w4a16,
41 use_int8_w8a16=use_int8_w8a16,
42 + use_mxfp4_w4a4=use_mxfp4_w4a4,
43 per_act_token_quant=per_act_token_quant,
44 block_shape=block_shape,
45 )
vllm/model_executor/layers/fused_moe/utils.py
+22 -2
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from math import prod
4 -from typing import Optional
4 +from typing import Optional, Union
5
6 import torch
7
@@ -10,6 +10,9 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
10 per_token_group_quant_fp8)
11 from vllm.model_executor.layers.quantization.utils.int8_utils import (
12 per_token_group_quant_int8, per_token_quant_int8)
13 +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
14 + quant_dequant_mxfp4)
15 +from vllm.platforms import current_platform
16 from vllm.utils import cdiv
17
18
@@ -74,10 +77,25 @@ def _int8_quantize(
77 return A, A_scale
78
79
80 +def _mxfp4_quantize(
81 + A: torch.Tensor,
82 + A_scale: Optional[torch.Tensor],
83 + per_act_token_quant: bool,
84 + block_shape: Optional[list[int]] = None,
85 +) -> tuple[torch.Tensor, None]:
86 + assert block_shape is None
87 + if not current_platform.supports_mx():
88 + A = quant_dequant_mxfp4(A)
89 + else:
90 + raise NotImplementedError()
91 +
92 + return A, None
93 +
94 +
95 def moe_kernel_quantize_input(
96 A: torch.Tensor,
97 A_scale: Optional[torch.Tensor],
80 - quant_dtype: Optional[torch.dtype],
98 + quant_dtype: Union[None, torch.dtype, str],
99 per_act_token_quant: bool,
100 block_shape: Optional[list[int]] = None,
101 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
@@ -85,6 +103,8 @@ def moe_kernel_quantize_input(
103 return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
104 elif quant_dtype == torch.int8:
105 return _int8_quantize(A, A_scale, per_act_token_quant, block_shape)
106 + elif quant_dtype == "mxfp4":
107 + return _mxfp4_quantize(A, A_scale, per_act_token_quant, block_shape)
108 else:
109 return A, A_scale
110
vllm/model_executor/layers/quantization/quark/quark.py
-6
@@ -237,12 +237,6 @@ class QuarkConfig(QuantizationConfig):
237 "Quark model is not in MX-FP4 format: not group_size=32")
238 return False
239
240 - # Weights need to use static quantization.
241 - if weight_quant.get("is_dynamic") is True:
242 - logger.debug(
243 - "Quark model is not in MX-FP4 format: not weight static")
244 - return False
245 -
240 # Activations need to use dynamic quantization.
241 if input_quant.get("is_dynamic") is False:
242 logger.debug(
vllm/model_executor/layers/quantization/quark/quark_moe.py
+167 -3
@@ -5,11 +5,12 @@ from typing import Any, Callable, Optional
5
6 import torch
7
8 -import vllm.model_executor.layers.fused_moe # noqa
8 from vllm import _custom_ops as ops
9 from vllm.logger import init_logger
10 from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEMethodBase,
11 FusedMoeWeightScaleSupported)
12 +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
13 + OCP_MX_BLOCK_SIZE)
14 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
15 all_close_1d, normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize)
16 from vllm.model_executor.utils import set_weight_attrs
@@ -17,7 +18,9 @@ from vllm.platforms import current_platform
18
19 logger = init_logger(__name__)
20
20 -__all__ = ["QuarkMoEMethod", "QuarkW8A8Fp8MoEMethod"]
21 +__all__ = [
22 + "QuarkMoEMethod", "QuarkW8A8Fp8MoEMethod", "QuarkW4A4MXFp4MoEMethod"
23 +]
24
25
26 class QuarkMoEMethod(FusedMoEMethodBase):
@@ -40,6 +43,8 @@ class QuarkMoEMethod(FusedMoEMethodBase):
43
44 if quant_config._is_fp8_w8a8(weight_config, input_config):
45 return QuarkW8A8Fp8MoEMethod(weight_config, input_config)
46 + elif quant_config._is_mx_fp4(weight_config, input_config):
47 + return QuarkW4A4MXFp4MoEMethod(weight_config, input_config)
48 else:
49 raise RuntimeError("Unsupported FusedMoe scheme")
50
@@ -242,4 +247,163 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
247 w1_scale=layer.w13_weight_scale,
248 w2_scale=layer.w2_weight_scale,
249 a1_scale=layer.w13_input_scale,
245 - a2_scale=layer.w2_input_scale)
250 + a2_scale=layer.w2_input_scale,
251 + activation=activation)
252 +
253 +
254 +class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
255 +
256 + def __init__(self, weight_config: dict[str, Any], input_config: dict[str,
257 + Any]):
258 + self.weight_quant = weight_config
259 + self.input_quant = input_config
260 +
261 + weight_qscheme = self.weight_quant.get("qscheme")
262 + input_qscheme = self.input_quant.get("qscheme")
263 + if not (weight_qscheme == "per_group"
264 + and input_qscheme == "per_group"):
265 + raise ValueError(
266 + "For MX(FP4) Fused MoE layers, only per-group scales "
267 + "for weights and activations are supported. Found "
268 + f"{weight_qscheme}, {input_qscheme}") # noqa E501
269 +
270 + self.static_input_scales = not self.input_quant.get("is_dynamic")
271 +
272 + if self.static_input_scales:
273 + raise NotImplementedError(
274 + "QuarkW4A4MXFp4MoEMethod with static input scales is currently "
275 + "not implemented. Please open an issue.")
276 +
277 + if not current_platform.supports_mx():
278 + self.emulate = True
279 + logger.warning_once(
280 + "The current platform does not support native MXFP4 "
281 + "computation. Simulated weight dequantization and activation "
282 + "QDQ (quantize and dequantize) will be used, with the linear "
283 + "layers computed in high precision.")
284 + else:
285 + self.emulate = True
286 + logger.warning_once(
287 + "The current platform supports native MXFP4 "
288 + "computation, but kernels are not yet integrated in vLLM. "
289 + "Simulated weight dequantization and activation "
290 + "QDQ (quantize and dequantize) will be used, with the linear "
291 + "layers computed in high precision.")
292 +
293 + def create_weights(self, layer: torch.nn.Module, num_experts: int,
294 + hidden_size: int, intermediate_size_per_partition: int,
295 + params_dtype: torch.dtype, **extra_weight_attrs):
296 +
297 + # Add the quantization method used (per tensor/grouped/channel)
298 + # to ensure the weight scales are loaded in properly
299 + extra_weight_attrs.update(
300 + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value})
301 +
302 + params_dtype = torch.uint8
303 +
304 + # WEIGHTS
305 + w13_weight = torch.nn.Parameter(torch.empty(
306 + num_experts,
307 + 2 * intermediate_size_per_partition,
308 + hidden_size // 2,
309 + dtype=params_dtype),
310 + requires_grad=False)
311 + layer.register_parameter("w13_weight", w13_weight)
312 +
313 + set_weight_attrs(w13_weight, extra_weight_attrs)
314 +
315 + w2_weight = torch.nn.Parameter(torch.empty(
316 + num_experts,
317 + hidden_size,
318 + intermediate_size_per_partition // 2,
319 + dtype=params_dtype),
320 + requires_grad=False)
321 + layer.register_parameter("w2_weight", w2_weight)
322 +
323 + set_weight_attrs(w2_weight, extra_weight_attrs)
324 +
325 + # WEIGHT_SCALES
326 + w13_weight_scale = torch.nn.Parameter(
327 + torch.ones(
328 + num_experts,
329 + 2 * intermediate_size_per_partition,
330 + hidden_size // OCP_MX_BLOCK_SIZE,
331 + dtype=params_dtype,
332 + ),
333 + requires_grad=False,
334 + )
335 + w2_weight_scale = torch.nn.Parameter(
336 + torch.ones(
337 + num_experts,
338 + hidden_size,
339 + intermediate_size_per_partition // OCP_MX_BLOCK_SIZE,
340 + dtype=params_dtype,
341 + ),
342 + requires_grad=False,
343 + )
344 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
345 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
346 +
347 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
348 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
349 +
350 + def apply(
351 + self,
352 + layer: torch.nn.Module,
353 + x: torch.Tensor,
354 + router_logits: torch.Tensor,
355 + top_k: int,
356 + renormalize: bool,
357 + use_grouped_topk: bool = False,
358 + topk_group: Optional[int] = None,
359 + num_expert_group: Optional[int] = None,
360 + global_num_experts: int = -1,
361 + expert_map: Optional[torch.Tensor] = None,
362 + custom_routing_function: Optional[Callable] = None,
363 + scoring_func: str = "softmax",
364 + e_score_correction_bias: Optional[torch.Tensor] = None,
365 + apply_router_weight_on_input: bool = False,
366 + activation: str = "silu",
367 + enable_eplb: bool = False,
368 + expert_load_view: Optional[torch.Tensor] = None,
369 + logical_to_physical_map: Optional[torch.Tensor] = None,
370 + logical_replica_count: Optional[torch.Tensor] = None,
371 + ) -> torch.Tensor:
372 +
373 + if enable_eplb:
374 + raise NotImplementedError(
375 + "EPLB not supported for `QuarkW4A4MXFp4MoEMethod` yet.")
376 +
377 + from vllm.model_executor.layers.fused_moe import fused_experts
378 +
379 + topk_weights, topk_ids = FusedMoE.select_experts(
380 + hidden_states=x,
381 + router_logits=router_logits,
382 + use_grouped_topk=use_grouped_topk,
383 + top_k=top_k,
384 + renormalize=renormalize,
385 + topk_group=topk_group,
386 + num_expert_group=num_expert_group,
387 + custom_routing_function=custom_routing_function,
388 + scoring_func=scoring_func,
389 + e_score_correction_bias=e_score_correction_bias)
390 +
391 + out = fused_experts(
392 + x,
393 + layer.w13_weight,
394 + layer.w2_weight,
395 + topk_weights=topk_weights,
396 + topk_ids=topk_ids,
397 + inplace=True,
398 + use_mxfp4_w4a4=True,
399 + global_num_experts=global_num_experts,
400 + apply_router_weight_on_input=apply_router_weight_on_input,
401 + expert_map=expert_map,
402 + w1_scale=layer.w13_weight_scale,
403 + w2_scale=layer.w2_weight_scale,
404 + a1_scale=None,
405 + a2_scale=None,
406 + block_shape=None,
407 + activation=activation,
408 + )
409 + return out
vllm/model_executor/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py
+32 -46
@@ -6,14 +6,16 @@ from typing import Any, Callable, Optional
6 import torch
7 import torch.nn.functional as F
8
9 -import vllm.envs as envs
9 +from vllm.logger import init_logger
10 from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
11 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
12 - OCP_MX_BLOCK_SIZE, per_token_group_quant_mxfp4)
12 + OCP_MX_BLOCK_SIZE, dequant_mxfp4, quant_dequant_mxfp4)
13 from vllm.model_executor.parameter import (GroupQuantScaleParameter,
14 PackedvLLMParameter)
15 from vllm.platforms import current_platform
16
17 +logger = init_logger(__name__)
18 +
19 __all__ = ["QuarkW4A4MXFP4"]
20
21
@@ -25,7 +27,29 @@ class QuarkW4A4MXFP4(QuarkScheme):
27 self.qscheme = "per_group"
28 self.weight_quant_spec = weight_quant_spec
29 self.input_quant_spec = input_quant_spec
28 - self.emulate = not current_platform.supports_mx()
30 +
31 + self.static_input_scales = not input_quant_spec.get("is_dynamic")
32 +
33 + if self.static_input_scales:
34 + raise NotImplementedError(
35 + "QuarkW4A4MXFP4 with static input scales is currently not "
36 + "implemented. Please open an issue.")
37 +
38 + if not current_platform.supports_mx():
39 + self.emulate = True
40 + logger.warning_once(
41 + "The current platform does not support native MXFP4 "
42 + "computation. Simulated weight dequantization and activation "
43 + "QDQ (quantize and dequantize) will be used, with the linear "
44 + "layers computed in high precision.")
45 + else:
46 + self.emulate = True
47 + logger.warning_once(
48 + "The current platform supports native MXFP4 "
49 + "computation, but kernels are not yet integrated in vLLM. "
50 + "Simulated weight dequantization and activation "
51 + "QDQ (quantize and dequantize) will be used, with the linear "
52 + "layers computed in high precision.")
53
54 @classmethod
55 def get_min_capability(cls) -> int:
@@ -37,43 +61,6 @@ class QuarkW4A4MXFP4(QuarkScheme):
61 layer.weight_scale = torch.nn.Parameter(layer.weight_scale.data,
62 requires_grad=False)
63
40 - if self.emulate:
41 - try:
42 - from quark.torch.export.nn.modules import realquantizer
43 - from quark.torch.quantization.config.config import (
44 - QuantizationSpec)
45 - except ImportError as err:
46 - raise ImportError(
47 - "The package `amd-quark` is required to use AMD Quark "
48 - "MX-FP4 models. Please install it with `pip install "
49 - "amd-quark`.") from err
50 -
51 - weight_quant_spec = QuantizationSpec.from_dict(
52 - self.weight_quant_spec)
53 -
54 - weight_quantizer = realquantizer.get_real_quantizer(
55 - qspec=weight_quant_spec,
56 - quantizer=None,
57 - real_quantized=True,
58 - reorder=False,
59 - float_dtype=self.out_dtype,
60 - scale_shape=layer.weight_scale.shape,
61 - zero_point_shape=None,
62 - )
63 - weight_quantizer.scale.data = layer.weight_scale.data
64 -
65 - if not envs.VLLM_QUARK_EMU_MEM_OPT:
66 - layer.weight = torch.nn.Parameter(
67 - weight_quantizer(layer.weight.data).to(self.out_dtype),
68 - requires_grad=False,
69 - )
70 - else:
71 - self.weight_quantizer = weight_quantizer
72 - layer.weight_scale = None
73 -
74 - # This call is necessary to release the scales memory.
75 - torch.cuda.empty_cache()
76 -
64 def create_weights(self, layer: torch.nn.Module,
65 output_partition_sizes: list[int],
66 input_size_per_partition: int,
@@ -116,11 +103,10 @@ class QuarkW4A4MXFP4(QuarkScheme):
103 bias: Optional[torch.Tensor] = None) -> torch.Tensor:
104
105 if self.emulate:
119 - if envs.VLLM_QUARK_EMU_MEM_OPT:
120 - dq_w = self.weight_quantizer(layer.weight).to(self.out_dtype)
121 - else:
122 - dq_w = layer.weight
123 - qdq_x, _ = per_token_group_quant_mxfp4(x, OCP_MX_BLOCK_SIZE)
124 - return F.linear(qdq_x, dq_w, bias)
106 + dq_w = dequant_mxfp4(layer.weight, layer.weight_scale, x.dtype)
107 +
108 + x = quant_dequant_mxfp4(x)
109 +
110 + return F.linear(x, dq_w, bias)
111 else:
112 raise NotImplementedError()
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+52 -30
@@ -1,45 +1,67 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
3 import torch
4
5 +from vllm.utils import direct_register_custom_op
6 +
7 OCP_MX_BLOCK_SIZE = 32
8
9
9 -def per_token_group_quant_mxfp4(x: torch.Tensor,
10 - block_k: int,
11 - scale_calculation_mode: str = "even"
12 - ) -> tuple[torch.Tensor, torch.Tensor]:
10 +def _dequant_mxfp4(x: torch.Tensor, scale: torch.Tensor,
11 + float_dtype: torch.dtype) -> torch.Tensor:
12 try:
14 - from quark.torch.kernel.hw_emulation.hw_emulation_interface import (
15 - fake_quantize_fp4_fp6_per_group_with_scale)
16 - from quark.torch.quantization.utils import (even_round,
17 - reshape_to_blocks)
13 + from quark.torch.kernel import mx
14 except ImportError as err:
15 raise ImportError("The package `amd-quark` is required to use "
16 "MX-FP4 models. Please install it with `pip install "
17 "amd-quark`.") from err
18
23 - axis = -1
24 - block_x = reshape_to_blocks(x, block_k, axis)
25 - amax, _ = torch.max(torch.abs(block_x), dim=-1, keepdim=True)
26 - amax = amax.squeeze(-1)
27 -
28 - # TODO: there are other rounding strategies supported in quark and in the
29 - # config.json that we do not check for here!
30 - if scale_calculation_mode != "even":
31 - raise NotImplementedError(
32 - f"Scale calculation mode {scale_calculation_mode} is not yet "
33 - "supported in MX-FP4 quantization")
34 - scale = even_round(amax, "fp4")
35 -
36 - # Apply dequantize(quantize(x)).
37 - x = fake_quantize_fp4_fp6_per_group_with_scale(
38 - x,
39 - scale.to(x.device),
40 - axis=axis,
41 - group_size=block_k,
42 - quant_dtype="fp4",
19 + return mx.dq_mxfp4(x, scale, float_dtype)
20 +
21 +
22 +def _dequant_mxfp4_fake(x: torch.Tensor, scale: torch.Tensor,
23 + float_dtype: torch.dtype) -> torch.Tensor:
24 + return torch.empty((*x.shape[:-1], x.shape[-1] * 2),
25 + dtype=float_dtype,
26 + device=x.device)
27 +
28 +
29 +def _quant_dequant_mxfp4(x: torch.Tensor,
30 + scale_calculation_mode: str = "even") -> torch.Tensor:
31 + try:
32 + from quark.torch.kernel import mx
33 + except ImportError as err:
34 + raise ImportError("The package `amd-quark` is required to use "
35 + "MX-FP4 models. Please install it with `pip install "
36 + "amd-quark`.") from err
37 +
38 + return mx.qdq_mxfp4(x, scale_calculation_mode)
39 +
40 +
41 +def _quant_dequant_mxfp4_fake(x: torch.Tensor,
42 + scale_calculation_mode: str = "even"
43 + ) -> torch.Tensor:
44 + return torch.empty_like(x)
45 +
46 +
47 +try:
48 + direct_register_custom_op(
49 + op_name="dequant_mxfp4",
50 + op_func=_dequant_mxfp4,
51 + mutates_args=[],
52 + fake_impl=_dequant_mxfp4_fake,
53 )
54 + dequant_mxfp4 = torch.ops.vllm.dequant_mxfp4
55 +except AttributeError as error:
56 + raise error
57
45 - return x, scale
58 +try:
59 + direct_register_custom_op(
60 + op_name="quant_dequant_mxfp4",
61 + op_func=_quant_dequant_mxfp4,
62 + mutates_args=[],
63 + fake_impl=_quant_dequant_mxfp4_fake,
64 + )
65 + quant_dequant_mxfp4 = torch.ops.vllm.quant_dequant_mxfp4
66 +except AttributeError as error:
67 + raise error