| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from fractions import Fraction |
| 4 | |
| 5 | import pytest |
| 6 | import pytest_asyncio |
| 7 | from mistral_common.protocol.transcription.request import ( |
| 8 | StreamingMode, |
| 9 | TranscriptionRequest, |
| 10 | ) |
| 11 | from mistral_common.tokens.tokenizers.audio import Audio |
| 12 | from mistral_common.tokens.tokenizers.mistral import MistralTokenizer |
| 13 | from mistral_common.tokens.tokenizers.tekken import SpecialTokenPolicy |
| 14 | |
| 15 | from vllm import LLM, SamplingParams |
| 16 | from vllm.assets.audio import AudioAsset |
| 17 | from vllm.config import CUDAGraphMode |
| 18 | from vllm.engine.arg_utils import AsyncEngineArgs |
| 19 | from vllm.engine.protocol import StreamingInput |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.renderers.inputs.preprocess import parse_model_prompt |
| 22 | from vllm.utils.math_utils import cdiv |
| 23 | from vllm.v1.engine.async_llm import AsyncLLM |
| 24 | from vllm.v1.kv_cache_interface import SlidingWindowSpec |
| 25 | |
| 26 | from ....utils import ROCM_ENGINE_KWARGS, create_new_process_for_each_test |
| 27 | |
| 28 | MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602" |
| 29 | AUDIO_LAYER_NAME = "whisper_encoder.whisper_encoder.layers.0.layers.self_attn.attn" |
| 30 | ENGINE_CONFIG = { |
| 31 | "model": MODEL_NAME, |
| 32 | "max_model_len": 8192, |
| 33 | "max_num_seqs": 4, |
| 34 | "limit_mm_per_prompt": {"audio": 1}, |
| 35 | "config_format": "mistral", |
| 36 | "load_format": "mistral", |
| 37 | "tokenizer_mode": "mistral", |
| 38 | "enforce_eager": True, |
| 39 | "gpu_memory_utilization": 0.9, |
| 40 | **ROCM_ENGINE_KWARGS, |
| 41 | } |
| 42 | |
| 43 | |
| 44 | EXPECTED_TEXT = [ |
| 45 | ( |
| 46 | " First words I spoke in the original phonograph. " |
| 47 | "A little piece of practical poetry. Mary had a little lamb," |
| 48 | " its fleece was quite a slow, and everywhere that Mary went, " |
| 49 | "the lamb was sure to go." |
| 50 | ), |
| 51 | ( |
| 52 | " And the 0-1 pitch on the way to Edgar Martinez. Swung on" |
| 53 | " the line. Down the left field line for OBS. Here comes Joy. " |
| 54 | "Here is Junior to third base. They're going to wave him in. " |
| 55 | "The throw to the plate will be late. The Mariners are going" |
| 56 | " to play. For the American League Championship, " |
| 57 | "I don't believe it. It just continues. My, oh, my." |
| 58 | ), |
| 59 | ] |
| 60 | |
| 61 | |
| 62 | def _normalize(texts: list[str]) -> list[str]: |
| 63 | # The model occasionally transcribes "OBS" as "a base hit" and |
| 64 | # "oh, my" as "oh my", but both are acoustically valid. Normalise so |
| 65 | # the assertion is stable across runs and hardware. |
| 66 | texts[1] = texts[1].replace("a base hit", "OBS").replace("oh my", "oh, my") |
| 67 | return texts |
| 68 | |
| 69 | |
| 70 | def assert_encoder_kv_cache_spec(engine: LLM) -> None: |
| 71 | vllm_config = engine.llm_engine.vllm_config |
| 72 | audio_config = vllm_config.model_config.hf_config.audio_config |
| 73 | kv_cache_specs_per_rank = engine.llm_engine.model_executor.get_kv_cache_specs() |
| 74 | |
| 75 | assert len(kv_cache_specs_per_rank) == 1 |
| 76 | kv_cache_specs = kv_cache_specs_per_rank[0] |
| 77 | assert AUDIO_LAYER_NAME in kv_cache_specs, kv_cache_specs.keys() |
| 78 | spec = kv_cache_specs[AUDIO_LAYER_NAME] |
| 79 | |
| 80 | assert audio_config.sliding_window == 750 |
| 81 | assert audio_config.block_pool_size == 4 |
| 82 | assert isinstance(spec, SlidingWindowSpec) |
| 83 | assert spec.block_size == 16 |
| 84 | assert spec.num_kv_heads == 32 |
| 85 | assert spec.tokens_per_state == Fraction(1, 4) |
| 86 | # cdiv(750, 4) == 188 pooled tokens cover the model's window; the extra |
| 87 | # +1 is an eviction margin (see whisper_causal.py get_kv_cache_spec). |
| 88 | assert spec.sliding_window == cdiv(750, 4) + 1 == 189 |
| 89 | assert ( |
| 90 | spec.max_admission_blocks_per_request( |
| 91 | max_in_flight_tokens=1, |
| 92 | max_model_len=vllm_config.model_config.max_model_len, |
| 93 | ) |
| 94 | == 13 |
| 95 | ) |
| 96 | |
| 97 | |
| 98 | @pytest.fixture |
| 99 | def audio_assets() -> list[AudioAsset]: |
| 100 | return [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")] |
| 101 | |
| 102 | |
| 103 | @pytest.fixture |
| 104 | def tokenizer() -> MistralTokenizer: |
| 105 | return MistralTokenizer.from_hf_hub(MODEL_NAME) |
| 106 | |
| 107 | |
| 108 | @pytest_asyncio.fixture |
| 109 | async def async_engine(): |
| 110 | gpu_memory_utilization = ENGINE_CONFIG.get("gpu_memory_utilization", 0.9) |
| 111 | from vllm.platforms import current_platform |
| 112 | |
| 113 | if current_platform.is_rocm(): |
| 114 | from tests.utils import wait_for_memory_to_settle |
| 115 | |
| 116 | wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) |
| 117 | |
| 118 | engine_args = AsyncEngineArgs(**ENGINE_CONFIG) |
| 119 | llm = AsyncLLM.from_engine_args(engine_args) |
| 120 | try: |
| 121 | yield llm |
| 122 | finally: |
| 123 | shutdown_timeout = 60.0 if current_platform.is_rocm() else None |
| 124 | llm.shutdown(timeout=shutdown_timeout) |
| 125 | del llm |
| 126 | import torch |
| 127 | |
| 128 | torch._dynamo.reset() |
| 129 | from vllm.distributed import cleanup_dist_env_and_memory |
| 130 | |
| 131 | cleanup_dist_env_and_memory() |
| 132 | from tests.utils import wait_for_memory_to_settle |
| 133 | |
| 134 | wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) |
| 135 | |
| 136 | |
| 137 | def _build_inputs( |
| 138 | audio_assets: list[AudioAsset], tokenizer: MistralTokenizer |
| 139 | ) -> tuple[list[dict], list[SamplingParams]]: |
| 140 | audio_config = tokenizer.instruct_tokenizer.tokenizer.audio |
| 141 | inputs = [] |
| 142 | sampling_params = [] |
| 143 | for audio_asset in audio_assets: |
| 144 | audio = Audio.from_file(audio_asset.get_local_path(), strict=False) |
| 145 | req = TranscriptionRequest( |
| 146 | audio=audio.to_base64(audio.format), |
| 147 | streaming=StreamingMode.OFFLINE, |
| 148 | language=None, |
| 149 | ) |
| 150 | tokenized = tokenizer.instruct_tokenizer.encode_transcription(req) |
| 151 | tokens = tokenized.tokens |
| 152 | audio_array = tokenized.audios[0].audio_array |
| 153 | |
| 154 | num_samples = audio_array.shape[0] |
| 155 | max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1 |
| 156 | sampling_params.append(SamplingParams(temperature=0.0, max_tokens=max_tokens)) |
| 157 | inputs.append( |
| 158 | { |
| 159 | "multi_modal_data": {"audio": [(audio_array, None)]}, |
| 160 | "prompt_token_ids": tokens, |
| 161 | } |
| 162 | ) |
| 163 | return inputs, sampling_params |
| 164 | |
| 165 | |
| 166 | def _assert_expected_text(outputs) -> None: |
| 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}" |
| 173 | ) |
| 174 | |
| 175 | |
| 176 | def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypatch): |
| 177 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 178 | |
| 179 | vllm_kwargs = {**ENGINE_CONFIG} |
| 180 | vllm_kwargs["model_name"] = vllm_kwargs.pop("model") |
| 181 | |
| 182 | with vllm_runner(**vllm_kwargs) as vllm_model: |
| 183 | assert_encoder_kv_cache_spec(vllm_model.llm) |
| 184 | inputs, sampling_params = _build_inputs(audio_assets, tokenizer) |
| 185 | |
| 186 | outputs = vllm_model.llm.generate(inputs, sampling_params=sampling_params) |
| 187 | _assert_expected_text(outputs) |
| 188 | |
| 189 | |
| 190 | @pytest.mark.skipif( |
| 191 | not current_platform.is_cuda(), reason="Full cudagraphs require CUDA" |
| 192 | ) |
| 193 | @pytest.mark.parametrize( |
| 194 | "cudagraph_mode", |
| 195 | [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE], |
| 196 | ) |
| 197 | @create_new_process_for_each_test(method="spawn") |
| 198 | def test_voxtral_realtime_cudagraph( |
| 199 | audio_assets, tokenizer, vllm_runner, monkeypatch, cudagraph_mode |
| 200 | ): |
| 201 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 202 | |
| 203 | vllm_kwargs = {**ENGINE_CONFIG} |
| 204 | vllm_kwargs["model_name"] = vllm_kwargs.pop("model") |
| 205 | vllm_kwargs["enforce_eager"] = False |
| 206 | vllm_kwargs["compilation_config"] = { |
| 207 | "cudagraph_mode": cudagraph_mode, |
| 208 | # max_num_seqs is 4, so this is the only decode size worth capturing. |
| 209 | "cudagraph_capture_sizes": [4], |
| 210 | } |
| 211 | |
| 212 | with vllm_runner(**vllm_kwargs) as vllm_model: |
| 213 | inputs, sampling_params = _build_inputs(audio_assets, tokenizer) |
| 214 | |
| 215 | outputs = vllm_model.llm.generate(inputs, sampling_params=sampling_params) |
| 216 | _assert_expected_text(outputs) |
| 217 | |
| 218 | |
| 219 | @pytest.mark.asyncio |
| 220 | async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine): |
| 221 | # Lazy import to avoid CUDA-reinitialization error |
| 222 | from vllm.model_executor.models.voxtral_realtime import VoxtralRealtimeBuffer |
| 223 | |
| 224 | sampling_params = SamplingParams(temperature=0.0, max_tokens=1) |
| 225 | audio_config = tokenizer.instruct_tokenizer.audio_encoder.audio_config |
| 226 | |
| 227 | async def input_stream(audio_buffer): |
| 228 | async for prompt in audio_buffer.get_input_stream(): |
| 229 | parsed_prompt = parse_model_prompt(async_engine.model_config, prompt) |
| 230 | (engine_input,) = await async_engine.renderer.render_cmpl_async( |
| 231 | [parsed_prompt] |
| 232 | ) |
| 233 | yield StreamingInput(prompt=engine_input) |
| 234 | |
| 235 | output_tokens_list = [] |
| 236 | for i, audio_asset in enumerate(audio_assets): |
| 237 | output_tokens = [] |
| 238 | audio = Audio.from_file(audio_asset.get_local_path(), strict=False) |
| 239 | |
| 240 | req = TranscriptionRequest( |
| 241 | streaming=StreamingMode.OFFLINE, |
| 242 | audio=audio.to_base64(audio.format), |
| 243 | language=None, |
| 244 | ) |
| 245 | audio_enc = tokenizer.encode_transcription(req) |
| 246 | |
| 247 | buffer = VoxtralRealtimeBuffer(audio_config, audio_enc.tokens) |
| 248 | await buffer.append_audio(audio_enc.audios[0].audio_array) |
| 249 | await buffer.append_audio(None) |
| 250 | |
| 251 | request_id = f"session-{i}" |
| 252 | |
| 253 | async for resp in async_engine.generate( |
| 254 | prompt=input_stream(buffer), |
| 255 | sampling_params=sampling_params, |
| 256 | request_id=request_id, |
| 257 | ): |
| 258 | tokens = resp.outputs[0].token_ids[-1:] |
| 259 | output_tokens.extend(tokens) |
| 260 | await buffer.append_tokens(tokens) |
| 261 | |
| 262 | output_tokens_list.append(output_tokens) |
| 263 | |
| 264 | texts = _normalize( |
| 265 | [ |
| 266 | tokenizer.decode( |
| 267 | output_tokens, special_token_policy=SpecialTokenPolicy.IGNORE |
| 268 | ) |
| 269 | for output_tokens in output_tokens_list |
| 270 | ] |
| 271 | ) |
| 272 | for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)): |
| 273 | assert got == expected, ( |
| 274 | f"Output mismatch at index {i}:\n" |
| 275 | f" got: {got!r}\n" |
| 276 | f" expected: {expected!r}" |
| 277 | ) |