diff --git a/vllm/v1/worker/gpu/model_states/__init__.py b/vllm/v1/worker/gpu/model_states/__init__.py index e24c7e9b1c..dc52dc4ee5 100644 --- a/vllm/v1/worker/gpu/model_states/__init__.py +++ b/vllm/v1/worker/gpu/model_states/__init__.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn from vllm.config import VllmConfig +from vllm.model_executor.layers.attention import CrossAttention from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache @@ -18,13 +19,13 @@ def init_model_state( cls = model.get_model_state_cls() return cls(vllm_config, model, encoder_cache, device) - if ( - "WhisperForConditionalGeneration" in vllm_config.model_config.architectures - or "CohereAsrForConditionalGeneration" in vllm_config.model_config.architectures - ): - from vllm.v1.worker.gpu.model_states.whisper import WhisperModelState + # Cross-attention encoder-decoder models (Whisper, CohereASR, NemotronParse, ...) + if any(isinstance(m, CrossAttention) for m in model.modules()): + from vllm.v1.worker.gpu.model_states.encoder_decoder import ( + EncoderDecoderModelState, + ) - return WhisperModelState(vllm_config, model, encoder_cache, device) + return EncoderDecoderModelState(vllm_config, model, encoder_cache, device) if vllm_config.model_config.is_hybrid: from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index ee5d9384fa..4d71b4d5b3 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -7,7 +7,6 @@ import torch.nn as nn from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode -from vllm.tasks import GenerationTask from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_attn_metadata @@ -62,28 +61,6 @@ class DefaultModelState(ModelState): device=self.device, ) - def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]: - from vllm.model_executor.models.interfaces import ( - supports_realtime, - supports_transcription, - ) - from vllm.model_executor.models.interfaces_base import is_text_generation_model - - supported_tasks = list[GenerationTask]() - - if is_text_generation_model(self.model): - supported_tasks.append("generate") - - if supports_transcription(self.model): - if self.model.supports_transcription_only: - return ("transcription",) - supported_tasks.append("transcription") - - if supports_realtime(self.model): - supported_tasks.append("realtime") - - return tuple(supported_tasks) - def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: if self.rope_state is not None: assert new_req_data.prefill_token_ids is not None diff --git a/vllm/v1/worker/gpu/model_states/whisper.py b/vllm/v1/worker/gpu/model_states/encoder_decoder.py similarity index 90% rename from vllm/v1/worker/gpu/model_states/whisper.py rename to vllm/v1/worker/gpu/model_states/encoder_decoder.py index df432d82ce..2a07e4a1a3 100644 --- a/vllm/v1/worker/gpu/model_states/whisper.py +++ b/vllm/v1/worker/gpu/model_states/encoder_decoder.py @@ -23,7 +23,7 @@ from vllm.v1.worker.utils import AttentionGroup @dataclass -class WhisperAttnMetadata(ModelSpecificAttnMetadata): +class EncoderDecoderAttnMetadata(ModelSpecificAttnMetadata): encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]] def get_extra_common_attn_kwargs( @@ -41,7 +41,11 @@ class WhisperAttnMetadata(ModelSpecificAttnMetadata): } -class WhisperModelState(ModelState): +class EncoderDecoderModelState(ModelState): + """ModelState for cross-attention encoder-decoder models + (Whisper, CohereASR, NemotronParse, FireRedLID, ...) + """ + def __init__( self, vllm_config: VllmConfig, @@ -80,9 +84,6 @@ class WhisperModelState(ModelState): self.encoder_outputs: list[torch.Tensor] = [] - def get_supported_generation_tasks(self): - return ("transcription",) - def get_mm_embeddings( self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch ) -> None: @@ -94,11 +95,11 @@ class WhisperModelState(ModelState): encoder_inputs[req_id] = req_encoder_inputs _, mm_kwargs = self.encoder_runner.prepare_mm_inputs(encoder_inputs) if mm_kwargs: - # Whisper consumes encoder outputs through `encoder_outputs`, not - # `inputs_embeds`. Single modality (audio) so execute_mm_encoder - # preserves request order; use its return value directly. - # No need to store in encoder_cache: cross-attention K/V are written - # to the KV cache on the first step; decode steps use the cache. + # Encoder-decoder models consume encoder outputs through the + # `encoder_outputs` forward kwarg, not `inputs_embeds`. Single modality + # so execute_mm_encoder preserves request order; use its return value + # directly. No need to store in encoder_cache: cross-attention K/V are + # written to the KV cache on the first step; decode steps use the cache. self.encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs) else: # Decode steps: encoder K/V are in cross-attention KV cache. @@ -131,7 +132,7 @@ class WhisperModelState(ModelState): else: num_reqs = input_batch.num_reqs num_tokens = input_batch.num_tokens - whisper_attn_metadata = WhisperAttnMetadata( + enc_dec_attn_metadata = EncoderDecoderAttnMetadata( self._get_encoder_seq_lens( input_batch.req_ids, attn_groups, for_capture, num_reqs ) @@ -158,7 +159,7 @@ class WhisperModelState(ModelState): kv_cache_config=kv_cache_config, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, dcp_local_seq_lens=input_batch.dcp_local_seq_lens, - model_specific_attn_metadata=whisper_attn_metadata, + model_specific_attn_metadata=enc_dec_attn_metadata, for_cudagraph_capture=for_capture, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 86f28e08ea..be631bb94b 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -46,9 +46,25 @@ class ModelState(ABC): ) -> None: raise NotImplementedError - @abstractmethod + model: nn.Module + def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]: - raise NotImplementedError + from vllm.model_executor.models.interfaces import ( + supports_realtime, + supports_transcription, + ) + from vllm.model_executor.models.interfaces_base import is_text_generation_model + + supported_tasks = list[GenerationTask]() + if is_text_generation_model(self.model): + supported_tasks.append("generate") + if supports_transcription(self.model): + if self.model.supports_transcription_only: + return ("transcription",) + supported_tasks.append("transcription") + if supports_realtime(self.model): + supported_tasks.append("realtime") + return tuple(supported_tasks) def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: return None