[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>

tests/quantization/test_quark_maybe_update_config.py

@@ -1,63 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for QuarkConfig.maybe_update_config.
-
-Fetches real HF configs (metadata only, no model weights) to verify
-that dynamic_mxfp4_quant is only enabled for DeepSeek-V3-family models.
-
-Run: pytest tests/quantization/test_quark_maybe_update_config.py -v
-"""
-
-import pytest
-from transformers import AutoConfig
-
-from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
-
-
-def _make_quark_config() -> QuarkConfig:
-    """Create a minimal QuarkConfig for testing."""
-    return QuarkConfig(quant_config={}, kv_cache_group=[], pack_method="reorder")
-
-
-# ---------------------------------------------------------------------------
-# Non-deepseek models must not flip dynamic_mxfp4_quant
-# ---------------------------------------------------------------------------
-@pytest.mark.parametrize(
-    "model_name",
-    ["amd/MiniMax-M2.1-MXFP4"],
-)
-def test_non_deepseek_model_stays_false(model_name: str):
-    """Non-deepseek_v3 models must not enable dynamic_mxfp4_quant."""
-    hf_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
-    qcfg = _make_quark_config()
-
-    qcfg.maybe_update_config(model_name, hf_config=hf_config)
-
-    assert qcfg.dynamic_mxfp4_quant is False
-
-
-# ---------------------------------------------------------------------------
-# DeepSeek-V3 family + fp4 must enable dynamic_mxfp4_quant
-# ---------------------------------------------------------------------------
-@pytest.mark.parametrize(
-    "model_name",
-    ["amd/DeepSeek-R1-MXFP4-ASQ"],
-)
-def test_deepseek_family_fp4_enables_flag(model_name: str):
-    hf_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
-    qcfg = _make_quark_config()
-
-    qcfg.maybe_update_config(model_name, hf_config=hf_config)
-
-    assert qcfg.dynamic_mxfp4_quant is True
-
-
-# ---------------------------------------------------------------------------
-# Missing hf_config → warn and stay False
-# ---------------------------------------------------------------------------
-def test_missing_hf_config_stays_false():
-    qcfg = _make_quark_config()
-
-    qcfg.maybe_update_config("some/model")
-
-    assert qcfg.dynamic_mxfp4_quant is False

vllm/_aiter_ops.py

@@ -1355,6 +1355,13 @@ class rocm_aiter_ops:
     def is_triton_gemm_enabled(cls) -> bool:
         return cls._AITER_ENABLED and cls._TRITON_UNQUANT_GEMM
 
+    @classmethod
+    @if_aiter_supported
+    def is_tgemm_enabled(cls) -> bool:
+        from vllm.platforms.rocm import on_gfx950
+
+        return cls.is_linear_enabled() and on_gfx950()
+
     @staticmethod
     @if_aiter_supported
     def register_ops_once() -> None:

vllm/envs.py

@@ -999,6 +999,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # use aiter linear op if aiter ops are enabled
     # The following list of related ops
     # - scaled_mm (per-tensor / rowwise)
+    # - use aiter tuned gemms for unquantized gemms
     "VLLM_ROCM_USE_AITER_LINEAR": lambda: (
         os.getenv("VLLM_ROCM_USE_AITER_LINEAR", "True").lower() in ("true", "1")
     ),

vllm/model_executor/layers/quantization/quark/quark.py

@@ -5,7 +5,6 @@ import fnmatch
 from typing import TYPE_CHECKING, Any, cast
 
 import torch
-from transformers import PretrainedConfig
 
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
@@ -45,10 +44,6 @@ __all__ = ["QuarkLinearMethod"]
 
 logger = init_logger(__name__)
 
-# model_type values that use dynamic MXFP4 re-quantization for
-# OCP MX fp4 Quark checkpoints
-_DEEPSEEK_V3_FAMILY_MODEL_TYPES = frozenset({"deepseek_v3"})
-
 
 class QuarkConfig(QuantizationConfig):
     def __init__(
@@ -65,32 +60,12 @@ class QuarkConfig(QuantizationConfig):
         self.kv_cache_group = kv_cache_group
         self.kv_cache_config = kv_cache_config
         self.pack_method = pack_method
+        # Note : this flag is kept disabled because the overhead of
+        # dynamic mxfp4 quantization negates the performance gains
+        # that come from shifting to mxfp4. It is left here in case
+        # we want to re-enable it in the future.
         self.dynamic_mxfp4_quant = False
 
-    def maybe_update_config(
-        self,
-        model_name: str,
-        hf_config: PretrainedConfig | None = None,
-        revision: str | None = None,
-    ):
-        """Enable dynamic MXFP4 only for DeepSeek-V3-family + fp4 Quark checkpoints."""
-
-        if (
-            getattr(hf_config, "model_type", None)
-            not in _DEEPSEEK_V3_FAMILY_MODEL_TYPES
-        ):
-            return
-
-        quant_config = getattr(hf_config, "quantization_config", None)
-        if quant_config is not None:
-            quant_dtype = (
-                quant_config.get("global_quant_config", {})
-                .get("weight", {})
-                .get("dtype")
-            )
-            if quant_dtype == "fp4":
-                self.dynamic_mxfp4_quant = True
-
     def get_linear_method(self) -> "QuarkLinearMethod":
         return QuarkLinearMethod(self)
 

vllm/model_executor/layers/utils.py

@@ -159,6 +159,7 @@ def rocm_unquantized_gemm_impl(
             and weight.is_contiguous()
         )
     )
+
     if use_skinny_reduce_counting:
         return ops.wvSplitKrc(x, weight, cu_count, bias)
 
@@ -174,17 +175,21 @@ def rocm_unquantized_gemm_impl(
         and k % 8 == 0
     )
 
-    if not use_skinny:
-        return torch.nn.functional.linear(x, weight, bias)
+    if use_skinny:
+        x_view = x.reshape(-1, x.size(-1))
+        if m > 8 and 0 < n <= 4:
+            cu_count = num_compute_units()
+            out = ops.wvSplitK(weight, x_view, cu_count, bias)
+            return out.reshape(*x.shape[:-1], weight.shape[0])
+        elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None:
+            out = ops.LLMM1(weight, x_view, 4)
+            return out.reshape(*x.shape[:-1], weight.shape[0])
+
+    if rocm_aiter_ops.is_tgemm_enabled():
+        from aiter.tuned_gemm import tgemm
+
+        return tgemm.mm(x, weight, bias)
 
-    x_view = x.reshape(-1, x.size(-1))
-    if m > 8 and 0 < n <= 4:
-        cu_count = num_compute_units()
-        out = ops.wvSplitK(weight, x_view, cu_count, bias)
-        return out.reshape(*x.shape[:-1], weight.shape[0])
-    elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None:
-        out = ops.LLMM1(weight, x_view, 4)
-        return out.reshape(*x.shape[:-1], weight.shape[0])
     return torch.nn.functional.linear(x, weight, bias)