[Model] Voxtral Realtime: add support for `CUDAGraphMode.FULL_DECODE_ONLY` (#51167)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: Jorge Gallego Feliciano <jorge.gallego@mistral.ai>
Nicolò Lucchesi committed
Sep 14, 2026 at 21:16 UTC
995e8581f462a13e32f30cfba946c63d48cf31d7
3 files changed
+174
-62
tests/models/multimodal/generation/test_voxtral_realtime.py
+67
-39
@@ -14,7 +14,9 @@ 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.platforms import current_platform
20
from vllm.utils.math_utils import cdiv
21
from vllm.v1.engine.async_llm import AsyncLLM
22
from vllm.v1.kv_cache_interface import SlidingWindowSpec
@@ -130,6 +132,45 @@ async def async_engine():
132
wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization)
133
134
135
+def _build_inputs(
136
+ audio_assets: list[AudioAsset], tokenizer: MistralTokenizer
137
+) -> tuple[list[dict], list[SamplingParams]]:
138
+ audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
139
+ inputs = []
140
+ sampling_params = []
141
+ for audio_asset in audio_assets:
142
+ audio = Audio.from_file(audio_asset.get_local_path(), strict=False)
143
+ req = TranscriptionRequest(
144
+ audio=audio.to_base64(audio.format),
145
+ streaming=StreamingMode.OFFLINE,
146
+ language=None,
147
+ )
148
+ tokenized = tokenizer.instruct_tokenizer.encode_transcription(req)
149
+ tokens = tokenized.tokens
150
+ audio_array = tokenized.audios[0].audio_array
151
+
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))
155
+ inputs.append(
156
+ {
157
+ "multi_modal_data": {"audio": [(audio_array, None)]},
158
+ "prompt_token_ids": tokens,
159
+ }
160
+ )
161
+ return inputs, sampling_params
162
+
163
+
164
+def _assert_expected_text(outputs) -> None:
165
+ texts = _normalize([out.outputs[0].text for out in outputs])
166
+ for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
167
+ assert got == expected, (
168
+ f"Output mismatch at index {i}:\n"
169
+ f" got: {got!r}\n"
170
+ f" expected: {expected!r}"
171
+ )
172
+
173
+
174
def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypatch):
175
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
176
@@ -138,51 +179,38 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypa
179
180
with vllm_runner(**vllm_kwargs) as vllm_model:
181
assert_encoder_kv_cache_spec(vllm_model.llm)
141
- audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
142
-
143
- def from_file(file_path: str):
144
- audio = Audio.from_file(file_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)
182
+ inputs, sampling_params = _build_inputs(audio_assets, tokenizer)
183
152
- return (tokenized.tokens, tokenized.audios[0].audio_array)
153
-
154
- tokenized_list = [
155
- from_file(audio_asset.get_local_path()) for audio_asset in audio_assets
156
- ]
184
+ outputs = vllm_model.llm.generate(inputs, sampling_params=sampling_params)
185
+ _assert_expected_text(outputs)
186
158
- inputs = []
159
- sampling_params = []
187
161
- for tokens, audio_array in tokenized_list:
162
- num_samples = audio_array.shape[0]
163
- max_tokens = audio_config.num_audio_tokens(num_samples) - len(tokens) - 1
164
- sampling_params.append(
165
- SamplingParams(temperature=0.0, max_tokens=max_tokens)
166
- )
188
+@pytest.mark.skipif(
189
+ not current_platform.is_cuda(), reason="Full cudagraphs require CUDA"
190
+)
191
+@pytest.mark.parametrize(
192
+ "cudagraph_mode",
193
+ [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE],
194
+)
195
+def test_voxtral_realtime_cudagraph(
196
+ audio_assets, tokenizer, vllm_runner, monkeypatch, cudagraph_mode
197
+):
198
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
199
168
- input_dict = {
169
- "multi_modal_data": {"audio": [(audio_array, None)]},
170
- "prompt_token_ids": tokens,
171
- }
172
- inputs.append(input_dict)
200
+ vllm_kwargs = {**ENGINE_CONFIG}
201
+ vllm_kwargs["model_name"] = vllm_kwargs.pop("model")
202
+ vllm_kwargs["enforce_eager"] = False
203
+ vllm_kwargs["compilation_config"] = {
204
+ "cudagraph_mode": cudagraph_mode,
205
+ # max_num_seqs is 4, so this is the only decode size worth capturing.
206
+ "cudagraph_capture_sizes": [4],
207
+ }
208
174
- outputs = vllm_model.llm.generate(
175
- inputs,
176
- sampling_params=sampling_params,
177
- )
209
+ with vllm_runner(**vllm_kwargs) as vllm_model:
210
+ inputs, sampling_params = _build_inputs(audio_assets, tokenizer)
211
179
- texts = _normalize([out.outputs[0].text for out in outputs])
180
- for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
181
- assert got == expected, (
182
- f"Output mismatch at index {i}:\n"
183
- f" got: {got!r}\n"
184
- f" expected: {expected!r}"
185
- )
212
+ outputs = vllm_model.llm.generate(inputs, sampling_params=sampling_params)
213
+ _assert_expected_text(outputs)
214
215
216
@pytest.mark.asyncio
vllm/model_executor/models/voxtral_realtime.py
+9
-3
@@ -217,9 +217,15 @@ class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtim
217
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
218
super().__init__(vllm_config=vllm_config, prefix=prefix)
219
220
- assert (
221
- not vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs()
222
- ), "Voxtral realtime doesn't support full cudagraphs yet. Please use PIECEWISE."
220
+ # Full cudagraphs are supported for decode-only batches (the encoder's
221
+ # block-pooling attention builder is capture-safe for uniform
222
+ # single-token decode). Mixed/prefill batches are not, so block only
223
+ # pure FULL; FULL_DECODE_ONLY and FULL_AND_PIECEWISE are allowed.
224
+ cudagraph_mode = vllm_config.compilation_config.cudagraph_mode
225
+ assert not cudagraph_mode.mixed_mode().has_full_cudagraphs(), (
226
+ "Voxtral realtime supports full cudagraphs for decode-only batches. "
227
+ "Use cudagraph_mode=FULL_DECODE_ONLY or FULL_AND_PIECEWISE (not FULL)."
228
+ )
229
230
self.time_embedding: TimeEmbedding = TimeEmbedding(
231
dim=self.config.text_config.hidden_size
vllm/model_executor/models/whisper_causal.py
+98
-20
@@ -1,12 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import copy
3
import functools
4
import logging
5
import math
6
from dataclasses import replace
7
from fractions import Fraction
8
from functools import partial
9
+from typing import ClassVar
10
11
import torch
12
import torch.nn.functional as F
@@ -27,6 +27,7 @@ from vllm.model_executor.models.whisper import WhisperPosEmbedType
27
from vllm.utils.math_utils import cdiv
28
from vllm.v1.attention.backend import (
29
AttentionBackend,
30
+ AttentionCGSupport,
31
AttentionMetadata,
32
AttentionType,
33
CommonAttentionMetadata,
@@ -123,6 +124,16 @@ def create_whisper_attention_backend_with_block_pooling(
124
underlying_impl = underlying_attn_backend.get_impl_cls()
125
126
class WhisperCausalAttentionWithBlockPoolingBuilder(underlying_builder): # type: ignore
127
+ # Full cudagraphs only for uniform single-token decode: capture bakes in
128
+ # tensor addresses, so `build` writes metadata into persistent buffers
129
+ # (below). Prefill/mixed batches fall back to piecewise.
130
+ _cudagraph_support: ClassVar[AttentionCGSupport] = (
131
+ AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
132
+ )
133
+ # Disable the update_block_table fast path: it would bypass block pooling
134
+ # by splicing the raw slot_mapping/block_table into cached metadata.
135
+ supports_update_block_table: bool = False
136
+
137
def __init__(
138
self,
139
kv_cache_spec: AttentionSpec,
@@ -150,6 +161,38 @@ def create_whisper_attention_backend_with_block_pooling(
161
# num_heads_q for the encoder is the same as num_kv_heads
162
# (no GQA in whisper encoder)
163
self.num_heads_q = kv_cache_spec.num_kv_heads
164
+ # Persistent capture-safe buffers for pooled metadata; lazily sized
165
+ # on first `build` to match the incoming tensors' dtype/device.
166
+ self._pool_buffers_ready = False
167
+
168
+ def _maybe_init_pool_buffers(
169
+ self, common_attn_metadata: CommonAttentionMetadata
170
+ ) -> None:
171
+ if self._pool_buffers_ready:
172
+ return
173
+ device = common_attn_metadata.slot_mapping.device
174
+ sched = self.vllm_config.scheduler_config
175
+ self._qsl_buf = torch.empty(
176
+ sched.max_num_seqs + 1,
177
+ dtype=common_attn_metadata.query_start_loc.dtype,
178
+ device=device,
179
+ )
180
+ self._seq_buf = torch.empty(
181
+ sched.max_num_seqs,
182
+ dtype=common_attn_metadata.seq_lens.dtype,
183
+ device=device,
184
+ )
185
+ self._slot_buf = torch.empty(
186
+ sched.max_num_batched_tokens * block_pool_size,
187
+ dtype=common_attn_metadata.slot_mapping.dtype,
188
+ device=device,
189
+ )
190
+ self._pool_arange = torch.arange(
191
+ block_pool_size,
192
+ dtype=common_attn_metadata.slot_mapping.dtype,
193
+ device=device,
194
+ )
195
+ self._pool_buffers_ready = True
196
197
def build(
198
self,
@@ -157,27 +200,62 @@ def create_whisper_attention_backend_with_block_pooling(
200
common_attn_metadata: CommonAttentionMetadata,
201
fast_build: bool = False,
202
) -> AttentionMetadata:
160
- new_common_attn_metadata = copy.deepcopy(common_attn_metadata)
161
- new_common_attn_metadata.query_start_loc *= block_pool_size
162
- new_common_attn_metadata.query_start_loc_cpu *= block_pool_size
163
- new_common_attn_metadata.seq_lens *= block_pool_size
164
- if new_common_attn_metadata.seq_lens_cpu_upper_bound is not None:
165
- new_common_attn_metadata.seq_lens_cpu_upper_bound *= block_pool_size
166
- new_common_attn_metadata.num_actual_tokens *= block_pool_size
167
- new_common_attn_metadata.max_query_len *= block_pool_size
168
- new_common_attn_metadata.max_seq_len *= block_pool_size
169
- original_slot_mapping = common_attn_metadata.slot_mapping
170
- common_prefix_len *= block_pool_size
171
- new_common_attn_metadata.slot_mapping = (
172
- (
173
- original_slot_mapping.unsqueeze(1) * block_pool_size
174
- + torch.arange(block_pool_size, device=original_slot_mapping.device)
175
- )
176
- .flatten()
177
- .clamp(min=-1)
203
+ cm = common_attn_metadata
204
+ self._maybe_init_pool_buffers(cm)
205
+
206
+ # Scale query_start_loc / seq_lens into the persistent buffers.
207
+ n = cm.query_start_loc.numel()
208
+ assert n <= self._qsl_buf.numel(), (
209
+ f"{n=} exceeds pooled query_start_loc buffer {self._qsl_buf.numel()}"
210
+ )
211
+ qsl_buf = self._qsl_buf[:n]
212
+ qsl_buf.copy_(cm.query_start_loc)
213
+ qsl_buf.mul_(block_pool_size)
214
+
215
+ r = cm.seq_lens.numel()
216
+ assert r <= self._seq_buf.numel(), (
217
+ f"{r=} exceeds pooled seq_lens buffer {self._seq_buf.numel()}"
218
+ )
219
+ seq_buf = self._seq_buf[:r]
220
+ seq_buf.copy_(cm.seq_lens)
221
+ seq_buf.mul_(block_pool_size)
222
+
223
+ # Expand each slot into block_pool_size contiguous slots. Padding
224
+ # slots (-1) stay -1 after clamp so the KV write is skipped.
225
+ m = cm.slot_mapping.numel()
226
+ assert m * block_pool_size <= self._slot_buf.numel(), (
227
+ f"pooled slot_mapping {m * block_pool_size} exceeds buffer "
228
+ f"{self._slot_buf.numel()}"
229
+ )
230
+ slot_buf = self._slot_buf[: m * block_pool_size]
231
+ slot_buf.view(m, block_pool_size).copy_(
232
+ cm.slot_mapping.unsqueeze(1) * block_pool_size + self._pool_arange
233
+ )
234
+ slot_buf.clamp_(min=-1)
235
+
236
+ # CPU tensors aren't captured; fresh scaled copies avoid mutating
237
+ # the runner's shared buffers.
238
+ qsl_cpu = cm.query_start_loc_cpu * block_pool_size
239
+ seq_lens_cpu_upper_bound = (
240
+ cm.seq_lens_cpu_upper_bound * block_pool_size
241
+ if cm.seq_lens_cpu_upper_bound is not None
242
+ else None
243
+ )
244
+
245
+ new_common_attn_metadata = cm.replace(
246
+ query_start_loc=qsl_buf,
247
+ query_start_loc_cpu=qsl_cpu,
248
+ seq_lens=seq_buf,
249
+ seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
250
+ num_actual_tokens=cm.num_actual_tokens * block_pool_size,
251
+ max_query_len=cm.max_query_len * block_pool_size,
252
+ max_seq_len=cm.max_seq_len * block_pool_size,
253
+ slot_mapping=slot_buf,
254
)
255
return super().build(
180
- common_prefix_len, new_common_attn_metadata, fast_build
256
+ common_prefix_len * block_pool_size,
257
+ new_common_attn_metadata,
258
+ fast_build,
259
)
260
261
# NOTE: We need a custom impl so we can use the transformed slot_mapping