File size: 13,776 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 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 | """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) # device copy: graph input
self.first_valid = torch.zeros(batch, dtype=torch.long, device=device)
self.padded = False
self.n = 0 # host copy of pos, for bounds and bucket choice (no sync)
self.window: int | None = None # None = prefill mode
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: # eager prefill from position 0
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]) # (B, W)
self.attn_kwargs = {"attn_mask": valid[:, None, None, :]} # (B, 1, 1, W)
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) # captured into the graph; the host copy moves in step()
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]:
# Warm up on a side stream (lazy cuBLAS / SDPA setup must not happen during capture).
# The warm-up runs the real step: its K/V writes are the ones replay will make, and
# the position it advances is put back.
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]:
# The warm-up runs the real step (its K/V writes are the ones replay will make) and
# advances the lengths; capture runs the host half again. Put both back each time.
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
|