[ROCm][CI][Bugfix] Use VllmRunner for `voxtral_realtime` tests to avoid OOM on AMD GPU (#47536)
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Shanshan Shen committed
Jul 4, 2026 at 12:26 UTC
f329ce405b12623fb8b1cf1830f12e5a712523be
1 file changed
+64
-55
tests/models/multimodal/generation/test_voxtral_realtime.py
+64
-55
@@ -1,7 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import contextlib
4
-
3
import pytest
4
import pytest_asyncio
5
from mistral_common.protocol.transcription.request import (
@@ -12,7 +10,7 @@ from mistral_common.tokens.tokenizers.audio import Audio
10
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
11
from mistral_common.tokens.tokenizers.tekken import SpecialTokenPolicy
12
15
-from vllm import LLM, EngineArgs, SamplingParams
13
+from vllm import LLM, SamplingParams
14
from vllm.assets.audio import AudioAsset
15
from vllm.engine.arg_utils import AsyncEngineArgs
16
from vllm.utils.math_utils import cdiv
@@ -100,78 +98,89 @@ def tokenizer() -> MistralTokenizer:
98
return MistralTokenizer.from_hf_hub(MODEL_NAME)
99
100
103
-@pytest.fixture
104
-def engine(monkeypatch: pytest.MonkeyPatch):
105
- # Disable multiprocessing allows us to access model executor from LLM engine
106
- monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
107
- engine_args = EngineArgs(**ENGINE_CONFIG)
108
- llm = LLM.from_engine_args(engine_args)
109
- try:
110
- yield llm
111
- finally:
112
- with contextlib.suppress(Exception):
113
- llm.llm_engine.engine_core.shutdown()
114
- import torch
101
+@pytest_asyncio.fixture
102
+async def async_engine():
103
+ gpu_memory_utilization = ENGINE_CONFIG.get("gpu_memory_utilization", 0.9)
104
+ from vllm.platforms import current_platform
105
116
- torch.accelerator.empty_cache()
106
+ if current_platform.is_rocm():
107
+ from tests.utils import wait_for_rocm_memory_to_settle
108
109
+ wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
110
119
-@pytest_asyncio.fixture
120
-async def async_engine():
111
engine_args = AsyncEngineArgs(**ENGINE_CONFIG)
112
llm = AsyncLLM.from_engine_args(engine_args)
113
try:
114
yield llm
115
finally:
126
- llm.shutdown()
116
+ shutdown_timeout = 60.0 if current_platform.is_rocm() else None
117
+ llm.shutdown(timeout=shutdown_timeout)
118
+ del llm
119
+ import torch
120
121
+ torch._dynamo.reset()
122
+ from vllm.distributed import cleanup_dist_env_and_memory
123
129
-def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
130
- assert_encoder_kv_cache_spec(engine)
131
- audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
124
+ cleanup_dist_env_and_memory()
125
+ from tests.utils import wait_for_rocm_memory_to_settle
126
133
- def from_file(file_path: str):
134
- audio = Audio.from_file(file_path, strict=False)
135
- req = TranscriptionRequest(
136
- audio=audio.to_base64(audio.format),
137
- streaming=StreamingMode.OFFLINE,
138
- language=None,
139
- )
140
- tokenized = tokenizer.instruct_tokenizer.encode_transcription(req)
127
+ wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
128
142
- return (tokenized.tokens, tokenized.audios[0].audio_array)
129
144
- tokenized_list = [
145
- from_file(audio_asset.get_local_path()) for audio_asset in audio_assets
146
- ]
130
+def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypatch):
131
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
132
148
- inputs = []
149
- sampling_params = []
133
+ vllm_kwargs = {**ENGINE_CONFIG}
134
+ vllm_kwargs["model_name"] = vllm_kwargs.pop("model")
135
151
- for tokens, audio_array in tokenized_list:
152
- num_samples = audio_array.shape[0]
153
- max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1
154
- sampling_params.append(SamplingParams(temperature=0.0, max_tokens=max_tokens))
136
+ with vllm_runner(**vllm_kwargs) as vllm_model:
137
+ assert_encoder_kv_cache_spec(vllm_model.llm)
138
+ audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
139
156
- input_dict = {
157
- "multi_modal_data": {"audio": [(audio_array, None)]},
158
- "prompt_token_ids": tokens,
159
- }
160
- inputs.append(input_dict)
140
+ def from_file(file_path: str):
141
+ audio = Audio.from_file(file_path, strict=False)
142
+ req = TranscriptionRequest(
143
+ audio=audio.to_base64(audio.format),
144
+ streaming=StreamingMode.OFFLINE,
145
+ language=None,
146
+ )
147
+ tokenized = tokenizer.instruct_tokenizer.encode_transcription(req)
148
162
- outputs = engine.generate(
163
- inputs,
164
- sampling_params=sampling_params,
165
- )
149
+ return (tokenized.tokens, tokenized.audios[0].audio_array)
150
167
- texts = _normalize([out.outputs[0].text for out in outputs])
168
- for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
169
- assert got == expected, (
170
- f"Output mismatch at index {i}:\n"
171
- f" got: {got!r}\n"
172
- f" expected: {expected!r}"
151
+ tokenized_list = [
152
+ from_file(audio_asset.get_local_path()) for audio_asset in audio_assets
153
+ ]
154
+
155
+ inputs = []
156
+ sampling_params = []
157
+
158
+ for tokens, audio_array in tokenized_list:
159
+ num_samples = audio_array.shape[0]
160
+ max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1
161
+ sampling_params.append(
162
+ SamplingParams(temperature=0.0, max_tokens=max_tokens)
163
+ )
164
+
165
+ input_dict = {
166
+ "multi_modal_data": {"audio": [(audio_array, None)]},
167
+ "prompt_token_ids": tokens,
168
+ }
169
+ inputs.append(input_dict)
170
+
171
+ outputs = vllm_model.llm.generate(
172
+ inputs,
173
+ sampling_params=sampling_params,
174
)
175
176
+ texts = _normalize([out.outputs[0].text for out in outputs])
177
+ for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
178
+ assert got == expected, (
179
+ f"Output mismatch at index {i}:\n"
180
+ f" got: {got!r}\n"
181
+ f" expected: {expected!r}"
182
+ )
183
+
184
185
@pytest.mark.asyncio
186
async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine):