A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 16:30 UTC
[ROCm] Add dynamic mxfp4 quantization for DeepSeek V2 projection layers (#34157)
Signed-off-by: Doug Lehr <douglehr@amd.com> Signed-off-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Doug Lehr <douglehr@amd.com> Co-authored-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com> Co-authored-by: Gregory Shtrasberg <156009573+gshtras@users.noreply.github.com>
Douglas Lehr committed
Feb 26, 2026 at 10:00 UTC
ec8ab9d254d3b2e6b919a55277da599a7b9ab146
2 files changed
+106
-39
vllm/model_executor/layers/quantization/quark/quark.py
+43
-5
@@ -35,6 +35,7 @@ from vllm.model_executor.layers.quantization.quark.utils import (
35
)
36
from vllm.model_executor.models.utils import WeightsMapper
37
from vllm.platforms import current_platform
38
+from vllm.transformers_utils.config import get_config
39
40
if TYPE_CHECKING:
41
from vllm.model_executor.models.utils import WeightsMapper
@@ -59,6 +60,22 @@ class QuarkConfig(QuantizationConfig):
60
self.kv_cache_group = kv_cache_group
61
self.kv_cache_config = kv_cache_config
62
self.pack_method = pack_method
63
+ self.dynamic_mxfp4_quant = False
64
+
65
+ def maybe_update_config(self, model_name: str, revision: str | None = None):
66
+ self.hf_config = get_config(
67
+ model=model_name,
68
+ trust_remote_code=False, # or get from model_config if available
69
+ revision=revision,
70
+ config_format="auto",
71
+ )
72
+
73
+ quant_config = getattr(self.hf_config, "quantization_config", None)
74
+ if quant_config is not None:
75
+ quant_dtype = quant_config["global_quant_config"]["weight"]["dtype"]
76
+ model_type = self.hf_config.model_type
77
+ if quant_dtype == "fp4" and model_type == "deepseek_v3":
78
+ self.dynamic_mxfp4_quant = True
79
80
def get_linear_method(self) -> "QuarkLinearMethod":
81
return QuarkLinearMethod(self)
@@ -108,7 +125,20 @@ class QuarkConfig(QuantizationConfig):
125
if should_ignore_layer(
126
prefix, ignore=exclude_layers, fused_mapping=self.packed_modules_mapping
127
):
111
- return UnquantizedLinearMethod()
128
+ if (
129
+ "self_attn" not in prefix # only quantize attention projections
130
+ or not getattr(self, "dynamic_mxfp4_quant", False)
131
+ or not isinstance(layer, LinearBase) # Ignore other methods
132
+ ):
133
+ return UnquantizedLinearMethod()
134
+
135
+ scheme = self.get_scheme(
136
+ layer=layer,
137
+ layer_name=prefix,
138
+ dynamic_mxfp4_quant=True,
139
+ )
140
+ layer.scheme = scheme
141
+ return QuarkLinearMethod(self)
142
if isinstance(layer, LinearBase):
143
scheme = self.get_scheme(layer=layer, layer_name=prefix)
144
layer.scheme = scheme
@@ -450,7 +480,9 @@ class QuarkConfig(QuantizationConfig):
480
)
481
return global_quant_config
482
453
- def _get_scheme_from_config(self, config: dict[str, Any]) -> "QuarkScheme":
483
+ def _get_scheme_from_config(
484
+ self, config: dict[str, Any], dynamic_mxfp4_quant: bool = False
485
+ ) -> "QuarkScheme":
486
if config.get("output_tensors") or config.get("bias"):
487
raise NotImplementedError(
488
"Currently, Quark models with output_tensors "
@@ -473,7 +505,9 @@ class QuarkConfig(QuantizationConfig):
505
input_symmetric=input_config.get("symmetric"),
506
)
507
elif self._is_w_ocp_mx_a_x(weight_config, input_config):
476
- return QuarkOCP_MX(weight_config, input_config)
508
+ return QuarkOCP_MX(
509
+ weight_config, input_config, dynamic_mxfp4_quant=dynamic_mxfp4_quant
510
+ )
511
512
raise NotImplementedError(
513
"No quark compatible scheme was found. "
@@ -481,11 +515,15 @@ class QuarkConfig(QuantizationConfig):
515
f"Input config: {input_config}"
516
)
517
484
- def get_scheme(self, layer: torch.nn.Module, layer_name: str) -> "QuarkScheme":
518
+ def get_scheme(
519
+ self, layer: torch.nn.Module, layer_name: str, dynamic_mxfp4_quant: bool = False
520
+ ) -> "QuarkScheme":
521
layer_quant_config = self._find_matched_config(layer_name, layer)
522
523
# Find the quant_scheme
488
- scheme = self._get_scheme_from_config(layer_quant_config)
524
+ scheme = self._get_scheme_from_config(
525
+ layer_quant_config, dynamic_mxfp4_quant=dynamic_mxfp4_quant
526
+ )
527
# Raise error if device does not support the scheme
528
# (e.g. fp8 needs ada lovelace)
529
self._check_scheme_supported(scheme.get_min_capability())
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+63
-34
@@ -24,7 +24,12 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
24
OCP_MX_BLOCK_SIZE,
25
OCP_MX_Scheme,
26
)
27
-from vllm.model_executor.parameter import GroupQuantScaleParameter, PackedvLLMParameter
27
+from vllm.model_executor.parameter import (
28
+ GroupQuantScaleParameter,
29
+ ModelWeightParameter,
30
+ PackedvLLMParameter,
31
+)
32
+from vllm.model_executor.utils import set_weight_attrs
33
from vllm.platforms import current_platform
34
35
from .quark_scheme import QuarkScheme
@@ -169,13 +174,16 @@ except (ImportError, AttributeError, RuntimeError):
174
175
class QuarkOCP_MX(QuarkScheme):
176
def __init__(
172
- self, weight_quant_spec: dict[str, Any], input_quant_spec: dict[str, Any]
177
+ self,
178
+ weight_quant_spec: dict[str, Any],
179
+ input_quant_spec: dict[str, Any],
180
+ dynamic_mxfp4_quant: bool = False,
181
):
182
self.out_dtype = torch.get_default_dtype()
183
self.qscheme = "per_group"
184
self.weight_quant_spec = weight_quant_spec
185
self.input_quant_spec = input_quant_spec
178
-
186
+ self.dynamic_mxfp4_quant = dynamic_mxfp4_quant
187
self.weight_dtype = weight_quant_spec["dtype"].replace("fp", "mxfp")
188
self.input_dtype = input_quant_spec["dtype"].replace("fp", "mxfp")
189
@@ -269,7 +277,13 @@ class QuarkOCP_MX(QuarkScheme):
277
layer.weight_scale.data, requires_grad=False
278
)
279
else:
272
- if self.rocm_use_aiter_fp4_asm_gemm:
280
+ if self.dynamic_mxfp4_quant:
281
+ w_q, w_s = dynamic_mxfp4_quant(layer.weight)
282
+ layer.weight_scale = torch.nn.Parameter(
283
+ w_s.T.contiguous(), requires_grad=False
284
+ )
285
+ layer.weight = torch.nn.Parameter(w_q, requires_grad=False)
286
+ elif self.rocm_use_aiter_fp4_asm_gemm:
287
# shuffle weight scale
288
weight_scale_shuffle = layer.weight_scale.data
289
sm, sn = weight_scale_shuffle.shape
@@ -302,36 +316,51 @@ class QuarkOCP_MX(QuarkScheme):
316
weight_loader: Callable,
317
**kwargs,
318
):
305
- output_size_per_partition = sum(output_partition_sizes)
306
- layer.logical_widths = output_partition_sizes
307
-
308
- # WEIGHT
309
- weight = PackedvLLMParameter(
310
- data=torch.empty(
311
- output_size_per_partition,
312
- self.get_packed_dim(input_size_per_partition, self.weight_dtype),
313
- dtype=torch.uint8,
314
- ),
315
- input_dim=1,
316
- output_dim=0,
317
- packed_dim=1,
318
- packed_factor=self.packed_factor,
319
- weight_loader=weight_loader,
320
- )
321
- layer.register_parameter("weight", weight)
322
-
323
- # WEIGHT SCALE
324
- weight_scale = GroupQuantScaleParameter(
325
- data=torch.empty(
326
- output_size_per_partition,
327
- input_size_per_partition // OCP_MX_BLOCK_SIZE,
328
- dtype=torch.uint8,
329
- ),
330
- input_dim=1,
331
- output_dim=0,
332
- weight_loader=weight_loader,
333
- )
334
- layer.register_parameter("weight_scale", weight_scale)
319
+ if self.dynamic_mxfp4_quant:
320
+ weight = ModelWeightParameter(
321
+ data=torch.empty(
322
+ sum(output_partition_sizes),
323
+ input_size_per_partition,
324
+ dtype=params_dtype,
325
+ ),
326
+ input_dim=1,
327
+ output_dim=0,
328
+ weight_loader=weight_loader,
329
+ )
330
+
331
+ layer.register_parameter("weight", weight)
332
+ set_weight_attrs(weight, kwargs)
333
+ else:
334
+ output_size_per_partition = sum(output_partition_sizes)
335
+ layer.logical_widths = output_partition_sizes
336
+
337
+ # WEIGHT
338
+ weight = PackedvLLMParameter(
339
+ data=torch.empty(
340
+ output_size_per_partition,
341
+ self.get_packed_dim(input_size_per_partition, self.weight_dtype),
342
+ dtype=torch.uint8,
343
+ ),
344
+ input_dim=1,
345
+ output_dim=0,
346
+ packed_dim=1,
347
+ packed_factor=self.packed_factor,
348
+ weight_loader=weight_loader,
349
+ )
350
+ layer.register_parameter("weight", weight)
351
+
352
+ # WEIGHT SCALE
353
+ weight_scale = GroupQuantScaleParameter(
354
+ data=torch.empty(
355
+ output_size_per_partition,
356
+ input_size_per_partition // OCP_MX_BLOCK_SIZE,
357
+ dtype=torch.uint8,
358
+ ),
359
+ input_dim=1,
360
+ output_dim=0,
361
+ weight_loader=weight_loader,
362
+ )
363
+ layer.register_parameter("weight_scale", weight_scale)
364
365
def apply_weights(
366
self,