| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import Any |
| 4 | |
| 5 | import torch |
| 6 | import torch.nn as nn |
| 7 | |
| 8 | from vllm.config import VllmConfig, get_layers_from_vllm_config |
| 9 | from vllm.config.compilation import CUDAGraphMode |
| 10 | from vllm.model_executor.layers.attention import Attention |
| 11 | from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE |
| 12 | from vllm.v1.attention.backend import ( |
| 13 | AttentionCGSupport, |
| 14 | AttentionType, |
| 15 | CommonAttentionMetadata, |
| 16 | ) |
| 17 | from vllm.v1.core.sched.output import NewRequestData |
| 18 | from vllm.v1.kv_cache_interface import EncoderOnlyAttentionSpec, KVCacheConfig |
| 19 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 20 | from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache |
| 21 | from vllm.v1.worker.gpu.model_states.default import DefaultModelState |
| 22 | from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata |
| 23 | from vllm.v1.worker.gpu.states import RequestState |
| 24 | from vllm.v1.worker.utils import AttentionGroup |
| 25 | |
| 26 | |
| 27 | class EncoderOnlyModelState(DefaultModelState): |
| 28 | """ModelState for encoder-only (BERT/RoBERTa) models. |
| 29 | |
| 30 | Encoder attention needs no KV cache: it runs full, bidirectional |
| 31 | self-attention over each request's tokens in a single forward. Such layers |
| 32 | return no ``get_kv_cache_spec`` and therefore never enter |
| 33 | ``kv_cache_config.kv_cache_groups``, so the KV-backed attention path never |
| 34 | builds their metadata. We build their (non-causal) metadata here, keeping |
| 35 | the normal KV-backed path untouched. |
| 36 | """ |
| 37 | |
| 38 | # The V2 pooling path is not wired for prompt embeds. |
| 39 | supports_prompt_embeds = False |
| 40 | |
| 41 | def __init__( |
| 42 | self, |
| 43 | vllm_config: VllmConfig, |
| 44 | model: nn.Module, |
| 45 | encoder_cache: EncoderCache | None, |
| 46 | device: torch.device, |
| 47 | ): |
| 48 | super().__init__(vllm_config, model, encoder_cache, device) |
| 49 | |
| 50 | cache_config = vllm_config.cache_config |
| 51 | if cache_config.cache_dtype == "auto": |
| 52 | kv_cache_dtype = self.dtype |
| 53 | else: |
| 54 | kv_cache_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype] |
| 55 | |
| 56 | # Build an attention group (and its non-causal metadata builder) for the |
| 57 | # encoder-only layers, grouped by backend + query/KV head config. Models |
| 58 | # are typically uniform, yielding a single group. |
| 59 | attn_layers = get_layers_from_vllm_config(vllm_config, Attention) |
| 60 | groups: dict[tuple[tuple[str, str], int, int, int], AttentionGroup] = {} |
| 61 | for name, layer in attn_layers.items(): |
| 62 | if layer.attn_type != AttentionType.ENCODER_ONLY: |
| 63 | continue |
| 64 | # No KV cache is bound for these layers; give them an (empty) |
| 65 | # device tensor so the attention forward context is well-formed. |
| 66 | layer.kv_cache = torch.empty(0, dtype=kv_cache_dtype, device=device) |
| 67 | backend = layer.get_attn_backend() |
| 68 | # Metadata builders require a uniform query-head count per group. |
| 69 | key = ( |
| 70 | backend.full_cls_name(), |
| 71 | layer.num_heads, |
| 72 | layer.num_kv_heads, |
| 73 | layer.head_size, |
| 74 | ) |
| 75 | group = groups.get(key) |
| 76 | if group is None: |
| 77 | spec = EncoderOnlyAttentionSpec( |
| 78 | block_size=cache_config.block_size, |
| 79 | num_kv_heads=layer.num_kv_heads, |
| 80 | head_size=layer.head_size, |
| 81 | dtype=kv_cache_dtype, |
| 82 | ) |
| 83 | group = AttentionGroup(backend, [], spec, kv_cache_group_id=len(groups)) |
| 84 | groups[key] = group |
| 85 | group.layer_names.append(name) |
| 86 | |
| 87 | self.encoder_attn_groups = list(groups.values()) |
| 88 | for group in self.encoder_attn_groups: |
| 89 | group.create_metadata_builders(vllm_config, device) |
| 90 | |
| 91 | # Encoder attention reads neither the block table nor the slot mapping |
| 92 | # (full varlen self-attention over the batch's q/k/v), but the metadata |
| 93 | # builder expects both tensors to be present. |
| 94 | self._dummy_block_table = torch.zeros( |
| 95 | self.max_num_reqs, 1, dtype=torch.int32, device=device |
| 96 | ) |
| 97 | self._dummy_slot_mapping = torch.zeros( |
| 98 | self.max_num_tokens, dtype=torch.int64, device=device |
| 99 | ) |
| 100 | self.token_type_ids: dict[str, torch.Tensor] = {} |
| 101 | |
| 102 | def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: |
| 103 | super().add_request(req_index, new_req_data) |
| 104 | pooling_params = new_req_data.pooling_params |
| 105 | if pooling_params is None or pooling_params.extra_kwargs is None: |
| 106 | return |
| 107 | |
| 108 | token_type_start = pooling_params.extra_kwargs.get("compressed_token_type_ids") |
| 109 | if token_type_start is not None: |
| 110 | assert new_req_data.prompt_token_ids is not None |
| 111 | self.token_type_ids[new_req_data.req_id] = ( |
| 112 | torch.arange(len(new_req_data.prompt_token_ids), dtype=torch.int32) |
| 113 | >= token_type_start |
| 114 | ).to(torch.int32) |
| 115 | |
| 116 | def remove_request(self, req_id: str) -> None: |
| 117 | super().remove_request(req_id) |
| 118 | self.token_type_ids.pop(req_id, None) |
| 119 | |
| 120 | def prepare_inputs( |
| 121 | self, input_batch: InputBatch, req_states: RequestState |
| 122 | ) -> dict[str, torch.Tensor | None]: |
| 123 | model_inputs = super().prepare_inputs(input_batch, req_states) |
| 124 | if not self.token_type_ids: |
| 125 | return model_inputs |
| 126 | |
| 127 | token_type_ids_cpu = torch.zeros( |
| 128 | input_batch.num_tokens_after_padding, |
| 129 | dtype=torch.int32, |
| 130 | pin_memory=PIN_MEMORY, |
| 131 | ) |
| 132 | offset = 0 |
| 133 | for i, req_id in enumerate(input_batch.req_ids): |
| 134 | num_tokens = int(input_batch.num_scheduled_tokens[i]) |
| 135 | request_token_type_ids = self.token_type_ids.get(req_id) |
| 136 | if request_token_type_ids is not None: |
| 137 | start = int(input_batch.num_computed_tokens_np[i]) |
| 138 | token_type_ids_cpu[offset : offset + num_tokens].copy_( |
| 139 | request_token_type_ids[start : start + num_tokens] |
| 140 | ) |
| 141 | offset += num_tokens |
| 142 | model_inputs["token_type_ids"] = token_type_ids_cpu.to( |
| 143 | self.device, non_blocking=True |
| 144 | ) |
| 145 | return model_inputs |
| 146 | |
| 147 | def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]: |
| 148 | # Encoder groups are built here rather than in init_attn_backend, so |
| 149 | # their cudagraph support must be surfaced to the runner separately. |
| 150 | support = AttentionCGSupport.ALWAYS |
| 151 | backend: str | None = None |
| 152 | for group in self.encoder_attn_groups: |
| 153 | builder = group.get_metadata_builder(0) |
| 154 | cg_support = builder.get_cudagraph_support( |
| 155 | self.vllm_config, group.kv_cache_spec |
| 156 | ) |
| 157 | if cg_support.value < support.value: |
| 158 | support = cg_support |
| 159 | backend = group.backend.__name__ |
| 160 | return support, backend |
| 161 | |
| 162 | def prepare_attn( |
| 163 | self, |
| 164 | input_batch: InputBatch, |
| 165 | cudagraph_mode: CUDAGraphMode, |
| 166 | block_tables: tuple[torch.Tensor, ...], |
| 167 | slot_mappings: torch.Tensor, |
| 168 | attn_groups: list[list[AttentionGroup]], |
| 169 | kv_cache_config: KVCacheConfig, |
| 170 | for_capture: bool = False, |
| 171 | ubatch_idx: int = 0, |
| 172 | model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, |
| 173 | ) -> dict[str, Any]: |
| 174 | assert ubatch_idx == 0, "DBO is not supported" |
| 175 | attn_metadata = super().prepare_attn( |
| 176 | input_batch, |
| 177 | cudagraph_mode, |
| 178 | block_tables, |
| 179 | slot_mappings, |
| 180 | attn_groups, |
| 181 | kv_cache_config, |
| 182 | for_capture, |
| 183 | model_specific_attn_metadata=model_specific_attn_metadata, |
| 184 | ) |
| 185 | attn_metadata.update( |
| 186 | self._build_encoder_attn_metadata(input_batch, cudagraph_mode, for_capture) |
| 187 | ) |
| 188 | return attn_metadata |
| 189 | |
| 190 | def _build_encoder_attn_metadata( |
| 191 | self, input_batch: InputBatch, cudagraph_mode: CUDAGraphMode, for_capture: bool |
| 192 | ) -> dict[str, Any]: |
| 193 | if cudagraph_mode == CUDAGraphMode.FULL: |
| 194 | num_reqs = input_batch.num_reqs_after_padding |
| 195 | num_tokens = input_batch.num_tokens_after_padding |
| 196 | else: |
| 197 | num_reqs = input_batch.num_reqs |
| 198 | num_tokens = input_batch.num_tokens |
| 199 | max_query_len = int(input_batch.num_scheduled_tokens.max()) |
| 200 | if for_capture: |
| 201 | max_seq_len = self.max_model_len |
| 202 | else: |
| 203 | max_seq_len = int(input_batch.seq_lens_cpu_upper_bound[:num_reqs].max()) |
| 204 | |
| 205 | # The encoder builder forces ``causal=False`` regardless of this value. |
| 206 | common_attn_metadata = CommonAttentionMetadata( |
| 207 | query_start_loc=input_batch.query_start_loc, |
| 208 | query_start_loc_cpu=torch.from_numpy(input_batch.query_start_loc_np), |
| 209 | seq_lens=input_batch.seq_lens[:num_reqs], |
| 210 | num_reqs=num_reqs, |
| 211 | num_actual_tokens=num_tokens, |
| 212 | max_query_len=max_query_len, |
| 213 | max_seq_len=max_seq_len, |
| 214 | block_table_tensor=self._dummy_block_table[:num_reqs], |
| 215 | slot_mapping=self._dummy_slot_mapping[:num_tokens], |
| 216 | seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound[:num_reqs], |
| 217 | positions=input_batch.positions, |
| 218 | ) |
| 219 | |
| 220 | attn_metadata: dict[str, Any] = {} |
| 221 | for group in self.encoder_attn_groups: |
| 222 | builder = group.get_metadata_builder(0) |
| 223 | if for_capture: |
| 224 | metadata = builder.build_for_cudagraph_capture(common_attn_metadata) |
| 225 | else: |
| 226 | metadata = builder.build( |
| 227 | common_prefix_len=0, common_attn_metadata=common_attn_metadata |
| 228 | ) |
| 229 | for name in group.layer_names: |
| 230 | attn_metadata[name] = metadata |
| 231 | return attn_metadata |