[MRV2] Generalize use of `WhisperModelState` (#46096)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jun 22, 2026 at 12:40 UTC
fbf520cf3aab26f97337c6d6e299c66b96d0fcb9
4 files changed
+38
-43
vllm/v1/worker/gpu/model_states/__init__.py
+7
-6
@@ -4,6 +4,7 @@ import torch
4
import torch.nn as nn
5
6
from vllm.config import VllmConfig
7
+from vllm.model_executor.layers.attention import CrossAttention
8
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
9
10
@@ -18,13 +19,13 @@ def init_model_state(
19
cls = model.get_model_state_cls()
20
return cls(vllm_config, model, encoder_cache, device)
21
21
- if (
22
- "WhisperForConditionalGeneration" in vllm_config.model_config.architectures
23
- or "CohereAsrForConditionalGeneration" in vllm_config.model_config.architectures
24
- ):
25
- from vllm.v1.worker.gpu.model_states.whisper import WhisperModelState
22
+ # Cross-attention encoder-decoder models (Whisper, CohereASR, NemotronParse, ...)
23
+ if any(isinstance(m, CrossAttention) for m in model.modules()):
24
+ from vllm.v1.worker.gpu.model_states.encoder_decoder import (
25
+ EncoderDecoderModelState,
26
+ )
27
27
- return WhisperModelState(vllm_config, model, encoder_cache, device)
28
+ return EncoderDecoderModelState(vllm_config, model, encoder_cache, device)
29
30
if vllm_config.model_config.is_hybrid:
31
from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
vllm/v1/worker/gpu/model_states/default.py
-23
@@ -7,7 +7,6 @@ import torch.nn as nn
7
8
from vllm.config import VllmConfig
9
from vllm.config.compilation import CUDAGraphMode
10
-from vllm.tasks import GenerationTask
10
from vllm.v1.core.sched.output import NewRequestData
11
from vllm.v1.kv_cache_interface import KVCacheConfig
12
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
@@ -62,28 +61,6 @@ class DefaultModelState(ModelState):
61
device=self.device,
62
)
63
65
- def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
66
- from vllm.model_executor.models.interfaces import (
67
- supports_realtime,
68
- supports_transcription,
69
- )
70
- from vllm.model_executor.models.interfaces_base import is_text_generation_model
71
-
72
- supported_tasks = list[GenerationTask]()
73
-
74
- if is_text_generation_model(self.model):
75
- supported_tasks.append("generate")
76
-
77
- if supports_transcription(self.model):
78
- if self.model.supports_transcription_only:
79
- return ("transcription",)
80
- supported_tasks.append("transcription")
81
-
82
- if supports_realtime(self.model):
83
- supported_tasks.append("realtime")
84
-
85
- return tuple(supported_tasks)
86
-
64
def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
65
if self.rope_state is not None:
66
assert new_req_data.prefill_token_ids is not None
vllm/v1/worker/gpu/model_states/encoder_decoder.py
renamed
+13
-12
@@ -23,7 +23,7 @@ from vllm.v1.worker.utils import AttentionGroup
23
24
25
@dataclass
26
-class WhisperAttnMetadata(ModelSpecificAttnMetadata):
26
+class EncoderDecoderAttnMetadata(ModelSpecificAttnMetadata):
27
encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]]
28
29
def get_extra_common_attn_kwargs(
@@ -41,7 +41,11 @@ class WhisperAttnMetadata(ModelSpecificAttnMetadata):
41
}
42
43
44
-class WhisperModelState(ModelState):
44
+class EncoderDecoderModelState(ModelState):
45
+ """ModelState for cross-attention encoder-decoder models
46
+ (Whisper, CohereASR, NemotronParse, FireRedLID, ...)
47
+ """
48
+
49
def __init__(
50
self,
51
vllm_config: VllmConfig,
@@ -80,9 +84,6 @@ class WhisperModelState(ModelState):
84
85
self.encoder_outputs: list[torch.Tensor] = []
86
83
- def get_supported_generation_tasks(self):
84
- return ("transcription",)
85
-
87
def get_mm_embeddings(
88
self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch
89
) -> None:
@@ -94,11 +95,11 @@ class WhisperModelState(ModelState):
95
encoder_inputs[req_id] = req_encoder_inputs
96
_, mm_kwargs = self.encoder_runner.prepare_mm_inputs(encoder_inputs)
97
if mm_kwargs:
97
- # Whisper consumes encoder outputs through `encoder_outputs`, not
98
- # `inputs_embeds`. Single modality (audio) so execute_mm_encoder
99
- # preserves request order; use its return value directly.
100
- # No need to store in encoder_cache: cross-attention K/V are written
101
- # to the KV cache on the first step; decode steps use the cache.
98
+ # Encoder-decoder models consume encoder outputs through the
99
+ # `encoder_outputs` forward kwarg, not `inputs_embeds`. Single modality
100
+ # so execute_mm_encoder preserves request order; use its return value
101
+ # directly. No need to store in encoder_cache: cross-attention K/V are
102
+ # written to the KV cache on the first step; decode steps use the cache.
103
self.encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
104
else:
105
# Decode steps: encoder K/V are in cross-attention KV cache.
@@ -131,7 +132,7 @@ class WhisperModelState(ModelState):
132
else:
133
num_reqs = input_batch.num_reqs
134
num_tokens = input_batch.num_tokens
134
- whisper_attn_metadata = WhisperAttnMetadata(
135
+ enc_dec_attn_metadata = EncoderDecoderAttnMetadata(
136
self._get_encoder_seq_lens(
137
input_batch.req_ids, attn_groups, for_capture, num_reqs
138
)
@@ -158,7 +159,7 @@ class WhisperModelState(ModelState):
159
kv_cache_config=kv_cache_config,
160
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
161
dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
161
- model_specific_attn_metadata=whisper_attn_metadata,
162
+ model_specific_attn_metadata=enc_dec_attn_metadata,
163
for_cudagraph_capture=for_capture,
164
)
165
return attn_metadata
vllm/v1/worker/gpu/model_states/interface.py
+18
-2
@@ -46,9 +46,25 @@ class ModelState(ABC):
46
) -> None:
47
raise NotImplementedError
48
49
- @abstractmethod
49
+ model: nn.Module
50
+
51
def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
51
- raise NotImplementedError
52
+ from vllm.model_executor.models.interfaces import (
53
+ supports_realtime,
54
+ supports_transcription,
55
+ )
56
+ from vllm.model_executor.models.interfaces_base import is_text_generation_model
57
+
58
+ supported_tasks = list[GenerationTask]()
59
+ if is_text_generation_model(self.model):
60
+ supported_tasks.append("generate")
61
+ if supports_transcription(self.model):
62
+ if self.model.supports_transcription_only:
63
+ return ("transcription",)
64
+ supported_tasks.append("transcription")
65
+ if supports_realtime(self.model):
66
+ supported_tasks.append("realtime")
67
+ return tuple(supported_tasks)
68
69
def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
70
return None