| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | from typing import Any |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | import torch.nn as nn |
| 9 | |
| 10 | from vllm.config import VllmConfig |
| 11 | from vllm.config.compilation import CUDAGraphMode |
| 12 | from vllm.utils.torch_utils import PIN_MEMORY |
| 13 | from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheConfig |
| 14 | from vllm.v1.worker.gpu.attn_utils import build_attn_metadata |
| 15 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 16 | from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache |
| 17 | from vllm.v1.worker.gpu.model_states.interface import ( |
| 18 | ModelSpecificAttnMetadata, |
| 19 | ModelState, |
| 20 | ) |
| 21 | from vllm.v1.worker.gpu.states import RequestState |
| 22 | from vllm.v1.worker.utils import AttentionGroup |
| 23 | |
| 24 | |
| 25 | @dataclass |
| 26 | class EncoderDecoderAttnMetadata(ModelSpecificAttnMetadata): |
| 27 | encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]] |
| 28 | |
| 29 | def get_extra_common_attn_kwargs( |
| 30 | self, kv_cache_group_id: int, num_reqs: int |
| 31 | ) -> dict[str, Any]: |
| 32 | encoder_seq_lens = self.encoder_seq_lens.get(kv_cache_group_id) |
| 33 | if encoder_seq_lens is None: |
| 34 | return {} |
| 35 | encoder_seq_lens_gpu, encoder_seq_lens_cpu = encoder_seq_lens |
| 36 | return { |
| 37 | "encoder_seq_lens": encoder_seq_lens_gpu[:num_reqs], |
| 38 | "encoder_seq_lens_cpu": encoder_seq_lens_cpu[:num_reqs], |
| 39 | } |
| 40 | |
| 41 | |
| 42 | class EncoderDecoderModelState(ModelState): |
| 43 | """ModelState for cross-attention encoder-decoder models |
| 44 | (Whisper, CohereASR, NemotronParse, ...) |
| 45 | """ |
| 46 | |
| 47 | def __init__( |
| 48 | self, |
| 49 | vllm_config: VllmConfig, |
| 50 | model: nn.Module, |
| 51 | encoder_cache: EncoderCache | None, |
| 52 | device: torch.device, |
| 53 | ) -> None: |
| 54 | assert encoder_cache is not None |
| 55 | if vllm_config.model_config.enable_prompt_embeds: |
| 56 | raise ValueError( |
| 57 | "--enable-prompt-embeds is not supported with encoder-decoder models." |
| 58 | ) |
| 59 | super().__init__(vllm_config, model, encoder_cache, device) |
| 60 | |
| 61 | self.max_encoder_len = getattr( |
| 62 | self.model_config.hf_config, |
| 63 | "max_source_positions", |
| 64 | self.max_model_len, |
| 65 | ) |
| 66 | self.encoder_seq_lens_gpu = torch.zeros( |
| 67 | self.max_num_reqs, dtype=torch.int32, device=self.device |
| 68 | ) |
| 69 | |
| 70 | self.encoder_outputs: list[torch.Tensor] = [] |
| 71 | |
| 72 | def prepare_inputs_embeds( |
| 73 | self, |
| 74 | scheduled_encoder_inputs: dict[str, list[int]], |
| 75 | input_batch: InputBatch, |
| 76 | req_states: RequestState, |
| 77 | ) -> None: |
| 78 | # Ensure encoder inputs are ordered consistently with input_batch.req_ids. |
| 79 | encoder_inputs: dict[str, list[int]] = {} |
| 80 | for req_id in input_batch.req_ids: |
| 81 | req_encoder_inputs = scheduled_encoder_inputs.get(req_id, []) |
| 82 | if req_encoder_inputs: |
| 83 | encoder_inputs[req_id] = req_encoder_inputs |
| 84 | _, mm_kwargs = self.encoder_runner.prepare_mm_inputs(encoder_inputs) |
| 85 | if mm_kwargs: |
| 86 | # Encoder-decoder models consume encoder outputs through the |
| 87 | # `encoder_outputs` forward kwarg, not `inputs_embeds`. Single modality |
| 88 | # so execute_mm_encoder preserves request order; use its return value |
| 89 | # directly. No need to store in encoder_cache: cross-attention K/V are |
| 90 | # written to the KV cache on the first step; decode steps use the cache. |
| 91 | with self.encoder_runner.timed_encoder_operation(encoder_inputs.keys()): |
| 92 | self.encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs) |
| 93 | else: |
| 94 | # Decode steps: encoder K/V are in cross-attention KV cache. |
| 95 | self.encoder_outputs = [] |
| 96 | return None |
| 97 | |
| 98 | def prepare_inputs( |
| 99 | self, input_batch: InputBatch, req_states: RequestState |
| 100 | ) -> dict[str, Any]: |
| 101 | model_inputs = {"encoder_outputs": self.encoder_outputs} |
| 102 | self.encoder_outputs = [] |
| 103 | return model_inputs |
| 104 | |
| 105 | def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any]: |
| 106 | return {"encoder_outputs": []} |
| 107 | |
| 108 | def prepare_attn( |
| 109 | self, |
| 110 | input_batch: InputBatch, |
| 111 | cudagraph_mode: CUDAGraphMode, |
| 112 | block_tables: tuple[torch.Tensor, ...], |
| 113 | slot_mappings: torch.Tensor, |
| 114 | attn_groups: list[list[AttentionGroup]], |
| 115 | kv_cache_config: KVCacheConfig, |
| 116 | for_capture: bool = False, |
| 117 | ubatch_idx: int = 0, |
| 118 | ) -> dict[str, Any]: |
| 119 | assert ubatch_idx == 0, "DBO is not supported" |
| 120 | if cudagraph_mode == CUDAGraphMode.FULL: |
| 121 | num_reqs = input_batch.num_reqs_after_padding |
| 122 | num_tokens = input_batch.num_tokens_after_padding |
| 123 | else: |
| 124 | num_reqs = input_batch.num_reqs |
| 125 | num_tokens = input_batch.num_tokens |
| 126 | enc_dec_attn_metadata = EncoderDecoderAttnMetadata( |
| 127 | self._get_encoder_seq_lens( |
| 128 | input_batch.req_ids, attn_groups, for_capture, num_reqs |
| 129 | ) |
| 130 | ) |
| 131 | |
| 132 | query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) |
| 133 | max_query_len = input_batch.num_scheduled_tokens.max().item() |
| 134 | seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound |
| 135 | if for_capture: |
| 136 | max_seq_len = self.max_model_len |
| 137 | else: |
| 138 | max_seq_len = int(seq_lens_cpu_upper_bound[:num_reqs].max().item()) |
| 139 | attn_metadata = build_attn_metadata( |
| 140 | attn_groups=attn_groups, |
| 141 | num_reqs=num_reqs, |
| 142 | num_tokens=num_tokens, |
| 143 | query_start_loc_gpu=input_batch.query_start_loc, |
| 144 | query_start_loc_cpu=query_start_loc_cpu, |
| 145 | max_query_len=max_query_len, |
| 146 | seq_lens=input_batch.seq_lens, |
| 147 | max_seq_len=max_seq_len, |
| 148 | block_tables=block_tables, |
| 149 | slot_mappings=slot_mappings, |
| 150 | kv_cache_config=kv_cache_config, |
| 151 | seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, |
| 152 | dcp_local_seq_lens=input_batch.dcp_local_seq_lens, |
| 153 | model_specific_attn_metadata=enc_dec_attn_metadata, |
| 154 | for_cudagraph_capture=for_capture, |
| 155 | rswa_prefix_lens=input_batch.prompt_lens, |
| 156 | ) |
| 157 | return attn_metadata |
| 158 | |
| 159 | def _get_encoder_seq_lens( |
| 160 | self, |
| 161 | req_ids: list[str], |
| 162 | attn_groups: list[list[AttentionGroup]], |
| 163 | for_capture: bool, |
| 164 | num_reqs: int, |
| 165 | ) -> dict[int, tuple[torch.Tensor, np.ndarray]]: |
| 166 | encoder_seq_lens = torch.zeros( |
| 167 | num_reqs, dtype=torch.int32, pin_memory=PIN_MEMORY |
| 168 | ) |
| 169 | encoder_seq_lens_np = encoder_seq_lens.numpy() |
| 170 | if not for_capture: |
| 171 | # During normal execution, use actual encoder lengths. |
| 172 | for i, req_id in enumerate(req_ids): |
| 173 | mm_features = self.encoder_cache.mm_features.get(req_id, []) |
| 174 | encoder_seq_lens_np[i] = sum( |
| 175 | feature.mm_position.get_num_embeds() for feature in mm_features |
| 176 | ) |
| 177 | else: |
| 178 | # During CUDA graph capture, use max encoder length so max_seqlen_k |
| 179 | # is captured with the correct value for cross-attention. |
| 180 | encoder_seq_lens_np[:] = self.max_encoder_len |
| 181 | |
| 182 | self.encoder_seq_lens_gpu[:num_reqs].copy_(encoder_seq_lens, non_blocking=True) |
| 183 | self.encoder_seq_lens_gpu[num_reqs:].fill_(0) |
| 184 | encoder_seq_lens_gpu = self.encoder_seq_lens_gpu[:num_reqs] |
| 185 | |
| 186 | seq_lens_by_group: dict[int, tuple[torch.Tensor, np.ndarray]] = {} |
| 187 | for kv_cache_group_idx, groups in enumerate(attn_groups): |
| 188 | has_cross_attn = any( |
| 189 | isinstance(attn_group.kv_cache_spec, CrossAttentionSpec) |
| 190 | for attn_group in groups |
| 191 | ) |
| 192 | if has_cross_attn: |
| 193 | seq_lens_by_group[kv_cache_group_idx] = ( |
| 194 | encoder_seq_lens_gpu, |
| 195 | encoder_seq_lens_np, |
| 196 | ) |
| 197 | return seq_lens_by_group |