diff --git a/tests/kernels/attention/test_flashinfer.py b/tests/kernels/attention/test_flashinfer.py index c56af3a6e852..176cf0b39f4a 100644 --- a/tests/kernels/attention/test_flashinfer.py +++ b/tests/kernels/attention/test_flashinfer.py @@ -705,6 +705,142 @@ def test_flashinfer_decode_with_paged_fp8_kv( ) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +def test_packed_flashinfer_indices_refresh_on_graph_replay(): + from vllm.v1.attention.backends.flashinfer import _copy_page_indices_kernel + from vllm.v1.attention.ops.flashinfer_repage import ( + PackedKVPageGeometry, + repage_block_table, + ) + + cache = torch.empty((4, 2, 2, 128, 128), device="cuda", dtype=torch.bfloat16)[:, 1] + geometry = PackedKVPageGeometry.from_cache(cache, 32) + storage = torch.full((3, 10), 123, device="cuda", dtype=torch.int32) + source = storage[1:, 1:9] + source[0] = torch.tensor([0, 1, 2, 3, 8, 9, -1, 11], device="cuda") + lengths = torch.tensor([256, 0], device="cuda", dtype=torch.int32) + destination = torch.empty((2, 9), device="cuda", dtype=torch.int32) + indptr = torch.tensor([0, 8, 8], device="cuda", dtype=torch.int32) + flat = torch.empty(8, device="cuda", dtype=torch.int32) + + def run(): + repage_block_table(source, lengths, destination, geometry) + _copy_page_indices_kernel[(2,)]( + flat, + source, + source.stride(0), + indptr, + BLOCK_SIZE=32, + PAGES_PER_BLOCK=geometry.pages_per_block, + BLOCK_STRIDE_PAGES=geometry.block_stride_pages, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + source[0, 4] = 4 + original = source.clone() + destination.fill_(-999) + flat.fill_(-999) + graph.replay() + expected = torch.tensor( + [0, 1, 2, 3, 16, 33, -1, 35], device="cuda", dtype=torch.int32 + ) + torch.testing.assert_close(flat, expected) + torch.testing.assert_close(destination[0, :8], flat) + assert torch.all(destination[1, :8] == 0) + torch.testing.assert_close(source, original) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +@pytest.mark.parametrize( + "backend", + [ + "fa2", + pytest.param( + "trtllm-gen", + marks=pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="Requires Blackwell", + ), + ), + ], +) +@torch.inference_mode() +def test_packed_flashinfer_matches_dense_pages(backend): + from vllm.v1.attention.ops.flashinfer_repage import ( + PackedKVPageGeometry, + repage_block_table, + ) + + set_random_seed(0) + manager_size, page_size, heads, dim = 592, 16, 2, 64 + cache = torch.randn( + (3, 2, heads, manager_size, 2 * dim), device="cuda", dtype=torch.bfloat16 + )[:, 1] + geometry = PackedKVPageGeometry.from_cache(cache, page_size) + packed = geometry.read_view(cache).split(dim, dim=-1) + ratio = manager_size // page_size + dense = ( + cache.reshape(3, heads, ratio, page_size, 2 * dim) + .permute(0, 2, 1, 3, 4) + .reshape(3 * ratio, heads, page_size, 2 * dim) + .contiguous() + .split(dim, dim=-1) + ) + lengths = [manager_size + 5, manager_size - 1] + indptr, _, last, source = _make_paged_kv_metadata(lengths, page_size, 3 * ratio) + source[:, 0] = 3 * ratio - 1 + seq_lens = torch.tensor(lengths, device="cuda", dtype=torch.int32) + mapped = repage_block_table(source, seq_lens, torch.empty_like(source), geometry) + query = torch.randn((6, 4 * heads, dim), device="cuda", dtype=torch.bfloat16) + workspace = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda") + + def attend(kv, table): + if backend == "trtllm-gen": + return flashinfer.decode.trtllm_batch_decode_with_kv_cache( + query, + kv, + workspace, + table, + seq_lens, + max(lengths), + bmm1_scale=dim**-0.5, + bmm2_scale=1.0, + backend=backend, + q_len_per_req=3, + ) + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace, "HND", backend=backend + ) + indices = torch.cat( + [ + table[i, : (length + page_size - 1) // page_size] + for i, length in enumerate(lengths) + ] + ) + wrapper.plan( + torch.tensor([0, 3, 6], device="cpu", dtype=torch.int32), + indptr, + indices, + last, + 4 * heads, + heads, + dim, + page_size, + causal=False, + sm_scale=dim**-0.5, + q_data_type=query.dtype, + kv_data_type=query.dtype, + ) + return wrapper.run(query, kv, return_lse=True) + + torch.testing.assert_close( + attend(packed, mapped), attend(dense, source), atol=1e-2, rtol=1e-2 + ) + + @pytest.mark.skipif( not any(current_platform.is_device_capability_family(f) for f in (80, 90, 120)), reason="NVFP4 fa2 path", diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 638deeedc579..c3b688684a97 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -38,6 +38,7 @@ MambaSpec, MLAAttentionSpec, SlidingWindowMLASpec, + SlidingWindowSpec, UniformTypeKVCacheSpecs, iter_layer_specs, ) @@ -427,6 +428,27 @@ def test_mamba_split_measures_the_block_the_other_groups_already_force(self, wid class TestSlidingWindowBucketCap: + @pytest.mark.parametrize("layout", ["BLHNC", "BLNHC"]) + def test_sliding_window_keeps_manager_size_when_kernels_use_small_pages( + self, layout + ): + config = _mock_vllm_config(layout) + config.speculative_config = None + specs = { + "target": replace(_full(), block_size=640, head_size=256, head_size_v=256), + "draft.sliding": SlidingWindowSpec( + block_size=64, + num_kv_heads=8, + head_size=128, + dtype=torch.float16, + sliding_window=4096, + ), + } + groups = _get_packed_kv_cache_groups(config, specs) + sliding = next(g for g in groups if "draft.sliding" in g.layer_names) + assert sliding.kv_cache_spec.block_size == 640 + assert specs["draft.sliding"].block_size == 64 + def test_sliding_window_bucket_is_capped_at_the_main_page(self): """DeepSeek-V4.1 shape: 43 SlidingWindowMLASpec SWA caches beside an unbalanced 8-layer paged MLA bucket. Left whole, the SWA bucket would diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 9c9c26c3c2c0..a3c149a704ef 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/tests/v1/worker/test_attn_utils.py @@ -7,8 +7,10 @@ never addressed by the logical view. """ +from dataclasses import replace from types import SimpleNamespace from typing import Any +from weakref import ref import numpy as np import pytest @@ -41,6 +43,7 @@ from vllm.v1.worker.utils import ( AttentionGroup, allocate_kv_cache, + clear_layer_kv_caches, copy_kv_cache_blocks_inplace, ) @@ -568,6 +571,61 @@ def test_copy_kv_cache_blocks_shared_storage(layout: KVCacheLayout): torch.testing.assert_close(cache[1], expected[layer_idx][1]) +@pytest.mark.parametrize("padded", [False, True]) +def test_packed_allocation_retains_manager_blocks(padded): + class Backend(AttentionBackend): + @classmethod + def get_kv_cache_view_block_size(cls, spec, kernel_block_size, vllm_config): + return spec.block_size + + spec = FullAttentionSpec( + block_size=128, num_kv_heads=2, head_size=64, dtype=torch.bfloat16 + ) + block_stride = 2 * spec.page_size_bytes + layers = ["layer.0"] if padded else ["layer.0", "layer.1"] + if padded: + spec = replace(spec, page_size_padded=block_stride) + config = KVCacheConfig( + num_blocks=3, + kv_cache_tensors=[ + KVCacheTensor( + size=3 * block_stride, + layers=layers, + layer_stride=spec.page_size_bytes, + block_stride=block_stride, + ) + ], + kv_cache_groups=[KVCacheGroupSpec(layers, spec)], + ) + context = { + name: SimpleNamespace(get_attn_backend=lambda: Backend) for name in layers + } + caches = allocate_kv_cache( + config, + torch.device("cpu"), + KVCacheLayout.BLHNC, + [32], + vllm_config=SimpleNamespace( + compilation_config=SimpleNamespace(static_forward_context=context) + ), + ) + for cache in caches.values(): + assert cache.shape == (3, 2, 128, 128) + assert cache.stride(0) * cache.element_size() == block_stride + + +def test_clear_layer_kv_caches_releases_flashinfer_read_views(): + cache = torch.empty((2, 1, 128, 128), dtype=torch.bfloat16, device="cpu") + cache_ref = ref(cache) + layer = SimpleNamespace( + kv_cache=cache, + impl=SimpleNamespace(_repage_source=cache, _repage_view=cache[:, :, :32]), + ) + del cache + clear_layer_kv_caches([layer]) + assert cache_ref() is None + + def test_fixed_block_stride_propagates_outward_in_lhbnc(): num_blocks = 3 num_layers = 2 diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 3387a31794c9..50449fb86c71 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -373,6 +373,20 @@ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...] | None: None when the kernels consume any layout and express no preference.""" return None + @classmethod + def get_kv_cache_view_block_size( + cls, + spec: "KVCacheSpec", + kernel_block_size: int, + vllm_config: "VllmConfig", + ) -> int: + """Physical view size when a non-dense manager block needs splitting. + + Backends retaining manager blocks must adapt their read indices and + views themselves. The common block table still uses kernel blocks. + """ + return kernel_block_size + @classmethod def is_ssm(cls) -> bool: return False diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 85b7456dd889..fe9063c397bd 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, replace from enum import Enum from functools import partial -from typing import ClassVar +from typing import ClassVar, cast import numpy as np import torch @@ -86,6 +86,11 @@ cp_lse_ag_out_rs, dcp_a2a_lse_reduce, ) +from vllm.v1.attention.ops.flashinfer_repage import ( + PackedKVPageGeometry, + remap_page_ids, + repage_block_table, +) from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.kv_cache_interface import ( AttentionSpec, @@ -515,6 +520,33 @@ def get_impl_cls() -> type["FlashInferImpl"]: def get_builder_cls() -> type["FlashInferMetadataBuilder"]: return FlashInferMetadataBuilder + @classmethod + def get_kv_cache_view_block_size( + cls, + spec: KVCacheSpec, + kernel_block_size: int, + vllm_config: VllmConfig, + ) -> int: + spec = cast(AttentionSpec, spec) + if ( + vllm_config.cache_config.get_resolved_kv_cache_layout() + != KVCacheLayout.BLHNC + or spec.kv_quant_mode != KVQuantMode.NONE + ): + return kernel_block_size + transfer = vllm_config.kv_transfer_config + if transfer is not None and any( + transfer.has_connector(name) + for name in ("NixlConnector", "NixlPullConnector", "NixlPushConnector") + ): + # NIXL infers physical block splitting from the kernel block size. + # Re-paging retains manager-sized physical blocks instead. + raise NotImplementedError( + "NIXL's physical block splitting must be adapted before it can " + "transfer FlashInfer manager-sized re-paged caches." + ) + return spec.block_size + @staticmethod def get_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype: if kv_cache_dtype in ("fp8", "fp8_e4m3"): @@ -741,6 +773,7 @@ class FlashInferMetadata: """ cascade_wrapper: MultiLevelCascadeAttentionWrapper | None + repage_geometry: PackedKVPageGeometry | None = None class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): @@ -760,6 +793,9 @@ def __init__( self.attention_config = vllm_config.attention_config self._num_speculative_tokens = vllm_config.num_speculative_tokens self._workspace_buffer = None + self._repage_source: torch.Tensor | None = None + self._repage_geometry: PackedKVPageGeometry | None = None + self._repage_block_tables: torch.Tensor | None = None self._prefill_wrapper: ( BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper | None ) = None # Wrapper for prefill/append @@ -1002,6 +1038,29 @@ def __init__( def kv_cache_layout(self) -> KVCacheLayout: return self.cache_config.get_resolved_kv_cache_layout() + def _resolve_repage_geometry( + self, block_table: torch.Tensor + ) -> PackedKVPageGeometry | None: + if self.kv_cache_layout != KVCacheLayout.BLHNC: + return None + context = self.vllm_config.compilation_config.static_forward_context + cache = context[self.layer_names[0]].kv_cache + if cache is self._repage_source: + return self._repage_geometry + geometry = None + if cache.shape[2] != self.page_size: + geometry = PackedKVPageGeometry.from_cache(cache, self.page_size) + # The runner also pads the table's width for token alignment. + width = block_table.shape[1] + self._repage_block_tables = torch.empty( + (self.max_num_reqs, width), dtype=torch.int32, device=self.device + ) + else: + self._repage_block_tables = None + self._repage_source = cache + self._repage_geometry = geometry + return geometry + # Keep SM90 prefill/decode Q dtype selection in one place. def get_q_data_type(self, is_prefill: bool) -> torch.dtype: # The user sets --attention-config.disable_flashinfer_q_quantization @@ -1375,6 +1434,7 @@ def _compute_flashinfer_kv_metadata( block_table_tensor: torch.Tensor, num_reqs: int, page_size: int, + repage_geometry: PackedKVPageGeometry | None = None, ) -> torch.Tensor: """Compute paged_kv_indptr, paged_kv_indices and paged_kv_last_page_len. @@ -1398,6 +1458,10 @@ def _compute_flashinfer_kv_metadata( block_table_tensor.stride(0), paged_kv_indptr, BLOCK_SIZE=1024, + PAGES_PER_BLOCK=(repage_geometry.pages_per_block if repage_geometry else 1), + BLOCK_STRIDE_PAGES=( + repage_geometry.block_stride_pages if repage_geometry else 1 + ), ) # write self.paged_kv_last_page_len_cpu inplace @@ -1448,6 +1512,9 @@ def build( # - Prefill (FI native or TRTLLM) # - Decode (FI native, XQA, or trtllm-gen) use_cascade = common_prefix_len > 0 + repage_geometry = self._resolve_repage_geometry(block_table_tensor) + if repage_geometry is not None and use_cascade: + raise NotImplementedError("Cascade attention cannot consume re-paged KV.") uses_spec_reorder = self.reorder_batch_threshold > 1 # Page sizes >= 128 must use trtllm-gen; force it for prefill too. prefill_force_trtllm = ( @@ -1524,6 +1591,7 @@ def build( prefill=None, decode=None, cascade_wrapper=None, + repage_geometry=repage_geometry, ) # Guard access to seq_lens_cpu, which may not always be needed @@ -1603,6 +1671,7 @@ def build( block_table_tensor, num_reqs, page_size, + repage_geometry, ) else: paged_kv_indices = None @@ -1659,6 +1728,23 @@ def build( return attn_metadata # Step 3: Handle prefill and decode pathways case by case + needs_trtllm_block_tables = (num_prefills > 0 and prefill_use_trtllm) or ( + num_decodes > 0 and decode_with_flashinfer_trtllm_api + ) + trtllm_seq_lens = seq_lens + if self.use_dcp and needs_trtllm_block_tables: + assert common_attn_metadata.dcp_local_seq_lens is not None + trtllm_seq_lens = common_attn_metadata.dcp_local_seq_lens + + if repage_geometry is not None and needs_trtllm_block_tables: + assert self._repage_block_tables is not None + block_table_tensor = repage_block_table( + block_table_tensor, + trtllm_seq_lens, + self._repage_block_tables, + repage_geometry, + ) + ## PREFILL PATHWAY if num_prefills > 0: # Slices for shared prefill metadata @@ -1817,12 +1903,7 @@ def build( if num_decodes > 0: if decode_with_flashinfer_trtllm_api: assert self.flashinfer_trtllm_api_decode_kernel is not None - seq_lens_decode = seq_lens[:num_decodes] - if self.use_dcp: - assert common_attn_metadata.dcp_local_seq_lens is not None - seq_lens_decode = common_attn_metadata.dcp_local_seq_lens[ - :num_decodes - ] + seq_lens_decode = trtllm_seq_lens[:num_decodes] q_len_per_req, q_cu_seq_lens, ragged_q_lens = ( self._compute_decode_query_lens( qo_indptr, qo_indptr_cpu, num_decodes, num_decode_tokens @@ -1953,6 +2034,9 @@ def __init__( sinks: torch.Tensor | None = None, ) -> None: self.num_heads = num_heads + self._repage_source: torch.Tensor | None = None + self._repage_view_geometry: PackedKVPageGeometry | None = None + self._repage_view: torch.Tensor | None = None self.head_size = head_size self.scale = float(scale) self.num_kv_heads = num_kv_heads @@ -2203,6 +2287,18 @@ def forward( num_actual_tokens = attn_metadata.num_actual_tokens kv_cache_layout = attn_metadata.kv_cache_layout + if attn_metadata.repage_geometry is not None: + geometry = attn_metadata.repage_geometry + if ( + self._repage_source is not kv_cache + or self._repage_view_geometry != geometry + ): + self._repage_source = kv_cache + self._repage_view_geometry = geometry + self._repage_view = geometry.read_view(kv_cache) + assert self._repage_view is not None + kv_cache = self._repage_view + # FlashInfer treats uint8 KV cache as NVFP4. vLLM stores FP8 KV cache # as uint8 bytes, so pass FP8 caches with their logical dtype. if not self.is_kvcache_nvfp4 and kv_cache.dtype == torch.uint8: @@ -2888,6 +2984,8 @@ def _copy_page_indices_kernel( block_table_stride, cu_num_blocks, BLOCK_SIZE: tl.constexpr, + PAGES_PER_BLOCK: tl.constexpr = 1, + BLOCK_STRIDE_PAGES: tl.constexpr = 1, ): req_idx = tl.program_id(0) row_ptr = block_table + req_idx * block_table_stride @@ -2898,6 +2996,8 @@ def _copy_page_indices_kernel( offset = tl.arange(0, BLOCK_SIZE) for i in tl.range(0, num_blocks, BLOCK_SIZE): block_ids = tl.load(row_ptr + i + offset, mask=i + offset < num_blocks) + if PAGES_PER_BLOCK != BLOCK_STRIDE_PAGES: + block_ids = remap_page_ids(block_ids, PAGES_PER_BLOCK, BLOCK_STRIDE_PAGES) tl.store( page_indices + start_idx + i + offset, block_ids, diff --git a/vllm/v1/attention/ops/flashinfer_repage.py b/vllm/v1/attention/ops/flashinfer_repage.py new file mode 100644 index 000000000000..60e5225372aa --- /dev/null +++ b/vllm/v1/attention/ops/flashinfer_repage.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Read-only kernel pages over a packed BF16/FP16 HND cache.""" + +from dataclasses import dataclass + +import torch + +from vllm.triton_utils import tl, triton + + +@dataclass(frozen=True) +class PackedKVPageGeometry: + shape: tuple[int, ...] + strides: tuple[int, ...] + page_size: int + pages_per_block: int + block_stride_pages: int + num_pages: int + + @classmethod + def from_cache(cls, cache: torch.Tensor, page_size: int) -> "PackedKVPageGeometry": + """Derive kernel pages from an allocator-owned [B,H,M,2D] cache.""" + num_blocks, _, block_size, _ = cache.shape + block_stride, head_stride, token_stride, _ = cache.stride() + if block_size % page_size: + raise ValueError("Kernel page size must divide the manager block size.") + page_stride = page_size * token_stride + if block_stride % page_stride: + raise ValueError("Packed KV block stride must contain whole kernel pages.") + pages_per_block = block_size // page_size + block_stride_pages = block_stride // page_stride + num_pages = (num_blocks - 1) * block_stride_pages + pages_per_block + if max(num_pages, head_stride, page_stride) > 2**31 - 1: + raise ValueError("Packed KV geometry exceeds FlashInfer's int32 limits.") + return cls( + shape=tuple(cache.shape), + strides=cache.stride(), + page_size=page_size, + pages_per_block=pages_per_block, + block_stride_pages=block_stride_pages, + num_pages=num_pages, + ) + + def read_view(self, cache: torch.Tensor) -> torch.Tensor: + _, heads, _, channels = self.shape + _, head_stride, token_stride, channel_stride = self.strides + # Only referenced pages are valid. Never register this view for writes, + # block copies, zeroing or transfer: its virtual holes alias other layers. + return cache.as_strided( + (self.num_pages, heads, self.page_size, channels), + (self.page_size * token_stride, head_stride, token_stride, channel_stride), + ) + + +@triton.jit +def remap_page_ids( + ids, PAGES_PER_BLOCK: tl.constexpr, BLOCK_STRIDE_PAGES: tl.constexpr +): + physical = ( + ids.to(tl.int64) // PAGES_PER_BLOCK * BLOCK_STRIDE_PAGES + + ids.to(tl.int64) % PAGES_PER_BLOCK + ) + return tl.where(ids >= 0, physical, ids).to(tl.int32) + + +@triton.jit(do_not_specialize=["num_cols"]) +def _repage_block_table_kernel( + block_table, + seq_lens, + out, + num_cols, + block_table_stride, + out_stride, + PAGE_SIZE: tl.constexpr, + PAGES_PER_BLOCK: tl.constexpr, + BLOCK_STRIDE_PAGES: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + tiles = tl.cdiv(num_cols, BLOCK_SIZE) + row = tl.program_id(0) // tiles + cols = (tl.program_id(0) % tiles) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + length = tl.load(seq_lens + row) + active = cols < tl.cdiv(length, PAGE_SIZE) + ids = tl.load( + block_table + row.to(tl.int64) * block_table_stride + cols, + (cols < num_cols) & active, + other=0, + ) + ids = remap_page_ids(ids, PAGES_PER_BLOCK, BLOCK_STRIDE_PAGES) + tl.store( + out + row.to(tl.int64) * out_stride + cols, + tl.where(active, ids, 0), + cols < num_cols, + ) + + +def repage_block_table( + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + geometry: PackedKVPageGeometry, +) -> torch.Tensor: + """Remap into the builder's separate, preallocated int32 scratch buffer.""" + rows, cols = block_table.shape + if rows == 0 or cols == 0: + return out[:rows, :cols] + _repage_block_table_kernel[(rows * triton.cdiv(cols, 256),)]( + block_table, + seq_lens, + out, + cols, + block_table.stride(0), + out.stride(0), + geometry.page_size, + geometry.pages_per_block, + geometry.block_stride_pages, + BLOCK_SIZE=256, + ) + return out[:rows, :cols] diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 53a01a5cabe2..cf286088e93f 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -2124,6 +2124,31 @@ def _get_packed_kv_cache_groups( if not layout.is_block_outermost or len(page_sizes) <= 1: return None + primary_block_size = max( + ( + spec.block_size + for spec in kv_cache_spec.values() + if isinstance(spec, FullAttentionSpec) + ), + default=0, + ) + # Each SWA manager block occupies a whole packed block. Kernel page + # splitting is selected later, independently of this physical size. + kv_cache_spec = { + name: replace(spec, block_size=primary_block_size) + if ( + type(spec) is SlidingWindowSpec + and spec.page_size_padded is None + and spec.block_size < primary_block_size + and primary_block_size % spec.block_size == 0 + ) + else spec + for name, spec in kv_cache_spec.items() + } + page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()} + if len(page_sizes) <= 1: + return None + buckets: list[dict[str, KVCacheSpec]] = [] for name, spec in kv_cache_spec.items(): for bucket in buckets: diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 8aa9ba3cb8b2..19b965de6095 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -401,6 +401,7 @@ def init_kv_cache( device, vllm_config.cache_config.get_resolved_kv_cache_layout(), kernel_block_sizes, + vllm_config=vllm_config, ) for layer_name, target in get_shared_kv_cache_layers(vllm_config).items(): kv_caches[layer_name] = kv_caches[target] diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index c02bdc74e22b..7db0d43c9129 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6532,6 +6532,9 @@ def _cleanup_profiling_kv_cache(self) -> None: layer.impl._k_scale_cache = None if hasattr(layer.impl, "_v_scale_cache"): layer.impl._v_scale_cache = None + if hasattr(layer.impl, "_repage_source"): + layer.impl._repage_source = None + layer.impl._repage_view = None gc.collect() torch.accelerator.empty_cache() @@ -7257,6 +7260,7 @@ def initialize_kv_cache_tensors( self.device, self.cache_config.get_resolved_kv_cache_layout(), kernel_block_sizes, + vllm_config=self.vllm_config, ) # Set up cross-layer KV cache sharing diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index db0f81b81c36..4293f7b2fd9a 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -10,7 +10,7 @@ import numpy as np import torch -from vllm.config import CacheConfig +from vllm.config import CacheConfig, VllmConfig from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.mamba.mamba_mixer2 import share_replayssm_ring_trackers @@ -37,6 +37,7 @@ MambaSpec, MLAAttentionSpec, UniformTypeKVCacheSpecs, + compute_layer_kv_cache_shape_bytes, create_kv_cache_views, ) from vllm.v1.worker.block_table import get_block_table_width @@ -401,6 +402,8 @@ def allocate_kv_cache( device: torch.device, layout: KVCacheLayout, kernel_block_sizes: list[int] | None = None, + *, + vllm_config: VllmConfig | None = None, ) -> dict[str, torch.Tensor]: """Allocate the KV cache and view it as ``[B, H, N, C]`` per layer. @@ -450,8 +453,28 @@ def allocate_kv_cache( kernel_block_size = None if kernel_block_sizes is not None and group_id < len(kernel_block_sizes): kernel_block_size = kernel_block_sizes[group_id] + view_block_size = kernel_block_size + dense_page_size = math.prod(compute_layer_kv_cache_shape_bytes(spec, 1)[1:]) if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None: - kernel_block_size = spec.storage_block_size + view_block_size = spec.storage_block_size + elif ( + vllm_config is not None + and kernel_block_size is not None + and kernel_block_size != spec.block_size + and tensor.block_stride != dense_page_size + ): + context = vllm_config.compilation_config.static_forward_context + view_block_sizes = { + context[name] + .get_attn_backend() + .get_kv_cache_view_block_size(spec, kernel_block_size, vllm_config) + for name in tensor.layers + } + if len(view_block_sizes) != 1: + raise ValueError( + "Layers sharing a KV tensor require the same view size." + ) + view_block_size = view_block_sizes.pop() views = create_kv_cache_views( buf, @@ -459,7 +482,7 @@ def allocate_kv_cache( num_blocks, layout, tensor, - kernel_block_size=kernel_block_size, + kernel_block_size=view_block_size, ) kv_caches.update(zip(tensor.layers, views)) return kv_caches @@ -715,7 +738,11 @@ def bind_kv_cache_to_layers( # splits conv/ssm), so the kv_caches dict can hold a single tensor per # layer for the KV connector to register. for layer_name, kv_cache in kv_caches.items(): - forward_context[layer_name].bind_kv_cache(kv_cache) + layer = forward_context[layer_name] + if hasattr(layer, "impl") and hasattr(layer.impl, "_repage_source"): + layer.impl._repage_source = None + layer.impl._repage_view = None + layer.bind_kv_cache(kv_cache) ordered_layer_names = sorted( kv_caches, key=lambda name: extract_layer_index(name, num_attn_module) @@ -742,6 +769,9 @@ def clear_layer_kv_caches(layers: Iterable[Any]) -> None: layer.impl._k_scale_cache = None if hasattr(layer.impl, "_v_scale_cache"): layer.impl._v_scale_cache = None + if hasattr(layer.impl, "_repage_source"): + layer.impl._repage_source = None + layer.impl._repage_view = None def copy_kv_cache_blocks_inplace(