| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | import torch.nn as nn |
| 5 | |
| 6 | from vllm.config import VllmConfig, get_layers_from_vllm_config |
| 7 | from vllm.model_executor.layers.attention import Attention, CrossAttention |
| 8 | from vllm.v1.attention.backend import AttentionType |
| 9 | from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache |
| 10 | from vllm.v1.worker.gpu.model_states.interface import ModelState |
| 11 | |
| 12 | |
| 13 | def init_model_state( |
| 14 | vllm_config: VllmConfig, |
| 15 | model: nn.Module, |
| 16 | encoder_cache: EncoderCache | None, |
| 17 | device: torch.device, |
| 18 | ) -> ModelState: |
| 19 | cls = resolve_model_state_cls(vllm_config, model) |
| 20 | |
| 21 | # Reject enable_prompt_embeds for states that would silently ignore it. |
| 22 | if vllm_config.model_config.enable_prompt_embeds and not cls.supports_prompt_embeds: |
| 23 | raise ValueError(f"--enable-prompt-embeds not supported with {cls.__name__}.") |
| 24 | |
| 25 | return cls(vllm_config, model, encoder_cache, device) |
| 26 | |
| 27 | |
| 28 | def resolve_model_state_cls( |
| 29 | vllm_config: VllmConfig, model: nn.Module |
| 30 | ) -> type[ModelState]: |
| 31 | # Let the model provide its own ModelState if it defines one. |
| 32 | if hasattr(model, "get_model_state_cls"): |
| 33 | return model.get_model_state_cls() |
| 34 | |
| 35 | # Cross-attention encoder-decoder models (Whisper, CohereASR, NemotronParse, ...) |
| 36 | if any(isinstance(m, CrossAttention) for m in model.modules()): |
| 37 | from vllm.v1.worker.gpu.model_states.encoder_decoder import ( |
| 38 | EncoderDecoderModelState, |
| 39 | ) |
| 40 | |
| 41 | return EncoderDecoderModelState |
| 42 | |
| 43 | # Encoder-only attention is non-causal and needs no KV cache. |
| 44 | if any( |
| 45 | layer.attn_type == AttentionType.ENCODER_ONLY |
| 46 | for layer in get_layers_from_vllm_config(vllm_config, Attention).values() |
| 47 | ): |
| 48 | from vllm.v1.worker.gpu.model_states.encoder_only import EncoderOnlyModelState |
| 49 | |
| 50 | return EncoderOnlyModelState |
| 51 | |
| 52 | if vllm_config.model_config.is_hybrid or vllm_config.model_config.is_attention_free: |
| 53 | from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState |
| 54 | |
| 55 | return MambaHybridModelState |
| 56 | |
| 57 | from vllm.v1.worker.gpu.model_states.default import DefaultModelState |
| 58 | |
| 59 | return DefaultModelState |