# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, replace from math import prod from typing import Any, cast import torch from vllm.config import ( VllmConfig, get_layers_from_vllm_config, set_current_vllm_config, ) from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.utils.torch_utils import get_dtype_size from vllm.v1.attention.backend import ( AttentionCGSupport, CommonAttentionMetadata, ) from vllm.v1.kv_cache_interface import ( AttentionSpec, KVCacheConfig, KVCacheSpec, KVQuantMode, MambaSpec, TQFullAttentionSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata from vllm.v1.worker.utils import ( AttentionGroup, add_kv_sharing_layers_to_kv_cache_groups, bind_kv_cache, prepare_kernel_block_sizes, ) logger = init_logger(__name__) @dataclass(frozen=True) class AttentionCGSupportInfo: min_cg_support: AttentionCGSupport = AttentionCGSupport.ALWAYS min_cg_attn_backend: str | None = None def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]: kv_cache_spec: dict[str, KVCacheSpec] = {} layer_type = cast(type[Any], AttentionLayerBase) attn_layers = get_layers_from_vllm_config(vllm_config, layer_type) for layer_name, attn_module in attn_layers.items(): if getattr(attn_module, "kv_sharing_target_layer_name", None): # This layer will use KV cache of the sharing target layer. continue # Skip modules that don't need KV cache (eg encoder-only attention) if spec := attn_module.get_kv_cache_spec(vllm_config): if isinstance(spec, AttentionSpec): backend = attn_module.get_attn_backend() # indexes_kv_by_block_stride() -> get_kv_cache_stride_order() -> # get_kv_cache_layout() needs the current vLLM config. with set_current_vllm_config(vllm_config): indexes = backend.indexes_kv_by_block_stride() spec = replace(spec, indexes_kv_by_block_stride=indexes) kv_cache_spec[layer_name] = spec return kv_cache_spec def get_shared_kv_cache_layers(vllm_config: VllmConfig): attn_layers = get_layers_from_vllm_config(vllm_config, Attention) return { layer_name: kv_tgt_layer for layer_name, attn_module in attn_layers.items() if (kv_tgt_layer := attn_module.kv_sharing_target_layer_name) } def init_attn_backend( kv_cache_config: KVCacheConfig, vllm_config: VllmConfig, device: torch.device, active_layer_names: set[str] | None = None, ) -> tuple[list[list[AttentionGroup]], AttentionCGSupportInfo, list[int]]: # Phase 1: discover attention groups for each kv cache group. attn_groups: list[list[AttentionGroup]] = [] # Add KV-sharing layers to their target's kv cache group so they are # discovered alongside the target layer in Phase 1 below. add_kv_sharing_layers_to_kv_cache_groups( get_shared_kv_cache_layers(vllm_config), kv_cache_config.kv_cache_groups ) # Phase 1: discover attention groups for each kv cache group. for kv_cache_group_id, kv_cache_group_spec in enumerate( kv_cache_config.kv_cache_groups ): layer_names = kv_cache_group_spec.layer_names if active_layer_names is not None: layer_names = list(active_layer_names.intersection(layer_names)) layer_type = cast(type[Any], AttentionLayerBase) attn_layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names) group_map: dict[tuple[tuple[str, str], KVCacheSpec, int], AttentionGroup] = {} group_order: list[tuple[tuple[str, str], KVCacheSpec, int]] = [] for layer_name in layer_names: attn_backend = attn_layers[layer_name].get_attn_backend() layer_kv_cache_spec: KVCacheSpec = kv_cache_group_spec.kv_cache_spec if isinstance(layer_kv_cache_spec, UniformTypeKVCacheSpecs): layer_kv_cache_spec = layer_kv_cache_spec.kv_cache_specs[layer_name] # Split on per-rank num_heads_q so layers with different Q-head # counts (e.g. a spec-decode draft head and its target) get separate # metadata builders. num_heads_q = getattr(attn_layers[layer_name], "num_heads", 0) key = (attn_backend.full_cls_name(), layer_kv_cache_spec, num_heads_q) if key not in group_map: group_map[key] = AttentionGroup( attn_backend, [layer_name], layer_kv_cache_spec, kv_cache_group_id ) group_order.append(key) else: group_map[key].layer_names.append(layer_name) attn_groups.append([group_map[key] for key in group_order]) # Phase 2: pick a kernel block size per kv cache group that is supported # by all backends within that group. kernel_block_sizes = prepare_kernel_block_sizes(kv_cache_config, attn_groups) # Phase 3: create metadata builders and determine cudagraph support. attn_backend_workspace: torch.Tensor | None = None min_cg_support = AttentionCGSupport.ALWAYS min_cg_attn_backend = None for kv_cache_group_id, groups in enumerate(attn_groups): kernel_block_size = None if kv_cache_group_id < len(kernel_block_sizes): kernel_block_size = kernel_block_sizes[kv_cache_group_id] for group in groups: group.create_metadata_builders( vllm_config=vllm_config, device=device, kernel_block_size=kernel_block_size, num_metadata_builders=1, ) builder = group.get_metadata_builder(0) if attn_backend_workspace is None: if hasattr(builder, "_get_workspace_buffer"): attn_backend_workspace = builder._get_workspace_buffer() else: if hasattr(builder, "set_workspace_buffer"): builder.set_workspace_buffer(attn_backend_workspace) # Check cudagraph support for the attention backend cg_support = builder.get_cudagraph_support( vllm_config, cast(AttentionSpec, group.kv_cache_spec), ) if cg_support.value < min_cg_support.value: min_cg_support = cg_support min_cg_attn_backend = group.backend.__name__ attn_cg_support_info = AttentionCGSupportInfo( min_cg_support=min_cg_support, min_cg_attn_backend=min_cg_attn_backend ) return attn_groups, attn_cg_support_info, kernel_block_sizes def _allocate_kv_cache( kv_cache_config: KVCacheConfig, shared_layers: dict[str, str], device: torch.device ): kv_cache_raw_tensors: dict[str, torch.Tensor] = {} packed_backing: torch.Tensor | None = None for kv_cache_tensor in kv_cache_config.kv_cache_tensors: if kv_cache_tensor.block_stride > 0: # Allocate once; all packed tensors alias the same backing. if packed_backing is None: packed_backing = torch.zeros( kv_cache_tensor.size, dtype=torch.int8, device=device ) tensor = packed_backing else: tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device) for layer_name in kv_cache_tensor.shared_by: kv_cache_raw_tensors[layer_name] = tensor layer_names = set() for group in kv_cache_config.kv_cache_groups: for layer_name in group.layer_names: layer_names.add(layer_name) assert layer_names == (kv_cache_raw_tensors.keys() | shared_layers.keys()), ( "Some layers are not correctly initialized" ) return kv_cache_raw_tensors def _reshape_attention_kv_cache( kv_raw_tensor: torch.Tensor, kv_cache_spec: AttentionSpec, kv_cache_shape: tuple[int, ...], kv_cache_stride_order: tuple[int, ...], num_blocks: int, packing: tuple[int, int] | None, ) -> torch.Tensor: permuted_kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order) inv_order = [ kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order)) ] dtype = kv_cache_spec.dtype if packing is not None: offset, block_stride = packing assert inv_order[0] == 0 page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype) kv_cache = ( kv_raw_tensor.view(-1, block_stride)[:, offset : offset + page_bytes] .view(dtype) .view(permuted_kv_cache_shape) ) elif kv_cache_spec.page_size_padded is not None: # Use a strided view to skip the padding between physical pages. # # Only num-blocks-first layouts are supported (the block dimension is # dim 0 of the unpermuted shape). kv-first layouts such as ROCm's # ``(2, num_blocks, ...)`` are intentionally not supported here. For a # num-blocks-first layout the only stride that must change is the block # stride: every other (contiguous) stride already steps within the # unpadded region of a page, so no further adjustment is needed. assert kv_cache_shape[0] == num_blocks, ( "Padded KV pages require a num-blocks-first KV cache layout (got " f"shape {kv_cache_shape} with num_blocks={num_blocks}); " "kv-first layouts are not supported." ) dtype_size = get_dtype_size(kv_cache_spec.dtype) page_stride = kv_cache_spec.page_size_bytes // dtype_size num_blocks_dim = inv_order[0] strides = list(torch.empty(permuted_kv_cache_shape, device="meta").stride()) strides[num_blocks_dim] = page_stride kv_cache = torch.as_strided( kv_raw_tensor.view(dtype), size=permuted_kv_cache_shape, stride=tuple(strides), ) else: # No padding — safe to use a contiguous view. kv_cache = kv_raw_tensor.view(dtype).view(permuted_kv_cache_shape) return kv_cache.permute(*inv_order) def _reshape_kv_cache( attn_groups: Sequence[AttentionGroup], kv_cache_raw_tensors: dict[str, torch.Tensor], cache_dtype: str, kernel_block_sizes: list[int], shared_kv_cache_layers: dict[str, str], kv_cache_config: "KVCacheConfig | None" = None, ) -> dict[str, Any]: kv_caches: dict[str, Any] = {} has_attn, has_mamba = False, False layer_packing: dict[str, tuple[int, int]] = {} if kv_cache_config is not None: for kv_tensor in kv_cache_config.kv_cache_tensors: if kv_tensor.block_stride > 0: for ln in kv_tensor.shared_by: layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride) for group in attn_groups: if group.kv_cache_group_id >= len(kernel_block_sizes): continue kv_cache_spec = group.kv_cache_spec if kv_cache_spec.storage_block_size != kv_cache_spec.block_size: # use storage_block_size as the kernel block size for groups # that apply a compression on block size (eg. DeepSeek V4). kernel_block_size = kv_cache_spec.storage_block_size else: kernel_block_size = kernel_block_sizes[group.kv_cache_group_id] for layer_name in group.layer_names: if layer_name in shared_kv_cache_layers: # Shared layer — tensor will be aliased to its target later. continue kv_raw_tensor = kv_cache_raw_tensors[layer_name] packing = layer_packing.get(layer_name) if packing is not None: _, blk_stride = packing num_blocks = kv_raw_tensor.numel() // blk_stride else: assert kv_raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0 num_blocks = kv_raw_tensor.numel() // kv_cache_spec.page_size_bytes if isinstance(kv_cache_spec, AttentionSpec): has_attn = True # Use storage_block_size: it equals block_size for uncompressed # specs but is smaller for compressed ones (DeepSeek V4), which # store block_size tokens in block_size // compress_ratio slots. num_blocks_per_kv_block = ( kv_cache_spec.storage_block_size // kernel_block_size ) kernel_num_blocks = num_blocks * num_blocks_per_kv_block # Skipped layers (--kv-cache-dtype-skip-layers) keep the # unquantized shape; only the quantized primary uses the # quantized cache dtype's (possibly packed) layout. layer_cache_dtype = ( "auto" if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE and not isinstance(kv_cache_spec, TQFullAttentionSpec) else cache_dtype ) kv_cache_shape = group.backend.get_kv_cache_shape( kernel_num_blocks, kernel_block_size, kv_cache_spec.num_kv_heads, kv_cache_spec.head_size, cache_dtype_str=layer_cache_dtype, ) # FIXME(woosuk): Add kv_cache_stride_order to all attention backends. try: kv_cache_stride_order = group.backend.get_kv_cache_stride_order() assert len(kv_cache_stride_order) == len(kv_cache_shape) except (AttributeError, NotImplementedError): kv_cache_stride_order = tuple(range(len(kv_cache_shape))) kv_caches[layer_name] = _reshape_attention_kv_cache( kv_raw_tensor, kv_cache_spec, kv_cache_shape, kv_cache_stride_order, kernel_num_blocks, packing, ) elif isinstance(kv_cache_spec, MambaSpec): has_mamba = True state_tensors = [] storage_offset_bytes = 0 for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes): dtype_size = get_dtype_size(dtype) num_element_per_page = kv_cache_spec.page_size_bytes // dtype_size target_shape = (num_blocks, *shape) stride = torch.empty(target_shape).stride() target_stride = (num_element_per_page, *stride[1:]) assert storage_offset_bytes % dtype_size == 0 tensor = torch.as_strided( kv_raw_tensor.view(dtype), size=target_shape, stride=target_stride, storage_offset=storage_offset_bytes // dtype_size, ) state_tensors.append(tensor) storage_offset_bytes += stride[0] * dtype_size kv_caches[layer_name] = state_tensors else: raise NotImplementedError( f"Unsupported KV cache spec type: {type(kv_cache_spec)}" ) if has_attn and has_mamba: _update_hybrid_attention_layout( attn_groups=attn_groups, kv_caches=kv_caches, kernel_block_sizes=kernel_block_sizes, cache_dtype=cache_dtype, ) elif has_attn and kv_cache_config is not None: _align_mixed_attention_kv_cache_views( attn_groups=attn_groups, kv_caches=kv_caches, kernel_block_sizes=kernel_block_sizes, cache_dtype=cache_dtype, kv_cache_config=kv_cache_config, ) # Map any sharing layers to their target layer's KV cache. for layer_name, target_layer_name in shared_kv_cache_layers.items(): kv_caches[layer_name] = kv_caches[target_layer_name] return kv_caches def _align_mixed_attention_kv_cache_views( attn_groups: Iterable[AttentionGroup], kv_caches: dict[str, Any], kernel_block_sizes: list[int], cache_dtype: str, kv_cache_config: KVCacheConfig, ) -> None: """Align shared attention KV views when backends disagree on layout. Encoder-decoder models can share one raw allocation between decoder self-attention (K/V-first ROCM_ATTN, block dim 1) and cross-attention (blocks-first backends, block dim 0). Keep the physical storage in the K/V-first layout expected by ROCM_ATTN, and restride the blocks-first logical views so block IDs address the same bytes. """ block_dims_by_layer: dict[str, int] = {} for group in attn_groups: kv_cache_spec = group.kv_cache_spec if not isinstance(kv_cache_spec, AttentionSpec): continue if group.kv_cache_group_id >= len(kernel_block_sizes): continue block_dim = group.backend.get_kv_cache_block_dim( kernel_block_sizes[group.kv_cache_group_id], kv_cache_spec.num_kv_heads, kv_cache_spec.head_size, cache_dtype_str=cache_dtype, ) for layer_name in group.layer_names: if layer_name in kv_caches: block_dims_by_layer[layer_name] = block_dim for kv_tensor in kv_cache_config.kv_cache_tensors: if kv_tensor.block_stride > 0: continue shared_block_dims = { block_dims_by_layer[layer_name] for layer_name in kv_tensor.shared_by if layer_name in block_dims_by_layer } if 0 not in shared_block_dims or 1 not in shared_block_dims: continue for layer_name in kv_tensor.shared_by: if block_dims_by_layer.get(layer_name) == 0: _restride_blocks_first_kv_cache_to_kv_first_storage( kv_caches[layer_name] ) def _restride_blocks_first_kv_cache_to_kv_first_storage( kv_cache: torch.Tensor, ) -> None: assert kv_cache.ndim >= 3 assert kv_cache.shape[1] == 2 page_size = kv_cache.shape[2:].numel() num_blocks = kv_cache.shape[0] expected_tail_stride = torch.empty(kv_cache.shape[2:]).stride() if kv_cache.stride()[2:] != expected_tail_stride: logger.warning_once( "Skipping mixed KV-cache layout alignment for a non-NHD " "blocks-first attention view with stride %s.", kv_cache.stride(), ) return kv_cache.as_strided_( size=kv_cache.shape, stride=(page_size, num_blocks * page_size, *expected_tail_stride), ) def _update_hybrid_attention_layout( attn_groups: Iterable[AttentionGroup], kv_caches: dict[str, Any], kernel_block_sizes: list[int], cache_dtype: str, ) -> None: for group in attn_groups: if group.kv_cache_group_id >= len(kernel_block_sizes): continue kv_cache_spec = group.kv_cache_spec if not isinstance(kv_cache_spec, AttentionSpec): continue # Mirror the per-layer dtype selection used when building the shape # above. The block-dim index is dtype-independent for current backends # (quantization only changes the last dim), so this is a no-op today, # but it keeps both call sites consistent for skip layers. layer_cache_dtype = ( "auto" if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE and not isinstance(kv_cache_spec, TQFullAttentionSpec) else cache_dtype ) block_dim = group.backend.get_kv_cache_block_dim( kernel_block_sizes[group.kv_cache_group_id], kv_cache_spec.num_kv_heads, kv_cache_spec.head_size, cache_dtype_str=layer_cache_dtype, ) # if the first dim of the kvcache's layout is already num_blocks, continue if block_dim == 0: continue assert block_dim == 1, ( "Expected the dim `num_blocks` at the second dim when updating" " the kvcache's layout of full attention layer" ) for layer_name in group.layer_names: if layer_name not in kv_caches: # Shared layer — will be aliased to its target after this pass. continue kv_cache = kv_caches[layer_name] if kv_cache.shape[0] == 2: assert kv_cache.shape[1] != 2, ( f"Cannot determine layout for tensor of shape {kv_cache.shape}" ) hidden_size = kv_cache.shape[2:].numel() kv_cache.as_strided_( size=kv_cache.shape, stride=( hidden_size, 2 * hidden_size, *kv_cache.stride()[2:], ), ) def init_kv_cache( runner_kv_caches: list[torch.Tensor | list[torch.Tensor]], forward_context: dict[str, Any], kv_cache_config: KVCacheConfig, attn_groups: list[list[AttentionGroup]], device: torch.device, cache_dtype: str, kernel_block_sizes: list[int], vllm_config: VllmConfig, ) -> dict[str, Any]: shared_kv_cache_layers = get_shared_kv_cache_layers(vllm_config) kv_cache_raw_tensors = _allocate_kv_cache( kv_cache_config, shared_kv_cache_layers, device ) flattened_attn_groups = list(group for groups in attn_groups for group in groups) kv_caches = _reshape_kv_cache( attn_groups=flattened_attn_groups, kv_cache_raw_tensors=kv_cache_raw_tensors, kernel_block_sizes=kernel_block_sizes, cache_dtype=cache_dtype, shared_kv_cache_layers=shared_kv_cache_layers, kv_cache_config=kv_cache_config, ) # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per # decoder layer, so a layer name carries two integers (layer + module index). num_attn_module = ( 2 if vllm_config.model_config.hf_config.model_type in ("longcat_flash", "longcat_flash_ngram") else 1 ) bind_kv_cache(kv_caches, forward_context, runner_kv_caches, num_attn_module) return kv_caches def build_slot_mappings_by_layer( slot_mappings: torch.Tensor, kv_cache_config: KVCacheConfig ) -> dict[str, torch.Tensor]: slot_mappings_by_layer: dict[str, torch.Tensor] = {} kv_cache_groups = kv_cache_config.kv_cache_groups for slot_mapping, kv_cache_group in zip(slot_mappings, kv_cache_groups): for layer_name in kv_cache_group.layer_names: slot_mappings_by_layer[layer_name] = slot_mapping return slot_mappings_by_layer def build_attn_metadata( attn_groups: list[list[AttentionGroup]], num_reqs: int, num_tokens: int, query_start_loc_gpu: torch.Tensor, query_start_loc_cpu: torch.Tensor, max_query_len: int, seq_lens: torch.Tensor, max_seq_len: int, block_tables: Sequence[torch.Tensor], slot_mappings: torch.Tensor, kv_cache_config: KVCacheConfig, seq_lens_cpu_upper_bound: torch.Tensor | None = None, dcp_local_seq_lens: torch.Tensor | None = None, positions: torch.Tensor | None = None, mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None, model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, for_cudagraph_capture: bool = False, causal: bool | torch.Tensor | Mapping[int, bool] = True, rswa_prefix_lens: torch.Tensor | None = None, ) -> dict[str, Any]: seq_lens = seq_lens[:num_reqs] if dcp_local_seq_lens is not None: dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs] if seq_lens_cpu_upper_bound is not None: seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound[:num_reqs] attn_metadata: dict[str, Any] = {} num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) for i in range(num_kv_cache_groups): block_table = block_tables[i] slot_mapping = slot_mappings[i] # Per-group causal for hybrid drafters (mixed SWA/full attention). group_causal = ( causal if isinstance(causal, (bool, torch.Tensor)) else causal.get(i, True) ) common_attn_metadata_extra_kwargs = ( model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs) if model_specific_attn_metadata is not None else {} ) common_attn_metadata = CommonAttentionMetadata( query_start_loc=query_start_loc_gpu, query_start_loc_cpu=query_start_loc_cpu, seq_lens=seq_lens, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, max_seq_len=max_seq_len, num_reqs=num_reqs, num_actual_tokens=num_tokens, max_query_len=max_query_len, block_table_tensor=block_table, slot_mapping=slot_mapping, causal=group_causal, dcp_local_seq_lens=dcp_local_seq_lens, positions=positions, mm_req_doc_ranges=mm_req_doc_ranges, rswa_prefix_lens=rswa_prefix_lens, **common_attn_metadata_extra_kwargs, ) for attn_group in attn_groups[i]: attn_metadata_builder = attn_group.get_metadata_builder(0) if for_cudagraph_capture: metadata = attn_metadata_builder.build_for_cudagraph_capture( common_attn_metadata ) else: attn_metadata_extra_kwargs = ( model_specific_attn_metadata.get_extra_attn_kwargs( attn_metadata_builder, num_reqs, ) if model_specific_attn_metadata is not None else {} ) metadata = attn_metadata_builder.build( common_prefix_len=0, common_attn_metadata=common_attn_metadata, **attn_metadata_extra_kwargs, ) for layer_name in attn_group.layer_names: attn_metadata[layer_name] = metadata return attn_metadata def compute_mm_prefix_ranges( req_ids: list[str], mm_features: dict[str, list[MultiModalFeatureSpec]], sliding_window: int | None = None, ) -> dict[int, list[tuple[int, int]]]: """Compute PrefixLM bidirectional ranges for multimodal tokens. Ranges exceeding sliding_window are skipped to prevent early tokens from attending across the entire image span. """ req_doc_ranges: dict[int, list[tuple[int, int]]] = {} for req_idx, req_id in enumerate(req_ids): image_doc_ranges = [] for mm_feature in mm_features.get(req_id, ()): if mm_feature.modality not in ("image", "video"): continue for r in mm_feature.mm_position.extract_embeds_range(): if sliding_window is not None and (r[1] - r[0] + 1) > sliding_window: continue image_doc_ranges.append(r) req_doc_ranges[req_idx] = image_doc_ranges return req_doc_ranges