[MRV2] Generalize use of WhisperModelState (#46096)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

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

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

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

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