# SPDX-License-Identifier: Apache-2.0 """Generic paging mechanics + chunk-window eviction spec for the AR-Diffusion engine. Engine-generic, model-agnostic primitives — the layer a second model (e.g. the Cosmos port) reuses unchanged. Three concerns live here: * **Slot mapping** — absolute token positions → physical KV-cache slots, the standard PagedAttention layout ``slot(pos) = block_id(pos) * block_size + (pos % block_size)``. * **Pool I/O** — allocate per-layer FlashAttention-compatible paged K/V pools plus flat compatibility views for direct slot writes. * **Chunk-window eviction** — a ``SlidingWindowSpec`` subclass whose unit is a *chunk* (``sliding_window = window_chunks * chunk_size``) plus a manager that evicts at chunk boundaries. Memory policy / refcounting / ``null_block`` replacement stay in vLLM's ``BlockPool``; only the token-skip math is here. """ from __future__ import annotations from collections.abc import Sequence from dataclasses import dataclass import torch from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager from vllm.v1.kv_cache_interface import SlidingWindowSpec from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec # ── Slot mapping ──────────────────────────────────────────────────────────── def compute_slot_mapping( block_ids: Sequence[int], positions: torch.Tensor | Sequence[int], block_size: int, ) -> torch.Tensor: """Map absolute token positions to physical KV-cache slots. Args: block_ids: physical block id per block index (the request's block table). positions: absolute token positions to map (1-D). block_size: tokens per block. Returns: ``LongTensor`` of physical slots, one per position. """ if block_size <= 0: raise ValueError(f"block_size must be positive, got {block_size}") table = torch.as_tensor(block_ids, dtype=torch.long) pos = torch.as_tensor(positions, dtype=torch.long) block_index = torch.div(pos, block_size, rounding_mode="floor") offset = pos % block_size return table[block_index] * block_size + offset def chunk_slot_mapping( block_ids: Sequence[int], num_computed_tokens: int, chunk_size: int, block_size: int, ) -> torch.Tensor: """Slot mapping for the in-flight chunk's tokens (the commit write target). The chunk occupies absolute positions ``[num_computed_tokens, num_computed_tokens + chunk_size)``. """ positions = torch.arange( num_computed_tokens, num_computed_tokens + chunk_size, dtype=torch.long, ) return compute_slot_mapping(block_ids, positions, block_size) def resident_block_ids(block_ids: Sequence[int], null_block_id: int) -> list[int]: """Real (non-null) blocks currently resident, in table order. These are the blocks the read path gathers the attention window from; out-of-window positions are the shared ``null_block`` and are excluded. """ return [int(b) for b in block_ids if int(b) != null_block_id] # ── Pool I/O (allocate / write) ───────────────────────────────────────────── def allocate_kv_pool_with_views( num_blocks: int, block_size: int, num_layers: int, num_kv_heads: int, head_dim: int, dtype: torch.dtype, device: torch.device, ) -> tuple[list[list[torch.Tensor]], list[torch.Tensor], list[torch.Tensor]]: """Allocate vLLM-style paged KV pools plus flat slot-write views. Each cache keeps FlashAttention's block-table layout, ``(num_blocks, block_size, num_kv_heads, head_dim)``. The flat K/V views keep ``slot = block_id * block_size + offset`` writes simple without changing the kernel-facing layout, and ``kv_pools[layer][0]`` / ``[1]`` still address K and V. K and V are allocated separately rather than as the two halves of one ``(2, ...)`` tensor. Sharing one allocation makes them views of the same storage, and the paged-write custom op declares both as mutated. Inductor's reinplace pass can then re-inplace only the first of the two -- the second hits the "mutated arg aliases a graph input" case in ``_inductor/fx_passes/reinplace.py`` -- so ``auto_functionalized_v2``'s clone of the whole V pool survives into the compiled graph. That clone is one full pool per compiled region per denoising step. """ kv_pools: list[list[torch.Tensor]] = [] k_pools: list[torch.Tensor] = [] v_pools: list[torch.Tensor] = [] cache_shape = (num_blocks, block_size, num_kv_heads, head_dim) flat_shape = (num_blocks * block_size, num_kv_heads, head_dim) for _ in range(num_layers): k = torch.empty(cache_shape, dtype=dtype, device=device) v = torch.empty(cache_shape, dtype=dtype, device=device) kv_pools.append([k, v]) k_pools.append(k.reshape(flat_shape)) v_pools.append(v.reshape(flat_shape)) return kv_pools, k_pools, v_pools def pool_write_chunk( k_pool: torch.Tensor, v_pool: torch.Tensor, new_k: torch.Tensor, new_v: torch.Tensor, slot_mapping: torch.Tensor, ) -> None: """Write the committed chunk's K/V into pool slots (in place). ``new_k`` / ``new_v`` shape: ``(batch, chunk_size, num_kv_heads, head_dim)``. ``slot_mapping`` shape: ``(chunk_size,)`` — one physical slot per token position. The batch dim is written to the same slots (multi-batch writes the same positions; the caller is responsible for per-sequence bookkeeping). This is the *commit* write — called once per chunk, after the last denoise step. """ k_pool[slot_mapping] = new_k[0] # batch index 0 v_pool[slot_mapping] = new_v[0] def _drop_batch(t: torch.Tensor) -> torch.Tensor: """``(1, L, n_heads, head_dim)`` -> ``(L, n_heads, head_dim)``; pass 3-D through.""" if t.dim() == 4 and t.shape[0] == 1: return t[0] return t # ── Chunk-window eviction (spec + manager) ────────────────────────────────── def chunk_window_skipped_tokens( num_computed_tokens: int, *, chunk_size: int, sliding_window: int, sink_chunks: int, reset_at_boundary: bool, ) -> int: """Tokens outside the resident chunk window, snapped to a chunk boundary. Pure function so the eviction policy is unit-testable without constructing a manager. Two strategies: - ``reset_at_boundary``: at each chunk boundary everything past the sink is dropped. - otherwise (VGGT-style sliding replace): keep the last ``window`` tokens (plus the sink); the skip count snaps down to a chunk boundary so a chunk is never half-evicted. """ sink = sink_chunks * chunk_size if reset_at_boundary: completed = (num_computed_tokens // chunk_size) * chunk_size return max(0, completed - sink) skipped = max(0, num_computed_tokens - sliding_window - sink) return (skipped // chunk_size) * chunk_size class ChunkWindowManager(SlidingWindowManager): """``SlidingWindowManager`` that evicts at chunk boundaries. ``self.sliding_window`` is set by the base ``__init__``; the chunk fields are read from ``self.kv_cache_spec`` (a :class:`ChunkWindowSpec`). """ def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: spec = self.kv_cache_spec return chunk_window_skipped_tokens( num_computed_tokens, chunk_size=spec.chunk_size, sliding_window=self.sliding_window, sink_chunks=spec.sink_chunks, reset_at_boundary=spec.reset_at_boundary, ) def remove_skipped_blocks( self, request_id: str, total_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: """Free the middle gap while preserving the leading attention sink.""" del num_prompt_tokens num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens) if num_skipped_tokens <= 0: return sink_blocks = self.kv_cache_spec.sink_chunks num_skipped_blocks = num_skipped_tokens // self.block_size self._remove_blocks_in_range( request_id, sink_blocks, sink_blocks + num_skipped_blocks, ) # Register so KVCacheManager resolves ChunkWindowSpec to ChunkWindowManager. # Dispatch walks the spec's MRO, so without explicit registration the subclass # would silently fall back to the parent SlidingWindowManager (override ignored). # uniform_type_base_spec=None => its own KV cache group. @register_kv_cache_spec(manager_class=ChunkWindowManager, uniform_type_base_spec=None) @dataclass(frozen=True, kw_only=True) class ChunkWindowSpec(SlidingWindowSpec): # sliding_window (inherited) MUST equal window_chunks * chunk_size. chunk_size: int window_chunks: int sink_chunks: int = 0 reset_at_boundary: bool = False def __post_init__(self): super().__post_init__() if self.sliding_window != self.window_chunks * self.chunk_size: raise ValueError( "ChunkWindowSpec.sliding_window must equal " f"window_chunks * chunk_size ({self.window_chunks} * " f"{self.chunk_size}), got {self.sliding_window}" )