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 9, 02:30 UTC
fix[ROCm]: Remove unconditional aiter import (#32902)
Signed-off-by: rabi <ramishra@redhat.com>
Rabi Mishra committed
Feb 2, 2026 at 19:40 UTC
9eb58f8cf123dde43a97216cd7fa0f9c7af81043
3 files changed
+18
-7
vllm/_aiter_ops.py
+12
-4
@@ -31,7 +31,12 @@ IS_AITER_FOUND = is_aiter_found()
31
32
33
def is_aiter_found_and_supported() -> bool:
34
- if current_platform.is_rocm() and IS_AITER_FOUND:
34
+ """Check if AITER is available AND enabled via environment variable.
35
+
36
+ Checks: platform (ROCm), device arch (gfx9), library existence,
37
+ and VLLM_ROCM_USE_AITER env variable.
38
+ """
39
+ if current_platform.is_rocm() and IS_AITER_FOUND and envs.VLLM_ROCM_USE_AITER:
40
from vllm.platforms.rocm import on_gfx9
41
42
return on_gfx9()
@@ -40,13 +45,11 @@ def is_aiter_found_and_supported() -> bool:
45
46
def if_aiter_supported(func: Callable) -> Callable:
47
"""Decorator that only executes the function if
43
- ROCm AITER package is supported on gfx9 archs.
48
+ ROCm AITER package is supported and enabled on gfx9 archs.
49
"""
50
51
@functools.wraps(func)
52
def wrapper(*args, **kwargs):
48
- # checks the platform, device arch and aiter library existence.
49
-
53
if is_aiter_found_and_supported():
54
return func(*args, **kwargs)
55
@@ -63,6 +66,11 @@ if is_aiter_found_and_supported():
66
from aiter import dtypes
67
68
AITER_FP8_DTYPE = dtypes.fp8
69
+else:
70
+ # Placeholder when AITER is disabled - prevents NameError during module load.
71
+ # Note: When AITER is disabled, ops are not registered, so fake implementations
72
+ # referencing this variable won't actually be called at runtime.
73
+ AITER_FP8_DTYPE = _FP8_DTYPE
74
75
76
def _rocm_aiter_fused_moe_impl(
vllm/v1/attention/backends/rocm_aiter_fa.py
+3
-2
@@ -32,10 +32,11 @@ from vllm.v1.kv_cache_interface import AttentionSpec
32
_PARTITION_SIZE_ROCM = 256
33
_CP_TOKENS_PER_ITER_ROCM = 32 * 1024
34
if current_platform.is_rocm():
35
- import aiter
36
-
35
from vllm.triton_utils import tl, triton
36
37
+ if rocm_aiter_ops.is_enabled():
38
+ import aiter
39
+
40
def block_size(x, head_dim):
41
return min(65536 // x.element_size(), triton.next_power_of_2(head_dim))
42
vllm/v1/spec_decode/eagle.py
+3
-1
@@ -183,7 +183,9 @@ class SpecDecodeBaseProposer:
183
RocmAttentionMetadata,
184
]
185
# ROCM_AITER_FA is an optional backend
186
- if find_spec(
186
+ from vllm._aiter_ops import rocm_aiter_ops
187
+
188
+ if rocm_aiter_ops.is_enabled() and find_spec(
189
AttentionBackendEnum.ROCM_AITER_FA.get_path(include_classname=False)
190
):
191
from vllm.v1.attention.backends.rocm_aiter_fa import (