| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Mapping, Sequence |
| 4 | from contextlib import AbstractContextManager, nullcontext |
| 5 | from dataclasses import dataclass |
| 6 | from typing import TYPE_CHECKING, Any, cast |
| 7 | |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | |
| 11 | from vllm.config import ( |
| 12 | VllmConfig, |
| 13 | get_layers_from_vllm_config, |
| 14 | ) |
| 15 | from vllm.config.compilation import CUDAGraphMode |
| 16 | from vllm.model_executor.layers.attention import Attention |
| 17 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 18 | from vllm.multimodal.inputs import MultiModalFeatureSpec |
| 19 | from vllm.v1.attention.backend import ( |
| 20 | AttentionCGSupport, |
| 21 | CommonAttentionMetadata, |
| 22 | ) |
| 23 | from vllm.v1.attention.backends.utils import create_fast_prefill_custom_backend |
| 24 | from vllm.v1.hisparse.binding import ( |
| 25 | init_hisparse_kv_cache, |
| 26 | resolve_hisparse_specs, |
| 27 | ) |
| 28 | from vllm.v1.kv_cache_interface import ( |
| 29 | AttentionSpec, |
| 30 | KVCacheConfig, |
| 31 | KVCacheSpec, |
| 32 | MambaSpec, |
| 33 | UniformTypeKVCacheSpecs, |
| 34 | ) |
| 35 | from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata |
| 36 | from vllm.v1.worker.ubatch_utils import get_num_ubatches |
| 37 | from vllm.v1.worker.utils import ( |
| 38 | AttentionGroup, |
| 39 | add_kv_sharing_layers_to_kv_cache_groups, |
| 40 | allocate_kv_cache, |
| 41 | bind_kv_cache_to_layers, |
| 42 | prepare_kernel_block_sizes, |
| 43 | ) |
| 44 | |
| 45 | if TYPE_CHECKING: |
| 46 | from vllm.v1.worker.gpu.block_table import BlockTables |
| 47 | from vllm.v1.worker.gpu.cudagraph_utils import ( |
| 48 | BatchExecutionDescriptor, |
| 49 | CudaGraphManager, |
| 50 | ) |
| 51 | |
| 52 | |
| 53 | @dataclass(frozen=True) |
| 54 | class AttentionCGSupportInfo: |
| 55 | min_cg_support: AttentionCGSupport = AttentionCGSupport.ALWAYS |
| 56 | min_cg_attn_backend: str | None = None |
| 57 | |
| 58 | def narrow( |
| 59 | self, support: AttentionCGSupport, backend: str | None |
| 60 | ) -> "AttentionCGSupportInfo": |
| 61 | """Return an info tightened by ``support`` if it is more restrictive. |
| 62 | |
| 63 | Lets attention groups built outside ``init_attn_backend`` (e.g. |
| 64 | encoder-only layers) contribute to the runner's cudagraph decision. |
| 65 | """ |
| 66 | if support.value < self.min_cg_support.value: |
| 67 | return AttentionCGSupportInfo(support, backend) |
| 68 | return self |
| 69 | |
| 70 | |
| 71 | @dataclass(frozen=True) |
| 72 | class FastPrefillBatchMetadata: |
| 73 | """Per-step inputs for the KV-sharing fast prefill path.""" |
| 74 | |
| 75 | logits_indices_padded: torch.Tensor |
| 76 | num_logits_indices: int |
| 77 | max_logits_per_req: int |
| 78 | |
| 79 | |
| 80 | class FastPrefillHelper: |
| 81 | """Decides per step whether to arm the KV-sharing fast prefill path, and |
| 82 | stages the logits indices it runs on. |
| 83 | """ |
| 84 | |
| 85 | def __init__(self, cudagraph_manager: "CudaGraphManager", max_num_tokens: int): |
| 86 | self.max_num_tokens = max_num_tokens |
| 87 | self.cudagraph_manager = cudagraph_manager |
| 88 | self.logits_indices_buf = torch.zeros( |
| 89 | max_num_tokens, dtype=torch.int32, device=cudagraph_manager.device |
| 90 | ) |
| 91 | |
| 92 | def prepare( |
| 93 | self, |
| 94 | logits_indices: torch.Tensor, |
| 95 | num_reqs: int, |
| 96 | cu_num_logits_np: np.ndarray, |
| 97 | has_prefill: bool, |
| 98 | batch_desc: "BatchExecutionDescriptor", |
| 99 | num_active_loras: int = 0, |
| 100 | ) -> FastPrefillBatchMetadata | None: |
| 101 | if ( |
| 102 | not has_prefill |
| 103 | or batch_desc.cg_mode == CUDAGraphMode.FULL |
| 104 | or batch_desc.num_ubatches > 1 |
| 105 | ): |
| 106 | return None |
| 107 | buf = self.logits_indices_buf |
| 108 | num_logits = logits_indices.shape[0] |
| 109 | assert num_logits > 0 |
| 110 | buf[:num_logits].copy_(logits_indices) |
| 111 | # There might be leftover indices in buf[num_logits:] from previous |
| 112 | # iterations. Broadcast the scalar GPU-side to keep them valid. |
| 113 | buf[num_logits:] = logits_indices[-1] |
| 114 | # Pad so the model's KV-sharing layers run their reduced (logits-only) |
| 115 | # batch at a captured piecewise cudagraph size. Pre-capture, or with |
| 116 | # cudagraphs off, dispatch returns the unpadded count. |
| 117 | desc = self.cudagraph_manager.dispatch( |
| 118 | num_reqs=num_reqs, |
| 119 | num_tokens=num_logits, |
| 120 | uniform_token_count=None, |
| 121 | num_active_loras=num_active_loras, |
| 122 | ) |
| 123 | num_logits_padded = min(desc.num_tokens, self.max_num_tokens) |
| 124 | return FastPrefillBatchMetadata( |
| 125 | logits_indices_padded=buf[:num_logits_padded], |
| 126 | num_logits_indices=num_logits, |
| 127 | # Largest per-request logits count, known on the host. |
| 128 | max_logits_per_req=int(np.diff(cu_num_logits_np).max()), |
| 129 | ) |
| 130 | |
| 131 | |
| 132 | def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: |
| 133 | kv_cache_spec: dict[str, KVCacheSpec] = {} |
| 134 | layer_type = cast(type[Any], AttentionLayerBase) |
| 135 | attn_layers = get_layers_from_vllm_config(vllm_config, layer_type) |
| 136 | for layer_name, attn_module in attn_layers.items(): |
| 137 | if getattr(attn_module, "kv_sharing_target_layer_name", None): |
| 138 | # This layer will use KV cache of the sharing target layer. |
| 139 | continue |
| 140 | # Skip modules that don't need KV cache (eg encoder-only attention) |
| 141 | if spec := attn_module.get_kv_cache_spec(vllm_config): |
| 142 | if isinstance(spec, AttentionSpec): |
| 143 | spec = attn_module.get_attn_backend().customize_spec(spec) |
| 144 | kv_cache_spec[layer_name] = spec |
| 145 | if vllm_config.attention_config.hisparse_config is not None: |
| 146 | kv_cache_spec = resolve_hisparse_specs(vllm_config, kv_cache_spec, attn_layers) |
| 147 | return kv_cache_spec |
| 148 | |
| 149 | |
| 150 | def get_shared_kv_cache_layers(vllm_config: VllmConfig): |
| 151 | attn_layers = get_layers_from_vllm_config(vllm_config, Attention) |
| 152 | return { |
| 153 | layer_name: kv_tgt_layer |
| 154 | for layer_name, attn_module in attn_layers.items() |
| 155 | if (kv_tgt_layer := attn_module.kv_sharing_target_layer_name) |
| 156 | } |
| 157 | |
| 158 | |
| 159 | def get_kv_sharing_fast_prefill_eligible_layers( |
| 160 | vllm_config: VllmConfig, draft_layer_names: set[str] | None = None |
| 161 | ) -> set[str]: |
| 162 | """Trailing run of KV-sharing layers, eligible for fast prefill. |
| 163 | |
| 164 | In You Only Cache Once (https://arxiv.org/abs/2405.05254) or other similar |
| 165 | KV sharing setups, only the layers that generate KV caches are involved in |
| 166 | the prefill phase, enabling prefill to early exit. Layers are registered in |
| 167 | execution order, so the eligible layers are the contiguous suffix of |
| 168 | KV-sharing layers. |
| 169 | |
| 170 | Speculator draft layers register after the target model's layers (and may |
| 171 | themselves share KV), so they are excluded from the walk. |
| 172 | """ |
| 173 | if not vllm_config.cache_config.kv_sharing_fast_prefill: |
| 174 | return set() |
| 175 | shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config) |
| 176 | if not shared_kv_cache_layers: |
| 177 | return set() |
| 178 | eligible_layers: set[str] = set() |
| 179 | attn_layers = get_layers_from_vllm_config(vllm_config, Attention) |
| 180 | for layer_name in reversed(attn_layers): |
| 181 | if draft_layer_names is not None and layer_name in draft_layer_names: |
| 182 | continue |
| 183 | if layer_name not in shared_kv_cache_layers: |
| 184 | break |
| 185 | eligible_layers.add(layer_name) |
| 186 | return eligible_layers |
| 187 | |
| 188 | |
| 189 | def init_attn_backend( |
| 190 | kv_cache_config: KVCacheConfig, |
| 191 | vllm_config: VllmConfig, |
| 192 | device: torch.device, |
| 193 | active_layer_names: set[str] | None = None, |
| 194 | draft_layer_names: set[str] | None = None, |
| 195 | ) -> tuple[list[list[AttentionGroup]], AttentionCGSupportInfo, list[int]]: |
| 196 | # Phase 1: discover attention groups for each kv cache group. |
| 197 | attn_groups: list[list[AttentionGroup]] = [] |
| 198 | |
| 199 | # Add KV-sharing layers to their target's kv cache group so they are |
| 200 | # discovered alongside the target layer in Phase 1 below. |
| 201 | add_kv_sharing_layers_to_kv_cache_groups( |
| 202 | get_shared_kv_cache_layers(vllm_config), kv_cache_config.kv_cache_groups |
| 203 | ) |
| 204 | fast_prefill_eligible_layers = get_kv_sharing_fast_prefill_eligible_layers( |
| 205 | vllm_config, draft_layer_names |
| 206 | ) |
| 207 | |
| 208 | # Phase 1: discover attention groups for each kv cache group. |
| 209 | for kv_cache_group_id, kv_cache_group_spec in enumerate( |
| 210 | kv_cache_config.kv_cache_groups |
| 211 | ): |
| 212 | layer_names = kv_cache_group_spec.layer_names |
| 213 | if not kv_cache_group_spec.kv_cache_spec.has_layer_views: |
| 214 | attn_groups.append([]) |
| 215 | continue |
| 216 | if active_layer_names is not None: |
| 217 | layer_names = list(active_layer_names.intersection(layer_names)) |
| 218 | |
| 219 | layer_type = cast(type[Any], AttentionLayerBase) |
| 220 | attn_layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names) |
| 221 | |
| 222 | group_map: dict[tuple[tuple[str, str], KVCacheSpec, int], AttentionGroup] = {} |
| 223 | group_order: list[tuple[tuple[str, str], KVCacheSpec, int]] = [] |
| 224 | |
| 225 | for layer_name in layer_names: |
| 226 | attn_backend = attn_layers[layer_name].get_attn_backend() |
| 227 | if layer_name in fast_prefill_eligible_layers: |
| 228 | attn_backend = create_fast_prefill_custom_backend( |
| 229 | "FastPrefill", attn_backend |
| 230 | ) |
| 231 | |
| 232 | layer_kv_cache_spec: KVCacheSpec = kv_cache_group_spec.kv_cache_spec |
| 233 | if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs): |
| 234 | layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[layer_name] |
| 235 | |
| 236 | # Split on per-rank num_heads_q so layers with different Q-head |
| 237 | # counts (e.g. a spec-decode draft head and its target) get separate |
| 238 | # metadata builders. |
| 239 | num_heads_q = getattr(attn_layers[layer_name], "num_heads", 0) |
| 240 | key = (attn_backend.full_cls_name(), layer_kv_cache_spec, num_heads_q) |
| 241 | if key not in group_map: |
| 242 | group_map[key] = AttentionGroup( |
| 243 | attn_backend, [layer_name], layer_kv_cache_spec, kv_cache_group_id |
| 244 | ) |
| 245 | group_order.append(key) |
| 246 | else: |
| 247 | group_map[key].layer_names.append(layer_name) |
| 248 | |
| 249 | attn_groups.append([group_map[key] for key in group_order]) |
| 250 | |
| 251 | # Phase 2: pick a kernel block size per kv cache group that is supported |
| 252 | # by all backends within that group. |
| 253 | kernel_block_sizes = prepare_kernel_block_sizes(kv_cache_config, attn_groups) |
| 254 | |
| 255 | # Phase 3: create metadata builders and determine cudagraph support. |
| 256 | attn_backend_workspace: torch.Tensor | None = None |
| 257 | for kv_cache_group_id, groups in enumerate(attn_groups): |
| 258 | kernel_block_size = None |
| 259 | if kv_cache_group_id < len(kernel_block_sizes): |
| 260 | kernel_block_size = kernel_block_sizes[kv_cache_group_id] |
| 261 | for group in groups: |
| 262 | group.create_metadata_builders( |
| 263 | vllm_config=vllm_config, |
| 264 | device=device, |
| 265 | kernel_block_size=kernel_block_size, |
| 266 | # Microbatches build attention metadata concurrently, and some |
| 267 | # builders keep the prepared metadata on themselves (MLA stores |
| 268 | # it on the prefill backend), so each ubatch needs its own. |
| 269 | num_metadata_builders=get_num_ubatches(vllm_config.parallel_config), |
| 270 | ) |
| 271 | # The microbatches' builders share the workspace: they all issue |
| 272 | # attention on the one compute stream the threads hand off, so the |
| 273 | # buffer is written serially, as it already is across steps. |
| 274 | for builder in group.metadata_builders: |
| 275 | if attn_backend_workspace is None: |
| 276 | if hasattr(builder, "_get_workspace_buffer"): |
| 277 | attn_backend_workspace = builder._get_workspace_buffer() |
| 278 | elif hasattr(builder, "set_workspace_buffer"): |
| 279 | builder.set_workspace_buffer(attn_backend_workspace) |
| 280 | attn_cg_support_info = get_attn_cg_support(attn_groups, vllm_config) |
| 281 | return attn_groups, attn_cg_support_info, kernel_block_sizes |
| 282 | |
| 283 | |
| 284 | def get_attn_cg_support( |
| 285 | attn_groups: list[list[AttentionGroup]], |
| 286 | vllm_config: VllmConfig, |
| 287 | checked_layer_names: set[str] | None = None, |
| 288 | ) -> AttentionCGSupportInfo: |
| 289 | """Return the weakest CUDA graph support among the checked layers.""" |
| 290 | min_cg_support = AttentionCGSupport.ALWAYS |
| 291 | min_cg_attn_backend = None |
| 292 | for groups in attn_groups: |
| 293 | for group in groups: |
| 294 | if checked_layer_names is not None and checked_layer_names.isdisjoint( |
| 295 | group.layer_names |
| 296 | ): |
| 297 | continue |
| 298 | builder = group.get_metadata_builder(0) |
| 299 | cg_support = builder.get_cudagraph_support( |
| 300 | vllm_config, |
| 301 | group.kv_cache_spec, |
| 302 | ) |
| 303 | if cg_support.value < min_cg_support.value: |
| 304 | min_cg_support = cg_support |
| 305 | min_cg_attn_backend = group.backend.__name__ |
| 306 | return AttentionCGSupportInfo( |
| 307 | min_cg_support=min_cg_support, |
| 308 | min_cg_attn_backend=min_cg_attn_backend, |
| 309 | ) |
| 310 | |
| 311 | |
| 312 | def get_query_lens_mismatch_unsupported_backend( |
| 313 | attn_groups: list[list[AttentionGroup]], |
| 314 | checked_layer_names: set[str] | None = None, |
| 315 | ) -> str | None: |
| 316 | """Name the first backend needing the CPU query lengths to be exact, if any. |
| 317 | |
| 318 | The attention selector already excludes these when adaptive verification is |
| 319 | enabled, but models that hard-wire their backend never consult it. See |
| 320 | AttentionBackend.supports_device_cpu_query_lens_mismatch(). |
| 321 | """ |
| 322 | for groups in attn_groups: |
| 323 | for group in groups: |
| 324 | if checked_layer_names is not None and checked_layer_names.isdisjoint( |
| 325 | group.layer_names |
| 326 | ): |
| 327 | continue |
| 328 | if not group.backend.supports_device_cpu_query_lens_mismatch(): |
| 329 | return group.backend.__name__ |
| 330 | return None |
| 331 | |
| 332 | |
| 333 | def get_varlen_cudagraph_unsupported_backend( |
| 334 | attn_groups: list[list[AttentionGroup]], |
| 335 | vllm_config: VllmConfig, |
| 336 | max_query_len: int, |
| 337 | checked_layer_names: set[str] | None = None, |
| 338 | ) -> tuple[str, int | None] | None: |
| 339 | """Name the first backend whose FULL cudagraphs cannot replay decode batches |
| 340 | of 1 to max_query_len tokens per request, with its bound, if any. |
| 341 | |
| 342 | ALWAYS builders replay any batch; the others need a bound of at least |
| 343 | max_query_len. See AttentionMetadataBuilder.get_varlen_cudagraph_max_query_len(). |
| 344 | """ |
| 345 | for groups in attn_groups: |
| 346 | for group in groups: |
| 347 | if checked_layer_names is not None and checked_layer_names.isdisjoint( |
| 348 | group.layer_names |
| 349 | ): |
| 350 | continue |
| 351 | builder = group.get_metadata_builder(0) |
| 352 | spec = group.kv_cache_spec |
| 353 | support = builder.get_cudagraph_support(vllm_config, spec) |
| 354 | if support == AttentionCGSupport.ALWAYS: |
| 355 | continue |
| 356 | bound = ( |
| 357 | # NEVER means never, whatever an inherited override says. |
| 358 | None |
| 359 | if support == AttentionCGSupport.NEVER |
| 360 | else builder.get_varlen_cudagraph_max_query_len(vllm_config, spec) |
| 361 | ) |
| 362 | if bound is None or bound < max_query_len: |
| 363 | return group.backend.__name__, bound |
| 364 | return None |
| 365 | |
| 366 | |
| 367 | def init_kv_cache( |
| 368 | forward_context: dict[str, Any], |
| 369 | kv_cache_config: KVCacheConfig, |
| 370 | device: torch.device, |
| 371 | kernel_block_sizes: list[int], |
| 372 | vllm_config: VllmConfig, |
| 373 | kv_cache_allocation_context: AbstractContextManager | None = None, |
| 374 | *, |
| 375 | block_tables: "BlockTables | None" = None, |
| 376 | ) -> dict[str, Any]: |
| 377 | allocation_context = kv_cache_allocation_context or nullcontext() |
| 378 | with allocation_context: |
| 379 | if vllm_config.attention_config.hisparse_config is not None: |
| 380 | assert block_tables is not None |
| 381 | kv_caches = init_hisparse_kv_cache( |
| 382 | kv_cache_config, |
| 383 | device, |
| 384 | kernel_block_sizes, |
| 385 | vllm_config, |
| 386 | forward_context, |
| 387 | block_tables, |
| 388 | ) |
| 389 | else: |
| 390 | kv_caches = allocate_kv_cache( |
| 391 | kv_cache_config, |
| 392 | device, |
| 393 | vllm_config.cache_config.get_resolved_kv_cache_layout(), |
| 394 | kernel_block_sizes, |
| 395 | ) |
| 396 | for layer_name, target in get_shared_kv_cache_layers(vllm_config).items(): |
| 397 | kv_caches[layer_name] = kv_caches[target] |
| 398 | # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per |
| 399 | # decoder layer, so a layer name carries two integers (layer + module index). |
| 400 | num_attn_module = ( |
| 401 | 2 |
| 402 | if vllm_config.model_config.hf_config.model_type |
| 403 | in ("longcat_flash", "longcat_flash_ngram") |
| 404 | else 1 |
| 405 | ) |
| 406 | bindable_caches = { |
| 407 | name: cache for name, cache in kv_caches.items() if name in forward_context |
| 408 | } |
| 409 | bind_kv_cache_to_layers( |
| 410 | bindable_caches, |
| 411 | forward_context, |
| 412 | num_attn_module, |
| 413 | kv_cache_groups=kv_cache_config.kv_cache_groups, |
| 414 | ) |
| 415 | return kv_caches |
| 416 | |
| 417 | |
| 418 | def build_slot_mappings_by_layer( |
| 419 | slot_mappings: torch.Tensor, kv_cache_config: KVCacheConfig |
| 420 | ) -> dict[str, torch.Tensor]: |
| 421 | slot_mappings_by_layer: dict[str, torch.Tensor] = {} |
| 422 | kv_cache_groups = kv_cache_config.kv_cache_groups |
| 423 | for slot_mapping, kv_cache_group in zip(slot_mappings, kv_cache_groups): |
| 424 | for layer_name in kv_cache_group.layer_names: |
| 425 | slot_mappings_by_layer[layer_name] = slot_mapping |
| 426 | return slot_mappings_by_layer |
| 427 | |
| 428 | |
| 429 | def build_attn_metadata( |
| 430 | attn_groups: list[list[AttentionGroup]], |
| 431 | num_reqs: int, |
| 432 | num_tokens: int, |
| 433 | query_start_loc_gpu: torch.Tensor, |
| 434 | query_start_loc_cpu: torch.Tensor, |
| 435 | max_query_len: int, |
| 436 | seq_lens: torch.Tensor, |
| 437 | max_seq_len: int, |
| 438 | block_tables: Sequence[torch.Tensor], |
| 439 | slot_mappings: torch.Tensor, |
| 440 | kv_cache_config: KVCacheConfig, |
| 441 | seq_lens_cpu_upper_bound: torch.Tensor | None = None, |
| 442 | dcp_local_seq_lens: torch.Tensor | None = None, |
| 443 | dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None, |
| 444 | positions: torch.Tensor | None = None, |
| 445 | is_prefilling: torch.Tensor | None = None, |
| 446 | mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None, |
| 447 | model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, |
| 448 | for_cudagraph_capture: bool = False, |
| 449 | causal: bool | torch.Tensor | Mapping[int, bool] = True, |
| 450 | rswa_prefix_lens: torch.Tensor | None = None, |
| 451 | ubatch_idx: int = 0, |
| 452 | fast_prefill: FastPrefillBatchMetadata | None = None, |
| 453 | req_idx: np.ndarray | None = None, |
| 454 | ) -> dict[str, Any]: |
| 455 | seq_lens = seq_lens[:num_reqs] |
| 456 | if dcp_local_seq_lens is not None: |
| 457 | dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs] |
| 458 | if dcp_local_seq_lens_cpu_upper_bound is not None: |
| 459 | dcp_local_seq_lens_cpu_upper_bound = dcp_local_seq_lens_cpu_upper_bound[ |
| 460 | :num_reqs |
| 461 | ] |
| 462 | if seq_lens_cpu_upper_bound is not None: |
| 463 | seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound[:num_reqs] |
| 464 | |
| 465 | attn_metadata: dict[str, Any] = {} |
| 466 | token_to_req_indices: torch.Tensor | None = None |
| 467 | # Mamba groups with the same spec and builder differ only in their state |
| 468 | # indices, so later groups re-gather those from the first group's metadata. |
| 469 | # Also at capture, so FULL graphs share the batch-level buffers. |
| 470 | cached_metadata: dict[tuple[KVCacheSpec, type], Any] = {} |
| 471 | num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) |
| 472 | for i in range(num_kv_cache_groups): |
| 473 | if not attn_groups[i]: |
| 474 | continue |
| 475 | block_table = block_tables[i] |
| 476 | slot_mapping = slot_mappings[i] |
| 477 | # Per-group causal for hybrid drafters (mixed SWA/full attention). |
| 478 | group_causal = ( |
| 479 | causal if isinstance(causal, (bool, torch.Tensor)) else causal.get(i, True) |
| 480 | ) |
| 481 | |
| 482 | common_attn_metadata_extra_kwargs = ( |
| 483 | model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs) |
| 484 | if model_specific_attn_metadata is not None |
| 485 | else {} |
| 486 | ) |
| 487 | # Model-specific metadata (e.g. Mamba hybrid) may supply its own |
| 488 | # padding-aware is_prefilling, which takes precedence over the default. |
| 489 | group_is_prefilling = common_attn_metadata_extra_kwargs.pop( |
| 490 | "is_prefilling", is_prefilling |
| 491 | ) |
| 492 | if fast_prefill is not None: |
| 493 | common_attn_metadata_extra_kwargs.update( |
| 494 | logits_indices_padded=fast_prefill.logits_indices_padded, |
| 495 | num_logits_indices=fast_prefill.num_logits_indices, |
| 496 | max_logits_per_req=fast_prefill.max_logits_per_req, |
| 497 | ) |
| 498 | common_attn_metadata = CommonAttentionMetadata( |
| 499 | query_start_loc=query_start_loc_gpu, |
| 500 | query_start_loc_cpu=query_start_loc_cpu, |
| 501 | seq_lens=seq_lens, |
| 502 | seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, |
| 503 | max_seq_len=max_seq_len, |
| 504 | num_reqs=num_reqs, |
| 505 | num_actual_tokens=num_tokens, |
| 506 | max_query_len=max_query_len, |
| 507 | block_table_tensor=block_table, |
| 508 | slot_mapping=slot_mapping, |
| 509 | causal=group_causal, |
| 510 | dcp_local_seq_lens=dcp_local_seq_lens, |
| 511 | dcp_local_seq_lens_cpu_upper_bound=dcp_local_seq_lens_cpu_upper_bound, |
| 512 | positions=positions, |
| 513 | is_prefilling=group_is_prefilling, |
| 514 | mm_req_doc_ranges=mm_req_doc_ranges, |
| 515 | rswa_prefix_lens=rswa_prefix_lens, |
| 516 | req_idx=req_idx, |
| 517 | _token_to_req_indices_cache=token_to_req_indices, |
| 518 | **common_attn_metadata_extra_kwargs, |
| 519 | ) |
| 520 | |
| 521 | for attn_group in attn_groups[i]: |
| 522 | attn_metadata_builder = attn_group.get_metadata_builder(ubatch_idx) |
| 523 | reuse_key = None |
| 524 | if attn_metadata_builder.supports_update_block_table and isinstance( |
| 525 | attn_group.kv_cache_spec, MambaSpec |
| 526 | ): |
| 527 | reuse_key = (attn_group.kv_cache_spec, type(attn_metadata_builder)) |
| 528 | if reuse_key in cached_metadata: |
| 529 | metadata = attn_metadata_builder.update_block_table( |
| 530 | cached_metadata[reuse_key], block_table, slot_mapping |
| 531 | ) |
| 532 | elif for_cudagraph_capture: |
| 533 | metadata = attn_metadata_builder.build_for_cudagraph_capture( |
| 534 | common_attn_metadata |
| 535 | ) |
| 536 | else: |
| 537 | attn_metadata_extra_kwargs = ( |
| 538 | model_specific_attn_metadata.get_extra_attn_kwargs( |
| 539 | attn_metadata_builder, |
| 540 | num_reqs, |
| 541 | ) |
| 542 | if model_specific_attn_metadata is not None |
| 543 | else {} |
| 544 | ) |
| 545 | metadata = attn_metadata_builder.build( |
| 546 | common_prefix_len=0, |
| 547 | common_attn_metadata=common_attn_metadata, |
| 548 | **attn_metadata_extra_kwargs, |
| 549 | ) |
| 550 | if reuse_key is not None and reuse_key not in cached_metadata: |
| 551 | cached_metadata[reuse_key] = metadata |
| 552 | for layer_name in attn_group.layer_names: |
| 553 | attn_metadata[layer_name] = metadata |
| 554 | token_to_req_indices = common_attn_metadata._token_to_req_indices_cache |
| 555 | return attn_metadata |
| 556 | |
| 557 | |
| 558 | def compute_mm_prefix_ranges( |
| 559 | req_ids: list[str], |
| 560 | mm_features: dict[str, list[MultiModalFeatureSpec]], |
| 561 | sliding_window: int | None = None, |
| 562 | ) -> dict[int, list[tuple[int, int]]]: |
| 563 | """Compute PrefixLM bidirectional ranges for multimodal tokens. |
| 564 | |
| 565 | Ranges exceeding sliding_window are skipped to prevent early tokens |
| 566 | from attending across the entire image span. |
| 567 | """ |
| 568 | req_doc_ranges: dict[int, list[tuple[int, int]]] = {} |
| 569 | for req_idx, req_id in enumerate(req_ids): |
| 570 | image_doc_ranges = [] |
| 571 | for mm_feature in mm_features.get(req_id, ()): |
| 572 | if mm_feature.modality not in ("image", "video"): |
| 573 | continue |
| 574 | for r in mm_feature.mm_position.extract_embeds_range(): |
| 575 | if sliding_window is not None and (r[1] - r[0] + 1) > sliding_window: |
| 576 | continue |
| 577 | image_doc_ranges.append(r) |
| 578 | req_doc_ranges[req_idx] = image_doc_ranges |
| 579 | return req_doc_ranges |