diff --git a/vllm/v1/worker/gpu/model_states/__init__.py b/vllm/v1/worker/gpu/model_states/__init__.py index 06b5a92c395..b096fcaf5e6 100644 --- a/vllm/v1/worker/gpu/model_states/__init__.py +++ b/vllm/v1/worker/gpu/model_states/__init__.py @@ -13,7 +13,10 @@ def init_model_state( encoder_cache: EncoderCache | None, device: torch.device, ): - if "WhisperForConditionalGeneration" in vllm_config.model_config.architectures: + 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 return WhisperModelState(vllm_config, model, encoder_cache, device)