@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 63988f3c2d

[Quantization][Refactor][1/N] Adopt `QuantKey` in `QuarkConfig` and methods, relying on `weight_quant_key`, `act_quant_key` for quant method dispatch (#52958)

Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

fxmarty-amd committed Sep 1, 2026 at 09:55 UTC 63988f3c2d79db505ca2ac24dc1ed51f4b63249a
13 files changed +1573 -475
tests/quantization/test_online_mxfp4.py
+8 -4
@@ -35,6 +35,10 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
35 mxfp4_quantize,
36 quant_dequant_mxfp4,
37 )
38 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
39 + kMxfp4Dynamic,
40 + kMxfp4Static,
41 +)
42 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
43 from vllm.model_executor.utils import replace_parameter
44 from vllm.platforms import PlatformEnum, current_platform
@@ -346,9 +350,9 @@ def test_online_mxfp4_moe_matches_quark(
350 online_layer = make_layer("online_layer")
351
352 checkpoint_method = QuarkOCP_MX_MoEMethod(
349 - weight_config={"qscheme": "per_group", "dtype": "fp4"},
350 - input_config={"dtype": "fp4", "is_dynamic": True},
353 moe=checkpoint_layer.moe_config,
354 + weight_quant_key=kMxfp4Static,
355 + activation_quant_key=kMxfp4Dynamic,
356 )
357 online_method = Mxfp4OnlineMoEMethod(layer=online_layer)
358
@@ -566,8 +570,8 @@ def test_online_mxfp4_dense_matches_quark(
570 checkpoint_weight, checkpoint_weight_scale = mxfp4_quantize(weight)
571
572 checkpoint_scheme = QuarkOCP_MX(
569 - weight_quant_spec={"qscheme": "per_group", "dtype": "fp4"},
570 - input_quant_spec={"dtype": "fp4", "is_dynamic": True},
573 + weight_quant_key=kMxfp4Static,
574 + activation_quant_key=kMxfp4Dynamic,
575 )
576 checkpoint_scheme.create_weights(
577 layer=checkpoint_layer,
tests/quantization/test_quark.py
+792 -67
@@ -10,6 +10,7 @@ See also `tests/kernels/moe/test_ocp_mx_moe.py`.
10 import importlib.metadata
11 from dataclasses import dataclass
12 from importlib.util import find_spec
13 +from types import SimpleNamespace
14
15 import huggingface_hub
16 import lm_eval
@@ -18,22 +19,62 @@ import torch
19 from packaging import version
20
21 from vllm._aiter_ops import is_aiter_found_and_supported
22 +from vllm.model_executor.layers.fused_moe import (
23 + RoutedExperts,
24 + UnquantizedFusedMoEMethod,
25 +)
26 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
27 +from vllm.model_executor.layers.fused_moe.config import (
28 + FusedMoEConfig,
29 + FusedMoEParallelConfig,
30 + RoutingMethodType,
31 +)
32 +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
33 from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501
34 QuarkConfig,
35 QuarkLinearMethod,
36 + QuarkNVFP4,
37 + QuarkOCP_MX,
38 QuarkW8A8Fp8,
39 QuarkW8A8Fp8PerBlock,
40 QuarkW8A8Int8,
41 )
42 from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E501
43 + QuarkMoEMethod,
44 QuarkW4A8Fp8MoEMethod,
45 QuarkW8A8Int8MoEMethod,
46 )
47 +from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
48 +from vllm.model_executor.layers.quantization.quark.utils import QuarkQTensorHint
49 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
50 quant_dequant_mxfp4,
51 )
52 from vllm.model_executor.layers.quantization.utils.quant_utils import (
53 + QuantKey,
54 is_layer_skipped,
55 + kFp8Dynamic128Sym,
56 + kFp8DynamicTensorSym,
57 + kFp8DynamicTokenSym,
58 + kFp8Static128BlockE8M0Sym,
59 + kFp8Static128BlockSym,
60 + kFp8StaticChannelSym,
61 + kFp8StaticTensorSym,
62 + kInt4W4A8StaticChannelSym,
63 + kInt8DynamicTensorAsym,
64 + kInt8DynamicTensorSym,
65 + kInt8DynamicTokenAsym,
66 + kInt8DynamicTokenSym,
67 + kInt8StaticChannelSym,
68 + kInt8StaticTensorAsym,
69 + kInt8StaticTensorSym,
70 + kMxfp4Dynamic,
71 + kMxfp4Static,
72 + kMxfp6E2M3Dynamic,
73 + kMxfp6E2M3Static,
74 + kMxfp6E3M2Dynamic,
75 + kMxfp6E3M2Static,
76 + kNvfp4Dynamic,
77 + kNvfp4Static,
78 )
79 from vllm.platforms import current_platform
80 from vllm.transformers_utils.repo_utils import hf_api
@@ -62,6 +103,576 @@ AITER_AVAILABLE = is_aiter_found_and_supported()
103
104 DEVICE_TYPE = current_platform.device_type
105
106 +
107 +@dataclass(frozen=True)
108 +class QTensorConfig:
109 + name: str
110 + weight: QuarkQTensorHint
111 + input_tensors: QuarkQTensorHint
112 + weight_quant_key: QuantKey | None = None
113 + act_quant_key: QuantKey | None = None
114 + dispatch_cls: type[QuarkScheme] | type[QuarkMoEMethod] | None = None
115 + expected_error: tuple[type[Exception], str] | None = None
116 +
117 +
118 +QTENSOR_CONFIGS = [
119 + QTensorConfig(
120 + name="fp8_w8a8_static_tensor",
121 + weight={"dtype": "fp8_e4m3", "qscheme": "per_tensor", "is_dynamic": False},
122 + input_tensors={
123 + "dtype": "fp8_e4m3",
124 + "qscheme": "per_tensor",
125 + "is_dynamic": False,
126 + },
127 + weight_quant_key=kFp8StaticTensorSym,
128 + act_quant_key=kFp8StaticTensorSym,
129 + dispatch_cls=QuarkW8A8Fp8,
130 + ),
131 + QTensorConfig(
132 + name="fp8_w8a8_static_tensor_single_entry_lists",
133 + weight=[
134 + {
135 + "dtype": "fp8_e4m3",
136 + "qscheme": "per_tensor",
137 + "is_dynamic": False,
138 + }
139 + ],
140 + input_tensors=[
141 + {
142 + "dtype": "fp8_e4m3",
143 + "qscheme": "per_tensor",
144 + "is_dynamic": False,
145 + }
146 + ],
147 + weight_quant_key=kFp8StaticTensorSym,
148 + act_quant_key=kFp8StaticTensorSym,
149 + dispatch_cls=QuarkW8A8Fp8,
150 + ),
151 + QTensorConfig(
152 + name="fp8_w8a8_dynamic_tensor",
153 + weight={"dtype": "fp8_e4m3", "qscheme": "per_tensor", "is_dynamic": False},
154 + input_tensors={
155 + "dtype": "fp8_e4m3",
156 + "qscheme": "per_tensor",
157 + "is_dynamic": True,
158 + },
159 + weight_quant_key=kFp8StaticTensorSym,
160 + act_quant_key=kFp8DynamicTensorSym,
161 + dispatch_cls=QuarkW8A8Fp8,
162 + ),
163 + QTensorConfig(
164 + name="fp8_w8a8_dynamic_token",
165 + weight={"dtype": "fp8_e4m3", "qscheme": "per_channel", "is_dynamic": False},
166 + input_tensors={
167 + "dtype": "fp8_e4m3",
168 + "qscheme": "per_channel",
169 + "is_dynamic": True,
170 + },
171 + weight_quant_key=kFp8StaticChannelSym,
172 + act_quant_key=kFp8DynamicTokenSym,
173 + dispatch_cls=QuarkW8A8Fp8,
174 + ),
175 + QTensorConfig(
176 + name="fp8_w8a8_channel_static_tensor",
177 + weight={"dtype": "fp8_e4m3", "qscheme": "per_channel", "is_dynamic": False},
178 + input_tensors={
179 + "dtype": "fp8_e4m3",
180 + "qscheme": "per_tensor",
181 + "is_dynamic": False,
182 + },
183 + weight_quant_key=kFp8StaticChannelSym,
184 + act_quant_key=kFp8StaticTensorSym,
185 + dispatch_cls=QuarkW8A8Fp8,
186 + ),
187 + QTensorConfig(
188 + name="fp8_w8a8_tensor_dynamic_token",
189 + weight={"dtype": "fp8_e4m3", "qscheme": "per_tensor", "is_dynamic": False},
190 + input_tensors={
191 + "dtype": "fp8_e4m3",
192 + "qscheme": "per_channel",
193 + "is_dynamic": True,
194 + },
195 + weight_quant_key=kFp8StaticTensorSym,
196 + act_quant_key=kFp8DynamicTokenSym,
197 + dispatch_cls=QuarkW8A8Fp8,
198 + ),
199 + QTensorConfig(
200 + name="fp8_w8a8_dynamic_block_fp32",
201 + weight={
202 + "dtype": "fp8_e4m3",
203 + "qscheme": "per_block",
204 + "is_dynamic": False,
205 + "block_size": [128, 128],
206 + "symmetric": True,
207 + },
208 + input_tensors={
209 + "dtype": "fp8_e4m3",
210 + "qscheme": "per_group",
211 + "is_dynamic": True,
212 + "group_size": 128,
213 + "symmetric": True,
214 + },
215 + weight_quant_key=kFp8Static128BlockSym,
216 + act_quant_key=kFp8Dynamic128Sym,
217 + dispatch_cls=QuarkW8A8Fp8PerBlock,
218 + ),
219 + QTensorConfig(
220 + name="fp8_w8a8_dynamic_block_e8m0",
221 + weight={
222 + "dtype": "fp8_e4m3",
223 + "qscheme": "per_block",
224 + "is_dynamic": False,
225 + "block_size": [128, 128],
226 + "symmetric": True,
227 + "scale_type": "float8_e8m0fnu",
228 + },
229 + input_tensors={
230 + "dtype": "fp8_e4m3",
231 + "qscheme": "per_group",
232 + "is_dynamic": True,
233 + "group_size": 128,
234 + "symmetric": True,
235 + },
236 + weight_quant_key=kFp8Static128BlockE8M0Sym,
237 + act_quant_key=kFp8Dynamic128Sym,
238 + dispatch_cls=QuarkW8A8Fp8PerBlock,
239 + ),
240 + QTensorConfig(
241 + name="fp8_w8a8_block_static_input",
242 + weight={
243 + "dtype": "fp8_e4m3",
244 + "qscheme": "per_block",
245 + "is_dynamic": False,
246 + "block_size": [128, 128],
247 + "symmetric": True,
248 + },
249 + input_tensors={
250 + "dtype": "fp8_e4m3",
251 + "qscheme": "per_group",
252 + "is_dynamic": False,
253 + "group_size": 128,
254 + "symmetric": True,
255 + },
256 + expected_error=(NotImplementedError, "No quark compatible scheme"),
257 + ),
258 + QTensorConfig(
259 + name="fp8_w8a8_block_group_size_mismatch",
260 + weight={
261 + "dtype": "fp8_e4m3",
262 + "qscheme": "per_block",
263 + "is_dynamic": False,
264 + "block_size": [128, 128],
265 + "symmetric": True,
266 + },
267 + input_tensors={
268 + "dtype": "fp8_e4m3",
269 + "qscheme": "per_group",
270 + "is_dynamic": True,
271 + "group_size": 64,
272 + "symmetric": True,
273 + },
274 + expected_error=(NotImplementedError, "No quark compatible scheme"),
275 + ),
276 + QTensorConfig(
277 + name="fp8_w8a8_block_missing_block_size",
278 + weight={
279 + "dtype": "fp8_e4m3",
280 + "qscheme": "per_block",
281 + "is_dynamic": False,
282 + "symmetric": True,
283 + },
284 + input_tensors={
285 + "dtype": "fp8_e4m3",
286 + "qscheme": "per_group",
287 + "is_dynamic": True,
288 + "group_size": 128,
289 + "symmetric": True,
290 + },
291 + expected_error=(ValueError, "requires `block_size`"),
292 + ),
293 + QTensorConfig(
294 + name="int8_w8a8_static_symmetric",
295 + weight={
296 + "dtype": "int8",
297 + "qscheme": "per_tensor",
298 + "is_dynamic": False,
299 + "symmetric": True,
300 + },
301 + input_tensors={
302 + "dtype": "int8",
303 + "qscheme": "per_tensor",
304 + "is_dynamic": False,
305 + "symmetric": True,
306 + },
307 + weight_quant_key=kInt8StaticTensorSym,
308 + act_quant_key=kInt8StaticTensorSym,
309 + dispatch_cls=QuarkW8A8Int8,
310 + ),
311 + QTensorConfig(
312 + name="int8_w8a8_static_asymmetric",
313 + weight={
314 + "dtype": "int8",
315 + "qscheme": "per_tensor",
316 + "is_dynamic": False,
317 + "symmetric": True,
318 + },
319 + input_tensors={
320 + "dtype": "int8",
321 + "qscheme": "per_tensor",
322 + "is_dynamic": False,
323 + "symmetric": False,
324 + },
325 + weight_quant_key=kInt8StaticTensorSym,
326 + act_quant_key=kInt8StaticTensorAsym,
327 + dispatch_cls=QuarkW8A8Int8,
328 + ),
329 + QTensorConfig(
330 + name="int8_w8a8_channel_static_symmetric",
331 + weight={
332 + "dtype": "int8",
333 + "qscheme": "per_channel",
334 + "is_dynamic": False,
335 + "symmetric": True,
336 + },
337 + input_tensors={
338 + "dtype": "int8",
339 + "qscheme": "per_tensor",
340 + "is_dynamic": False,
341 + "symmetric": True,
342 + },
343 + weight_quant_key=kInt8StaticChannelSym,
344 + act_quant_key=kInt8StaticTensorSym,
345 + dispatch_cls=QuarkW8A8Int8,
346 + ),
347 + QTensorConfig(
348 + name="int8_w8a8_channel_static_asymmetric",
349 + weight={
350 + "dtype": "int8",
351 + "qscheme": "per_channel",
352 + "is_dynamic": False,
353 + "symmetric": True,
354 + },
355 + input_tensors={
356 + "dtype": "int8",
357 + "qscheme": "per_tensor",
358 + "is_dynamic": False,
359 + "symmetric": False,
360 + },
361 + weight_quant_key=kInt8StaticChannelSym,
362 + act_quant_key=kInt8StaticTensorAsym,
363 + dispatch_cls=QuarkW8A8Int8,
364 + ),
365 + QTensorConfig(
366 + name="int8_w8a8_dynamic_tensor_symmetric",
367 + weight={
368 + "dtype": "int8",
369 + "qscheme": "per_tensor",
370 + "is_dynamic": False,
371 + "symmetric": True,
372 + },
373 + input_tensors={
374 + "dtype": "int8",
375 + "qscheme": "per_channel",
376 + "is_dynamic": True,
377 + "symmetric": True,
378 + },
379 + weight_quant_key=kInt8StaticTensorSym,
380 + act_quant_key=kInt8DynamicTensorSym,
381 + dispatch_cls=QuarkW8A8Int8,
382 + ),
383 + QTensorConfig(
384 + name="int8_w8a8_dynamic_tensor_asymmetric",
385 + weight={
386 + "dtype": "int8",
387 + "qscheme": "per_tensor",
388 + "is_dynamic": False,
389 + "symmetric": True,
390 + },
391 + input_tensors={
392 + "dtype": "int8",
393 + "qscheme": "per_channel",
394 + "is_dynamic": True,
395 + "symmetric": False,
396 + },
397 + weight_quant_key=kInt8StaticTensorSym,
398 + act_quant_key=kInt8DynamicTensorAsym,
399 + dispatch_cls=QuarkW8A8Int8,
400 + ),
401 + QTensorConfig(
402 + name="int8_w8a8_dynamic_token",
403 + weight={
404 + "dtype": "int8",
405 + "qscheme": "per_channel",
406 + "is_dynamic": False,
407 + "symmetric": True,
408 + },
409 + input_tensors={
410 + "dtype": "int8",
411 + "qscheme": "per_channel",
412 + "is_dynamic": True,
413 + "symmetric": True,
414 + },
415 + weight_quant_key=kInt8StaticChannelSym,
416 + act_quant_key=kInt8DynamicTokenSym,
417 + dispatch_cls=QuarkW8A8Int8,
418 + ),
419 + QTensorConfig(
420 + name="int8_w8a8_dynamic_token_asymmetric",
421 + weight={
422 + "dtype": "int8",
423 + "qscheme": "per_channel",
424 + "is_dynamic": False,
425 + "symmetric": True,
426 + },
427 + input_tensors={
428 + "dtype": "int8",
429 + "qscheme": "per_channel",
430 + "is_dynamic": True,
431 + "symmetric": False,
432 + },
433 + weight_quant_key=kInt8StaticChannelSym,
434 + act_quant_key=kInt8DynamicTokenAsym,
435 + dispatch_cls=QuarkW8A8Int8,
436 + ),
437 + QTensorConfig(
438 + name="ocp_mx_mxfp4_weight_only",
439 + weight={
440 + "dtype": "fp4",
441 + "qscheme": "per_group",
442 + "group_size": 32,
443 + "scale_format": "e8m0",
444 + "is_dynamic": False,
445 + },
446 + input_tensors=None,
447 + weight_quant_key=kMxfp4Static,
448 + act_quant_key=None,
449 + dispatch_cls=QuarkOCP_MX,
450 + ),
451 + QTensorConfig(
452 + name="ocp_mx_mxfp4_activation",
453 + weight={
454 + "dtype": "fp4",
455 + "qscheme": "per_group",
456 + "group_size": 32,
457 + "scale_format": "e8m0",
458 + "is_dynamic": False,
459 + },
460 + input_tensors={
461 + "dtype": "fp4",
462 + "qscheme": "per_group",
463 + "group_size": 32,
464 + "scale_format": "e8m0",
465 + "is_dynamic": True,
466 + },
467 + weight_quant_key=kMxfp4Static,
468 + act_quant_key=kMxfp4Dynamic,
469 + dispatch_cls=QuarkOCP_MX,
470 + ),
471 + QTensorConfig(
472 + name="ocp_mx_mxfp6_e3m2",
473 + weight={
474 + "dtype": "fp6_e3m2",
475 + "qscheme": "per_group",
476 + "group_size": 32,
477 + "scale_format": "e8m0",
478 + "is_dynamic": False,
479 + },
480 + input_tensors={
481 + "dtype": "fp6_e3m2",
482 + "qscheme": "per_group",
483 + "group_size": 32,
484 + "scale_format": "e8m0",
485 + "is_dynamic": True,
486 + },
487 + weight_quant_key=kMxfp6E3M2Static,
488 + act_quant_key=kMxfp6E3M2Dynamic,
489 + dispatch_cls=QuarkOCP_MX,
490 + ),
491 + QTensorConfig(
492 + name="ocp_mx_mxfp4_mxfp6_e3m2_activation",
493 + weight={
494 + "dtype": "fp4",
495 + "qscheme": "per_group",
496 + "group_size": 32,
497 + "scale_format": "e8m0",
498 + "is_dynamic": False,
499 + },
500 + input_tensors={
501 + "dtype": "fp6_e3m2",
502 + "qscheme": "per_group",
503 + "group_size": 32,
504 + "scale_format": "e8m0",
505 + "is_dynamic": True,
506 + },
507 + weight_quant_key=kMxfp4Static,
508 + act_quant_key=kMxfp6E3M2Dynamic,
509 + dispatch_cls=QuarkOCP_MX,
510 + ),
511 + QTensorConfig(
512 + name="ocp_mx_mxfp4_mxfp6_e2m3_activation",
513 + weight={
514 + "dtype": "fp4",
515 + "qscheme": "per_group",
516 + "group_size": 32,
517 + "scale_format": "e8m0",
518 + "is_dynamic": False,
519 + },
520 + input_tensors={
521 + "dtype": "fp6_e2m3",
522 + "qscheme": "per_group",
523 + "group_size": 32,
524 + "scale_format": "e8m0",
525 + "is_dynamic": True,
526 + },
527 + weight_quant_key=kMxfp4Static,
528 + act_quant_key=kMxfp6E2M3Dynamic,
529 + dispatch_cls=QuarkOCP_MX,
530 + ),
531 + QTensorConfig(
532 + name="ocp_mx_mxfp6_e2m3",
533 + weight={
534 + "dtype": "fp6_e2m3",
535 + "qscheme": "per_group",
536 + "group_size": 32,
537 + "scale_format": "e8m0",
538 + "is_dynamic": False,
539 + },
540 + input_tensors={
541 + "dtype": "fp6_e2m3",
542 + "qscheme": "per_group",
543 + "group_size": 32,
544 + "scale_format": "e8m0",
545 + "is_dynamic": True,
546 + },
547 + weight_quant_key=kMxfp6E2M3Static,
548 + act_quant_key=kMxfp6E2M3Dynamic,
549 + dispatch_cls=QuarkOCP_MX,
550 + ),
551 + QTensorConfig(
552 + name="nvfp4",
553 + weight=[
554 + {
555 + "dtype": "fp4",
556 + "qscheme": "per_group",
557 + "group_size": 16,
558 + "is_dynamic": False,
559 + },
560 + {"dtype": "fp8_e4m3", "qscheme": "per_tensor", "is_dynamic": False},
561 + ],
562 + input_tensors=[
563 + {
564 + "dtype": "fp4",
565 + "qscheme": "per_group",
566 + "group_size": 16,
567 + "is_dynamic": True,
568 + },
569 + {"dtype": "fp8_e4m3", "qscheme": "per_tensor", "is_dynamic": False},
570 + ],
571 + weight_quant_key=kNvfp4Static,
572 + act_quant_key=kNvfp4Dynamic,
573 + dispatch_cls=QuarkNVFP4,
574 + ),
575 + QTensorConfig(
576 + name="w4a8_fp8_static",
577 + weight=[
578 + {"dtype": "fp8_e4m3", "qscheme": "per_tensor", "is_dynamic": False},
579 + {
580 + "dtype": "int4",
581 + "qscheme": "per_channel",
582 + "is_dynamic": False,
583 + "symmetric": True,
584 + "ch_axis": 0,
585 + },
586 + ],
587 + input_tensors={
588 + "dtype": "fp8_e4m3",
589 + "qscheme": "per_tensor",
590 + "is_dynamic": False,
591 + },
592 + weight_quant_key=kInt4W4A8StaticChannelSym,
593 + act_quant_key=kFp8StaticTensorSym,
594 + dispatch_cls=QuarkW4A8Fp8MoEMethod,
595 + ),
596 + QTensorConfig(
597 + name="w4a8_fp8_dynamic",
598 + weight=[
599 + {"dtype": "fp8_e4m3", "qscheme": "per_tensor", "is_dynamic": False},
600 + {
601 + "dtype": "int4",
602 + "qscheme": "per_channel",
603 + "is_dynamic": False,
604 + "symmetric": True,
605 + "ch_axis": 0,
606 + },
607 + ],
608 + input_tensors={
609 + "dtype": "fp8_e4m3",
610 + "qscheme": "per_channel",
611 + "is_dynamic": True,
612 + },
613 + weight_quant_key=kInt4W4A8StaticChannelSym,
614 + act_quant_key=kFp8DynamicTokenSym,
615 + dispatch_cls=QuarkW4A8Fp8MoEMethod,
616 + ),
617 + QTensorConfig(
618 + name="w4a8_fp8_static_single_entry_input",
619 + weight=[
620 + {"dtype": "fp8_e4m3", "qscheme": "per_tensor", "is_dynamic": False},
621 + {
622 + "dtype": "int4",
623 + "qscheme": "per_channel",
624 + "is_dynamic": False,
625 + "symmetric": True,
626 + "ch_axis": 0,
627 + },
628 + ],
629 + input_tensors=[
630 + {
631 + "dtype": "fp8_e4m3",
632 + "qscheme": "per_tensor",
633 + "is_dynamic": False,
634 + }
635 + ],
636 + weight_quant_key=kInt4W4A8StaticChannelSym,
637 + act_quant_key=kFp8StaticTensorSym,
638 + dispatch_cls=QuarkW4A8Fp8MoEMethod,
639 + ),
640 +]
641 +
642 +
643 +def _make_qtensor_config(
644 + weight: QuarkQTensorHint,
645 + input_tensors: QuarkQTensorHint,
646 + exclude: list[str] | None = None,
647 +) -> QuarkConfig:
648 + return QuarkConfig(
649 + {
650 + "global_quant_config": {
651 + "weight": weight,
652 + "input_tensors": input_tensors,
653 + },
654 + "layer_type_quant_config": {},
655 + "exclude": exclude or [],
656 + }
657 + )
658 +
659 +
660 +def _make_test_moe_config() -> FusedMoEConfig:
661 + return FusedMoEConfig(
662 + num_experts=8,
663 + experts_per_token=2,
664 + hidden_dim=256,
665 + intermediate_size=256,
666 + num_local_experts=8,
667 + num_logical_experts=8,
668 + activation=MoEActivation.SILU,
669 + device=current_platform.device_type,
670 + routing_method=RoutingMethodType.Renormalize,
671 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
672 + in_dtype=torch.bfloat16,
673 + )
674 +
675 +
676 if QUARK_MXFP4_AVAILABLE:
677 from quark.torch.export.nn.modules.realquantizer import StaticScaledRealQuantizer
678 from quark.torch.kernel import mx as mx_kernel
@@ -98,83 +709,197 @@ def test_quark_config_preserves_existing_packed_modules_mapping():
709 assert config.packed_modules_mapping["custom_proj"] == ["a", "b"]
710
711
101 -def test_quark_fp8_w8a8_detects_per_block_config():
102 - config = QuarkConfig({})
103 - weight_config = {
104 - "dtype": "fp8_e4m3",
105 - "qscheme": "per_block",
106 - "is_dynamic": False,
107 - "block_size": [128, 128],
108 - "symmetric": True,
109 - }
110 - input_config = {
111 - "dtype": "fp8_e4m3",
112 - "qscheme": "per_group",
113 - "is_dynamic": True,
114 - "group_size": 128,
115 - "symmetric": True,
116 - }
712 +def test_quant_method_dispatch_ignored(default_vllm_config):
713 + config = _make_qtensor_config(None, None, exclude=["linear", "experts"])
714
118 - assert config._is_fp8_w8a8(weight_config, input_config)
715 + class TestLinear(LinearBase):
716 + def __init__(self):
717 + torch.nn.Module.__init__(self)
718
719 + class TestRoutedExperts(RoutedExperts):
720 + def __init__(self):
721 + torch.nn.Module.__init__(self)
722 + self.moe_config = _make_test_moe_config()
723
121 -def test_quark_fp8_w8a8_rejects_per_block_static_input():
122 - config = QuarkConfig({})
123 - weight_config = {
124 - "dtype": "fp8_e4m3",
125 - "qscheme": "per_block",
126 - "is_dynamic": False,
127 - "block_size": [128, 128],
128 - "symmetric": True,
129 - }
130 - input_config = {
131 - "dtype": "fp8_e4m3",
132 - "qscheme": "per_group",
133 - "is_dynamic": False,
134 - "group_size": 128,
135 - "symmetric": True,
136 - }
724 + assert config.get_quant_method_target("linear", LinearBase) == (
725 + None,
726 + None,
727 + UnquantizedLinearMethod,
728 + )
729 + assert isinstance(
730 + config.get_quant_method(TestLinear(), "linear"), UnquantizedLinearMethod
731 + )
732
138 - assert not config._is_fp8_w8a8(weight_config, input_config)
733 + assert config.get_quant_method_target("experts", RoutedExperts) == (
734 + None,
735 + None,
736 + UnquantizedFusedMoEMethod,
737 + )
738 + assert isinstance(
739 + config.get_quant_method(TestRoutedExperts(), "experts"),
740 + UnquantizedFusedMoEMethod,
741 + )
742
743 + dynamic_mxfp4_config = _make_qtensor_config(
744 + {
745 + "dtype": "fp4",
746 + "qscheme": "per_group",
747 + "group_size": 32,
748 + "scale_format": "e8m0",
749 + "is_dynamic": False,
750 + },
751 + None,
752 + exclude=["self_attn.q_proj", "mlp.down_proj"],
753 + )
754 + dynamic_mxfp4_config.dynamic_mxfp4_quant = True
755 +
756 + assert dynamic_mxfp4_config.get_quant_method_target(
757 + "self_attn.q_proj", LinearBase
758 + ) == (kMxfp4Static, None, QuarkLinearMethod)
759 + attention_proj = TestLinear()
760 + assert isinstance(
761 + dynamic_mxfp4_config.get_quant_method(attention_proj, "self_attn.q_proj"),
762 + QuarkLinearMethod,
763 + )
764 + assert isinstance(attention_proj.scheme, QuarkOCP_MX)
765 + assert attention_proj.scheme.dynamic_mxfp4_quant
766 +
767 + assert dynamic_mxfp4_config.get_quant_method_target(
768 + "mlp.down_proj", LinearBase
769 + ) == (None, None, UnquantizedLinearMethod)
770 + assert isinstance(
771 + dynamic_mxfp4_config.get_quant_method(TestLinear(), "mlp.down_proj"),
772 + UnquantizedLinearMethod,
773 + )
774
141 -def test_quark_fp8_w8a8_rejects_per_block_group_size_mismatch():
142 - config = QuarkConfig({})
143 - weight_config = {
144 - "dtype": "fp8_e4m3",
145 - "qscheme": "per_block",
146 - "is_dynamic": False,
147 - "block_size": [128, 128],
148 - "symmetric": True,
149 - }
150 - input_config = {
151 - "dtype": "fp8_e4m3",
152 - "qscheme": "per_group",
153 - "is_dynamic": True,
154 - "group_size": 64,
155 - "symmetric": True,
156 - }
775
158 - assert not config._is_fp8_w8a8(weight_config, input_config)
776 +@pytest.mark.parametrize("case", QTENSOR_CONFIGS, ids=lambda case: case.name)
777 +def test_quant_method_dispatch_target(case):
778 + config = _make_qtensor_config(case.weight, case.input_tensors)
779 + if case.expected_error is not None:
780 + error_type, error_message = case.expected_error
781 + with pytest.raises(error_type, match=error_message):
782 + config.get_quant_method_target("linear", LinearBase)
783 + return
784
785 + assert case.dispatch_cls is not None
786 + is_linear = issubclass(case.dispatch_cls, QuarkScheme)
787
161 -def test_quark_w8a8_fp8_per_block_requires_block_size():
162 - weight_config = {
163 - "dtype": "fp8_e4m3",
164 - "qscheme": "per_block",
165 - "is_dynamic": False,
166 - "symmetric": True,
167 - }
168 - input_config = {
169 - "dtype": "fp8_e4m3",
170 - "qscheme": "per_group",
171 - "is_dynamic": True,
172 - "group_size": 128,
173 - "symmetric": True,
174 - }
788 + weight_quant_key, act_quant_key, method_cls = config.get_quant_method_target(
789 + "linear" if is_linear else "experts",
790 + LinearBase if is_linear else RoutedExperts,
791 + )
792 +
793 + assert weight_quant_key == case.weight_quant_key
794 + assert act_quant_key == case.act_quant_key
795 + assert method_cls is (QuarkLinearMethod if is_linear else case.dispatch_cls)
796 +
797 +
798 +@pytest.mark.parametrize(
799 + ("weight", "input_tensors"),
800 + [
801 + pytest.param(
802 + {
803 + "dtype": "int8",
804 + "qscheme": "per_group",
805 + "is_dynamic": False,
806 + "symmetric": True,
807 + },
808 + {
809 + "dtype": "int8",
810 + "qscheme": "per_tensor",
811 + "is_dynamic": False,
812 + "symmetric": True,
813 + },
814 + id="single_entry",
815 + ),
816 + pytest.param(
817 + [
818 + {"dtype": "int8", "qscheme": "per_tensor"},
819 + {"dtype": "int8", "qscheme": "per_tensor"},
820 + ],
821 + [
822 + {"dtype": "int8", "qscheme": "per_tensor"},
823 + {"dtype": "int8", "qscheme": "per_tensor"},
824 + ],
825 + id="multi_entry",
826 + ),
827 + ],
828 +)
829 +def test_quant_method_dispatch_unsupported(weight, input_tensors):
830 + config = _make_qtensor_config(weight, input_tensors)
831 +
832 + class TestRoutedExperts(RoutedExperts):
833 + def __init__(self):
834 + torch.nn.Module.__init__(self)
835 +
836 + with pytest.raises(RuntimeError, match="^Unsupported FusedMoe scheme$"):
837 + config.get_quant_method_target("experts", RoutedExperts)
838 +
839 + with pytest.raises(RuntimeError, match="^Unsupported FusedMoe scheme$"):
840 + config.get_quant_method(TestRoutedExperts(), "experts")
841 +
842 +
843 +@pytest.mark.parametrize(
844 + "case",
845 + [case for case in QTENSOR_CONFIGS if case.expected_error is None],
846 + ids=lambda case: case.name,
847 +)
848 +def test_quant_method_dispatch_instantiation(case, monkeypatch, default_vllm_config):
849 + config = _make_qtensor_config(case.weight, case.input_tensors)
850 + assert case.dispatch_cls is not None
851 + if issubclass(case.dispatch_cls, QuarkScheme):
852 +
853 + class TestLinear(LinearBase):
854 + def __init__(self):
855 + torch.nn.Module.__init__(self)
856 +
857 + monkeypatch.setattr(
858 + "vllm.model_executor.layers.quantization.quark.schemes."
859 + "quark_w8a8_fp8.get_current_vllm_config",
860 + lambda: SimpleNamespace(model_config=SimpleNamespace(dtype=torch.bfloat16)),
861 + )
862 + layer = TestLinear()
863 + method = config.get_quant_method(layer, "linear")
864 +
865 + assert isinstance(method, QuarkLinearMethod)
866 + assert isinstance(layer.scheme, case.dispatch_cls)
867 + if case.weight_quant_key == kFp8Static128BlockE8M0Sym:
868 + # TODO: Remove once E8M0 quant key is properly handled in oracle
869 + assert layer.scheme.weight_quant_key == kFp8Static128BlockSym
870 + else:
871 + assert layer.scheme.weight_quant_key == case.weight_quant_key
872 + assert layer.scheme.activation_quant_key == case.act_quant_key
873 + else:
874 +
875 + class TestRoutedExperts(RoutedExperts):
876 + def __init__(self):
877 + torch.nn.Module.__init__(self)
878 + self.moe_config = _make_test_moe_config()
879 +
880 + for target in (
881 + "select_fp8_moe_backend",
882 + "select_int8_moe_backend",
883 + "select_mxfp4_moe_backend",
884 + "backend_to_kernel_cls",
885 + "select_nvfp4_moe_backend",
886 + ):
887 + monkeypatch.setattr(
888 + f"vllm.model_executor.layers.quantization.quark.quark_moe.{target}",
889 + lambda *args, **kwargs: (object(), object()),
890 + )
891 +
892 + # AssertionError: W4A8 FP8 MoE requires ROCm AITER fused MoE support
893 + monkeypatch.setattr(
894 + "vllm.model_executor.layers.quantization.quark.quark_moe."
895 + "rocm_aiter_ops.is_fused_moe_enabled",
896 + lambda: True,
897 + )
898 +
899 + layer = TestRoutedExperts()
900 + method = config.get_quant_method(layer, "experts")
901
176 - with pytest.raises(ValueError, match="requires `block_size`"):
177 - QuarkW8A8Fp8PerBlock(weight_config, input_config)
902 + assert isinstance(method, case.dispatch_cls)
903
904
905 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
vllm/model_executor/layers/quantization/quark/quark.py
+386 -161
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import fnmatch
5 -from typing import TYPE_CHECKING, Any, cast
5 +from typing import TYPE_CHECKING, Any, NamedTuple, cast
6
7 import torch
8 from transformers import PretrainedConfig
@@ -37,9 +37,35 @@ from vllm.model_executor.layers.quantization.quark.schemes import (
37 QuarkW8A8Int8,
38 )
39 from vllm.model_executor.layers.quantization.quark.utils import (
40 + QuarkQTensorHint,
41 deep_compare,
42 should_ignore_layer,
43 )
44 +from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
45 + _ACTIVATION_QUANT_KEY_MAP,
46 + _WEIGHT_QUANT_KEY_MAP,
47 +)
48 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
49 + QuantKey,
50 + kFp8Dynamic128Sym,
51 + kFp8DynamicTensorSym,
52 + kFp8DynamicTokenSym,
53 + kFp8Static128BlockE8M0Sym,
54 + kFp8Static128BlockSym,
55 + kFp8StaticChannelSym,
56 + kFp8StaticTensorSym,
57 + kInt4W4A8StaticChannelSym,
58 + kInt8DynamicTensorAsym,
59 + kInt8DynamicTensorSym,
60 + kInt8DynamicTokenAsym,
61 + kInt8DynamicTokenSym,
62 + kInt8StaticChannelSym,
63 + kInt8StaticTensorAsym,
64 + kInt8StaticTensorSym,
65 + kMxfp4Static,
66 + kNvfp4Dynamic,
67 + kNvfp4Static,
68 +)
69 from vllm.model_executor.models.utils import WeightsMapper
70 from vllm.platforms import current_platform
71
@@ -55,6 +81,15 @@ logger = init_logger(__name__)
81 _DEEPSEEK_V3_FAMILY_MODEL_TYPES = frozenset({"deepseek_v3", "deepseek_v32"})
82
83
84 +class QuantKeyMatch(NamedTuple):
85 + matches: bool
86 + activation_quant_key: QuantKey | None
87 + weight_quant_key: QuantKey | None
88 +
89 + def __bool__(self) -> bool:
90 + return self.matches
91 +
92 +
93 class QuarkConfig(QuantizationConfig):
94 def __init__(
95 self,
@@ -155,41 +190,102 @@ class QuarkConfig(QuantizationConfig):
190 def get_quant_method(
191 self, layer: torch.nn.Module, prefix: str
192 ) -> "QuantizeMethodBase | None":
158 - # Check if the layer is skipped for quantization.
193 + weight_quant_key, activation_quant_key, method_cls = (
194 + self.get_quant_method_target(prefix, type(layer))
195 + )
196 +
197 exclude_layers = cast(list[str], self.quant_config.get("exclude"))
160 - if should_ignore_layer(
198 + is_ignored = should_ignore_layer(
199 prefix,
200 ignore=exclude_layers,
201 fused_mapping=self.packed_modules_mapping,
202 check_children=isinstance(layer, RoutedExperts),
165 - ):
166 - if isinstance(layer, RoutedExperts):
167 - return UnquantizedFusedMoEMethod(layer.moe_config)
168 - if (
169 - "self_attn" not in prefix # only quantize attention projections
170 - or not getattr(self, "dynamic_mxfp4_quant", False)
171 - or not isinstance(layer, LinearBase) # Ignore other methods
172 - ):
173 - return UnquantizedLinearMethod()
203 + )
204 + dynamic_mxfp4_quant = (
205 + is_ignored and "self_attn" in prefix and self.dynamic_mxfp4_quant
206 + )
207
175 - scheme = self.get_scheme(
176 - layer=layer,
177 - layer_name=prefix,
178 - dynamic_mxfp4_quant=True,
208 + if method_cls is UnquantizedFusedMoEMethod:
209 + return UnquantizedFusedMoEMethod(layer.moe_config)
210 + if method_cls is UnquantizedLinearMethod:
211 + return UnquantizedLinearMethod()
212 + if method_cls is QuarkLinearMethod:
213 + _, _, scheme_cls = self.get_scheme_cls(type(layer), prefix)
214 + scheme = self.init_scheme(
215 + scheme_cls,
216 + weight_quant_key=weight_quant_key,
217 + activation_quant_key=activation_quant_key,
218 + dynamic_mxfp4_quant=dynamic_mxfp4_quant,
219 )
220 layer.scheme = scheme
221 return QuarkLinearMethod(self)
182 - if isinstance(layer, LinearBase):
183 - scheme = self.get_scheme(layer=layer, layer_name=prefix)
184 - layer.scheme = scheme
185 - return QuarkLinearMethod(self)
186 - if isinstance(layer, Attention):
222 + if method_cls is QuarkKVCacheMethod:
223 return QuarkKVCacheMethod(self)
224 + if method_cls is not None and issubclass(method_cls, QuarkMoEMethod):
225 + return QuarkMoEMethod.get_moe_method(
226 + self,
227 + module=layer,
228 + method_cls=method_cls,
229 + weight_quant_key=weight_quant_key,
230 + activation_quant_key=activation_quant_key,
231 + )
232
189 - if isinstance(layer, RoutedExperts):
190 - return QuarkMoEMethod.get_moe_method(self, module=layer, layer_name=prefix)
233 return None
234
235 + def get_quant_method_target(
236 + self, prefix: str, layer_type: type[torch.nn.Module]
237 + ) -> tuple[
238 + QuantKey | None,
239 + QuantKey | None,
240 + type[QuantizeMethodBase]
241 + | type[UnquantizedLinearMethod]
242 + | type[UnquantizedFusedMoEMethod]
243 + | None,
244 + ]:
245 + """Return weight key, activation key, and quant method class for
246 + the given ``prefix`` and ``layer_type``.
247 +
248 + This is the counterpart of ``get_quant_method`` without quant method
249 + instantiation.
250 +
251 + TODO: integrate as part of base QuantizationConfig.get_quant_method_target
252 + in the future.
253 +
254 + ``None`` denotes an unquantized/BF16 activation.
255 + """
256 + is_routed_experts = issubclass(layer_type, RoutedExperts)
257 + if should_ignore_layer(
258 + prefix,
259 + ignore=cast(list[str], self.quant_config.get("exclude")),
260 + fused_mapping=self.packed_modules_mapping,
261 + check_children=is_routed_experts,
262 + ):
263 + if is_routed_experts:
264 + return None, None, UnquantizedFusedMoEMethod
265 + if issubclass(layer_type, LinearBase):
266 + if "self_attn" in prefix and self.dynamic_mxfp4_quant:
267 + weight_quant_key, activation_key, _ = self.get_scheme_cls(
268 + layer_type, prefix
269 + )
270 + return weight_quant_key, activation_key, QuarkLinearMethod
271 + return None, None, UnquantizedLinearMethod
272 + return None, None, None
273 + if issubclass(layer_type, LinearBase):
274 + weight_quant_key, activation_key, _ = self.get_scheme_cls(
275 + layer_type, prefix
276 + )
277 + return weight_quant_key, activation_key, QuarkLinearMethod
278 + if issubclass(layer_type, Attention):
279 + return None, None, QuarkKVCacheMethod
280 + if is_routed_experts:
281 + weight_quant_key, activation_quant_key, quant_method_cls = (
282 + QuarkMoEMethod.get_moe_method_target(self, layer_type, prefix)
283 + )
284 + if quant_method_cls is None:
285 + return None, None, None
286 + return weight_quant_key, activation_quant_key, quant_method_cls
287 + return None, None, None
288 +
289 @classmethod
290 def from_config(cls, config: dict[str, Any]) -> "QuarkConfig":
291 export_config = config.get("export")
@@ -290,15 +386,20 @@ class QuarkConfig(QuantizationConfig):
386
387 def _is_fp8_w4a8(
388 self,
293 - weight_quant: list[dict[str, Any]] | None,
294 - input_quant: dict[str, Any] | None,
295 - ) -> bool:
389 + weight_quant: QuarkQTensorHint,
390 + input_quant: QuarkQTensorHint,
391 + ) -> QuantKeyMatch:
392 # Confirm weights and input quantized.
297 - if weight_quant is None or input_quant is None:
298 - return False
299 -
300 - if not isinstance(weight_quant, list) or len(weight_quant) != 2:
301 - return False
393 + if isinstance(input_quant, list):
394 + if len(input_quant) != 1:
395 + return QuantKeyMatch(False, None, None)
396 + input_quant = input_quant[0]
397 + if (
398 + not isinstance(weight_quant, list)
399 + or len(weight_quant) != 2
400 + or not isinstance(input_quant, dict)
401 + ):
402 + return QuantKeyMatch(False, None, None)
403
404 # Confirm weight scheme is supported
405 is_w4a8_dtype = (
@@ -321,24 +422,26 @@ class QuarkConfig(QuantizationConfig):
422 and is_static_weight
423 and is_per_tensor_fp8_and_per_channel_int4_weight
424 ):
324 - return False
425 + return QuantKeyMatch(False, None, None)
426
427 # Dynamic quantization is always supported if weights supported.
428 if input_quant.get("is_dynamic"):
328 - return True
429 + return QuantKeyMatch(True, kFp8DynamicTokenSym, kInt4W4A8StaticChannelSym)
430
431 # Confirm activation scheme is supported.
432 is_per_tensor_activation = input_quant.get("qscheme") == "per_tensor"
332 - return is_per_tensor_activation
433 + if not is_per_tensor_activation:
434 + return QuantKeyMatch(False, None, None)
435 + return QuantKeyMatch(True, kFp8StaticTensorSym, kInt4W4A8StaticChannelSym)
436
437 def _is_fp8_w8a8(
438 self,
439 weight_quant: dict[str, Any] | None,
440 input_quant: dict[str, Any] | None,
338 - ) -> bool:
441 + ) -> QuantKeyMatch:
442 # Confirm weights and input quantized.
340 - if weight_quant is None or input_quant is None:
341 - return False
443 + if not isinstance(weight_quant, dict) or not isinstance(input_quant, dict):
444 + return QuantKeyMatch(False, None, None)
445
446 # Confirm weight scheme is supported
447 is_fp8_dtype = (
@@ -358,64 +461,125 @@ class QuarkConfig(QuantizationConfig):
461 )
462
463 if not is_fp8_dtype or not is_static_weight:
361 - return False
464 + return QuantKeyMatch(False, None, None)
465
466 if is_per_block_weight:
467 if not input_quant.get("is_dynamic"):
365 - return False
366 - return (
468 + return QuantKeyMatch(False, None, None)
469 + matches = (
470 input_quant.get("qscheme") == "per_group"
471 and input_quant.get("group_size") == block_size[1]
472 and input_quant.get("symmetric") is True
473 )
474 + if not matches:
475 + return QuantKeyMatch(False, None, None)
476 + weight_quant_key = (
477 + kFp8Static128BlockE8M0Sym
478 + if weight_quant.get("scale_type") == "float8_e8m0fnu"
479 + else kFp8Static128BlockSym
480 + )
481 + return QuantKeyMatch(True, kFp8Dynamic128Sym, weight_quant_key)
482
483 if not is_per_tensor_or_channel_weight:
373 - return False
484 + return QuantKeyMatch(False, None, None)
485
486 # Dynamic quantization is always supported if tensor/channel weights
487 # are supported.
488 if input_quant.get("is_dynamic"):
378 - return True
489 + activation_quant_key = (
490 + kFp8DynamicTokenSym
491 + if input_quant.get("qscheme") == "per_channel"
492 + else kFp8DynamicTensorSym
493 + )
494 + weight_quant_key = (
495 + kFp8StaticChannelSym
496 + if weight_quant.get("qscheme") == "per_channel"
497 + else kFp8StaticTensorSym
498 + )
499 + return QuantKeyMatch(True, activation_quant_key, weight_quant_key)
500
501 # Confirm activation scheme is supported.
502 is_per_tensor_activation = input_quant.get("qscheme") == "per_tensor"
382 - return is_per_tensor_activation
503 + if not is_per_tensor_activation:
504 + return QuantKeyMatch(False, None, None)
505 + weight_quant_key = (
506 + kFp8StaticChannelSym
507 + if weight_quant.get("qscheme") == "per_channel"
508 + else kFp8StaticTensorSym
509 + )
510 + return QuantKeyMatch(True, kFp8StaticTensorSym, weight_quant_key)
511
384 - def _is_static_tensor_w8a8(
512 + def _is_w8a8_int8(
513 self,
514 weight_quant: dict[str, Any] | None,
515 input_quant: dict[str, Any] | None,
388 - ) -> bool:
516 + ) -> QuantKeyMatch:
517 # Confirm weights and input quantized.
390 - if weight_quant is None or input_quant is None:
391 - return False
518 + if not isinstance(weight_quant, dict) or not isinstance(input_quant, dict):
519 + return QuantKeyMatch(False, None, None)
520
521 is_int8_dtype = (
522 weight_quant.get("dtype") == "int8" and input_quant.get("dtype") == "int8"
523 )
524
397 - is_tensor = (
398 - weight_quant.get("qscheme") in ["per_tensor", "per_channel"]
399 - and input_quant.get("qscheme") == "per_tensor"
400 - )
401 -
402 - is_static = not weight_quant.get("is_dynamic") and not input_quant.get(
403 - "is_dynamic"
525 + is_valid_weight_scheme = weight_quant.get("qscheme") in [
526 + "per_tensor",
527 + "per_channel",
528 + ]
529 + is_static_input = input_quant.get(
530 + "qscheme"
531 + ) == "per_tensor" and not input_quant.get("is_dynamic")
532 + is_dynamic_input = (
533 + input_quant.get("qscheme") == "per_channel"
534 + and input_quant.get("is_dynamic") is True
535 )
405 -
536 is_weight_symmetric = weight_quant.get("symmetric") is True
537
538 # Both symmetric and asymmetric input quantization supported.
539 # Only symmetric weight quantization supported.
410 - return is_int8_dtype and is_tensor and is_weight_symmetric and is_static
540 + matches = (
541 + is_int8_dtype
542 + and is_valid_weight_scheme
543 + and is_weight_symmetric
544 + and (
545 + (is_static_input and not weight_quant.get("is_dynamic"))
546 + or is_dynamic_input
547 + )
548 + )
549 + if not matches:
550 + return QuantKeyMatch(False, None, None)
551 + weight_quant_key = (
552 + kInt8StaticChannelSym
553 + if weight_quant.get("qscheme") == "per_channel"
554 + else kInt8StaticTensorSym
555 + )
556 + if is_static_input:
557 + activation_quant_key = (
558 + kInt8StaticTensorSym
559 + if input_quant.get("symmetric") is True
560 + else kInt8StaticTensorAsym
561 + )
562 + elif weight_quant.get("qscheme") == "per_channel":
563 + activation_quant_key = (
564 + kInt8DynamicTokenSym
565 + if input_quant.get("symmetric") is True
566 + else kInt8DynamicTokenAsym
567 + )
568 + else:
569 + activation_quant_key = (
570 + kInt8DynamicTensorSym
571 + if input_quant.get("symmetric") is True
572 + else kInt8DynamicTensorAsym
573 + )
574 + return QuantKeyMatch(True, activation_quant_key, weight_quant_key)
575
576 def _is_w4a8_mxfp4_fp8(
577 self,
578 weight_quant: dict[str, Any] | None,
579 input_quant: dict[str, Any] | None,
416 - ) -> bool:
580 + ) -> QuantKeyMatch:
581 if weight_quant is None or input_quant is None:
418 - return False
582 + return QuantKeyMatch(False, None, None)
583
584 is_weight_mxfp4 = (
585 weight_quant.get("dtype") == "fp4"
@@ -432,53 +596,24 @@ class QuarkConfig(QuantizationConfig):
596 and input_quant.get("symmetric") is True # Symmetric quantization
597 )
598
435 - return is_weight_mxfp4 and is_input_fp8
436 -
437 - def _is_dynamic_per_token_w8a8(
438 - self,
439 - weight_quant: dict[str, Any] | None,
440 - input_quant: dict[str, Any] | None,
441 - ) -> bool:
442 - """Detect W8A8 INT8 with per-tensor or per-channel
443 - weights and dynamic per-token input."""
444 - if weight_quant is None or input_quant is None:
445 - return False
446 -
447 - is_int8_dtype = (
448 - weight_quant.get("dtype") == "int8" and input_quant.get("dtype") == "int8"
449 - )
450 -
451 - is_valid_weight_scheme = weight_quant.get("qscheme") in [
452 - "per_tensor",
453 - "per_channel",
454 - ]
455 - is_per_token_input = input_quant.get("qscheme") == "per_channel"
456 -
457 - is_dynamic_input = input_quant.get("is_dynamic") is True
458 - is_weight_symmetric = weight_quant.get("symmetric") is True
459 -
460 - return (
461 - is_int8_dtype
462 - and is_valid_weight_scheme
463 - and is_per_token_input
464 - and is_dynamic_input
465 - and is_weight_symmetric
466 - )
599 + if not (is_weight_mxfp4 and is_input_fp8):
600 + return QuantKeyMatch(False, None, None)
601 + return QuantKeyMatch(True, kFp8StaticTensorSym, kMxfp4Static)
602
603 def _is_nvfp4(
604 self,
470 - weight_quant: dict[str, Any] | list[dict[str, Any]] | None,
471 - input_quant: dict[str, Any] | list[dict[str, Any]] | None,
472 - ) -> bool:
605 + weight_quant: QuarkQTensorHint,
606 + input_quant: QuarkQTensorHint,
607 + ) -> QuantKeyMatch:
608 # Confirm weights and input quantized.
609 if weight_quant is None or input_quant is None:
475 - return False
610 + return QuantKeyMatch(False, None, None)
611
612 # Confirm both weight_quant and input_quant are lists with 2 elements
613 if not isinstance(weight_quant, list) or len(weight_quant) != 2:
479 - return False
614 + return QuantKeyMatch(False, None, None)
615 if not isinstance(input_quant, list) or len(input_quant) != 2:
481 - return False
616 + return QuantKeyMatch(False, None, None)
617
618 # First element should be fp4 with per_group quantization
619 is_fp4_per_group_weight = (
@@ -506,36 +641,35 @@ class QuarkConfig(QuantizationConfig):
641 and not input_quant[1].get("is_dynamic")
642 )
643
509 - return (
644 + matches = (
645 is_fp4_per_group_weight # type: ignore[return-value]
646 and is_fp4_per_group_input
647 and is_fp8_per_tensor_weight
648 and is_fp8_per_tensor_input
649 )
650 + if not matches:
651 + return QuantKeyMatch(False, None, None)
652 + return QuantKeyMatch(True, kNvfp4Dynamic, kNvfp4Static)
653
654 def _is_w_ocp_mx_a_x(
517 - self, weight_quant: dict[str, Any] | None, input_quant: dict[str, Any] | None
518 - ) -> bool:
655 + self,
656 + weight_quant: dict[str, Any] | None,
657 + input_quant: dict[str, Any] | None,
658 + allow_static_fp8: bool = False,
659 + ) -> QuantKeyMatch:
660 """
661 This check returns True only if it is an OCP-MX weight quantization.
521 - The activation can be any data type (e.g., FP16/BF16, FP8, or OCP-MX format).
662 + The activation can be FP16/BF16, OCP MXFP4, MXFP8, FP8.
663 The rationale for checking only the weight type is that
664 the model loading concept and process primarily concerns the weights themselves.
665 """
666 # Confirm weights quantized.
526 - if weight_quant is None:
667 + if not isinstance(weight_quant, dict):
668 logger.debug(
669 "Quark model's weight quantization is incompatible with OCP_MX format: "
529 - "weight_quant is not set."
670 + "weight_quant is not a dictionary."
671 )
531 - return False
532 -
533 - if isinstance(weight_quant, list):
534 - logger.debug(
535 - "Quark model's weight quantization is incompatible with OCP_MX format: "
536 - "weight_quant is a list (e.g. fp8_w4a8), OCP_MX requires a single dict."
537 - )
538 - return False
672 + return QuantKeyMatch(False, None, None)
673
674 # Input and weight qscheme needs to be per group.
675 if weight_quant.get("qscheme") != "per_group":
@@ -543,7 +677,7 @@ class QuarkConfig(QuantizationConfig):
677 "Quark model's weight quantization is incompatible with OCP MX format: "
678 "weight is not per_group."
679 )
546 - return False
680 + return QuantKeyMatch(False, None, None)
681
682 # Input and weight group size needs to be 32.
683 if weight_quant.get("group_size") != 32:
@@ -551,7 +685,7 @@ class QuarkConfig(QuantizationConfig):
685 "Quark model's weight quantization is incompatible with OCP MX format: "
686 "group_size of weight is not 32."
687 )
554 - return False
688 + return QuantKeyMatch(False, None, None)
689
690 # Activations and weight scales need to be in e8m0 format.
691 if weight_quant.get("scale_format") != "e8m0":
@@ -559,7 +693,7 @@ class QuarkConfig(QuantizationConfig):
693 "Quark model's weight quantization is incompatible with OCP MX format: "
694 "scale_format of weight is not e8m0."
695 )
562 - return False
696 + return QuantKeyMatch(False, None, None)
697
698 # Input and weight dtypes need to be any of fp4,
699 # fp6_e3m2 or fp6_e3m2, possibly mixed.
@@ -572,9 +706,48 @@ class QuarkConfig(QuantizationConfig):
706 "Quark model's weight quantization is incompatible with OCP MX format: "
707 "dtype is not in {fp4, fp6_e3m2, fp6_e2m3}."
708 )
575 - return False
709 + return QuantKeyMatch(False, None, None)
710 +
711 + weight_dtype = weight_quant["dtype"].replace("fp", "mxfp")
712 + weight_quant_key = _WEIGHT_QUANT_KEY_MAP[weight_dtype]
713 + if input_quant is None:
714 + activation_quant_key = None
715 + elif not input_quant.get("is_dynamic"):
716 + if (
717 + allow_static_fp8
718 + and input_quant.get("dtype") == "fp8_e4m3"
719 + and input_quant.get("qscheme") == "per_tensor"
720 + and input_quant.get("symmetric") is True
721 + ):
722 + activation_quant_key = kFp8StaticTensorSym
723 + else:
724 + logger.debug(
725 + "Quark model's OCP MX quantization is incompatible with static "
726 + "input scales."
727 + )
728 + return QuantKeyMatch(False, None, None)
729 + elif input_quant["dtype"] == "fp8_e4m3":
730 + activation_quant_key = kFp8DynamicTensorSym
731 + else:
732 + input_dtype = input_quant["dtype"].replace("fp", "mxfp")
733 + if input_dtype not in _ACTIVATION_QUANT_KEY_MAP:
734 + raise ValueError(
735 + f"Unsupported input_dtype={input_dtype} in Quark's vLLM "
736 + "integration. Supported activation dtypes are "
737 + f"{_ACTIVATION_QUANT_KEY_MAP.keys()}, or None for "
738 + "weight-only quantization."
739 + )
740 + activation_quant_key = _ACTIVATION_QUANT_KEY_MAP[input_dtype]
741 + return QuantKeyMatch(True, activation_quant_key, weight_quant_key)
742
577 - return True
743 + @staticmethod
744 + def _unwrap_single_quant_config(
745 + quant_config: QuarkQTensorHint,
746 + ) -> dict[str, Any] | None:
747 + """Unwrap one-entry quantization-config lists."""
748 + if isinstance(quant_config, list):
749 + return quant_config[0] if len(quant_config) == 1 else None
750 + return quant_config
751
752 def get_layer_quant_config_from_name(
753 self, layer_name: str
@@ -617,14 +790,14 @@ class QuarkConfig(QuantizationConfig):
790 return None
791
792 def _find_matched_config(
620 - self, layer_name: str, module: torch.nn.Module
793 + self, layer_name: str, module: torch.nn.Module | type[torch.nn.Module]
794 ) -> dict[str, Any]:
795 # Priority order:
796 # 1. layer_quant_config,
797 # 2. layer_type_quant_config,
798 # 3. global_quant_config.
799
627 - layer_type = cast(str, type(module))
800 + layer_type = cast(str, module if isinstance(module, type) else type(module))
801 layer_type_quant_config = cast(
802 dict[str, Any], self.quant_config.get("layer_type_quant_config")
803 )
@@ -662,51 +835,71 @@ class QuarkConfig(QuantizationConfig):
835 return layer_quant_config
836 return fallback_config
837
665 - def _get_scheme_from_config(
666 - self, config: dict[str, Any], dynamic_mxfp4_quant: bool = False
667 - ) -> "QuarkScheme":
838 + def _get_scheme_cls_from_config(
839 + self, config: dict[str, Any]
840 + ) -> tuple[QuantKey | None, QuantKey | None, type["QuarkScheme"]]:
841 if config.get("output_tensors") or config.get("bias"):
842 raise NotImplementedError(
843 "Currently, Quark models with output_tensors "
844 "and bias quantized are not supported"
845 )
673 - weight_config = cast(dict[str, Any], config.get("weight"))
674 - input_config = cast(dict[str, Any], config.get("input_tensors"))
675 -
676 - if self._is_nvfp4(weight_config, input_config):
677 - return QuarkNVFP4()
678 - elif self._is_fp8_w8a8(weight_config, input_config):
679 - is_fp8_w8a8_supported = self._check_scheme_supported(
680 - QuarkW8A8Fp8.get_min_capability(), error=False
681 - )
682 - if is_fp8_w8a8_supported:
683 - if weight_config.get("qscheme") == "per_block":
684 - return QuarkW8A8Fp8PerBlock(weight_config, input_config)
685 - return QuarkW8A8Fp8(weight_config, input_config)
686 - elif self._is_static_tensor_w8a8(weight_config, input_config):
687 - weight_qscheme = cast(str, weight_config.get("qscheme"))
688 - return QuarkW8A8Int8(
689 - qscheme=weight_qscheme,
690 - is_static_input_scheme=True,
691 - input_symmetric=input_config.get("symmetric"),
846 + weight_config = cast(QuarkQTensorHint, config.get("weight"))
847 + input_config = cast(QuarkQTensorHint, config.get("input_tensors"))
848 +
849 + if (
850 + isinstance(weight_config, list)
851 + and len(weight_config) > 1
852 + or isinstance(input_config, list)
853 + and len(input_config) > 1
854 + ):
855 + if match := self._is_nvfp4(weight_config, input_config):
856 + return (
857 + match.weight_quant_key,
858 + match.activation_quant_key,
859 + QuarkNVFP4,
860 + )
861 +
862 + raise NotImplementedError(
863 + "Multi-entry weight or activation quantization configs are only "
864 + "supported for NVFP4."
865 )
693 - elif self._is_w4a8_mxfp4_fp8(weight_config, input_config):
694 - is_w4a8_supported = self._check_scheme_supported(
695 - QuarkW4A8_MXFP4_FP8.get_min_capability(), error=False
866 +
867 + weight_config = self._unwrap_single_quant_config(weight_config)
868 + input_config = self._unwrap_single_quant_config(input_config)
869 +
870 + if (
871 + isinstance(weight_config, dict)
872 + and weight_config.get("qscheme") == "per_block"
873 + and not weight_config.get("block_size")
874 + ):
875 + raise ValueError(
876 + "Quark W8A8 FP8 per-block weight quantization requires "
877 + "`block_size` in the weight quantization config."
878 )
697 - if is_w4a8_supported:
698 - return QuarkW4A8_MXFP4_FP8(weight_config, input_config)
699 - elif self._is_dynamic_per_token_w8a8(weight_config, input_config):
700 - weight_qscheme = cast(str, weight_config.get("qscheme"))
701 - return QuarkW8A8Int8(
702 - qscheme=weight_qscheme,
703 - is_static_input_scheme=False,
704 - input_symmetric=input_config.get("symmetric"),
879 +
880 + if match := self._is_fp8_w8a8(weight_config, input_config):
881 + assert weight_config is not None
882 + if weight_config.get("qscheme") == "per_block":
883 + return (
884 + match.weight_quant_key,
885 + match.activation_quant_key,
886 + QuarkW8A8Fp8PerBlock,
887 + )
888 + return match.weight_quant_key, match.activation_quant_key, QuarkW8A8Fp8
889 + elif match := self._is_w8a8_int8(weight_config, input_config):
890 + return (
891 + match.weight_quant_key,
892 + match.activation_quant_key,
893 + QuarkW8A8Int8,
894 )
706 - elif self._is_w_ocp_mx_a_x(weight_config, input_config):
707 - return QuarkOCP_MX(
708 - weight_config, input_config, dynamic_mxfp4_quant=dynamic_mxfp4_quant
895 + elif match := self._is_w4a8_mxfp4_fp8(weight_config, input_config):
896 + return (
897 + match.weight_quant_key,
898 + match.activation_quant_key,
899 + QuarkW4A8_MXFP4_FP8,
900 )
901 + elif match := self._is_w_ocp_mx_a_x(weight_config, input_config):
902 + return match.weight_quant_key, match.activation_quant_key, QuarkOCP_MX
903
904 raise NotImplementedError(
905 "No quark compatible scheme was found. "
@@ -714,19 +907,51 @@ class QuarkConfig(QuantizationConfig):
907 f"Input config: {input_config}"
908 )
909
717 - def get_scheme(
718 - self, layer: torch.nn.Module, layer_name: str, dynamic_mxfp4_quant: bool = False
910 + def get_scheme_cls(
911 + self, layer_type: type[torch.nn.Module], layer_name: str
912 + ) -> tuple[QuantKey | None, QuantKey | None, type["QuarkScheme"]]:
913 + """Return quantization keys and scheme class without initializing it."""
914 + return self._get_scheme_cls_from_config(
915 + self._find_matched_config(layer_name, layer_type)
916 + )
917 +
918 + def init_scheme(
919 + self,
920 + scheme_cls: type["QuarkScheme"],
921 + weight_quant_key: QuantKey | None,
922 + activation_quant_key: QuantKey | None,
923 + dynamic_mxfp4_quant: bool = False,
924 ) -> "QuarkScheme":
720 - layer_quant_config = self._find_matched_config(layer_name, layer)
925 + """Construct a Quark scheme selected by get_scheme_cls."""
926 + if scheme_cls not in (
927 + QuarkW8A8Fp8,
928 + QuarkW8A8Fp8PerBlock,
929 + QuarkW8A8Int8,
930 + QuarkOCP_MX,
931 + QuarkNVFP4,
932 + QuarkW4A8_MXFP4_FP8,
933 + ):
934 + raise AssertionError(f"Unsupported Quark scheme class: {scheme_cls}")
935 +
936 + kwargs: dict[str, Any] = {"activation_quant_key": activation_quant_key}
937 +
938 + # These classes handle different possible `weight_quant_key`.
939 + if scheme_cls in (
940 + QuarkW8A8Fp8,
941 + QuarkW8A8Fp8PerBlock,
942 + QuarkW8A8Int8,
943 + QuarkOCP_MX,
944 + ):
945 + kwargs["weight_quant_key"] = weight_quant_key
946 +
947 + if scheme_cls is QuarkOCP_MX:
948 + kwargs["dynamic_mxfp4_quant"] = dynamic_mxfp4_quant
949 +
950 + scheme = scheme_cls(**kwargs)
951
722 - # Find the quant_scheme
723 - scheme = self._get_scheme_from_config(
724 - layer_quant_config, dynamic_mxfp4_quant=dynamic_mxfp4_quant
725 - )
952 # Raise error if device does not support the scheme
953 # (e.g. fp8 needs ada lovelace)
954 self._check_scheme_supported(scheme.get_min_capability())
729 -
955 return scheme
956
957 @staticmethod
vllm/model_executor/layers/quantization/quark/quark_moe.py
+208 -111
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import Any
4 +from typing import TYPE_CHECKING
5
6 import torch
7
@@ -56,19 +56,27 @@ from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
56 make_nvfp4_moe_quant_config,
57 select_nvfp4_moe_backend,
58 )
59 +from vllm.model_executor.layers.quantization.quark.utils import QuarkQTensorHint
60 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
61 + _ACTIVATION_QUANT_KEY_MAP,
62 + _WEIGHT_QUANT_KEY_MAP,
63 OCP_MX_BLOCK_SIZE,
64 OCP_MX_Scheme,
65 )
66 from vllm.model_executor.layers.quantization.utils.quant_utils import (
67 GroupShape,
68 + QuantKey,
69 kFp8DynamicTensorSym,
70 kFp8DynamicTokenSym,
71 kFp8StaticChannelSym,
72 kFp8StaticTensorSym,
73 + kInt4W4A8StaticChannelSym,
74 + kInt8DynamicTensorAsym,
75 kInt8DynamicTensorSym,
76 + kInt8DynamicTokenAsym,
77 kInt8DynamicTokenSym,
78 kInt8StaticChannelSym,
79 + kInt8StaticTensorAsym,
80 kInt8StaticTensorSym,
81 kMxfp4Dynamic,
82 kNvfp4Dynamic,
@@ -82,6 +90,9 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
90 from vllm.model_executor.utils import replace_parameter, set_weight_attrs
91 from vllm.platforms import current_platform
92
93 +if TYPE_CHECKING:
94 + from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
95 +
96 logger = init_logger(__name__)
97
98 __all__ = [
@@ -93,67 +104,154 @@ __all__ = [
104
105
106 class QuarkMoEMethod(FusedMoEMethodBase):
96 - def __init__(self, moe: FusedMoEConfig):
107 + supported_activation_quant_keys: list[QuantKey | None] = []
108 + supported_weight_quant_keys: list[QuantKey] = []
109 +
110 + def __init__(
111 + self,
112 + moe: FusedMoEConfig,
113 + weight_quant_key: QuantKey,
114 + activation_quant_key: QuantKey | None,
115 + ):
116 super().__init__(moe)
117 + if activation_quant_key not in self.supported_activation_quant_keys:
118 + raise ValueError(
119 + f"Unsupported activation quant key: {activation_quant_key}"
120 + )
121 + if weight_quant_key not in self.supported_weight_quant_keys:
122 + raise ValueError(f"Unsupported weight quant key: {weight_quant_key}")
123 + self.activation_quant_key = activation_quant_key
124 + self.weight_quant_key = weight_quant_key
125 self.has_bias = self.moe.has_bias
126
127 @staticmethod
101 - def get_moe_method(
102 - quant_config: "QuarkConfig", # type: ignore # noqa E501 # noqa F821
103 - module: RoutedExperts,
128 + def get_moe_method_target(
129 + quant_config: "QuarkConfig",
130 + layer_type: type[torch.nn.Module],
131 layer_name: str,
105 - ) -> "QuarkMoEMethod":
106 - layer_quant_config = quant_config._find_matched_config(layer_name, module)
107 -
132 + ) -> tuple[
133 + QuantKey | None,
134 + QuantKey | None,
135 + type["QuarkMoEMethod"] | None,
136 + ]:
137 + """Return the quantization target for a routed-experts layer."""
138 + layer_quant_config = quant_config._find_matched_config(layer_name, layer_type)
139 if layer_quant_config.get("output_tensors") or layer_quant_config.get("bias"):
140 raise NotImplementedError(
141 "Currently, Quark models with "
142 "output_tensors and bias "
143 "quantized are not supported"
144 )
145 + weight_config: QuarkQTensorHint = layer_quant_config.get("weight")
146 + input_config: QuarkQTensorHint = layer_quant_config.get("input_tensors")
147 +
148 + if (
149 + isinstance(weight_config, list)
150 + and len(weight_config) > 1
151 + or isinstance(input_config, list)
152 + and len(input_config) > 1
153 + ):
154 + if match := quant_config._is_fp8_w4a8(weight_config, input_config):
155 + return (
156 + match.weight_quant_key,
157 + match.activation_quant_key,
158 + QuarkW4A8Fp8MoEMethod,
159 + )
160 + if match := quant_config._is_nvfp4(weight_config, input_config):
161 + return (
162 + match.weight_quant_key,
163 + match.activation_quant_key,
164 + QuarkNvfp4MoEMethod,
165 + )
166 + raise RuntimeError("Unsupported FusedMoe scheme")
167 +
168 + weight_config = quant_config._unwrap_single_quant_config(weight_config)
169 + input_config = quant_config._unwrap_single_quant_config(input_config)
170 +
171 + if match := quant_config._is_fp8_w8a8(weight_config, input_config):
172 + return (
173 + match.weight_quant_key,
174 + match.activation_quant_key,
175 + QuarkW8A8Fp8MoEMethod,
176 + )
177 + if match := quant_config._is_w_ocp_mx_a_x(
178 + weight_config, input_config, allow_static_fp8=True
179 + ):
180 + return (
181 + match.weight_quant_key,
182 + match.activation_quant_key,
183 + QuarkOCP_MX_MoEMethod,
184 + )
185 + if match := quant_config._is_w8a8_int8(weight_config, input_config):
186 + return (
187 + match.weight_quant_key,
188 + match.activation_quant_key,
189 + QuarkW8A8Int8MoEMethod,
190 + )
191 + raise RuntimeError("Unsupported FusedMoe scheme")
192
115 - weight_config = layer_quant_config.get("weight")
116 - input_config = layer_quant_config.get("input_tensors")
193 + @staticmethod
194 + def get_moe_method(
195 + quant_config: "QuarkConfig", # type: ignore # noqa E501 # noqa F821
196 + module: RoutedExperts,
197 + method_cls: type["QuarkMoEMethod"],
198 + weight_quant_key: QuantKey | None,
199 + activation_quant_key: QuantKey | None,
200 + ) -> "QuarkMoEMethod":
201 + if weight_quant_key is None:
202 + raise RuntimeError("Unsupported FusedMoe scheme")
203
118 - if quant_config._is_fp8_w4a8(weight_config, input_config):
119 - return QuarkW4A8Fp8MoEMethod(weight_config, input_config, module.moe_config)
120 - elif quant_config._is_nvfp4(weight_config, input_config):
204 + if method_cls is QuarkW4A8Fp8MoEMethod:
205 + return QuarkW4A8Fp8MoEMethod(module.moe_config, activation_quant_key)
206 + elif method_cls is QuarkNvfp4MoEMethod:
207 return QuarkNvfp4MoEMethod(
122 - weight_config, input_config, module.moe_config, quant_config
208 + module.moe_config, quant_config, activation_quant_key
209 + )
210 + elif method_cls is QuarkW8A8Fp8MoEMethod:
211 + return QuarkW8A8Fp8MoEMethod(
212 + module.moe_config, weight_quant_key, activation_quant_key
213 )
124 - elif quant_config._is_fp8_w8a8(weight_config, input_config):
125 - return QuarkW8A8Fp8MoEMethod(weight_config, input_config, module.moe_config)
126 - elif quant_config._is_w_ocp_mx_a_x(weight_config, input_config):
214 + elif method_cls is QuarkOCP_MX_MoEMethod:
215 # All OCP MX schemes (W4A16, W4A8, etc.) handled by QuarkOCP_MX_MoEMethod
216 # Backend selection happens inside via oracle
129 - return QuarkOCP_MX_MoEMethod(weight_config, input_config, module.moe_config)
130 - elif quant_config._is_static_tensor_w8a8(
131 - weight_config, input_config
132 - ) or quant_config._is_dynamic_per_token_w8a8(weight_config, input_config):
217 + return QuarkOCP_MX_MoEMethod(
218 + module.moe_config, weight_quant_key, activation_quant_key
219 + )
220 + elif method_cls is QuarkW8A8Int8MoEMethod:
221 return QuarkW8A8Int8MoEMethod(
134 - weight_config, input_config, module.moe_config
222 + module.moe_config, weight_quant_key, activation_quant_key
223 )
224 else:
225 raise RuntimeError("Unsupported FusedMoe scheme")
226
227
228 class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
229 + supported_activation_quant_keys = [
230 + kFp8StaticTensorSym,
231 + kFp8DynamicTensorSym,
232 + kFp8DynamicTokenSym,
233 + ]
234 + supported_weight_quant_keys = [
235 + kFp8StaticChannelSym,
236 + kFp8StaticTensorSym,
237 + ]
238 +
239 def __init__(
240 self,
143 - weight_config: dict[str, Any],
144 - input_config: dict[str, Any],
241 moe: FusedMoEConfig,
242 + weight_quant_key: QuantKey,
243 + activation_quant_key: QuantKey | None,
244 ):
147 - super().__init__(moe)
148 - self.weight_quant = weight_config
149 - self.input_quant = input_config
150 -
151 - self.weight_qscheme = self.weight_quant.get("qscheme")
152 - self.input_qscheme = self.input_quant.get("qscheme")
153 - self.weight_dtype = self.weight_quant.get("dtype", "").replace(
154 - "fp8_e4m3", "fp8"
245 + super().__init__(moe, weight_quant_key, activation_quant_key)
246 + self.weight_qscheme = (
247 + "per_channel" if weight_quant_key == kFp8StaticChannelSym else "per_tensor"
248 )
156 - self.input_dtype = self.input_quant.get("dtype", "").replace("fp8_e4m3", "fp8")
249 + self.input_qscheme = (
250 + "per_channel"
251 + if activation_quant_key == kFp8DynamicTokenSym
252 + else "per_tensor"
253 + )
254 + self.static_input_scales = activation_quant_key == kFp8StaticTensorSym
255 per_tensor = (
256 self.weight_qscheme == "per_tensor" and self.input_qscheme == "per_tensor"
257 )
@@ -170,28 +268,10 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
268 f"{self.weight_qscheme}, {self.input_qscheme}"
269 ) # noqa E501
270
173 - self.static_input_scales = not self.input_quant.get("is_dynamic")
174 - if self.static_input_scales and per_channel:
175 - raise ValueError(
176 - "For FP8 Fused MoE layer, we require either per tensor or "
177 - "channelwise, dynamic per token quantization."
178 - )
179 -
180 - # Determine quant keys for oracle backend selection
181 - if per_channel:
182 - weight_key = kFp8StaticChannelSym
183 - activation_key = kFp8DynamicTokenSym
184 - elif self.static_input_scales:
185 - weight_key = kFp8StaticTensorSym
186 - activation_key = kFp8StaticTensorSym
187 - else:
188 - weight_key = kFp8StaticTensorSym
189 - activation_key = kFp8DynamicTensorSym
190 -
271 self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
272 config=moe,
193 - weight_key=weight_key,
194 - activation_key=activation_key,
273 + weight_key=self.weight_quant_key,
274 + activation_key=self.activation_quant_key,
275 )
276
277 self.model_type = getattr(
@@ -500,17 +580,41 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
580 class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
581 """Quark W8A8 INT8 MoE method."""
582
583 + supported_activation_quant_keys = [
584 + kInt8StaticTensorSym,
585 + kInt8StaticTensorAsym,
586 + kInt8DynamicTensorSym,
587 + kInt8DynamicTensorAsym,
588 + kInt8DynamicTokenSym,
589 + kInt8DynamicTokenAsym,
590 + ]
591 + supported_weight_quant_keys = [
592 + kInt8StaticChannelSym,
593 + kInt8StaticTensorSym,
594 + ]
595 +
596 def __init__(
597 self,
505 - weight_config: dict[str, Any],
506 - input_config: dict[str, Any],
598 moe: FusedMoEConfig,
599 + weight_quant_key: QuantKey,
600 + activation_quant_key: QuantKey | None,
601 ):
509 - super().__init__(moe)
510 - self.weight_quant = weight_config
511 - self.input_quant = input_config
512 - self.weight_qscheme = self.weight_quant.get("qscheme", "per_tensor")
513 - self.static_input_scales = not self.input_quant.get("is_dynamic", False)
602 + super().__init__(moe, weight_quant_key, activation_quant_key)
603 +
604 + # TODO: oracle needs to support asym int8.
605 + if self.activation_quant_key == kInt8StaticTensorAsym:
606 + self.activation_quant_key = kInt8StaticTensorSym
607 + if self.activation_quant_key == kInt8DynamicTensorAsym:
608 + self.activation_quant_key = kInt8DynamicTensorSym
609 + if self.activation_quant_key == kInt8DynamicTokenAsym:
610 + self.activation_quant_key = kInt8DynamicTokenSym
611 +
612 + self.weight_qscheme = (
613 + "per_channel" if weight_quant_key == kInt8StaticChannelSym else "per_tensor"
614 + )
615 +
616 + assert self.activation_quant_key is not None
617 + self.static_input_scales = self.activation_quant_key.scale.static
618
619 self.moe_quant_config: FusedMoEQuantConfig | None = None
620 self.moe_kernel: mk.FusedMoEKernel | None = None
@@ -528,17 +632,10 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
632 # Map the Quark weight scheme to oracle quant keys. Per-channel
633 # weights pair with dynamic per-token activations; per-tensor
634 # weights with dynamic per-tensor activations.
531 - if self.weight_qscheme == "per_channel":
532 - weight_key = kInt8StaticChannelSym
533 - activation_key = kInt8DynamicTokenSym
534 - else:
535 - weight_key = kInt8StaticTensorSym
536 - activation_key = kInt8DynamicTensorSym
537 -
635 self.int8_backend, self.experts_cls = select_int8_moe_backend(
636 config=moe,
540 - weight_key=weight_key,
541 - activation_key=activation_key,
637 + weight_key=self.weight_quant_key,
638 + activation_key=self.activation_quant_key,
639 )
640
641 def create_weights(
@@ -876,15 +973,18 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
973
974
975 class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
976 + supported_activation_quant_keys = [
977 + kFp8DynamicTokenSym,
978 + kFp8StaticTensorSym,
979 + ]
980 + supported_weight_quant_keys = [kInt4W4A8StaticChannelSym]
981 +
982 def __init__(
983 self,
881 - weight_config: dict[str, Any],
882 - input_config: dict[str, Any],
984 moe: FusedMoEConfig,
985 + activation_quant_key: QuantKey | None,
986 ):
885 - super().__init__(moe)
886 - self.weight_quant = weight_config
887 - self.input_quant = input_config
987 + super().__init__(moe, kInt4W4A8StaticChannelSym, activation_quant_key)
988
989 assert rocm_aiter_ops.is_fused_moe_enabled(), (
990 "W4A8 FP8 MoE requires ROCm AITER fused MoE support."
@@ -1032,37 +1132,36 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
1132
1133
1134 class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1135 + supported_activation_quant_keys = [
1136 + *_ACTIVATION_QUANT_KEY_MAP.values(),
1137 + kFp8DynamicTensorSym,
1138 + kFp8StaticTensorSym,
1139 + None,
1140 + ]
1141 + supported_weight_quant_keys = [*_WEIGHT_QUANT_KEY_MAP.values()]
1142 +
1143 def __init__(
1144 self,
1037 - weight_config: dict[str, Any],
1038 - input_config: dict[str, Any] | None,
1145 moe: FusedMoEConfig,
1146 + weight_quant_key: QuantKey,
1147 + activation_quant_key: QuantKey | None,
1148 ):
1041 - super().__init__(moe)
1042 - self.weight_quant = weight_config
1043 - self.input_quant = input_config
1044 -
1045 - weight_qscheme = self.weight_quant.get("qscheme")
1046 - if not weight_qscheme == "per_group":
1047 - raise ValueError(
1048 - "For MX(FP4) Fused MoE layers, only per-group scales "
1049 - f"for weights are supported. Found {weight_qscheme}."
1050 - ) # noqa E501
1051 -
1052 - self.weight_dtype = self.weight_quant["dtype"].replace("fp", "mxfp")
1053 - if self.input_quant is not None:
1054 - input_quant = self.input_quant["dtype"]
1055 - if input_quant in ["fp4", "fp6_e3m2", "fp6_e2m3"]:
1056 - self.input_dtype = input_quant.replace("fp", "mxfp")
1057 - elif input_quant == "fp8_e4m3":
1058 - self.input_dtype = input_quant.replace("fp8_e4m3", "fp8")
1059 - else:
1060 - raise NotImplementedError(
1061 - f"Current input dtype {input_quant} is not compatible \
1062 - with OCP MX (weight) MoE quantization. Please open an issue"
1063 - )
1064 - else:
1149 + super().__init__(moe, weight_quant_key, activation_quant_key)
1150 + self.weight_dtype = next(
1151 + dtype
1152 + for dtype, quant_key in _WEIGHT_QUANT_KEY_MAP.items()
1153 + if quant_key == weight_quant_key
1154 + )
1155 + if activation_quant_key in {kFp8DynamicTensorSym, kFp8StaticTensorSym}:
1156 + self.input_dtype: str | None = "fp8"
1157 + elif activation_quant_key is None:
1158 self.input_dtype = None
1159 + else:
1160 + self.input_dtype = next(
1161 + dtype
1162 + for dtype, quant_key in _ACTIVATION_QUANT_KEY_MAP.items()
1163 + if quant_key == activation_quant_key
1164 + )
1165
1166 self.ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype(
1167 self.input_dtype, self.weight_dtype
@@ -1085,10 +1184,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1184 self.w13_precision_config = None
1185 self.w2_precision_config = None
1186
1088 - if self.input_quant is not None:
1089 - self.static_input_scales = not self.input_quant.get("is_dynamic")
1090 - else:
1091 - self.static_input_scales = False
1187 + self.static_input_scales = activation_quant_key == kFp8StaticTensorSym
1188
1189 # Select backend based on OCP MX scheme
1190 if self.ocp_mx_scheme == "w_mxfp4":
@@ -1388,6 +1484,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1484 f"in vLLM for OCP MX scheme {self.ocp_mx_scheme}. Please open an issue."
1485 )
1486 else:
1487 + assert self.input_dtype is not None
1488 return ocp_mx_moe_quant_config(
1489 quant_dtype=self.input_dtype,
1490 weight_dtype=self.weight_dtype,
@@ -1459,24 +1556,24 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1556
1557
1558 class QuarkNvfp4MoEMethod(QuarkMoEMethod):
1559 + supported_activation_quant_keys = [kNvfp4Dynamic]
1560 + supported_weight_quant_keys = [kNvfp4Static]
1561 +
1562 def __init__(
1563 self,
1464 - weight_config: dict[str, Any],
1465 - input_config: dict[str, Any],
1564 moe: FusedMoEConfig,
1565 quant_config: "QuarkConfig", # type: ignore # noqa E501 # noqa F821
1566 + activation_quant_key: QuantKey | None,
1567 ):
1469 - super().__init__(moe)
1470 - self.weight_quant = weight_config
1471 - self.input_quant = input_config
1568 + super().__init__(moe, kNvfp4Static, activation_quant_key)
1569 self.quant_config = quant_config
1570 self.group_size = 16
1571
1572 # Select experts implementation.
1573 self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
1574 config=self.moe,
1478 - weight_key=kNvfp4Static,
1479 - activation_key=kNvfp4Dynamic,
1575 + weight_key=self.weight_quant_key,
1576 + activation_key=self.activation_quant_key,
1577 )
1578
1579 def create_weights(
vllm/model_executor/layers/quantization/quark/schemes/quark_nvfp4.py
+10
@@ -14,6 +14,11 @@ from vllm.model_executor.kernels.linear.nvfp4.emulation import (
14 from vllm.model_executor.layers.quantization.quark.schemes.quark_scheme import (
15 QuarkScheme,
16 )
17 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
18 + QuantKey,
19 + kNvfp4Dynamic,
20 + kNvfp4Static,
21 +)
22 from vllm.model_executor.parameter import (
23 GroupQuantScaleParameter,
24 ModelWeightParameter,
@@ -36,9 +41,14 @@ class QuarkNVFP4(QuarkScheme):
41 - input_scale_2: bfloat16/float32, scalar (global input scale)
42 """
43
44 + supported_activation_quant_keys = [kNvfp4Dynamic]
45 + supported_weight_quant_keys = [kNvfp4Static]
46 +
47 def __init__(
48 self,
49 + activation_quant_key: QuantKey | None,
50 ):
51 + super().__init__(kNvfp4Static, activation_quant_key)
52 self.kernel = init_nvfp4_linear_kernel()
53 self.group_size = 16
54
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+18 -44
@@ -3,7 +3,6 @@
3
4 from collections.abc import Callable
5 from fractions import Fraction
6 -from typing import Any
6
7 import torch
8
@@ -15,15 +14,14 @@ from vllm.model_executor.kernels.linear import (
14 init_mxfp6_linear_kernel,
15 )
16 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
17 + _ACTIVATION_QUANT_KEY_MAP,
18 + _WEIGHT_QUANT_KEY_MAP,
19 OCP_MX_BLOCK_SIZE,
20 )
21 from vllm.model_executor.layers.quantization.utils.quant_utils import (
22 QuantKey,
22 - kMxfp4Dynamic,
23 kMxfp4Static,
24 - kMxfp6E2M3Dynamic,
24 kMxfp6E2M3Static,
26 - kMxfp6E3M2Dynamic,
25 kMxfp6E3M2Static,
26 )
27 from vllm.model_executor.parameter import (
@@ -38,45 +36,32 @@ from .quark_scheme import QuarkScheme
36
37 logger = init_logger(__name__)
38
41 -_WEIGHT_QUANT_KEY_MAP: dict[str, QuantKey] = {
42 - "mxfp4": kMxfp4Static,
43 - "mxfp6_e3m2": kMxfp6E3M2Static,
44 - "mxfp6_e2m3": kMxfp6E2M3Static,
45 -}
46 -
47 -_ACTIVATION_QUANT_KEY_MAP: dict[str, QuantKey] = {
48 - "mxfp4": kMxfp4Dynamic,
49 - "mxfp6_e3m2": kMxfp6E3M2Dynamic,
50 - "mxfp6_e2m3": kMxfp6E2M3Dynamic,
51 -}
52 -
39
40 class QuarkOCP_MX(QuarkScheme):
41 ocp_mx_linear: MxFp6LinearKernel | MxFp4LinearKernel
42 + supported_activation_quant_keys = [*_ACTIVATION_QUANT_KEY_MAP.values(), None]
43 + supported_weight_quant_keys = [*_WEIGHT_QUANT_KEY_MAP.values()]
44
45 def __init__(
46 self,
59 - weight_quant_spec: dict[str, Any],
60 - input_quant_spec: dict[str, Any] | None,
47 + weight_quant_key: QuantKey,
48 + activation_quant_key: QuantKey | None,
49 dynamic_mxfp4_quant: bool = False,
50 ):
51 + super().__init__(weight_quant_key, activation_quant_key)
52 self.dynamic_mxfp4_quant = dynamic_mxfp4_quant
64 - self.weight_dtype = weight_quant_spec["dtype"].replace("fp", "mxfp")
65 - self.input_dtype: str | None = None
66 - if input_quant_spec is not None:
67 - self.input_dtype = input_quant_spec["dtype"].replace("fp", "mxfp")
68 -
69 - if self.input_dtype not in [None, *_ACTIVATION_QUANT_KEY_MAP]:
70 - raise ValueError(
71 - f"Unsupported input_dtype={self.input_dtype} for QuarkOCP_MX. "
72 - f"Supported activation dtypes are {_ACTIVATION_QUANT_KEY_MAP.keys()}, "
73 - "or None for weight-only quantization."
53 + self.weight_dtype = next(
54 + dtype
55 + for dtype, quant_key in _WEIGHT_QUANT_KEY_MAP.items()
56 + if quant_key == weight_quant_key
57 + )
58 + self.input_dtype = (
59 + next(
60 + dtype
61 + for dtype, quant_key in _ACTIVATION_QUANT_KEY_MAP.items()
62 + if quant_key == activation_quant_key
63 )
75 -
76 - self.weight_quant_key = _WEIGHT_QUANT_KEY_MAP[self.weight_dtype]
77 - self.activation_quant_key = (
78 - _ACTIVATION_QUANT_KEY_MAP[self.input_dtype]
79 - if self.input_dtype is not None
64 + if activation_quant_key is not None
65 else None
66 )
67
@@ -85,17 +70,6 @@ class QuarkOCP_MX(QuarkScheme):
70 else:
71 self.packed_factor = Fraction(numerator=8, denominator=6)
72
88 - if input_quant_spec is None:
89 - self.static_input_scales = False
90 - else:
91 - self.static_input_scales = not input_quant_spec.get("is_dynamic")
92 -
93 - if self.static_input_scales:
94 - raise NotImplementedError(
95 - "QuarkOCP_MX with static input scales is currently not "
96 - "implemented. Please open an issue."
97 - )
98 -
73 if not current_platform.supports_mx():
74 logger.warning_once(
75 "The current platform does not support native MXFP4/MXFP6 "
vllm/model_executor/layers/quantization/quark/schemes/quark_scheme.py
+19
@@ -5,6 +5,8 @@ from abc import ABC, abstractmethod
5
6 import torch
7
8 +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
9 +
10 __all__ = ["QuarkScheme"]
11
12
@@ -14,6 +16,23 @@ class QuarkScheme(ABC):
16 of different quantization schemes supported by Quark.
17 """
18
19 + supported_activation_quant_keys: list[QuantKey | None] = []
20 + supported_weight_quant_keys: list[QuantKey] = []
21 +
22 + def __init__(
23 + self,
24 + weight_quant_key: QuantKey,
25 + activation_quant_key: QuantKey | None,
26 + ):
27 + if activation_quant_key not in self.supported_activation_quant_keys:
28 + raise ValueError(
29 + f"Unsupported activation quant key: {activation_quant_key}"
30 + )
31 + if weight_quant_key not in self.supported_weight_quant_keys:
32 + raise ValueError(f"Unsupported weight quant key: {weight_quant_key}")
33 + self.activation_quant_key = activation_quant_key
34 + self.weight_quant_key = weight_quant_key
35 +
36 @classmethod
37 @abstractmethod
38 def get_min_capability(cls) -> int:
vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py
+27 -38
@@ -3,7 +3,6 @@
3
4 from collections.abc import Callable
5 from fractions import Fraction
6 -from typing import Any
6
7 import torch
8 import torch.nn.functional as F
@@ -11,7 +10,10 @@ import torch.nn.functional as F
10 from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
11 from vllm.logger import init_logger
12 from vllm.model_executor.layers.quantization.utils.quant_utils import (
13 + QuantKey,
14 get_fp8_min_max,
15 + kFp8StaticTensorSym,
16 + kMxfp4Static,
17 )
18 from vllm.model_executor.parameter import (
19 GroupQuantScaleParameter,
@@ -38,39 +40,28 @@ class QuarkW4A8_MXFP4_FP8(QuarkScheme):
40 Uses the AITER Triton kernel and falls back to emulation if AITER not available.
41 """
42
43 + supported_activation_quant_keys = [kFp8StaticTensorSym]
44 + supported_weight_quant_keys = [kMxfp4Static]
45 +
46 def __init__(
47 self,
43 - weight_quant_spec: dict[str, Any],
44 - input_quant_spec: dict[str, Any],
48 + activation_quant_key: QuantKey | None,
49 ):
50 + super().__init__(kMxfp4Static, activation_quant_key)
51 self.out_dtype = None
52
53 self.weight_dtype = "mxfp4"
54 self.packed_factor: Fraction = Fraction(2, 1) # 2 FP4 values per byte
55 self.weight_block_size = OCP_MX_BLOCK_SIZE
56
52 - self.is_static_input_scheme = not input_quant_spec.get("is_dynamic")
53 - self.input_qscheme = input_quant_spec.get("qscheme") # "per_tensor"
54 -
57 self.fp8_min, self.fp8_max = get_fp8_min_max()
58 self.fp8_dtype = current_platform.fp8_dtype()
59
58 - if not self.is_static_input_scheme:
59 - raise NotImplementedError(
60 - "Dynamic FP8 activation quantization is not yet supported "
61 - "for W4A8. The current implementation expects static per-tensor "
62 - "FP8 scales stored in the checkpoint."
63 - )
64 -
60 kernel_supported_gpu = False
61 if current_platform.is_rocm():
62 kernel_supported_gpu = current_platform.supports_mx()
63
69 - self.use_aiter_kernel = (
70 - is_aiter_found_and_supported()
71 - and self.is_static_input_scheme
72 - and kernel_supported_gpu
73 - )
64 + self.use_aiter_kernel = is_aiter_found_and_supported() and kernel_supported_gpu
65
66 if not self.use_aiter_kernel:
67 logger.warning_once(
@@ -128,17 +119,16 @@ class QuarkW4A8_MXFP4_FP8(QuarkScheme):
119 layer.register_parameter("weight_scale", weight_scale)
120
121 # INPUT SCALE (FP8 per-tensor static scale)
131 - if self.is_static_input_scheme:
132 - input_scale = PerTensorScaleParameter(
133 - data=torch.empty(
134 - len(output_partition_sizes),
135 - dtype=torch.float32,
136 - ),
137 - weight_loader=weight_loader,
138 - )
139 - # Initialize to avoid NaN
140 - input_scale[:] = torch.finfo(torch.float32).min
141 - layer.register_parameter("input_scale", input_scale)
122 + input_scale = PerTensorScaleParameter(
123 + data=torch.empty(
124 + len(output_partition_sizes),
125 + dtype=torch.float32,
126 + ),
127 + weight_loader=weight_loader,
128 + )
129 + # Initialize to avoid NaN
130 + input_scale[:] = torch.finfo(torch.float32).min
131 + layer.register_parameter("input_scale", input_scale)
132
133 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
134 # Ensuring weights & scales are non-trainable
@@ -147,16 +137,15 @@ class QuarkW4A8_MXFP4_FP8(QuarkScheme):
137 layer.weight_scale.data, requires_grad=False
138 )
139
150 - if self.is_static_input_scheme:
151 - input_scale = layer.input_scale.data
152 - # For fused modules (QKV), take the max scale
153 - if input_scale.numel() != 1:
154 - input_scale = input_scale.max()
140 + input_scale = layer.input_scale.data
141 + # For fused modules (QKV), take the max scale
142 + if input_scale.numel() != 1:
143 + input_scale = input_scale.max()
144
156 - layer.input_scale = torch.nn.Parameter(
157 - torch.tensor(input_scale, dtype=torch.float32),
158 - requires_grad=False,
159 - )
145 + layer.input_scale = torch.nn.Parameter(
146 + torch.tensor(input_scale, dtype=torch.float32),
147 + requires_grad=False,
148 + )
149
150 def apply_weights(
151 self,
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
+41 -44
@@ -2,7 +2,6 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable
5 -from typing import Any, cast
5
6 import torch
7 from torch.nn import Parameter
@@ -18,8 +17,11 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
17 )
18 from vllm.model_executor.layers.quantization.utils.quant_utils import (
19 GroupShape,
20 + QuantKey,
21 kFp8Dynamic128Sym,
22 + kFp8DynamicTensorSym,
23 kFp8DynamicTokenSym,
24 + kFp8Static128BlockE8M0Sym,
25 kFp8Static128BlockSym,
26 kFp8StaticChannelSym,
27 kFp8StaticTensorSym,
@@ -42,36 +44,26 @@ logger = init_logger(__name__)
44
45
46 class QuarkW8A8Fp8(QuarkScheme):
47 + supported_activation_quant_keys: list[QuantKey | None] = [
48 + kFp8DynamicTensorSym,
49 + kFp8DynamicTokenSym,
50 + kFp8StaticTensorSym,
51 + ]
52 + supported_weight_quant_keys: list[QuantKey] = [
53 + kFp8StaticChannelSym,
54 + kFp8StaticTensorSym,
55 + ]
56 +
57 def __init__(
46 - self, weight_config: dict[str, Any], input_config: dict[str, Any] | None
58 + self,
59 + weight_quant_key: QuantKey,
60 + activation_quant_key: QuantKey | None,
61 ):
48 - self.weight_qscheme = cast(str, weight_config.get("qscheme"))
49 - self.is_static_input_scheme: bool = False
50 - self.input_qscheme: str | None = None
51 - if input_config is not None:
52 - self.is_static_input_scheme = not cast(bool, input_config.get("is_dynamic"))
53 - self.input_qscheme = cast(str, input_config.get("qscheme"))
54 -
55 - if self.weight_qscheme == "per_block":
56 - raise ValueError(
57 - "Quark W8A8 FP8 per-block weight quantization should use "
58 - "QuarkW8A8Fp8PerBlock."
59 - )
60 - per_token_activation = (
61 - not self.is_static_input_scheme and self.input_qscheme == "per_channel"
62 - )
63 - per_channel_weight = self.weight_qscheme == "per_channel"
64 -
65 - self.activation_quant_key = (
66 - kFp8DynamicTokenSym if per_token_activation else kFp8StaticTensorSym
67 - )
68 - # A per-output-channel weight scale is one fp32 value per weight row
69 - # (length N). Tag it as ``GroupShape.PER_CHANNEL`` to match the
70 - # canonical compressed-tensors CHANNEL strategy, so kernel selection
71 - # (e.g. AITER's pre-shuffled FP8 GEMM) treats it uniformly.
72 - self.weight_quant_key = (
73 - kFp8StaticChannelSym if per_channel_weight else kFp8StaticTensorSym
62 + super().__init__(weight_quant_key, activation_quant_key)
63 + self.weight_qscheme = (
64 + "per_channel" if weight_quant_key == kFp8StaticChannelSym else "per_tensor"
65 )
66 + self.is_static_input_scheme = activation_quant_key == kFp8StaticTensorSym
67 self.out_dtype = torch.get_default_dtype()
68 self.input_dtype = get_current_vllm_config().model_config.dtype
69
@@ -122,6 +114,8 @@ class QuarkW8A8Fp8(QuarkScheme):
114 layer.input_scale = Parameter(input_scale, requires_grad=False)
115 else:
116 weight_scale = layer.weight_scale.data
117 +
118 + assert self.activation_quant_key is not None
119 if self.activation_quant_key.scale.group_shape == GroupShape.PER_TOKEN:
120 weight_scale = weight_scale.view(-1, 1)
121 layer.weight = Parameter(weight.t(), requires_grad=False)
@@ -191,6 +185,7 @@ class QuarkW8A8Fp8(QuarkScheme):
185 input_scale[:] = torch.finfo(torch.float32).min
186 layer.register_parameter("input_scale", input_scale)
187
188 + assert self.activation_quant_key is not None
189 self.fp8_linear = init_fp8_linear_kernel(
190 activation_quant_key=self.activation_quant_key,
191 weight_quant_key=self.weight_quant_key,
@@ -210,22 +205,23 @@ class QuarkW8A8Fp8(QuarkScheme):
205
206
207 class QuarkW8A8Fp8PerBlock(QuarkScheme):
208 + supported_activation_quant_keys = [kFp8Dynamic128Sym]
209 + supported_weight_quant_keys = [
210 + kFp8Static128BlockSym,
211 + kFp8Static128BlockE8M0Sym,
212 + ]
213 +
214 def __init__(
214 - self, weight_config: dict[str, Any], input_config: dict[str, Any] | None
215 + self,
216 + weight_quant_key: QuantKey,
217 + activation_quant_key: QuantKey | None,
218 ):
216 - self.weight_qscheme = cast(str, weight_config.get("qscheme"))
217 - if self.weight_qscheme != "per_block":
218 - raise ValueError("QuarkW8A8Fp8PerBlock requires per-block weights.")
219 - block_size = weight_config.get("block_size")
220 - if not block_size:
221 - raise ValueError(
222 - "Quark W8A8 FP8 per-block weight quantization requires "
223 - "`block_size` in the weight quantization config."
224 - )
225 - self.weight_config = weight_config
226 - self.weight_block_size = list(block_size)
227 - self.activation_quant_key = kFp8Dynamic128Sym
228 - self.weight_quant_key = kFp8Static128BlockSym
219 + super().__init__(weight_quant_key, activation_quant_key)
220 + self.weight_scale_dtype = weight_quant_key.scale.dtype
221 + if self.weight_quant_key == kFp8Static128BlockE8M0Sym:
222 + # TODO: oracle needs to support kFp8Static128BlockE8M0Sym.
223 + self.weight_quant_key = kFp8Static128BlockSym
224 + self.weight_block_size = list(weight_quant_key.scale.group_shape)
225 self.out_dtype = torch.get_default_dtype()
226 self.input_dtype = get_current_vllm_config().model_config.dtype
227
@@ -272,8 +268,8 @@ class QuarkW8A8Fp8PerBlock(QuarkScheme):
268 layer.register_parameter("weight", weight)
269
270 scale_dtype = (
275 - torch.float8_e8m0fnu
276 - if self.weight_config.get("scale_type") == "float8_e8m0fnu"
271 + self.weight_scale_dtype
272 + if self.weight_scale_dtype != torch.float32
273 else None
274 )
275 weight_scale = create_fp8_scale_parameter(
@@ -287,6 +283,7 @@ class QuarkW8A8Fp8PerBlock(QuarkScheme):
283 # DeepSeek V4 weight mappers route checkpoint ".scale" tensors here.
284 layer.register_parameter("weight_scale_inv", weight_scale)
285
286 + assert self.activation_quant_key is not None
287 self.fp8_linear = init_fp8_linear_kernel(
288 activation_quant_key=self.activation_quant_key,
289 weight_quant_key=self.weight_quant_key,
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py
+33 -6
@@ -10,6 +10,16 @@ from vllm.model_executor.kernels.linear import (
10 init_int8_linear_kernel,
11 )
12 from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
13 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
14 + QuantKey,
15 + kInt8DynamicTensorAsym,
16 + kInt8DynamicTensorSym,
17 + kInt8DynamicTokenAsym,
18 + kInt8DynamicTokenSym,
19 + kInt8StaticChannelSym,
20 + kInt8StaticTensorAsym,
21 + kInt8StaticTensorSym,
22 +)
23 from vllm.model_executor.parameter import (
24 BasevLLMParameter,
25 ChannelQuantScaleParameter,
@@ -21,15 +31,32 @@ logger = init_logger(__name__)
31
32
33 class QuarkW8A8Int8(QuarkScheme):
34 + supported_activation_quant_keys = [
35 + kInt8StaticTensorSym,
36 + kInt8StaticTensorAsym,
37 + kInt8DynamicTensorSym,
38 + kInt8DynamicTensorAsym,
39 + kInt8DynamicTokenSym,
40 + kInt8DynamicTokenAsym,
41 + ]
42 + supported_weight_quant_keys = [
43 + kInt8StaticChannelSym,
44 + kInt8StaticTensorSym,
45 + ]
46 +
47 def __init__(
48 self,
26 - qscheme: str,
27 - is_static_input_scheme: bool | None,
28 - input_symmetric: bool | None,
49 + weight_quant_key: QuantKey,
50 + activation_quant_key: QuantKey | None,
51 ):
30 - self.qscheme = qscheme
31 - self.is_static_input_scheme = is_static_input_scheme
32 - self.input_symmetric = input_symmetric
52 + super().__init__(weight_quant_key, activation_quant_key)
53 + self.qscheme = (
54 + "per_channel" if weight_quant_key == kInt8StaticChannelSym else "per_tensor"
55 + )
56 +
57 + assert activation_quant_key is not None
58 + self.is_static_input_scheme = activation_quant_key.scale.static
59 + self.input_symmetric = activation_quant_key.symmetric
60
61 @classmethod
62 def get_min_capability(cls) -> int:
vllm/model_executor/layers/quantization/quark/utils.py
+2
@@ -8,6 +8,8 @@ from typing import Any
8 import regex as re
9 import torch
10
11 +QuarkQTensorHint = dict[str, Any] | list[dict[str, Any]] | None
12 +
13
14 def deep_compare(dict1: Any, dict2: Any) -> bool:
15 if type(dict1) is not type(dict2):
vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py
+21
@@ -3,11 +3,32 @@
3 from enum import Enum
4
5 from vllm.logger import init_logger
6 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
7 + QuantKey,
8 + kMxfp4Dynamic,
9 + kMxfp4Static,
10 + kMxfp6E2M3Dynamic,
11 + kMxfp6E2M3Static,
12 + kMxfp6E3M2Dynamic,
13 + kMxfp6E3M2Static,
14 +)
15
16 logger = init_logger(__name__)
17
18 OCP_MX_BLOCK_SIZE = 32
19
20 +_WEIGHT_QUANT_KEY_MAP: dict[str, QuantKey] = {
21 + "mxfp4": kMxfp4Static,
22 + "mxfp6_e3m2": kMxfp6E3M2Static,
23 + "mxfp6_e2m3": kMxfp6E2M3Static,
24 +}
25 +
26 +_ACTIVATION_QUANT_KEY_MAP: dict[str, QuantKey] = {
27 + "mxfp4": kMxfp4Dynamic,
28 + "mxfp6_e3m2": kMxfp6E3M2Dynamic,
29 + "mxfp6_e2m3": kMxfp6E2M3Dynamic,
30 +}
31 +
32 OCP_MX_DTYPES = {
33 "mxfp4",
34 "mxfp6_e3m2",
vllm/model_executor/layers/quantization/utils/quant_utils.py
+8
@@ -222,6 +222,11 @@ kFp8Dynamic128Sym = QuantKey(FP8_DTYPE, kDynamic128Scale, symmetric=True)
222
223 kStatic128BlockScale = ScaleDesc(torch.float32, True, GroupShape(128, 128))
224 kFp8Static128BlockSym = QuantKey(FP8_DTYPE, kStatic128BlockScale, symmetric=True)
225 +kFp8Static128BlockE8M0Sym = QuantKey(
226 + FP8_DTYPE,
227 + ScaleDesc(torch.float8_e8m0fnu, True, GroupShape(128, 128)),
228 + symmetric=True,
229 +)
230
231 kMxfp8StaticScale = ScaleDesc(torch.uint8, True, GroupShape(1, 32))
232 kMxfp8Static = QuantKey(FP8_DTYPE, kMxfp8StaticScale, symmetric=True)
@@ -284,6 +289,9 @@ kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True
289 kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True)
290 kInt8StaticTensorSym = QuantKey(torch.int8, kStaticTensorScale, symmetric=True)
291 kInt8DynamicTensorSym = QuantKey(torch.int8, kDynamicTensorScale, symmetric=True)
292 +kInt8DynamicTokenAsym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=False)
293 +kInt8StaticTensorAsym = QuantKey(torch.int8, kStaticTensorScale, symmetric=False)
294 +kInt8DynamicTensorAsym = QuantKey(torch.int8, kDynamicTensorScale, symmetric=False)
295
296 # INT4 W4A8 quantization keys
297