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 10, 01:30 UTC
[ROCm][CI] Enable AITER Unified Attention On ROCm For gpt-oss Test (#32431)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
Micah Williamson committed
Jan 15, 2026 at 18:55 UTC
46f8a982b191e3a3d3a1eccaf18b184c391ac2ac
1 file changed
+15
-2
tests/entrypoints/openai/test_serving_chat.py
+15
-2
@@ -10,6 +10,7 @@ import pytest
10
import pytest_asyncio
11
from openai import OpenAI
12
13
+from vllm._aiter_ops import is_aiter_found_and_supported
14
from vllm.config.multimodal import MultiModalConfig
15
from vllm.entrypoints.openai.chat_completion.protocol import (
16
ChatCompletionRequest,
@@ -106,9 +107,21 @@ def gptoss_speculative_server(default_server_args: list[str]):
107
"--speculative-config",
108
f'{{"model": "{GPT_OSS_SPECULATOR_NAME}", '
109
f'"method": "eagle3", "num_speculative_tokens": 3}}',
109
- "--attention-backend=TRITON_ATTN",
110
+ f"--attention-backend={
111
+ 'TRITON_ATTN'
112
+ if not is_aiter_found_and_supported()
113
+ else 'ROCM_AITER_UNIFIED_ATTN'
114
+ }",
115
]
111
- with RemoteOpenAIServer(GPT_OSS_MODEL_NAME, server_args) as remote_server:
116
+ # gpt-oss requires AITER unified attention on ROCm
117
+ # TODO: Remove after fixing TRITON_ATTN issue on ROCm
118
+ # https://github.com/vllm-project/vllm/issues/32434
119
+ env_dict = None
120
+ if is_aiter_found_and_supported():
121
+ env_dict = {"VLLM_ROCM_USE_AITER": "1"}
122
+ with RemoteOpenAIServer(
123
+ GPT_OSS_MODEL_NAME, server_args, env_dict=env_dict
124
+ ) as remote_server:
125
yield remote_server
126
127