""" Key-Value Cache for autoregressive generation. During generation each new token only needs to attend to its own Q against the full history of K and V. Rather than recomputing K/V for the entire sequence every step, we cache them per layer and append the new token's K/V pair. This reduces generation from O(n^2) to O(n) in sequence length per step. This module provides TWO implementations of the same interface, kept side-by-side for benchmarking and educational comparison: 1. ``KVCacheNaive`` — the textbook implementation. Each ``append`` call does ``torch.cat([prev_k, k_new], dim=2)``, which allocates a fresh tensor and copies the entire history every step. Simple to read, but appends are O(n_history) so total generation cost is O(n^2) *just from cache management*, on top of the unavoidable O(n) per token from attention itself. 2. ``KVCachePrealloc`` — the optimized version. Per-layer fixed-shape buffers of capacity ``(batch, n_kv_heads, capacity, head_dim)`` are allocated lazily on first append. Subsequent appends are in-place slice writes, so they are O(n_new) regardless of how much history already exists. Buffers grow geometrically (double on overflow) when the optional ``max_seq_len`` is not pinned, giving amortized O(1) appends with no upfront waste. Truncation and reset become O(1) — they only adjust per-layer length counters; the underlying buffers are reused across runs. Both classes expose the same public API: append(layer_idx, k_new, v_new) -> (k_view, v_view) get(layer_idx) -> (k, v) | None update(layer_idx, (k, v)) truncate(layer_idx, length) truncate_all(length) trim_to_last(layer_idx, max_len) reset() advance(num_tokens) seq_len (property) length(layer_idx) so the kernel and model code is generic over which cache is in use. ``KVCache`` is aliased to ``KVCachePrealloc`` for the default import. """ import torch class KVCacheNaive: """ Educational baseline cache: stores ``(k, v)`` tuples per layer and grows them via ``torch.cat`` on every append. Kept in the codebase as the reference implementation that the pre-allocated version is benchmarked against — re-allocating and copying the full history each step makes its per-step cost grow linearly with the sequence length already cached. """ def __init__(self, n_layers, **_unused): self.n_layers = n_layers self.cache = [None] * n_layers self.current_pos = 0 # --- position tracking --- @property def seq_len(self): return self.current_pos def advance(self, num_tokens): self.current_pos += num_tokens def length(self, layer_idx): entry = self.cache[layer_idx] return 0 if entry is None else entry[0].size(2) # --- core append (O(n_history) — the whole point of the comparison) --- def append(self, layer_idx, k_new, v_new): prev = self.cache[layer_idx] if prev is None: k, v = k_new, v_new else: prev_k, prev_v = prev k = torch.cat([prev_k, k_new], dim=2) v = torch.cat([prev_v, v_new], dim=2) self.cache[layer_idx] = (k, v) return k, v # --- back-compat tuple API --- def get(self, layer_idx): return self.cache[layer_idx] def update(self, layer_idx, value): self.cache[layer_idx] = value def get_all(self): return self.cache # --- truncation / reset --- def truncate(self, layer_idx, length): if self.cache[layer_idx] is not None: k, v = self.cache[layer_idx] self.cache[layer_idx] = (k[:, :, :length, :], v[:, :, :length, :]) def truncate_all(self, length): for i in range(self.n_layers): self.truncate(i, length) self.current_pos = length def trim_to_last(self, layer_idx, max_len): entry = self.cache[layer_idx] if entry is None: return k, v = entry if k.size(2) <= max_len: return self.cache[layer_idx] = (k[:, :, -max_len:, :], v[:, :, -max_len:, :]) def reset(self): for i in range(self.n_layers): self.cache[i] = None self.current_pos = 0 class KVCachePrealloc: """ Pre-allocated KV cache: per-layer fixed buffers with in-place writes. Buffers are allocated lazily on the first ``append`` (the cache does not need to know dtype/shape/device up front) and grow geometrically when their capacity is exceeded — appends are amortized O(1). Truncation and reset only flip per-layer length counters, so speculative-decoding rollbacks become free and buffers are reused across generation runs. """ def __init__(self, n_layers, max_seq_len=None, initial_capacity=256): """ Args: n_layers: Number of transformer layers. max_seq_len: Optional fixed maximum capacity. If given, buffers are allocated once at this capacity and never grow; an append that would overflow raises ``RuntimeError``. initial_capacity: Starting capacity when ``max_seq_len`` is None. Buffers double in capacity on overflow. """ self.n_layers = n_layers self.max_seq_len = max_seq_len self.initial_capacity = initial_capacity self.k_buffers = [None] * n_layers self.v_buffers = [None] * n_layers self.layer_lens = [0] * n_layers self.current_pos = 0 # --- position tracking --- @property def seq_len(self): return self.current_pos def advance(self, num_tokens): self.current_pos += num_tokens def length(self, layer_idx): return self.layer_lens[layer_idx] # --- core append (in-place, O(n_new)) --- def append(self, layer_idx, k_new, v_new): """ Append new K/V tensors for ``layer_idx`` and return views over the full prefix (length = previous + new). Both ``k_new`` and ``v_new`` have shape ``(batch, n_kv_heads, n_new_tokens, head_dim)``. """ self._ensure_capacity(layer_idx, k_new, v_new) k_buf = self.k_buffers[layer_idx] v_buf = self.v_buffers[layer_idx] cur = self.layer_lens[layer_idx] n = k_new.size(2) k_buf[:, :, cur:cur + n, :] = k_new v_buf[:, :, cur:cur + n, :] = v_new new_len = cur + n self.layer_lens[layer_idx] = new_len return k_buf[:, :, :new_len, :], v_buf[:, :, :new_len, :] # --- back-compat tuple API --- def get(self, layer_idx): n = self.layer_lens[layer_idx] if n == 0 or self.k_buffers[layer_idx] is None: return None k_buf = self.k_buffers[layer_idx] v_buf = self.v_buffers[layer_idx] return k_buf[:, :, :n, :], v_buf[:, :, :n, :] def update(self, layer_idx, value): """ Replace the stored K/V for ``layer_idx`` with ``value`` (a ``(k, v)`` tuple). Copies the data into the pre-allocated buffer; reallocates the buffer if it does not exist or is too small. """ if value is None: self.layer_lens[layer_idx] = 0 return k_new, v_new = value n = k_new.size(2) # Reset length so the ensure-capacity check sizes for the full payload self.layer_lens[layer_idx] = 0 self._ensure_capacity(layer_idx, k_new, v_new) k_buf = self.k_buffers[layer_idx] v_buf = self.v_buffers[layer_idx] k_buf[:, :, :n, :] = k_new v_buf[:, :, :n, :] = v_new self.layer_lens[layer_idx] = n def get_all(self): return [self.get(i) for i in range(self.n_layers)] # --- truncation / reset (O(1) — just adjust length counters) --- def truncate(self, layer_idx, length): if self.k_buffers[layer_idx] is None: return self.layer_lens[layer_idx] = min(self.layer_lens[layer_idx], length) def truncate_all(self, length): for i in range(self.n_layers): self.truncate(i, length) self.current_pos = length def trim_to_last(self, layer_idx, max_len): """ Keep only the most recent ``max_len`` tokens for ``layer_idx``. Used by sliding-window attention: when the cached length exceeds the window, slide the kept slice to the front of the buffer in-place so the next ``append`` writes contiguously after it. """ cur = self.layer_lens[layer_idx] if cur <= max_len or self.k_buffers[layer_idx] is None: return k_buf = self.k_buffers[layer_idx] v_buf = self.v_buffers[layer_idx] # Clone the slice we want to keep before overwriting (avoid aliasing). keep_k = k_buf[:, :, cur - max_len:cur, :].clone() keep_v = v_buf[:, :, cur - max_len:cur, :].clone() k_buf[:, :, :max_len, :] = keep_k v_buf[:, :, :max_len, :] = keep_v self.layer_lens[layer_idx] = max_len def reset(self): """Clear lengths and global position. Buffers are kept for reuse.""" for i in range(self.n_layers): self.layer_lens[i] = 0 self.current_pos = 0 # --- internal: lazy buffer allocation and geometric growth --- def _ensure_capacity(self, layer_idx, k_new, v_new): n_new = k_new.size(2) cur = self.layer_lens[layer_idx] needed = cur + n_new k_buf = self.k_buffers[layer_idx] if k_buf is None: cap = self.max_seq_len if self.max_seq_len is not None else max( self.initial_capacity, n_new ) shape = (k_new.size(0), k_new.size(1), cap, k_new.size(3)) self.k_buffers[layer_idx] = torch.empty( shape, dtype=k_new.dtype, device=k_new.device ) self.v_buffers[layer_idx] = torch.empty( shape, dtype=v_new.dtype, device=v_new.device ) return if needed <= k_buf.size(2): return if self.max_seq_len is not None: raise RuntimeError( f"KVCachePrealloc layer {layer_idx} would exceed " f"max_seq_len={self.max_seq_len} (needed {needed})." ) # Grow geometrically — at least double, at least enough to fit `needed`. new_cap = max(needed, k_buf.size(2) * 2) shape = (k_buf.size(0), k_buf.size(1), new_cap, k_buf.size(3)) new_k = torch.empty(shape, dtype=k_buf.dtype, device=k_buf.device) new_v = torch.empty(shape, dtype=k_buf.dtype, device=k_buf.device) new_k[:, :, :cur, :] = k_buf[:, :, :cur, :] new_v[:, :, :cur, :] = self.v_buffers[layer_idx][:, :, :cur, :] self.k_buffers[layer_idx] = new_k self.v_buffers[layer_idx] = new_v # Default cache used by the model code. Aliased so existing imports # (`from src.utils.kv_cache import KVCache`) pick up the optimized version. KVCache = KVCachePrealloc