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, 11:15 UTC
[ROCm] Add env flags to disable dynamic MXFP4 quant and enable AITER tuned GEMMs for Attention Projection Layers (#39987)
Signed-off-by: Hemanth Acharya <heachary@amd.com>
Hemanth Acharya committed
Apr 30, 2026 at 04:37 UTC
6841f5dc77e9200a2fa45a4bf935b23bd843bf30
5 files changed
+28
-103
tests/quantization/test_quark_maybe_update_config.py
deleted
-63
@@ -1,63 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Unit tests for QuarkConfig.maybe_update_config.
4
-
5
-Fetches real HF configs (metadata only, no model weights) to verify
6
-that dynamic_mxfp4_quant is only enabled for DeepSeek-V3-family models.
7
-
8
-Run: pytest tests/quantization/test_quark_maybe_update_config.py -v
9
-"""
10
-
11
-import pytest
12
-from transformers import AutoConfig
13
-
14
-from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
15
-
16
-
17
-def _make_quark_config() -> QuarkConfig:
18
- """Create a minimal QuarkConfig for testing."""
19
- return QuarkConfig(quant_config={}, kv_cache_group=[], pack_method="reorder")
20
-
21
-
22
-# ---------------------------------------------------------------------------
23
-# Non-deepseek models must not flip dynamic_mxfp4_quant
24
-# ---------------------------------------------------------------------------
25
-@pytest.mark.parametrize(
26
- "model_name",
27
- ["amd/MiniMax-M2.1-MXFP4"],
28
-)
29
-def test_non_deepseek_model_stays_false(model_name: str):
30
- """Non-deepseek_v3 models must not enable dynamic_mxfp4_quant."""
31
- hf_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
32
- qcfg = _make_quark_config()
33
-
34
- qcfg.maybe_update_config(model_name, hf_config=hf_config)
35
-
36
- assert qcfg.dynamic_mxfp4_quant is False
37
-
38
-
39
-# ---------------------------------------------------------------------------
40
-# DeepSeek-V3 family + fp4 must enable dynamic_mxfp4_quant
41
-# ---------------------------------------------------------------------------
42
-@pytest.mark.parametrize(
43
- "model_name",
44
- ["amd/DeepSeek-R1-MXFP4-ASQ"],
45
-)
46
-def test_deepseek_family_fp4_enables_flag(model_name: str):
47
- hf_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
48
- qcfg = _make_quark_config()
49
-
50
- qcfg.maybe_update_config(model_name, hf_config=hf_config)
51
-
52
- assert qcfg.dynamic_mxfp4_quant is True
53
-
54
-
55
-# ---------------------------------------------------------------------------
56
-# Missing hf_config → warn and stay False
57
-# ---------------------------------------------------------------------------
58
-def test_missing_hf_config_stays_false():
59
- qcfg = _make_quark_config()
60
-
61
- qcfg.maybe_update_config("some/model")
62
-
63
- assert qcfg.dynamic_mxfp4_quant is False
vllm/_aiter_ops.py
+7
@@ -1355,6 +1355,13 @@ class rocm_aiter_ops:
1355
def is_triton_gemm_enabled(cls) -> bool:
1356
return cls._AITER_ENABLED and cls._TRITON_UNQUANT_GEMM
1357
1358
+ @classmethod
1359
+ @if_aiter_supported
1360
+ def is_tgemm_enabled(cls) -> bool:
1361
+ from vllm.platforms.rocm import on_gfx950
1362
+
1363
+ return cls.is_linear_enabled() and on_gfx950()
1364
+
1365
@staticmethod
1366
@if_aiter_supported
1367
def register_ops_once() -> None:
vllm/envs.py
+1
@@ -999,6 +999,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
999
# use aiter linear op if aiter ops are enabled
1000
# The following list of related ops
1001
# - scaled_mm (per-tensor / rowwise)
1002
+ # - use aiter tuned gemms for unquantized gemms
1003
"VLLM_ROCM_USE_AITER_LINEAR": lambda: (
1004
os.getenv("VLLM_ROCM_USE_AITER_LINEAR", "True").lower() in ("true", "1")
1005
),
vllm/model_executor/layers/quantization/quark/quark.py
+4
-29
@@ -5,7 +5,6 @@ import fnmatch
5
from typing import TYPE_CHECKING, Any, cast
6
7
import torch
8
-from transformers import PretrainedConfig
8
9
from vllm.logger import init_logger
10
from vllm.model_executor.layers.attention import Attention
@@ -45,10 +44,6 @@ __all__ = ["QuarkLinearMethod"]
44
45
logger = init_logger(__name__)
46
48
-# model_type values that use dynamic MXFP4 re-quantization for
49
-# OCP MX fp4 Quark checkpoints
50
-_DEEPSEEK_V3_FAMILY_MODEL_TYPES = frozenset({"deepseek_v3"})
51
-
47
48
class QuarkConfig(QuantizationConfig):
49
def __init__(
@@ -65,32 +60,12 @@ 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
+ # Note : this flag is kept disabled because the overhead of
64
+ # dynamic mxfp4 quantization negates the performance gains
65
+ # that come from shifting to mxfp4. It is left here in case
66
+ # we want to re-enable it in the future.
67
self.dynamic_mxfp4_quant = False
68
70
- def maybe_update_config(
71
- self,
72
- model_name: str,
73
- hf_config: PretrainedConfig | None = None,
74
- revision: str | None = None,
75
- ):
76
- """Enable dynamic MXFP4 only for DeepSeek-V3-family + fp4 Quark checkpoints."""
77
-
78
- if (
79
- getattr(hf_config, "model_type", None)
80
- not in _DEEPSEEK_V3_FAMILY_MODEL_TYPES
81
- ):
82
- return
83
-
84
- quant_config = getattr(hf_config, "quantization_config", None)
85
- if quant_config is not None:
86
- quant_dtype = (
87
- quant_config.get("global_quant_config", {})
88
- .get("weight", {})
89
- .get("dtype")
90
- )
91
- if quant_dtype == "fp4":
92
- self.dynamic_mxfp4_quant = True
93
-
69
def get_linear_method(self) -> "QuarkLinearMethod":
70
return QuarkLinearMethod(self)
71
vllm/model_executor/layers/utils.py
+16
-11
@@ -159,6 +159,7 @@ def rocm_unquantized_gemm_impl(
159
and weight.is_contiguous()
160
)
161
)
162
+
163
if use_skinny_reduce_counting:
164
return ops.wvSplitKrc(x, weight, cu_count, bias)
165
@@ -174,17 +175,21 @@ def rocm_unquantized_gemm_impl(
175
and k % 8 == 0
176
)
177
177
- if not use_skinny:
178
- return torch.nn.functional.linear(x, weight, bias)
179
-
180
- x_view = x.reshape(-1, x.size(-1))
181
- if m > 8 and 0 < n <= 4:
182
- cu_count = num_compute_units()
183
- out = ops.wvSplitK(weight, x_view, cu_count, bias)
184
- return out.reshape(*x.shape[:-1], weight.shape[0])
185
- elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None:
186
- out = ops.LLMM1(weight, x_view, 4)
187
- return out.reshape(*x.shape[:-1], weight.shape[0])
178
+ if use_skinny:
179
+ x_view = x.reshape(-1, x.size(-1))
180
+ if m > 8 and 0 < n <= 4:
181
+ cu_count = num_compute_units()
182
+ out = ops.wvSplitK(weight, x_view, cu_count, bias)
183
+ return out.reshape(*x.shape[:-1], weight.shape[0])
184
+ elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None:
185
+ out = ops.LLMM1(weight, x_view, 4)
186
+ return out.reshape(*x.shape[:-1], weight.shape[0])
187
+
188
+ if rocm_aiter_ops.is_tgemm_enabled():
189
+ from aiter.tuned_gemm import tgemm
190
+
191
+ return tgemm.mm(x, weight, bias)
192
+
193
return torch.nn.functional.linear(x, weight, bias)
194
195