From 6841f5dc77e9200a2fa45a4bf935b23bd843bf30 Mon Sep 17 00:00:00 2001 From: Hemanth Acharya Date: Thu, 30 Apr 2026 04:37:46 +0530 Subject: [PATCH] [ROCm] Add env flags to disable dynamic MXFP4 quant and enable AITER tuned GEMMs for Attention Projection Layers (#39987) Signed-off-by: Hemanth Acharya --- .../test_quark_maybe_update_config.py | 63 ------------------- vllm/_aiter_ops.py | 7 +++ vllm/envs.py | 1 + .../layers/quantization/quark/quark.py | 33 ++-------- vllm/model_executor/layers/utils.py | 25 +++++--- 5 files changed, 27 insertions(+), 102 deletions(-) delete mode 100644 tests/quantization/test_quark_maybe_update_config.py diff --git a/tests/quantization/test_quark_maybe_update_config.py b/tests/quantization/test_quark_maybe_update_config.py deleted file mode 100644 index 0142e869c2..0000000000 --- a/tests/quantization/test_quark_maybe_update_config.py +++ /dev/null @@ -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 diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 2681fd4a73..f702a025f5 100644 --- a/vllm/_aiter_ops.py +++ b/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: diff --git a/vllm/envs.py b/vllm/envs.py index 73278bf8dc..fa7a5fe70f 100755 --- a/vllm/envs.py +++ b/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") ), diff --git a/vllm/model_executor/layers/quantization/quark/quark.py b/vllm/model_executor/layers/quantization/quark/quark.py index 6aaf9a6458..dd7f33e261 100644 --- a/vllm/model_executor/layers/quantization/quark/quark.py +++ b/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) diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index e26b511de4..f7fb2ad047 100644 --- a/vllm/model_executor/layers/utils.py +++ b/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)