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 4, 15:00 UTC
[ROCm][CI] Add missing quantization methods and fix online quant test failures (#39801)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Apr 27, 2026 at 15:08 UTC
5e2c37facde9f5edd68a7de8293107089e9887d8
9 files changed
+50
-23
tests/quantization/test_configs.py
+2
-2
@@ -54,13 +54,13 @@ MODEL_ARG_EXPTYPES = [
54
(
55
"TheBloke/OpenHermes-2.5-Mistral-7B-AWQ",
56
None,
57
- "awq_marlin" if current_platform.is_cuda() else "awq",
57
+ "awq_marlin" if current_platform.is_cuda_alike() else "awq",
58
),
59
("TheBloke/OpenHermes-2.5-Mistral-7B-AWQ", "awq", "awq"),
60
(
61
"TheBloke/OpenHermes-2.5-Mistral-7B-AWQ",
62
"marlin",
63
- "awq_marlin" if current_platform.is_cuda() else "ERROR",
63
+ "awq_marlin" if current_platform.is_cuda_alike() else "ERROR",
64
),
65
("TheBloke/OpenHermes-2.5-Mistral-7B-AWQ", "gptq", "ERROR"),
66
]
tests/quantization/test_mixed_precision.py
+1
-1
@@ -43,7 +43,7 @@ TEST_CONFIGS = {
43
"amd/Qwen3-8B-WMXFP4FP8-AMXFP4FP8-AMP-KVFP8": {"arc_challenge": 0.52, "mmlu": 0.72},
44
# Non-mixed-precision (PTQ) model
45
# - Reference for pipeline compatibility verification -> No conflicts or breakings
46
- "amd/Llama-2-70b-chat-hf-FP8-MLPerf-fp8_attn_quark_format": {
46
+ "amd/Llama-2-70b-chat-hf_FP8_MLPerf_V2": {
47
"arc_challenge": 0.53,
48
"mmlu": 0.61,
49
},
tests/quantization/test_turboquant.py
+4
-1
@@ -352,7 +352,10 @@ def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Te
352
gen = torch.Generator(device="cpu")
353
gen.manual_seed(seed)
354
G = torch.randn(d, d, generator=gen, device="cpu", dtype=torch.float32)
355
- Q, R = torch.linalg.qr(G)
355
+ # torch.linalg.qr on CPU requires LAPACK, which some torch wheels
356
+ # (ROCm) ship without. Run QR on accelerator instead
357
+ qr_device = "cuda" if torch.cuda.is_available() else "cpu"
358
+ Q, R = torch.linalg.qr(G.to(qr_device))
359
diag_sign = torch.sign(torch.diag(R))
360
diag_sign[diag_sign == 0] = 1.0
361
Q = Q * diag_sign.unsqueeze(0)
tests/quantization/utils.py
+4
-5
@@ -80,14 +80,13 @@ def _test_online_quant_peak_mem_impl(
80
print(f"GPU memory used after loading weights: {model_memory_gib} GiB")
81
print(f"Peak GPU memory usage while loading weights: {peak_memory_gib} GiB")
82
83
- # model specific, allenai/OLMoE-1B-7B-0125-Instruct fp8 online quant
84
- # uses 6.65 GiB for weight loading (bf16 checkpoint is ~12.89 GiB)
83
expected_model_memory_gib = 6.7
84
85
# for allenai/OLMoE-1B-7B-0125-Instruct the number we see today is 9.06
88
- # GiB, which is 1.36x above model_memory_gib. A slightly higher number is
89
- # expected as when we load and quantize weights in a streaming fashion we
90
- # need to have individual weights in bf16 + fp8 alive at the same time.
86
+ # GiB on CUDA, which is 1.36x above model_memory_gib. A slightly higher
87
+ # number is expected as when we load and quantize weights in a streaming
88
+ # fashion we need to have individual weights in bf16 + fp8 alive at the
89
+ # same time.
90
expected_peak_memory_gib = expected_model_memory_gib * 1.4
91
92
assert model_memory_gib < expected_model_memory_gib, (
vllm/model_executor/kernels/linear/scaled_mm/pytorch.py
+9
-1
@@ -144,13 +144,21 @@ class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
144
# For CUDA platform please validate if the torch._scaled_mm supports
145
# rowwise scaled GEMM before using it
146
147
+ # torch._scaled_mm rowwise requires scale_a = (m, 1), scale_b = (1, n).
148
+ # CompressedTensors stores weight_scale as (n, 1), so `.t()` yields (1, n).
149
+ # ModelOpt FP8_PER_CHANNEL_PER_TOKEN stores it as 1-D (n,); reshape to
150
+ # (1, n) so both paths satisfy the rowwise contract.
151
+ scale_b = Bs.view(1, -1) if Bs.dim() == 1 else Bs.t()
152
+ if As.dim() == 1:
153
+ As = As.view(-1, 1)
154
+
155
# Fused GEMM_DQ Rowwise GEMM
156
output = torch._scaled_mm(
157
A,
158
B,
159
out_dtype=out_dtype,
160
scale_a=As,
153
- scale_b=Bs.t(),
161
+ scale_b=scale_b,
162
bias=bias,
163
)
164
vllm/model_executor/layers/quantization/quark/quark_moe.py
+7
-1
@@ -1025,9 +1025,15 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1025
get_current_vllm_config().model_config.hf_config, "model_type", None
1026
)
1027
1028
+ # TODO(aiter): extend once rocm_aiter_fused_experts gains dispatch
1029
+ # for the other OCP MX schemes. Today its CK MoE kernel only has an
1030
+ # entry for `w_mxfp4` (w4a16); mixed schemes like `w_mxfp4_a_mxfp6_*`
1031
+ # fall through to QuantMethod.NO and raise "Unsupported kernel config
1032
+ # for moe heuristic dispatch".
1033
+ _AITER_NATIVE_OCP_MX_SCHEMES = ("w_mxfp4",)
1034
self.emulate = (
1035
not current_platform.supports_mx()
1030
- or not self.ocp_mx_scheme.startswith("w_mxfp4")
1036
+ or self.ocp_mx_scheme not in _AITER_NATIVE_OCP_MX_SCHEMES
1037
) and (
1038
self.mxfp4_backend is Mxfp4MoeBackend.NONE or not self.use_rocm_aiter_moe
1039
)
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+12
-8
@@ -376,11 +376,15 @@ class QuarkOCP_MX(QuarkScheme):
376
dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype)
377
qdq_x = self.quant_dequant_func(x)
378
return F.linear(qdq_x, dq_w, bias)
379
- else:
380
- return torch.ops.vllm.gemm_with_dynamic_quant(
381
- x,
382
- layer.weight,
383
- layer.weight_scale,
384
- self.rocm_use_aiter_fp4_asm_gemm,
385
- self.out_dtype,
386
- )
379
+ y = torch.ops.vllm.gemm_with_dynamic_quant(
380
+ x,
381
+ layer.weight,
382
+ layer.weight_scale,
383
+ self.rocm_use_aiter_fp4_asm_gemm,
384
+ self.out_dtype,
385
+ )
386
+ # gemm_with_dynamic_quant has no bias argument; add it here so the
387
+ # native path matches F.linear (e.g. qkv_proj with qkv_bias=True).
388
+ if bias is not None:
389
+ y = y + bias
390
+ return y
vllm/model_executor/model_loader/base_loader.py
+1
-1
@@ -65,7 +65,7 @@ class BaseModelLoader(ABC):
65
66
# Log peak GPU memory after loading weights. This is needed
67
# to have test coverage on peak memory for online quantization.
68
- if current_platform.is_cuda():
68
+ if current_platform.is_cuda_alike():
69
peak_memory = torch.accelerator.max_memory_allocated()
70
logger.debug_once(
71
"Peak GPU memory after loading weights: %s GiB",
vllm/platforms/rocm.py
+10
-3
@@ -414,10 +414,17 @@ class RocmPlatform(Platform):
414
"gguf",
415
"quark",
416
"mxfp4",
417
- "gpt_oss_mxfp4",
417
+ "mxfp8",
418
"torchao",
419
"bitsandbytes",
420
+ "modelopt",
421
"modelopt_fp4",
422
+ "modelopt_mxfp8",
423
+ "modelopt_mixed",
424
+ "fp8_per_tensor",
425
+ "fp8_per_block",
426
+ "online",
427
+ "gpt_oss_mxfp4",
428
]
429
430
@classmethod
@@ -785,9 +792,9 @@ class RocmPlatform(Platform):
792
def get_current_memory_usage(
793
cls, device: torch.types.Device | None = None
794
) -> float:
795
+ torch.cuda.empty_cache()
796
torch.cuda.reset_peak_memory_stats(device)
789
- free_mem, total_mem = torch.cuda.mem_get_info(device)
790
- return total_mem - free_mem
797
+ return torch.cuda.max_memory_allocated(device)
798
799
@classmethod
800
def get_device_communicator_cls(cls) -> str: