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 6, 20:30 UTC
[Model Runner V2] Add WhisperModelState [6/N] (#35790)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Mar 11, 2026 at 14:20 UTC
55eed6b7a52463e0eecb5adc45710c61f546b1ec
8 files changed
+232
-20
.buildkite/test_areas/model_runner_v2.yaml
+1
-2
@@ -47,8 +47,7 @@ steps:
47
- python3 offline_inference/audio_language.py --seed 0
48
- python3 offline_inference/vision_language.py --seed 0
49
- python3 offline_inference/vision_language_multi_image.py --seed 0
50
- # TODO: uncomment once https://github.com/vllm-project/vllm/pull/35790 is merged.
51
- #- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0 # TODO
50
+ - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
51
# for pooling models
52
- python3 pooling/embed/vision_embedding_offline.py --seed 0
53
# for features demo
vllm/v1/worker/gpu/attn_utils.py
+6
@@ -3,6 +3,7 @@
3
from collections.abc import Sequence
4
from typing import Any, cast
5
6
+import numpy as np
7
import torch
8
9
from vllm.config import VllmConfig, get_layers_from_vllm_config
@@ -180,6 +181,7 @@ def build_attn_metadata(
181
slot_mappings: torch.Tensor,
182
kv_cache_config: KVCacheConfig,
183
dcp_local_seq_lens: torch.Tensor | None = None,
184
+ encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]] | None = None,
185
) -> dict[str, Any]:
186
seq_lens = seq_lens[:num_reqs]
187
if dcp_local_seq_lens is not None:
@@ -204,6 +206,10 @@ def build_attn_metadata(
206
causal=True,
207
dcp_local_seq_lens=dcp_local_seq_lens,
208
)
209
+ if encoder_seq_lens and i in encoder_seq_lens:
210
+ encoder_seq_lens_gpu, encoder_seq_lens_cpu = encoder_seq_lens[i]
211
+ common_attn_metadata.encoder_seq_lens = encoder_seq_lens_gpu
212
+ common_attn_metadata.encoder_seq_lens_cpu = encoder_seq_lens_cpu
213
214
for attn_group in attn_groups[i]:
215
attn_metadata_builder = attn_group.get_metadata_builder(0)
vllm/v1/worker/gpu/cudagraph_utils.py
+1
@@ -389,5 +389,6 @@ def prepare_inputs_to_capture(
389
slot_mappings,
390
attn_groups,
391
kv_cache_config,
392
+ for_capture=True,
393
)
394
return attn_metadata, slot_mappings_by_layer
vllm/v1/worker/gpu/model_runner.py
+32
-5
@@ -125,6 +125,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
125
self.max_model_len = self.model_config.max_model_len
126
self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
127
self.max_num_reqs = self.scheduler_config.max_num_seqs
128
+ self.is_encoder_decoder = self.model_config.is_encoder_decoder
129
130
self.use_async_scheduling = self.scheduler_config.async_scheduling
131
self.output_copy_stream = torch.cuda.Stream(self.device)
@@ -159,12 +160,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
160
if self.supports_mm_inputs and self.is_first_pp_rank:
161
self.encoder_cache = EncoderCache()
162
163
+ # Speculative decoding.
164
self.speculator = None
165
self.num_speculative_steps = 0
166
self.use_aux_hidden_state_outputs = False
167
use_strict_rejection_sampling = False
168
if self.speculative_config is not None:
169
self.num_speculative_steps = self.speculative_config.num_speculative_tokens
170
+ use_strict_rejection_sampling = (
171
+ self.speculative_config.rejection_sample_method == "strict"
172
+ )
173
+
174
if self.is_last_pp_rank:
175
self.speculator = init_speculator(self.vllm_config, self.device)
176
@@ -173,13 +179,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
179
self.use_aux_hidden_state_outputs = True
180
if self.pp_size > 1:
181
raise ValueError("EAGLE3 with pipeline parallel is not supported.")
176
- use_strict_rejection_sampling = (
177
- self.speculative_config.rejection_sample_method == "strict"
178
- )
182
183
# Draft tokens propagation - for spec-dec + struct outputs.
184
self.draft_tokens_handler = DraftTokensHandler(self.device)
185
186
+ # General request states.
187
self.req_states = RequestState(
188
max_num_reqs=self.max_num_reqs,
189
max_model_len=self.max_model_len,
@@ -243,7 +247,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
247
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
248
tasks: list[SupportedTask] = []
249
if self.model_config.runner_type == "generate":
246
- tasks.append("generate")
250
+ tasks.extend(self.model_state.get_supported_generation_tasks())
251
if self.pooling_runner is not None:
252
tasks.extend(self.pooling_runner.get_supported_pooling_tasks())
253
return tuple(tasks)
@@ -307,11 +311,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
311
for kv_cache_group in kv_cache_config.kv_cache_groups
312
]
313
314
+ block_table_max_model_len = self.max_model_len
315
+ if self.is_encoder_decoder:
316
+ # Cross-attention block tables need to index encoder tokens
317
+ # (e.g., Whisper ~1500), which can exceed decoder max_model_len.
318
+ block_table_max_model_len = max(
319
+ block_table_max_model_len,
320
+ getattr(self.model_config.hf_config, "max_source_positions", 0),
321
+ )
322
+
323
self.block_tables = BlockTables(
324
block_sizes=block_sizes,
325
max_num_reqs=self.max_num_reqs,
326
max_num_batched_tokens=self.max_num_tokens,
314
- max_model_len=self.max_model_len,
327
+ max_model_len=block_table_max_model_len,
328
device=self.device,
329
cp_size=self.dcp_size,
330
cp_rank=self.dcp_rank,
@@ -870,6 +883,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):
883
)
884
num_tokens_across_dp = None
885
886
+ skip_compiled = False
887
+ if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
888
+ # Encoder-decoder models such as Whisper should run eager/non-compiled
889
+ # when encoder inputs are scheduled, because this step updates
890
+ # cross-attention cache with dynamic encoder outputs.
891
+ # Override batch_desc to NONE.
892
+ skip_compiled = True
893
+ batch_desc = BatchExecutionDescriptor(
894
+ cg_mode=CUDAGraphMode.NONE,
895
+ num_tokens=num_toks,
896
+ num_reqs=num_reqs,
897
+ )
898
+
899
if self.dp_size > 1:
900
batch_desc, num_tokens_across_dp = sync_cudagraph_and_dp_padding(
901
self.cudagraph_manager,
@@ -984,6 +1010,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1010
num_tokens_across_dp=num_tokens_across_dp,
1011
batch_descriptor=batch_descriptor,
1012
slot_mapping=slot_mappings_by_layer,
1013
+ skip_compiled=skip_compiled,
1014
):
1015
self.kv_connector.pre_forward(scheduler_output)
1016
model_output = self.model(**model_inputs)
vllm/v1/worker/gpu/model_states/__init__.py
+5
@@ -13,6 +13,11 @@ def init_model_state(
13
encoder_cache: EncoderCache | None,
14
device: torch.device,
15
):
16
+ if "WhisperForConditionalGeneration" in vllm_config.model_config.architectures:
17
+ from vllm.v1.worker.gpu.model_states.whisper import WhisperModelState
18
+
19
+ return WhisperModelState(vllm_config, model, encoder_cache, device)
20
+
21
from vllm.v1.worker.gpu.model_states.default import DefaultModelState
22
23
return DefaultModelState(vllm_config, model, encoder_cache, device)
vllm/v1/worker/gpu/model_states/default.py
+3
-4
@@ -109,7 +109,7 @@ class DefaultModelState(ModelState):
109
110
def prepare_inputs(
111
self, input_batch: InputBatch, req_states: RequestState
112
- ) -> dict[str, torch.Tensor | None]:
112
+ ) -> dict[str, Any]:
113
if not self.uses_mrope:
114
# Common case (1D positions).
115
return {}
@@ -126,9 +126,7 @@ class DefaultModelState(ModelState):
126
]
127
return {"positions": mrope_positions}
128
129
- def prepare_dummy_inputs(
130
- self, num_reqs: int, num_tokens: int
131
- ) -> dict[str, torch.Tensor | None]:
129
+ def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
130
model_inputs = {}
131
if self.supports_mm_inputs:
132
inputs_embeds = self.encoder_runner.inputs_embeds[:num_tokens]
@@ -146,6 +144,7 @@ class DefaultModelState(ModelState):
144
slot_mappings: torch.Tensor,
145
attn_groups: list[list[AttentionGroup]],
146
kv_cache_config: KVCacheConfig,
147
+ for_capture: bool = False,
148
) -> dict[str, Any]:
149
if cudagraph_mode == CUDAGraphMode.FULL:
150
# Use padded sizes - padding is handled by model_runner.prepare_attn.
vllm/v1/worker/gpu/model_states/interface.py
+10
-9
@@ -8,6 +8,7 @@ import torch.nn as nn
8
9
from vllm.config import VllmConfig
10
from vllm.config.compilation import CUDAGraphMode
11
+from vllm.tasks import GenerationTask
12
from vllm.v1.core.sched.output import NewRequestData
13
from vllm.v1.kv_cache_interface import KVCacheConfig
14
from vllm.v1.worker.gpu.input_batch import InputBatch
@@ -27,13 +28,14 @@ class ModelState(ABC):
28
) -> None:
29
raise NotImplementedError
30
30
- @abstractmethod
31
+ def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
32
+ return ("generate",)
33
+
34
def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
32
- raise NotImplementedError
35
+ return None
36
34
- @abstractmethod
37
def apply_staged_writes(self) -> None:
36
- raise NotImplementedError
38
+ return None
39
40
@abstractmethod
41
def get_mm_embeddings(
@@ -41,19 +43,17 @@ class ModelState(ABC):
43
scheduled_encoder_inputs: dict[str, list[int]],
44
input_batch: InputBatch,
45
req_states: RequestState,
44
- ) -> torch.Tensor:
46
+ ) -> torch.Tensor | None:
47
raise NotImplementedError
48
49
@abstractmethod
50
def prepare_inputs(
51
self, input_batch: InputBatch, req_states: RequestState
50
- ) -> dict[str, torch.Tensor | None]:
52
+ ) -> dict[str, Any]:
53
raise NotImplementedError
54
55
@abstractmethod
54
- def prepare_dummy_inputs(
55
- self, num_reqs: int, num_tokens: int
56
- ) -> dict[str, torch.Tensor | None]:
56
+ def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
57
raise NotImplementedError
58
59
@abstractmethod
@@ -65,5 +65,6 @@ class ModelState(ABC):
65
slot_mappings: torch.Tensor,
66
attn_groups: list[list[AttentionGroup]],
67
kv_cache_config: KVCacheConfig,
68
+ for_capture: bool = False,
69
) -> dict[str, Any]:
70
raise NotImplementedError
vllm/v1/worker/gpu/model_states/whisper.py
new
+174
@@ -0,0 +1,174 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Any
4
+
5
+import numpy as np
6
+import torch
7
+import torch.nn as nn
8
+
9
+from vllm.config import VllmConfig
10
+from vllm.config.compilation import CUDAGraphMode
11
+from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheConfig
12
+from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
13
+from vllm.v1.worker.gpu.input_batch import InputBatch
14
+from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
15
+from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
16
+from vllm.v1.worker.gpu.model_states.interface import ModelState
17
+from vllm.v1.worker.gpu.states import RequestState
18
+from vllm.v1.worker.utils import AttentionGroup
19
+
20
+
21
+class WhisperModelState(ModelState):
22
+ def __init__(
23
+ self,
24
+ vllm_config: VllmConfig,
25
+ model: nn.Module,
26
+ encoder_cache: EncoderCache | None,
27
+ device: torch.device,
28
+ ) -> None:
29
+ self.vllm_config = vllm_config
30
+ self.model_config = vllm_config.model_config
31
+ self.scheduler_config = vllm_config.scheduler_config
32
+ self.model = model
33
+ self.max_num_reqs = vllm_config.scheduler_config.max_num_seqs
34
+ self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
35
+ self.max_model_len = self.model_config.max_model_len
36
+ self.device = device
37
+
38
+ assert encoder_cache is not None
39
+ self.encoder_cache = encoder_cache
40
+ self.encoder_runner = EncoderRunner(
41
+ model=self.model,
42
+ max_num_tokens=self.max_num_tokens,
43
+ hidden_size=self.model_config.get_inputs_embeds_size(),
44
+ encoder_cache=self.encoder_cache,
45
+ dtype=self.model_config.dtype,
46
+ device=self.device,
47
+ )
48
+
49
+ self.max_encoder_len = getattr(
50
+ self.model_config.hf_config,
51
+ "max_source_positions",
52
+ self.max_model_len,
53
+ )
54
+ self.encoder_seq_lens_gpu = torch.zeros(
55
+ self.max_num_reqs, dtype=torch.int32, device=self.device
56
+ )
57
+
58
+ self.encoder_outputs: list[torch.Tensor] = []
59
+
60
+ def get_supported_generation_tasks(self):
61
+ return ("transcription",)
62
+
63
+ def get_mm_embeddings(
64
+ self,
65
+ scheduled_encoder_inputs: dict[str, list[int]],
66
+ input_batch: InputBatch,
67
+ req_states: RequestState,
68
+ ) -> None:
69
+ # Ensure encoder inputs are ordered consistently with input_batch.req_ids.
70
+ encoder_inputs: dict[str, list[int]] = {}
71
+ for req_id in input_batch.req_ids:
72
+ req_encoder_inputs = scheduled_encoder_inputs.get(req_id, [])
73
+ if req_encoder_inputs:
74
+ encoder_inputs[req_id] = req_encoder_inputs
75
+ _, mm_kwargs = self.encoder_runner.prepare_mm_inputs(encoder_inputs)
76
+ if mm_kwargs:
77
+ # Whisper consumes encoder outputs through `encoder_outputs`, not
78
+ # `inputs_embeds`. Single modality (audio) so execute_mm_encoder
79
+ # preserves request order; use its return value directly.
80
+ # No need to store in encoder_cache: cross-attention K/V are written
81
+ # to the KV cache on the first step; decode steps use the cache.
82
+ self.encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
83
+ else:
84
+ # Decode steps: encoder K/V are in cross-attention KV cache.
85
+ self.encoder_outputs = []
86
+ return None
87
+
88
+ def prepare_inputs(
89
+ self, input_batch: InputBatch, req_states: RequestState
90
+ ) -> dict[str, Any]:
91
+ model_inputs = {"encoder_outputs": self.encoder_outputs}
92
+ self.encoder_outputs = []
93
+ return model_inputs
94
+
95
+ def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]:
96
+ return {"encoder_outputs": []}
97
+
98
+ def prepare_attn(
99
+ self,
100
+ input_batch: InputBatch,
101
+ cudagraph_mode: CUDAGraphMode,
102
+ block_tables: tuple[torch.Tensor, ...],
103
+ slot_mappings: torch.Tensor,
104
+ attn_groups: list[list[AttentionGroup]],
105
+ kv_cache_config: KVCacheConfig,
106
+ for_capture: bool = False,
107
+ ) -> dict[str, Any]:
108
+ if cudagraph_mode == CUDAGraphMode.FULL:
109
+ num_reqs = input_batch.num_reqs_after_padding
110
+ num_tokens = input_batch.num_tokens_after_padding
111
+ else:
112
+ num_reqs = input_batch.num_reqs
113
+ num_tokens = input_batch.num_tokens
114
+ encoder_seq_lens = self._get_encoder_seq_lens(
115
+ input_batch.req_ids, attn_groups, for_capture
116
+ )
117
+
118
+ query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
119
+ max_query_len = input_batch.num_scheduled_tokens.max().item()
120
+ attn_metadata = build_attn_metadata(
121
+ attn_groups=attn_groups,
122
+ num_reqs=num_reqs,
123
+ num_tokens=num_tokens,
124
+ query_start_loc_gpu=input_batch.query_start_loc,
125
+ query_start_loc_cpu=query_start_loc_cpu,
126
+ max_query_len=max_query_len,
127
+ seq_lens=input_batch.seq_lens,
128
+ max_seq_len=self.max_model_len,
129
+ block_tables=block_tables,
130
+ slot_mappings=slot_mappings,
131
+ kv_cache_config=kv_cache_config,
132
+ dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
133
+ encoder_seq_lens=encoder_seq_lens,
134
+ )
135
+ return attn_metadata
136
+
137
+ def _get_encoder_seq_lens(
138
+ self,
139
+ req_ids: list[str],
140
+ attn_groups: list[list[AttentionGroup]],
141
+ for_capture: bool,
142
+ ) -> dict[int, tuple[torch.Tensor, np.ndarray]]:
143
+ num_reqs = len(req_ids)
144
+ encoder_seq_lens_np = np.zeros(num_reqs, dtype=np.int32)
145
+ if not for_capture:
146
+ # During normal execution, use actual encoder lengths.
147
+ for i, req_id in enumerate(req_ids):
148
+ mm_features = self.encoder_cache.mm_features.get(req_id, [])
149
+ encoder_seq_lens_np[i] = sum(
150
+ feature.mm_position.get_num_embeds() for feature in mm_features
151
+ )
152
+ else:
153
+ # During CUDA graph capture, use max encoder length so max_seqlen_k
154
+ # is captured with the correct value for cross-attention.
155
+ encoder_seq_lens_np[:] = self.max_encoder_len
156
+
157
+ self.encoder_seq_lens_gpu[:num_reqs].copy_(
158
+ torch.from_numpy(encoder_seq_lens_np), non_blocking=True
159
+ )
160
+ self.encoder_seq_lens_gpu[num_reqs:].fill_(0)
161
+ encoder_seq_lens_gpu = self.encoder_seq_lens_gpu[:num_reqs]
162
+
163
+ seq_lens_by_group: dict[int, tuple[torch.Tensor, np.ndarray]] = {}
164
+ for kv_cache_group_idx, groups in enumerate(attn_groups):
165
+ has_cross_attn = any(
166
+ isinstance(attn_group.kv_cache_spec, CrossAttentionSpec)
167
+ for attn_group in groups
168
+ )
169
+ if has_cross_attn:
170
+ seq_lens_by_group[kv_cache_group_idx] = (
171
+ encoder_seq_lens_gpu,
172
+ encoder_seq_lens_np,
173
+ )
174
+ return seq_lens_by_group