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, 14:30 UTC
[ROCm] Add env to enable/disable aiter triton gemm (#28321)
Signed-off-by: Yong Hoon Shin <yhshin@meta.com>
Yong Hoon Shin committed
Nov 8, 2025 at 20:27 UTC
de2b78305f4e8fa07e0b8b2a4fb6ac03ab0dde20
2 files changed
+8
vllm/envs.py
+7
@@ -113,6 +113,7 @@ if TYPE_CHECKING:
113
VLLM_ROCM_USE_AITER_FP8BMM: bool = True
114
VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False
115
VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = True
116
+ VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True
117
VLLM_ROCM_USE_SKINNY_GEMM: bool = True
118
VLLM_ROCM_FP8_PADDING: bool = True
119
VLLM_ROCM_MOE_PADDING: bool = True
@@ -944,6 +945,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
945
os.getenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "True").lower()
946
in ("true", "1")
947
),
948
+ # Whether to use aiter triton kernels for gemm ops.
949
+ # By default is enabled.
950
+ "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: (
951
+ os.getenv("VLLM_ROCM_USE_AITER_TRITON_GEMM", "True").lower() in ("true", "1")
952
+ ),
953
# use rocm skinny gemms
954
"VLLM_ROCM_USE_SKINNY_GEMM": lambda: (
955
os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in ("true", "1")
@@ -1586,6 +1592,7 @@ def compute_hash() -> str:
1592
"VLLM_ROCM_USE_TRITON_ROPE",
1593
"VLLM_ROCM_USE_AITER_FP8BMM",
1594
"VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION",
1595
+ "VLLM_ROCM_USE_AITER_TRITON_GEMM",
1596
"VLLM_ROCM_USE_SKINNY_GEMM",
1597
"VLLM_ROCM_FP8_PADDING",
1598
"VLLM_ROCM_MOE_PADDING",
vllm/model_executor/layers/utils.py
+1
@@ -106,6 +106,7 @@ def default_unquantized_gemm(
106
def use_aiter_triton_gemm(n, m, k, dtype):
107
if (
108
envs.VLLM_ROCM_USE_AITER == 0
109
+ or envs.VLLM_ROCM_USE_AITER_TRITON_GEMM == 0
110
# MI300's - fp8nuz=True
111
or current_platform.is_fp8_fnuz()
112
or dtype not in [torch.float16, torch.bfloat16]