@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 05576df85c
Fork 0

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

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

Synced Nov 13, 12:00 UTC

[ROCm][Quantization] extend AMD Quark to support mixed-precision quantized model (#24239)

Signed-off-by: xuebwang-amd <xuebwang@amd.com> Co-authored-by: fxmarty-amd <felmarty@amd.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

xuebwang-amd committed Nov 12, 2025 at 01:05 UTC 05576df85c5274ee3045d90b0779d4adeecc09b9
3 files changed +127 -8
docs/features/quantization/quark.md
+33 -1
@@ -281,4 +281,36 @@ python quantize_quark.py --model_dir Qwen/Qwen1.5-MoE-A2.7B-Chat \
281 --group_size 32
282 ```
283
284 -The current integration supports [all combination of FP4, FP6_E3M2, FP6_E2M3](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py) used for either weights or activations. Eventually, some target hardware support mixed precision GEMM, as AMD Instinct MI350/MI355, for example using FP6 for activations and FP4 for weights.
284 +The current integration supports [all combination of FP4, FP6_E3M2, FP6_E2M3](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py) used for either weights or activations.
285 +
286 +## Using Quark Quantized layerwise Auto Mixed Precision (AMP) Models
287 +
288 +vLLM also supports loading layerwise mixed precision model quantized using AMD Quark. Currently, mixed scheme of {MXFP4, FP8} is supported, where FP8 here denotes for FP8 per-tensor scheme. More mixed precision schemes are planned to be supported in a near future, including
289 +
290 +- Unquantized Linear and/or MoE layer(s) as an option for each layer, i.e., mixed of {MXFP4, FP8, BF16/FP16}
291 +- MXFP6 quantization extension, i.e., {MXFP4, MXFP6, FP8, BF16/FP16}
292 +
293 +Although one can maximize serving throughput using the lowest precision supported on a given device (e.g. MXFP4 for AMD Instinct MI355, FP8 for AMD Instinct MI300), these aggressive schemes can be detrimental to accuracy recovering from quantization on target tasks. Mixed precision allows to strike a balance between maximizing accuracy and throughput.
294 +
295 +There are two steps to generate and deploy a mixed precision model quantized with AMD Quark, as shown below.
296 +
297 +### 1. Quantize a model using mixed precision in AMD Quark
298 +
299 +Firstly, the layerwise mixed-precision configuration for a given LLM model is searched and then quantized using AMD Quark. We will provide a detailed tutorial with Quark APIs later.
300 +
301 +As examples, we provide some ready-to-use quantized mixed precision model to show the usage in vLLM and the accuracy benifits. They are:
302 +
303 +- amd/Llama-2-70b-chat-hf-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8
304 +- amd/Mixtral-8x7B-Instruct-v0.1-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8
305 +- amd/Qwen3-8B-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8
306 +
307 +### 2. inference the quantized mixed precision model in vLLM
308 +
309 +Models quantized with AMD Quark using mixed precision can natively be reload in vLLM, and e.g. evaluated using lm-evaluation-harness as follow:
310 +
311 +```bash
312 +lm_eval --model vllm \
313 + --model_args pretrained=amd/Llama-2-70b-chat-hf-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8,tensor_parallel_size=4,dtype=auto,gpu_memory_utilization=0.8,trust_remote_code=False \
314 + --tasks mmlu \
315 + --batch_size auto
316 +```
tests/quantization/test_mixed_precision.py new
+69
@@ -0,0 +1,69 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Test quark-quantized {MXFP4, FP8} mixed precision models.
4 +
5 +Run `pytest tests/quantization/test_mixed_precision.py`.
6 +
7 +"""
8 +
9 +import importlib
10 +import importlib.metadata
11 +from dataclasses import dataclass
12 +
13 +import lm_eval
14 +import pytest
15 +from packaging import version
16 +
17 +QUARK_MXFP4_AVAILABLE = importlib.util.find_spec("quark") is not None and version.parse(
18 + importlib.metadata.version("amd-quark")
19 +) >= version.parse("0.8.99")
20 +
21 +
22 +@dataclass
23 +class ModelCase:
24 + model_id: str
25 + tp: int
26 +
27 +
28 +@dataclass
29 +class EvaluationConfig:
30 + model_name: str
31 +
32 + def get_model_args(self) -> str:
33 + return (
34 + f"pretrained={self.model_name},"
35 + "tensor_parallel_size=4,dtype=auto,gpu_memory_utilization=0.8,trust_remote_code=False"
36 + )
37 +
38 +
39 +TEST_CONFIGS = {
40 + # Mixed-precision (AMP) model
41 + # - Demonstrates end-to-end pipeline functionality
42 + "amd/Qwen3-8B-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8": {"arc_challenge": 0.52, "mmlu": 0.72},
43 + # Non-mixed-precision (PTQ) model
44 + # - Reference for pipeline compatibility verification -> No conflicts or breakings
45 + "amd/Llama-2-70b-chat-hf-FP8-MLPerf-fp8_attn_quark_format": {
46 + "arc_challenge": 0.53,
47 + "mmlu": 0.61,
48 + },
49 +}
50 +
51 +
52 +@pytest.mark.parametrize("model_name, accuracy_numbers", TEST_CONFIGS.items())
53 +@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
54 +def test_mixed_precision_model_accuracies(model_name: str, accuracy_numbers: dict):
55 + results = lm_eval.simple_evaluate(
56 + model="vllm",
57 + model_args=EvaluationConfig(model_name).get_model_args(),
58 + tasks=list(accuracy_numbers.keys()),
59 + batch_size=8,
60 + )
61 +
62 + rtol = 0.05
63 +
64 + for task, expect_accuracy in accuracy_numbers.items():
65 + measured_accuracy = results["results"][task]["acc,none"]
66 + assert (
67 + measured_accuracy - rtol < expect_accuracy
68 + and measured_accuracy + rtol > expect_accuracy
69 + ), f"Expected: {expect_accuracy} | Measured: {measured_accuracy}"
vllm/model_executor/layers/quantization/quark/quark.py
+25 -7
@@ -114,7 +114,14 @@ class QuarkConfig(QuantizationConfig):
114 layer_quant_names = list(layer_quant_config.keys())
115 layer_quant_set = set(layer_quant_names)
116
117 - if not kv_cache_set.issubset(layer_quant_set):
117 + if not (
118 + kv_cache_set.issubset(layer_quant_set)
119 + or any(
120 + fnmatch.fnmatchcase(layer_quant, pat)
121 + for layer_quant in list(layer_quant_set)
122 + for pat in list(kv_cache_set)
123 + )
124 + ):
125 raise ValueError(
126 "The Quark quantized model has the "
127 "kv_cache_group parameter setting, "
@@ -124,10 +131,15 @@ class QuarkConfig(QuantizationConfig):
131 )
132
133 q_configs = [
127 - cast(dict[str, Any], layer_quant_config.get(name))
128 - for name in kv_cache_group
134 + quant_cfg
135 + for name, quant_cfg in layer_quant_config.items()
136 + if any(fnmatch.fnmatchcase(name, pattern) for pattern in kv_cache_group)
137 ]
130 - if not all(deep_compare(q_config, q_configs[0]) for q_config in q_configs):
138 +
139 + if not all(
140 + deep_compare(q_config["output_tensors"], q_configs[0]["output_tensors"])
141 + for q_config in q_configs
142 + ):
143 raise ValueError(
144 "The quantization method used for kv_cache should "
145 "be the same, but the quantization method for the "
@@ -312,9 +324,15 @@ class QuarkConfig(QuantizationConfig):
324 layer_quant_config = cast(
325 dict[str, Any], self.quant_config.get("layer_quant_config")
326 )
315 - for name_pattern in layer_quant_config:
316 - if fnmatch.fnmatch(layer_name, name_pattern):
317 - return layer_quant_config[name_pattern]
327 +
328 + def _matches_pattern(layer_name, pattern):
329 + if "*" not in pattern:
330 + return layer_name in pattern
331 + return fnmatch.fnmatch(layer_name, pattern)
332 +
333 + for name_pattern, config in layer_quant_config.items():
334 + if _matches_pattern(layer_name, name_pattern):
335 + return config
336
337 layer_type = cast(str, type(module))
338 layer_type_quant_config = cast(