File size: 10,718 Bytes
d656a30 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 | """KV caches for incremental decoding (docs/INFERENCE.md).
The model talks to a cache through three calls per forward pass:
rope = cache.begin_step(T, rope_module) # positions of the T new tokens; builds the mask
view = cache.layer(i) # per layer: view.update(k, v), view.attn_kwargs
cache.end_step(T) # the new tokens are now stored
- KVCache (4.3, 4.4 static): one shared position for every row. Optional left padding
(`pad_lens`) lets prompts of different lengths share a batch: pad keys are masked.
- SlotKVCache (4.4 continuous): S independent slots, each with its own length, so
requests can join and leave between decode steps.
Buffers are [B, n_kv_heads, max_len, head_dim] per layer. No reference cycles: the
caches own plain tensor lists and `layer(i)` returns a short-lived view. (A cycle once
kept every finished cache's GPU memory alive until Python's cycle collector ran.)
Left padding needs no per-row RoPE positions: RoPE makes q·k depend only on the
distance between two positions, and padding shifts every real token of a row by the
same amount. Pad *queries* may attend to themselves, so no row is fully masked (a fully
masked softmax row is NaN, and NaN would leak through 0 × NaN in later layers).
"""
import torch
from torch import Tensor
def causal_kwargs(start: int, T: int, device) -> dict:
"""SDPA masking for T new queries at positions start..start+T-1 over start+T keys, no padding.
- start 0 (prefill / training): standard causal.
- one new token: it may see every cached key, so no mask.
- several new tokens on a non-empty cache (chunked prefill, speculative verify):
SDPA's is_causal aligns to the top-left, which is wrong here, so pass an explicit
mask where query i sees keys 0..start+i.
"""
if start == 0:
return {"is_causal": True}
if T == 1:
return {}
q_pos = torch.arange(start, start + T, device=device)[:, None]
k_pos = torch.arange(start + T, device=device)[None, :]
return {"attn_mask": k_pos <= q_pos}
def padded_causal_mask(start: int, T: int, first_valid: Tensor) -> Tensor:
"""(B, 1, T, start+T) bool: causal, keys before first_valid[b] hidden, diagonal always kept."""
dev = first_valid.device
q_pos = torch.arange(start, start + T, device=dev)[:, None] # (T, 1)
k_pos = torch.arange(start + T, device=dev)[None, :] # (1, K)
causal = k_pos <= q_pos
valid = k_pos[None] >= first_valid[:, None, None] # (B, 1, K)
return (causal[None] & (valid | (k_pos == q_pos)[None]))[:, None] # (B, 1, T, K)
class KVCache:
"""Shared-position cache. pad_lens (B,): left-pad count per row, or None for no padding."""
def __init__(self, n_layers: int, max_len: int, pad_lens: Tensor | None = None):
self.max_len = max_len
self.pos = 0
self.pad_lens = pad_lens
self.k: list[Tensor | None] = [None] * n_layers
self.v: list[Tensor | None] = [None] * n_layers
self.attn_kwargs: dict = {}
def begin_step(self, T: int, rope) -> tuple[Tensor, Tensor]:
cos, sin = rope(T, self.pos)
if self.pad_lens is None:
self.attn_kwargs = causal_kwargs(self.pos, T, cos.device)
else:
self.attn_kwargs = {"attn_mask": padded_causal_mask(self.pos, T, self.pad_lens)}
return cos, sin
def end_step(self, T: int) -> None:
self.pos += T
def layer(self, i: int) -> "LayerKV":
return LayerKV(self, i)
def nbytes(self) -> int:
return sum(t.numel() * t.element_size() for t in self.k + self.v if t is not None)
class LayerKV:
"""One layer's view of a KVCache."""
def __init__(self, cache: KVCache, layer_idx: int):
self.cache = cache
self.i = layer_idx
@property
def start(self) -> int:
"""Position of the first new token = number of tokens already cached."""
return self.cache.pos
@property
def attn_kwargs(self) -> dict:
return self.cache.attn_kwargs
def update(self, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
"""Store k, v (B, H, T, D) at positions start..start+T-1; return all cached K, V so far."""
c, i = self.cache, self.i
B, H, T, D = k.shape
if c.k[i] is None:
shape = (B, H, c.max_len, D)
c.k[i] = torch.zeros(shape, dtype=k.dtype, device=k.device)
c.v[i] = torch.zeros(shape, dtype=v.dtype, device=v.device)
start, end = self.start, self.start + T
if end > c.max_len:
raise ValueError(f"KV cache full: position {end} > max_len {c.max_len}")
c.k[i][:, :, start:end] = k
c.v[i][:, :, start:end] = v
return c.k[i][:, :, :end], c.v[i][:, :, :end]
class SlotKVCache:
"""S slots, each holding one sequence with its own length (continuous batching, 4.4).
Slot b holds keys at positions first[b] .. lens[b]-1 (positions before first[b] are
left padding from a batched prefill). A decode step feeds one token per slot, at
position lens[b]. Free slots still run (their output is ignored) so shapes stay fixed.
Every slot reserves max_len positions whatever its request needs: the waste that
paged KV (4.5) removes. `used_tokens()` vs `reserved_tokens()` measures it.
"""
def __init__(self, n_layers: int, n_slots: int, n_kv_heads: int, head_dim: int,
max_len: int, dtype: torch.dtype, device):
shape = (n_slots, n_kv_heads, max_len, head_dim)
self.k = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(n_layers)]
self.v = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(n_layers)]
self.n_slots, self.max_len = n_slots, max_len
# CPU copies drive scheduling without GPU syncs; device copies feed RoPE and masks.
self.first = [0] * n_slots
self.lens = [0] * n_slots
self.active = [False] * n_slots
self._lens_dev = torch.zeros(n_slots, dtype=torch.long, device=device)
self._first_dev = torch.zeros(n_slots, dtype=torch.long, device=device)
self._active_dev = torch.zeros(n_slots, dtype=torch.long, device=device)
self._rows = torch.arange(n_slots, device=device)
self.attn_kwargs: dict = {}
self._end = 0
# 4.11b: a fixed attention window (a graph bucket) instead of "as far as the longest
# slot", so the step's shapes don't change between steps. None = eager default.
self.window: int | None = None
# -- scheduling --------------------------------------------------------
def free_slots(self) -> list[int]:
return [s for s in range(self.n_slots) if not self.active[s]]
def admit(self, slots: list[int], prefill: KVCache, pad_lens: list[int]) -> None:
"""Copy a finished batched prefill (row j → slot slots[j]) into the slots."""
L = prefill.pos
if L > self.max_len:
raise ValueError(f"prompt length {L} > slot max_len {self.max_len}")
idx = torch.tensor(slots, device=self._lens_dev.device)
for i in range(len(self.k)):
self.k[i][idx, :, :L] = prefill.k[i][:, :, :L]
self.v[i][idx, :, :L] = prefill.v[i][:, :, :L]
for s, p in zip(slots, pad_lens):
self.first[s], self.lens[s], self.active[s] = p, L, True
self._sync()
def release(self, slot: int) -> None:
self.active[slot] = False
self.first[slot] = self.lens[slot] = 0
self._sync()
def _sync(self) -> None:
"""Push the CPU bookkeeping to the device (only on admit / release, not every step)."""
self._lens_dev.copy_(torch.tensor(self.lens))
self._first_dev.copy_(torch.tensor(self.first))
self._active_dev.copy_(torch.tensor(self.active, dtype=torch.long))
def can_admit(self, lengths: list[int]) -> bool:
return True # a free slot is all a request needs
def reserve_step(self) -> list[int]:
return [] # never preempts: capacity is per slot
def used_tokens(self) -> int:
return sum(n - f for n, f, a in zip(self.lens, self.first, self.active) if a)
def allocated_tokens(self) -> int:
"""KV positions claimed by running sequences (a whole slot each)."""
return sum(self.active) * self.max_len
def reserved_tokens(self) -> int:
return self.n_slots * self.max_len
def nbytes(self) -> int:
return sum(t.numel() * t.element_size() for t in self.k + self.v)
# -- model protocol (decode only: one token per slot) ------------------
def begin_step(self, T: int, rope) -> tuple[Tensor, Tensor]:
if T != 1:
raise ValueError("SlotKVCache decodes one token per slot; prefill goes through KVCache")
if max(self.lens) >= self.max_len:
raise ValueError(f"a slot is full (max_len {self.max_len})")
need = max(self.lens) + 1 # attend only as far as the longest slot
if self.window is not None and self.window < need:
raise ValueError(f"window {self.window} < {need} keys")
self._end = self.window or need
k_pos = torch.arange(self._end, device=self._lens_dev.device)
valid = (k_pos >= self._first_dev[:, None]) & (k_pos <= self._lens_dev[:, None])
self.attn_kwargs = {"attn_mask": valid[:, None, None, :]} # (S, 1, 1, end)
return rope.at(self._lens_dev[:, None])
def end_step(self, T: int) -> None:
self.advance_host()
self._lens_dev += self._active_dev # same update on the device, no copy
def advance_host(self) -> None:
"""The host half of end_step. A CUDA-graph replay runs only the device half, so the
graph decoder calls this after each replay."""
for s in range(self.n_slots):
if self.active[s]:
self.lens[s] += 1
def layer(self, i: int) -> "SlotLayerKV":
return SlotLayerKV(self, i)
class SlotLayerKV:
def __init__(self, cache: SlotKVCache, layer_idx: int):
self.cache = cache
self.i = layer_idx
@property
def attn_kwargs(self) -> dict:
return self.cache.attn_kwargs
def update(self, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
"""k, v (S, H, 1, D): write each slot's token at its own position lens[s]."""
c, i = self.cache, self.i
c.k[i][c._rows, :, c._lens_dev] = k[:, :, 0]
c.v[i][c._rows, :, c._lens_dev] = v[:, :, 0]
return c.k[i][:, :, :c._end], c.v[i][:, :, :c._end]
|