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 5, 18:30 UTC
[ROCm][CI] Fix Whisper translation test attention backend selection (#38508)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 31, 2026 at 00:21 UTC
b9cdc85207b3af83613fb45501b666e4e76c974d
1 file changed
+39
-6
tests/entrypoints/openai/speech_to_text/test_translation_validation.py
+39
-6
@@ -16,10 +16,41 @@ import soundfile as sf
16
17
from tests.entrypoints.openai.conftest import add_attention_backend
18
from tests.utils import RemoteOpenAIServer
19
+from vllm.logger import init_logger
20
+
21
+logger = init_logger(__name__)
22
23
SERVER_ARGS = ["--enforce-eager"]
24
25
26
+def _get_rocm_attention_config(model_name):
27
+ """Return appropriate ROCm attention config for the given model.
28
+
29
+ Whisper uses cross-attention (ENCODER_DECODER) which ROCM_AITER_FA does
30
+ not support. For Whisper we use ROCM_AITER_UNIFIED_ATTN (or TRITON_ATTN
31
+ as fallback); other models can use ROCM_AITER_FA.
32
+ """
33
+ from vllm.platforms import current_platform
34
+
35
+ if not current_platform.is_rocm():
36
+ return None
37
+
38
+ if "whisper" in model_name.lower():
39
+ try:
40
+ from vllm.platforms.rocm import _ON_MI3XX
41
+
42
+ if _ON_MI3XX:
43
+ return {"backend": "ROCM_AITER_UNIFIED_ATTN"}
44
+ except ImportError:
45
+ logger.warning(
46
+ "Could not import _ON_MI3XX from rocm platform, "
47
+ "falling back to TRITON_ATTN for Whisper."
48
+ )
49
+ return {"backend": "TRITON_ATTN"}
50
+
51
+ return {"backend": "ROCM_AITER_FA"}
52
+
53
+
54
def _get_server_args(attention_config):
55
"""Get server args with attention backend if specified."""
56
args = SERVER_ARGS.copy()
@@ -30,10 +61,11 @@ def _get_server_args(attention_config):
61
@pytest.fixture(
62
scope="module", params=["openai/whisper-small", "google/gemma-3n-E2B-it"]
63
)
33
-def server(request, rocm_aiter_fa_attention):
64
+def server(request):
65
# Parametrize over model name
66
+ attention_config = _get_rocm_attention_config(request.param)
67
with RemoteOpenAIServer(
36
- request.param, _get_server_args(rocm_aiter_fa_attention)
68
+ request.param, _get_server_args(attention_config)
69
) as remote_server:
70
yield remote_server, request.param
71
@@ -46,11 +78,12 @@ async def client_and_model(server):
78
79
80
@pytest.mark.asyncio
49
-async def test_non_asr_model(foscolo, rocm_aiter_fa_attention):
81
+async def test_non_asr_model(foscolo):
82
# text to text model
83
model_name = "JackFram/llama-68m"
84
+ attention_config = _get_rocm_attention_config(model_name)
85
with RemoteOpenAIServer(
53
- model_name, _get_server_args(rocm_aiter_fa_attention)
86
+ model_name, _get_server_args(attention_config)
87
) as remote_server:
88
client = remote_server.get_async_client()
89
@@ -61,7 +94,7 @@ async def test_non_asr_model(foscolo, rocm_aiter_fa_attention):
94
95
96
@pytest.mark.asyncio
64
-async def test_basic_audio_with_lora(mary_had_lamb, rocm_aiter_fa_attention):
97
+async def test_basic_audio_with_lora(mary_had_lamb):
98
"""Ensure STT (translate) requests can pass LoRA through to generate."""
99
# ROCm SPECIFIC CONFIGURATION:
100
# To ensure the test passes on ROCm, we modify the max model length to 512.
@@ -85,7 +118,7 @@ async def test_basic_audio_with_lora(mary_had_lamb, rocm_aiter_fa_attention):
118
"1",
119
]
120
88
- add_attention_backend(server_args, rocm_aiter_fa_attention)
121
+ add_attention_backend(server_args, _get_rocm_attention_config(model_name))
122
123
# Based on https://github.com/openai/openai-cookbook/blob/main/examples/Whisper_prompting_guide.ipynb.
124
with RemoteOpenAIServer(model_name, server_args) as remote_server: