| """CUDA-graph decode (4.11, docs/INFERENCE.md). |
| |
| A decode step of a small model is launch-bound: ~10 kernels per layer, each a few µs of |
| GPU work, and the CPU needs about as long to launch each one. A CUDA graph records the |
| whole step once and replays it with a single launch. |
| |
| A graph replays fixed kernels on fixed addresses, so the step must have static shapes: |
| - StaticKVCache keeps the position as a device tensor. K/V are written with index_copy_ |
| at that position, RoPE is gathered at it, and every decode step attends over a fixed |
| window of W keys, with keys past the position masked out. |
| - W comes from length buckets (128, 256, ..., max_len). Each bucket has its own graph; |
| the step uses the smallest bucket that holds pos + 1 keys. All graphs share one memory |
| pool and the same static buffers (cache, input token, position). |
| - Prefill stays eager: it runs once per request and has a different shape each time. |
| Sampling stays eager too. |
| |
| Not bit-identical to eager generate(): eager decode attends over exactly pos + 1 keys |
| with no mask; the graph attends over W keys with a mask, and SDPA picks a different |
| kernel for that. Graph replay *is* bit-identical to the same static step run eagerly. |
| |
| Autocast runs with cache_enabled=False: its weight-cast cache doesn't survive capture. |
| """ |
|
|
| import time |
|
|
| import torch |
| from torch import Tensor |
|
|
| from tinyllm.kv_cache import causal_kwargs, padded_causal_mask |
| from tinyllm.model import Transformer |
|
|
|
|
| def default_buckets(max_len: int, first: int = 128) -> tuple[int, ...]: |
| """first, 2*first, ... below max_len, then max_len itself.""" |
| out, b = [], first |
| while b < max_len: |
| out.append(b) |
| b *= 2 |
| return (*out, max_len) |
|
|
|
|
| class StaticKVCache: |
| """Shared-position cache with fixed shapes (B rows, max_len positions). |
| |
| Same model protocol as kv_cache.KVCache. Two modes: |
| - prefill (eager, position 0): as KVCache, writes and returns positions 0..T-1. |
| - decode (`window` set, T=1): position from the device tensor `pos`; returns the first |
| `window` keys with a mask, so shapes don't depend on the position. |
| first_valid (B,) is the left-pad count per row; zeros when unpadded. |
| """ |
|
|
| def __init__(self, n_layers: int, batch: int, max_len: int, device): |
| self.n_layers, self.batch, self.max_len = n_layers, batch, max_len |
| self.k: list[Tensor | None] = [None] * n_layers |
| self.v: list[Tensor | None] = [None] * n_layers |
| self.pos = torch.zeros(1, dtype=torch.long, device=device) |
| self.first_valid = torch.zeros(batch, dtype=torch.long, device=device) |
| self.padded = False |
| self.n = 0 |
| self.window: int | None = None |
| self.attn_kwargs: dict = {} |
|
|
| def reset(self, pad_lens: Tensor | None = None) -> None: |
| self.n, self.window = 0, None |
| self.pos.zero_() |
| self.padded = pad_lens is not None |
| self.first_valid.copy_(pad_lens if self.padded else torch.zeros_like(self.first_valid)) |
|
|
| def begin_step(self, T: int, rope) -> tuple[Tensor, Tensor]: |
| if self.window is None: |
| if self.n != 0: |
| raise ValueError("StaticKVCache prefills only from position 0") |
| cos, sin = rope(T, 0) |
| self.attn_kwargs = ({"attn_mask": padded_causal_mask(0, T, self.first_valid)} |
| if self.padded else causal_kwargs(0, T, cos.device)) |
| return cos, sin |
| if T != 1: |
| raise ValueError("decode mode takes one token per row") |
| k_pos = torch.arange(self.window, device=self.pos.device) |
| valid = (k_pos <= self.pos) & (k_pos >= self.first_valid[:, None]) |
| self.attn_kwargs = {"attn_mask": valid[:, None, None, :]} |
| return rope.cos.index_select(0, self.pos), rope.sin.index_select(0, self.pos) |
|
|
| def end_step(self, T: int) -> None: |
| if self.window is None: |
| self.n = T |
| self.pos.fill_(T) |
| else: |
| self.pos.add_(1) |
|
|
| def layer(self, i: int) -> "StaticLayerKV": |
| return StaticLayerKV(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 StaticLayerKV: |
| def __init__(self, cache: StaticKVCache, 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]: |
| c, i = self.cache, self.i |
| B, H, T, D = k.shape |
| if c.k[i] is None: |
| shape = (c.batch, 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) |
| elif c.k[i].dtype != k.dtype: |
| raise ValueError(f"cache holds {c.k[i].dtype}, got {k.dtype}") |
| if c.window is None: |
| if T > c.max_len: |
| raise ValueError(f"KV cache full: position {T} > max_len {c.max_len}") |
| c.k[i][:, :, :T] = k |
| c.v[i][:, :, :T] = v |
| return c.k[i][:, :, :T], c.v[i][:, :, :T] |
| c.k[i].index_copy_(2, c.pos, k) |
| c.v[i].index_copy_(2, c.pos, v) |
| return c.k[i][:, :, :c.window], c.v[i][:, :, :c.window] |
|
|
|
|
| class GraphDecoder: |
| """Prefill eagerly, then decode one token per row with one CUDA-graph replay per step. |
| |
| Reusable: prefill() resets the cache, and graphs captured for one request serve the |
| next (same batch size). use_graph=False runs the identical static step eagerly, the |
| reference that replay must match bit for bit. |
| """ |
|
|
| def __init__(self, model: Transformer, batch: int, max_len: int | None = None, |
| buckets: tuple[int, ...] | None = None, use_graph: bool = True): |
| device = next(model.parameters()).device |
| if device.type != "cuda": |
| raise ValueError("CUDA graphs need a CUDA model") |
| self.model = model |
| self.max_len = max_len or model.cfg.max_seq_len |
| if self.max_len > model.cfg.max_seq_len: |
| raise ValueError(f"max_len {self.max_len} > max_seq_len {model.cfg.max_seq_len}") |
| self.buckets = tuple(sorted(buckets)) if buckets else default_buckets(self.max_len) |
| if self.buckets[-1] < self.max_len: |
| raise ValueError(f"largest bucket {self.buckets[-1]} < max_len {self.max_len}") |
| self.cache = StaticKVCache(model.cfg.n_layers, batch, self.max_len, device) |
| self.tok = torch.zeros(batch, 1, dtype=torch.long, device=device) |
| self.use_graph = use_graph |
| self.graphs: dict[int, tuple[torch.cuda.CUDAGraph, Tensor]] = {} |
| self.pool = torch.cuda.graph_pool_handle() |
|
|
| def _autocast(self): |
| return torch.autocast("cuda", dtype=torch.bfloat16, cache_enabled=False) |
|
|
| @torch.no_grad() |
| def prefill(self, idx: Tensor, pad_lens: Tensor | None = None) -> Tensor: |
| """idx (B, T), left-padded by pad_lens (B,) if given. Returns logits (B, T, V).""" |
| if idx.shape[0] != self.cache.batch: |
| raise ValueError(f"batch {idx.shape[0]} != decoder batch {self.cache.batch}") |
| self.cache.reset(pad_lens) |
| with self._autocast(): |
| logits, _ = self.model(idx, cache=self.cache) |
| return logits |
|
|
| def _forward(self) -> Tensor: |
| with self._autocast(): |
| logits, _ = self.model(self.tok, cache=self.cache) |
| return logits |
|
|
| def _capture(self, window: int) -> tuple[torch.cuda.CUDAGraph, Tensor]: |
| |
| |
| |
| saved = self.cache.pos.clone() |
| side = torch.cuda.Stream() |
| side.wait_stream(torch.cuda.current_stream()) |
| with torch.cuda.stream(side): |
| self._forward() |
| torch.cuda.current_stream().wait_stream(side) |
| self.cache.pos.copy_(saved) |
| graph = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(graph, pool=self.pool): |
| out = self._forward() |
| return graph, out |
|
|
| def bucket(self, n_keys: int) -> int: |
| return next(b for b in self.buckets if b >= n_keys) |
|
|
| @torch.no_grad() |
| def step(self, tok: Tensor) -> Tensor: |
| """tok (B, 1) → logits (B, 1, V). With graphs, the result is a static buffer that |
| the next step overwrites: read it (or sample from it) first.""" |
| c = self.cache |
| if c.n >= self.max_len: |
| raise ValueError(f"KV cache full: position {c.n + 1} > max_len {self.max_len}") |
| c.window = self.bucket(c.n + 1) |
| self.tok.copy_(tok) |
| if not self.use_graph: |
| out = self._forward() |
| else: |
| if c.window not in self.graphs: |
| self.graphs[c.window] = self._capture(c.window) |
| graph, out = self.graphs[c.window] |
| graph.replay() |
| c.n += 1 |
| return out |
|
|
|
|
| class SlotGraphDecoder: |
| """The continuous engine's decode step (one token per SlotKVCache slot) as CUDA graphs |
| (4.11b). Admission, release and prefill stay eager in the engine. |
| |
| Why it can be a graph: the batch is always n_slots rows (free slots compute too), and |
| slot lengths / pad starts / active flags live in fixed device tensors that admit() and |
| release() update in place. The one changing shape is the attention window, so the cache |
| gets a fixed window from the length buckets, with keys past each slot's length masked. |
| |
| Each step writes the token into a static buffer and replays the graph for the smallest |
| bucket that holds the longest slot + 1. The replay advances the device lengths; the host |
| lengths move in advance_host(). use_graph=False runs the identical windowed step eagerly, |
| the reference that replay must match bit for bit. |
| """ |
|
|
| def __init__(self, model: Transformer, slots, buckets: tuple[int, ...] | None = None, |
| use_graph: bool = True): |
| device = next(model.parameters()).device |
| if device.type != "cuda": |
| raise ValueError("CUDA graphs need a CUDA model") |
| if not hasattr(slots, "advance_host"): |
| raise ValueError("SlotGraphDecoder supports the slot cache (kv='slots') only") |
| self.model, self.slots = model, slots |
| self.buckets = tuple(sorted(buckets)) if buckets else default_buckets(slots.max_len) |
| if self.buckets[-1] < slots.max_len: |
| raise ValueError(f"largest bucket {self.buckets[-1]} < max_len {slots.max_len}") |
| self.tok = torch.zeros(slots.n_slots, 1, dtype=torch.long, device=device) |
| self.use_graph = use_graph |
| self.graphs: dict[int, tuple[torch.cuda.CUDAGraph, Tensor]] = {} |
| self.pool = torch.cuda.graph_pool_handle() |
|
|
| def _forward(self) -> Tensor: |
| with torch.autocast("cuda", dtype=torch.bfloat16, cache_enabled=False): |
| logits, _ = self.model(self.tok, cache=self.slots) |
| return logits |
|
|
| def _capture(self) -> tuple[torch.cuda.CUDAGraph, Tensor]: |
| |
| |
| c = self.slots |
| lens, lens_dev = list(c.lens), c._lens_dev.clone() |
| side = torch.cuda.Stream() |
| side.wait_stream(torch.cuda.current_stream()) |
| with torch.cuda.stream(side): |
| self._forward() |
| torch.cuda.current_stream().wait_stream(side) |
| c.lens[:] = lens |
| c._lens_dev.copy_(lens_dev) |
| graph = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(graph, pool=self.pool): |
| out = self._forward() |
| c.lens[:] = lens |
| return graph, out |
|
|
| def bucket(self, n_keys: int) -> int: |
| return next(b for b in self.buckets if b >= n_keys) |
|
|
| @torch.no_grad() |
| def capture(self, windows: tuple[int, ...] | None = None) -> float: |
| """Capture graphs ahead of time (default: every bucket), so no timed step pays for it. |
| Safe between steps: the warm-up's K/V writes land at each slot's next position, which |
| the next real step overwrites. Returns the seconds it took.""" |
| t0 = time.perf_counter() |
| c = self.slots |
| try: |
| for w in windows or self.buckets: |
| if w not in self.graphs: |
| c.window = w |
| self.graphs[w] = self._capture() |
| finally: |
| c.window = None |
| torch.cuda.synchronize() |
| return time.perf_counter() - t0 |
|
|
| @torch.no_grad() |
| def step(self, tok: Tensor) -> Tensor: |
| """tok (n_slots, 1) → logits (n_slots, 1, V). With graphs, the result is a static |
| buffer that the next step overwrites: sample from it first.""" |
| c = self.slots |
| need = max(c.lens) + 1 |
| if need > c.max_len: |
| raise ValueError(f"a slot is full (max_len {c.max_len})") |
| c.window = self.bucket(need) |
| self.tok.copy_(tok) |
| try: |
| if not self.use_graph: |
| return self._forward() |
| if c.window not in self.graphs: |
| self.graphs[c.window] = self._capture() |
| graph, out = self.graphs[c.window] |
| graph.replay() |
| c.advance_host() |
| return out |
| finally: |
| c.window = None |
|
|
