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 6, 08:15 UTC
[ROCm][Quantization] make quark ocp mx dtype parser robust for weight-only quantization (#36232)
Signed-off-by: xuebwang-amd <xuebwang@amd.com>
xuebwang-amd committed
Mar 20, 2026 at 23:10 UTC
44eea10f68461852ff4467cd5b7924a46777c8c9
2 files changed
+19
-5
vllm/model_executor/layers/quantization/quark/quark_moe.py
+2
-1
@@ -92,7 +92,8 @@ class QuarkMoEMethod(FusedMoEMethodBase):
92
rocm_aiter_ops.is_fused_moe_enabled()
93
)
94
if (
95
- input_config.get("dtype") == "fp8_e4m3"
95
+ input_config is not None
96
+ and input_config.get("dtype") == "fp8_e4m3"
97
and not input_config.get("is_dynamic")
98
and not emulate
99
):
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+17
-4
@@ -176,7 +176,7 @@ class QuarkOCP_MX(QuarkScheme):
176
def __init__(
177
self,
178
weight_quant_spec: dict[str, Any],
179
- input_quant_spec: dict[str, Any],
179
+ input_quant_spec: dict[str, Any] | None,
180
dynamic_mxfp4_quant: bool = False,
181
):
182
self.out_dtype = torch.get_default_dtype()
@@ -185,7 +185,13 @@ class QuarkOCP_MX(QuarkScheme):
185
self.input_quant_spec = input_quant_spec
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")
188
+ self.input_dtype: str | None = None
189
+ if input_quant_spec is not None:
190
+ input_quant = input_quant_spec["dtype"]
191
+ if input_quant == "fp8_e4m3":
192
+ self.input_dtype = "fp8"
193
+ else:
194
+ self.input_dtype = input_quant.replace("fp", "mxfp")
195
196
self.ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype(
197
self.input_dtype, self.weight_dtype
@@ -200,14 +206,21 @@ class QuarkOCP_MX(QuarkScheme):
206
dequant_mxfp6, quant_dtype=self.weight_dtype.replace("mx", "")
207
)
208
203
- if self.input_dtype == "mxfp4":
209
+ if self.input_dtype is None:
210
+ self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = (
211
+ lambda x: x
212
+ ) # no input Q/DQ for weight-only
213
+ elif self.input_dtype == "mxfp4":
214
self.quant_dequant_func = quant_dequant_mxfp4
215
else:
216
self.quant_dequant_func = partial(
217
quant_dequant_mxfp6, quant_dtype=self.input_dtype.replace("mx", "")
218
)
219
210
- self.static_input_scales = not input_quant_spec.get("is_dynamic")
220
+ if input_quant_spec is None:
221
+ self.static_input_scales = False
222
+ else:
223
+ self.static_input_scales = not input_quant_spec.get("is_dynamic")
224
225
if self.static_input_scales:
226
raise NotImplementedError(