squad-0.8b 2026-10-26: retrained e2e expert (real web-app data, held-out tests 29/82 vs 17/82), prefix cache in the runtime
091795c verified | """The cross-request state cache: a text row continues from the deepest cached model state its tokens share. | |
| An agent asking one question at a time sends the same state again and again (the next question about the same page), | |
| or a state that only grows (one more recent action appended). A row here is [state prefix | question suffix]; its | |
| forward runs in chunks that continue from a saved model state, and the state at chosen positions of the prefix is | |
| kept on the GPU (LRU, bounded in MB) for later rows that share those tokens: | |
| - a checkpoint at position c: per Gated DeltaNet layer the recurrent state (fp32) and the last (kernel - 1) conv | |
| inputs before c; the full-attention layers' K / V of positions [0, c) (shared by every checkpoint of one row); | |
| the router's feature sum over [0, c) (specialists); and the route its slot layers ran with; | |
| - checkpoints are made at the state end p (the same state again: only the question is computed), at p - `tail` | |
| (a state that grows at its end: one more action appended), and at every multiple of `stride` (a state that | |
| changes in its middle, e.g. a new page after the same goal and steps); | |
| - a row restores the deepest checkpoint whose tokens it shares and runs the rest in chunks: the layers below the | |
| first slot layer (`split`; the router reads layer 5) over [c, end) first, which gives the route, then the layers | |
| from `split` on with that route's slot (graphs keyed by chunk length and attention range only: the slot index is a | |
| device tensor). The answer: the letter logits at the answer slots of the last chunk, masked, softmax. | |
| - each chunk is one CUDA graph replay: the conv runs on [saved conv inputs | chunk], the delta rule from the saved | |
| state (padding gets g = 0 and beta = 0, so the state stops at the real end), attention over [cached K / V | causal | |
| chunk] with rotary positions offset by the chunk start. | |
| Equal to the uncached forward up to bf16 round-off (other chunk boundaries in the delta rule and the conv, a masked | |
| SDPA instead of the causal kernel, the router's feature sum added in another order). A whole repeated request is | |
| answered by qev.py's memo before it gets here. | |
| """ | |
| from __future__ import annotations | |
| import time | |
| from collections import OrderedDict | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from . import model as M | |
| from .model import FULL, LINEAR, Batch | |
| from .prefix import PrefixQwen35 | |
| CHUNKS = (16, 32, 64, 128, 192, 256, 320, 384, 448, 512, 640, 768, 1024, 1536, 2048) | |
| SPANS = (1024, 2048, 4096, 8192) # attention key range of a chunk graph (chunk start + chunk length <= span) | |
| ATTN = {"kind": "varlen"} # chunk attention: "mask" (SDPA, boolean mask), "efficient" (SDPA mem-efficient, K / V | |
| # expanded to the query heads), "varlen" (torch's varlen FlashAttention, causal aligned to the chunk's end) | |
| def _bucket(n, buckets): | |
| for b in buckets: | |
| if n <= b: | |
| return b | |
| return None | |
| class CacheQwen35(PrefixQwen35): | |
| """PrefixQwen35 plus the chunk mode (b.ext): continue from / write to a StateCache's static buffers.""" | |
| def _gdn(self, i, x, b: Batch): | |
| ext = getattr(b, "ext", None) | |
| if ext is None: | |
| return super()._gdn(i, x, b) | |
| d = self.L[i] | |
| m = d["mod"] | |
| T = x.shape[0] | |
| kd, vd, nv = d["key_dim"], d["value_dim"], d["nv"] | |
| mixed, z, bb, a = torch.split(self.fused[i]["inp"](x), [2 * kd + vd, vd, nv, nv], dim=-1) | |
| st = ext["store"] | |
| li = st.gdn[i] | |
| Kc = st.Kc | |
| xin = torch.cat([st.conv[li], mixed], 0) # [saved conv inputs | chunk] | |
| y = M.fla_conv(xin[None], weight=d["conv_w"], bias=None, activation="silu") | |
| y = (y[0] if isinstance(y, tuple) else y)[:, Kc:] | |
| st.conv[li].copy_(xin.index_select(0, ext["tail_idx"])) # the last Kc real inputs (old ones if the chunk is short) | |
| q, k, v = torch.split(y[0], [kd, kd, vd], dim=-1) | |
| q = q.reshape(1, T, -1, d["hk"]) | |
| k = k.reshape(1, T, -1, d["hk"]) | |
| v = v.reshape(1, T, -1, d["hv"]) | |
| g, beta = M.gdn_gates(a, bb, m.A_log, m.dt_bias) | |
| g = g * ext["mask"][:, None] # padding: no decay, no write | |
| beta = beta * ext["mask"][:, None].to(beta.dtype) | |
| if d["nv"] // d["nk"] > 1: | |
| q = q.repeat_interleave(d["nv"] // d["nk"], dim=2) | |
| k = k.repeat_interleave(d["nv"] // d["nk"], dim=2) | |
| o, S = M.fla_gdr(q, k, v, g=g[None], beta=beta[None], initial_state=st.S[li][None], output_final_state=True, | |
| use_qk_l2norm_in_kernel=True) | |
| st.S[li].copy_(S[0]) | |
| o = M.gated_rmsnorm(o.reshape(-1, d["hv"]), z.reshape(-1, d["hv"]), m.norm.weight, m.norm.variance_epsilon) | |
| return o.reshape(T, -1) | |
| def _attn(self, i, x, b: Batch, rope): | |
| ext = getattr(b, "ext", None) | |
| if ext is None: | |
| return super()._attn(i, x, b, rope) | |
| d = self.L[i] | |
| m = d["mod"] | |
| T, hd, nh, nkv = x.shape[0], d["hd"], d["nh"], d["nkv"] | |
| qg, k, v = torch.split(self.fused[i]["inp"](x), [2 * nh * hd, nkv * hd, nkv * hd], dim=-1) | |
| q, gate = torch.chunk(qg.view(T, nh, 2 * hd), 2, dim=-1) | |
| gate = gate.reshape(T, nh * hd) | |
| q = M.rmsnorm(q, m.q_norm.weight, m.q_norm.eps) | |
| k = M.rmsnorm(k.view(T, nkv, hd), m.k_norm.weight, m.k_norm.eps) | |
| v = v.view(T, nkv, hd) | |
| cos, sin = rope | |
| q, k = M.apply_rope(q, cos, sin), M.apply_rope(k, cos, sin) | |
| st = ext["store"] | |
| la = st.att[i] | |
| P = ext["P"] | |
| st.K[la].index_copy_(0, ext["kv_idx"], k) | |
| st.V[la].index_copy_(0, ext["kv_idx"], v) | |
| kind = ATTN["kind"] | |
| if kind == "varlen": | |
| from torch.nn.attention.varlen import varlen_attn | |
| o = varlen_attn(q, st.K[la, :P], st.V[la, :P], ext["cu_q"], ext["cu_k"], T, P, scale=d["scale"], | |
| window_size=(-1, 0), enable_gqa=True) | |
| return o.reshape(T, nh * hd) * torch.sigmoid(gate) | |
| kk, vv = st.K[la, :P][None].transpose(1, 2), st.V[la, :P][None].transpose(1, 2) | |
| if kind == "efficient": | |
| from torch.nn.attention import SDPBackend, sdpa_kernel | |
| kk, vv = kk.repeat_interleave(nh // nkv, dim=1), vv.repeat_interleave(nh // nkv, dim=1) | |
| with sdpa_kernel([SDPBackend.EFFICIENT_ATTENTION]): | |
| o = F.scaled_dot_product_attention(q[None].transpose(1, 2), kk, vv, attn_mask=ext["amask"], | |
| scale=d["scale"]) | |
| else: | |
| o = F.scaled_dot_product_attention(q[None].transpose(1, 2), kk, vv, attn_mask=ext["amask"], | |
| scale=d["scale"], enable_gqa=True) | |
| return o[0].transpose(0, 1).reshape(T, nh * hd) * torch.sigmoid(gate) | |
| class _Group: | |
| """The checkpoints made by one row: its prefix tokens and the K / V of [0, p) they share.""" | |
| __slots__ = ("K", "V", "ckpts", "ids", "nbytes") | |
| def __init__(self, ids: np.ndarray): | |
| self.ids = ids | |
| self.K = self.V = None | |
| self.ckpts: list = [] | |
| self.nbytes = 0 | |
| class _Ckpt: | |
| __slots__ = ("fsum", "group", "lo", "nbytes", "pos", "route", "up") | |
| def __init__(self, group, pos, route, lo, up, fsum): | |
| self.group, self.pos, self.route, self.lo, self.up, self.fsum = group, pos, route, lo, up, fsum | |
| self.nbytes = sum(t.numel() * t.element_size() for t in (*lo, *up) if t is not None) + ( | |
| fsum.numel() * 4 if fsum is not None else 0) | |
| class StateCache: | |
| """See the module docstring. model: a CacheQwen35; slots: qev_fast.specialists.Specialists or None (the trunk).""" | |
| def __init__(self, model: CacheQwen35, pad_id: int, slots=None, mb: int = 0, stride: int = 2048, tail: int = 16, | |
| min_prefix: int = 256, max_slots: int = 4, align: int = 64, device="cuda"): | |
| self.model, self.pad_id, self.sl = model, pad_id, slots | |
| self.mb, self.stride, self.tail, self.min_prefix, self.max_slots = mb, stride, tail, min_prefix, max_slots | |
| self.align = align # checkpoints and chunk starts at multiples of this (fla's delta-rule chunk: 64) | |
| n = len(model.L) | |
| self.split = slots.b.layers[0] if slots is not None else n // 2 # the first slot layer | |
| self.router_layer = slots.router_layer if slots is not None else 0 | |
| ok_router = slots is None or "rmsnorm" in slots.router.features | |
| self.enabled_ok = ok_router and (slots is None or self.router_layer <= self.split) | |
| self.gdn = {i: j for j, i in enumerate(i for i, k in enumerate(model.kinds) if k == LINEAR)} | |
| self.att = {i: j for j, i in enumerate(i for i, k in enumerate(model.kinds) if k == FULL)} | |
| self.n_lo_g = sum(1 for i in self.gdn if i < self.split) | |
| self.n_lo_a = sum(1 for i in self.att if i < self.split) | |
| L0 = model.L[next(iter(self.gdn))] | |
| A0 = model.L[next(iter(self.att))] | |
| dt = model.embed.weight.dtype | |
| C = 2 * L0["key_dim"] + L0["value_dim"] | |
| self.Kc = L0["conv_w"].shape[-1] - 1 | |
| H = model.embed.weight.shape[1] | |
| self.maxp = SPANS[-1] | |
| self.conv = torch.zeros(len(self.gdn), self.Kc, C, device=device, dtype=dt) | |
| self.S = torch.zeros(len(self.gdn), L0["nv"], L0["hk"], L0["hv"], device=device, dtype=torch.float32) | |
| self.K = torch.zeros(len(self.att), self.maxp, A0["nkv"], A0["hd"], device=device, dtype=dt) | |
| self.V = torch.zeros_like(self.K) | |
| self.h = torch.zeros(self.maxp, H, device=device, dtype=dt) # the hidden state entering layer `split` | |
| self.Tc = torch.zeros((), dtype=torch.long, device=device) # chunk start | |
| self.Lr = torch.zeros((), dtype=torch.long, device=device) # real chunk length | |
| self.fout = torch.zeros(H, device=device, dtype=torch.float32) # the chunk's router feature sum | |
| self.slot_dev = torch.zeros(1, dtype=torch.long, device=device) # the loaded-slot index of the upper pass | |
| self.graphs: dict = {} | |
| self.pool = torch.cuda.graph_pool_handle() | |
| self.ckpts: OrderedDict = OrderedDict() # id -> _Ckpt, least recently used first | |
| self.groups: dict = {} # id -> _Group | |
| self.bytes = 0 | |
| try: # FlashAttention over [cached K / V | chunk] (torch's varlen kernel), else the mem-efficient SDPA | |
| from torch.nn.attention.varlen import varlen_attn # noqa: F401 | |
| except ImportError: | |
| ATTN["kind"] = "efficient" | |
| self.last_route = 0 # the route guessed for a row with no cached state | |
| self.cold_ok = None # callable: may a row with no cached state take this path now (FastQEV: no one waits) | |
| self.profile = False # diagnostics: per-phase milliseconds of the last row (synchronising) in last_timing | |
| self.last_timing: dict = {} | |
| self.stats = {"rows": 0, "hits": 0, "partial": 0, "cold": 0, "mispredict": 0, "tokens_reused": 0, | |
| "tokens_computed": 0, "ckpts_made": 0, "evicted": 0, "captures": 0, "capture_s": 0.0, | |
| "failed": []} | |
| # ---- graphs --------------------------------------------------------------------------------------------------- | |
| def _state(self, P, Lb, part): | |
| S = self.max_slots | |
| b = Batch([[self.pad_id] * Lb], [[0] * S], "cuda") | |
| st = {"b": b, "P": P, "Lb": Lb, "part": part, "nopts": torch.ones(S, dtype=torch.long, device="cuda")} | |
| return st | |
| def _ext(self, st): | |
| """The chunk tensors derived on the device from Tc / Lr (inside the graph).""" | |
| P, Lb = st["P"], st["Lb"] | |
| ar = torch.arange(Lb, device="cuda") | |
| cols = torch.arange(P, device="cuda") | |
| rows = ar[:, None] | |
| amask = (cols[None] < self.Tc) | ((cols[None] >= self.Tc) & (cols[None] - self.Tc <= rows)) | |
| cu_q = torch.stack([torch.zeros((), dtype=torch.int32, device="cuda"), | |
| torch.full((), Lb, dtype=torch.int32, device="cuda")]) | |
| cu_k = torch.stack([torch.zeros((), dtype=torch.int32, device="cuda"), (self.Tc + Lb).to(torch.int32)]) | |
| return {"store": self, "P": P, "mask": (ar < self.Lr).float(), "cu_q": cu_q, "cu_k": cu_k, "tail_idx": self.Lr + torch.arange(self.Kc, device="cuda"), | |
| "kv_idx": self.Tc + ar, "amask": amask, "pos": self.Tc + ar} | |
| def _forward(self, st): | |
| fm, b = self.model, st["b"] | |
| b.ext = self._ext(st) | |
| try: | |
| rope = M.rope_tables(b.ext["pos"], fm.inv_freq, fm.embed.weight.dtype) | |
| n = len(fm.L) | |
| part = st["part"] # "lo": layers < split; "up0" / "up1": the rest (trunk / slot_dev); "f0" / "f1": all | |
| if part in ("lo", "f0", "f1"): | |
| h = fm.embed(b.ids) | |
| for i in range(self.split): | |
| h = fm.layer(i, h, b, rope) | |
| if self.sl is not None and i == self.router_layer - 1: | |
| x = h.float() | |
| x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.sl.router.eps) | |
| self.fout.copy_((x * b.ext["mask"][:, None]).sum(0)) | |
| if part == "lo": | |
| self.h.index_copy_(0, b.ext["kv_idx"], h) | |
| return self.fout | |
| else: | |
| h = self.h.index_select(0, b.ext["kv_idx"]) | |
| ctx = self.sl.ctx if self.sl is not None else None | |
| if ctx is not None: | |
| ctx.clear() | |
| if part in ("up1", "f1"): | |
| ctx.dyn = self.slot_dev | |
| try: | |
| for i in range(self.split, n): | |
| h = fm.layer(i, h, b, rope, b.slots if i == n - 1 else None) | |
| finally: | |
| if ctx is not None: | |
| ctx.clear() | |
| hs = M.rmsnorm(h, fm.norm_w, fm.eps) | |
| W = fm.embed.weight.index_select(0, fm.letters) | |
| logits = F.linear(hs, W).float() | |
| logits = logits.masked_fill(torch.arange(logits.shape[1], device=logits.device)[None] >= st["nopts"][:, None], | |
| float("-inf")) | |
| return torch.softmax(logits, -1) | |
| finally: | |
| b.ext = None | |
| def _graph(self, part, P, Lb): | |
| key = (part, P, Lb) | |
| if key not in self.graphs: | |
| t0 = time.perf_counter() | |
| st = self._state(P, Lb, part) | |
| self.Tc.fill_(0) # in range for this shape (the warm-up runs read them; they clobber the Store too, so | |
| self.Lr.fill_(Lb) # a row captures what it needs before restoring anything: _ensure) | |
| s = torch.cuda.Stream() | |
| s.wait_stream(torch.cuda.current_stream()) | |
| with torch.cuda.stream(s), torch.no_grad(): | |
| for _ in range(2): | |
| self._forward(st) | |
| torch.cuda.current_stream().wait_stream(s) | |
| g = torch.cuda.CUDAGraph() | |
| try: | |
| with torch.cuda.graph(g, pool=self.pool), torch.no_grad(): | |
| out = self._forward(st) | |
| self.graphs[key] = (st, out, g) | |
| except Exception as e: # noqa: BLE001 - this shape stays uncached (the row takes the plain path) | |
| self.stats["failed"].append(f"{key}: {type(e).__name__}: {str(e)[:160]}") | |
| torch.cuda.synchronize() | |
| self.graphs[key] = None | |
| self.stats["captures"] += 1 | |
| self.stats["capture_s"] += time.perf_counter() - t0 | |
| return self.graphs[key] | |
| def parts(self, all_parts: bool = False): | |
| """The graph kinds precaptured: the fused ones (the two-pass ones are captured when a guess fails).""" | |
| if all_parts: | |
| return ("f0", "lo", "up0") + (("f1", "up1") if self.sl is not None else ()) | |
| return ("f0",) + (("f1",) if self.sl is not None else ()) | |
| def precapture(self, lock=None) -> dict: | |
| """Capture every chunk graph now (under `lock` per graph when given).""" | |
| t0 = time.perf_counter() | |
| for P in SPANS: | |
| for Lb in CHUNKS: | |
| if Lb > P: | |
| continue | |
| for part in self.parts(): | |
| if (part, P, Lb) in self.graphs: | |
| continue | |
| if lock is not None: | |
| with lock: | |
| self._graph(part, P, Lb) | |
| else: | |
| self._graph(part, P, Lb) | |
| torch.cuda.synchronize() | |
| return {"graphs": sum(v is not None for v in self.graphs.values()), "s": round(time.perf_counter() - t0, 1), | |
| "failed": self.stats["failed"][:3]} | |
| def _shape(self, s, e): | |
| Lb = _bucket(e - s, CHUNKS) | |
| P = _bucket(s + Lb, SPANS) if Lb is not None else None | |
| return P, Lb | |
| def _ensure(self, spans, parts): | |
| """Capture the graphs of these chunks [(s, e)] and kinds before a row touches the Store.""" | |
| for s, e in spans: | |
| P, Lb = self._shape(s, e) | |
| if P is None: | |
| raise _NotCached | |
| for part in parts: | |
| if self._graph(part, P, Lb) is None: | |
| raise _NotCached | |
| def _run(self, part, ids, s, e, slots=None, nopts=None): | |
| """One chunk [s, e) of ids through `part` -> its output tensor (lo: the feature sum; up: probabilities).""" | |
| Lr = e - s | |
| P, Lb = self._shape(s, e) | |
| if P is None: | |
| raise _NotCached | |
| hit = self._graph(part, P, Lb) | |
| if hit is None: | |
| raise _NotCached | |
| st, out, g = hit | |
| b = st["b"] | |
| buf = torch.full((Lb,), self.pad_id, dtype=torch.long) | |
| buf[:Lr] = torch.as_tensor(ids[s:e]) | |
| b.ids.copy_(buf, non_blocking=True) | |
| self.Tc.fill_(s) | |
| self.Lr.fill_(Lr) | |
| if slots is not None: | |
| sl = torch.zeros(self.max_slots, dtype=torch.long) | |
| no = torch.ones(self.max_slots, dtype=torch.long) | |
| sl[: len(slots)] = torch.as_tensor([x - s for x in slots]) | |
| no[: len(nopts)] = torch.as_tensor(nopts) | |
| b.slots.copy_(sl, non_blocking=True) | |
| st["nopts"].copy_(no, non_blocking=True) | |
| g.replay() | |
| return out | |
| # ---- checkpoints ---------------------------------------------------------------------------------------------- | |
| def recapture(self): | |
| """Drop the chunk graphs (they are captured again on use): after changing ATTN.""" | |
| self.graphs.clear() | |
| self.pool = torch.cuda.graph_pool_handle() # a fresh pool: capturing into one whose graphs are gone asserts | |
| self.clear() | |
| def clear(self): | |
| self.ckpts.clear() | |
| self.groups.clear() | |
| self.bytes = 0 | |
| def info(self) -> dict: | |
| return {**{k: v for k, v in self.stats.items() if k != "failed"}, "mb": self.mb, "used_mb": round(self.bytes / 2**20, 1), | |
| "checkpoints": len(self.ckpts), "groups": len(self.groups), "stride": self.stride, "tail": self.tail, | |
| "failed": self.stats["failed"][:3]} | |
| def _lookup(self, x: np.ndarray): | |
| """-> checkpoints usable by a row whose prefix tokens are x (pos <= the tokens it shares with their row).""" | |
| out = [] | |
| for g in self.groups.values(): | |
| m = min(len(g.ids), len(x)) | |
| if m == 0 or g.ids[0] != x[0]: | |
| continue | |
| ne = np.flatnonzero(g.ids[:m] != x[:m]) | |
| lcp = int(ne[0]) if len(ne) else m | |
| out += [c for c in g.ckpts if c.pos <= lcp] | |
| return out | |
| def _restore(self, c, lo=True, up=True): | |
| nl, na = self.n_lo_g, self.n_lo_a | |
| if c is None: | |
| if lo: | |
| self.conv[:nl].zero_() | |
| self.S[:nl].zero_() | |
| if up: | |
| self.conv[nl:].zero_() | |
| self.S[nl:].zero_() | |
| return | |
| if lo: | |
| self.conv[:nl].copy_(c.lo[0]) | |
| self.S[:nl].copy_(c.lo[1]) | |
| self.K[:na, : c.pos].copy_(c.group.K[:na, : c.pos]) | |
| self.V[:na, : c.pos].copy_(c.group.V[:na, : c.pos]) | |
| if up: | |
| self.conv[nl:].copy_(c.up[0]) | |
| self.S[nl:].copy_(c.up[1]) | |
| self.K[na:, : c.pos].copy_(c.group.K[na:, : c.pos]) | |
| self.V[na:, : c.pos].copy_(c.group.V[na:, : c.pos]) | |
| def _touch(self, c): | |
| self.ckpts.move_to_end(id(c)) | |
| def _evict(self): | |
| cap = self.mb * 2**20 | |
| while self.ckpts and self.bytes > cap: | |
| _, c = self.ckpts.popitem(last=False) | |
| self.bytes -= c.nbytes | |
| g = c.group | |
| g.ckpts.remove(c) | |
| self.stats["evicted"] += 1 | |
| if not g.ckpts: | |
| self.bytes -= g.nbytes | |
| self.groups.pop(id(g), None) | |
| def _down(self, x: int) -> int: | |
| return x - x % self.align if self.align else x | |
| def _bounds(self, c, p): | |
| """The chunk ends (and checkpoint positions) of a state of p tokens computed from c: the multiples of stride, | |
| p - tail and p (each rounded down to `align`); no chunk longer than the largest chunk graph.""" | |
| e = self._down(p) | |
| b = {k for k in range(self.stride, e, self.stride) if k > c} if self.stride else set() | |
| if self.tail and self._down(p - self.tail) > c: | |
| b.add(self._down(p - self.tail)) | |
| if e > c: | |
| b.add(e) | |
| out, s = [], c | |
| for x in sorted(b): | |
| while x - s > CHUNKS[-1]: | |
| s += CHUNKS[-1] | |
| out.append(s) | |
| out.append(x) | |
| s = x | |
| return out | |
| # ---- a row ------------------------------------------------------------------------------------------------------ | |
| def eligible(self, e: dict) -> bool: | |
| p = self._down(e.get("state_len", 0)) | |
| n = len(e["ids"]) | |
| return (self.mb > 0 and self.enabled_ok and not e.get("images") and p >= self.min_prefix and 0 < len(e["slots"]) | |
| <= self.max_slots and min(e["slots"]) >= p and n - p <= CHUNKS[-1] and p + _bucket(n - p, CHUNKS) <= self.maxp) | |
| def row(self, e: dict): | |
| """One eligible text row (encoder dict with "state_len") -> its slot probabilities (lists over each slot's | |
| options) and sets e["route"] (specialists); None when a chunk graph cannot be captured. The GPU lock must be | |
| held.""" | |
| try: | |
| return self._row(e) | |
| except _NotCached: | |
| return None | |
| def _tick(self, name): | |
| if self.profile: | |
| torch.cuda.synchronize() | |
| t = time.perf_counter() | |
| self.last_timing[name] = self.last_timing.get(name, 0.0) + 1000 * (t - self._t) | |
| self._t = t | |
| def _spans(self, c, bounds, p, n): | |
| return list(zip([c] + bounds[:-1], bounds)) + [(p, n)] | |
| def _part(self, kind, route): | |
| return kind + ("1" if route else "0") | |
| def _slot(self, route): | |
| if route: | |
| self.slot_dev.fill_(self.sl.loaded.index(route)) | |
| def _route(self, fsum, n): | |
| if self.sl is None: | |
| return 0 | |
| return int(self.sl.pick(self.sl.router.logits((fsum / n)[None]))[0]) | |
| def _row(self, e: dict): | |
| if self.profile: | |
| torch.cuda.synchronize() | |
| self.last_timing = {} | |
| self._t = time.perf_counter() | |
| ids, p_state = e["ids"], int(e["state_len"]) | |
| p = self._down(p_state) # the deepest checkpoint this row can make; [p, n) is computed every time | |
| n = len(ids) | |
| x = np.asarray(ids[:p], dtype=np.int64) | |
| cands = self._lookup(x) | |
| best = max(cands, key=lambda c: c.pos) if cands else None | |
| c0 = best.pos if best is not None else 0 | |
| if best is None and self.cold_ok is not None and not self.cold_ok(): | |
| return None # a new state while other calls wait: the plain path (no checkpoints) | |
| self.stats["rows"] += 1 | |
| self.stats["hits" if c0 == p else ("partial" if c0 > 0 else "cold")] += 1 | |
| # every layer at once, the slot layers with the route guessed (the reused state's, else the last one); the | |
| # route the row really takes comes from the router's feature sum after the last chunk | |
| guess = best.route if best is not None else self.last_route | |
| bounds = self._bounds(c0, p_state) if c0 < p else [] | |
| part = self._part("f", guess) | |
| self._ensure(self._spans(c0, bounds, p, n), (part,)) # captures (if any) before the restore | |
| self._tick("lookup") | |
| self._restore(best) | |
| self._slot(guess) | |
| fsum = best.fsum.clone() if (best is not None and best.fsum is not None) else torch.zeros_like(self.fout) | |
| snaps = {} | |
| s = c0 | |
| for b_ in bounds: | |
| self._run(part, ids, s, b_) | |
| fsum += self.fout | |
| snaps[b_] = self._snap(fsum) | |
| s = b_ | |
| self._tick("prefix") | |
| out = self._run(part, ids, p, n, e["slots"], e["nopts"]) | |
| fsum += self.fout | |
| route = self._route(fsum, n) | |
| self._tick("suffix") | |
| e["route"] = route | |
| if route != guess: # the slot layers ran with another slot: again with this one, from its deepest state | |
| self.stats["mispredict"] += 1 | |
| return self._two_pass(e, route, cands, ids, p_state, p, n, x) | |
| probs = out[: len(e["slots"])].cpu() | |
| self.last_route = route | |
| self._done(best, c0, n, x, p, bounds, snaps, route) | |
| return [probs[j, : e["nopts"][j]].tolist() for j in range(len(e["slots"]))] | |
| def _snap(self, fsum): | |
| nl = self.n_lo_g | |
| return ((self.conv[:nl].clone(), self.S[:nl].clone()), (self.conv[nl:].clone(), self.S[nl:].clone()), | |
| fsum.clone() if self.sl is not None else None) | |
| def _two_pass(self, e, route, cands, ids, p_state, p, n, x): | |
| """The row again with a known route: the layers below `split` from the deepest checkpoint of any route (they | |
| do not depend on it) are not reused here: both halves start from this route's deepest checkpoint.""" | |
| alt = [c for c in cands if c.route == route] | |
| start = max(alt, key=lambda c: c.pos) if alt else None | |
| c1 = start.pos if start is not None else 0 | |
| bounds = self._bounds(c1, p_state) if c1 < p else [] | |
| up = self._part("up", route) | |
| self._ensure(self._spans(c1, bounds, p, n), ("lo", up)) | |
| self._restore(start) | |
| fsum = start.fsum.clone() if (start is not None and start.fsum is not None) else torch.zeros_like(self.fout) | |
| lo_snap = {} | |
| s = c1 | |
| for b_ in bounds: # the lower layers: the hidden state entering `split` and their checkpoints | |
| self._run("lo", ids, s, b_) | |
| fsum += self.fout | |
| lo_snap[b_] = self._snap(fsum) | |
| s = b_ | |
| self._run("lo", ids, p, n) | |
| self._slot(route) | |
| snaps = {} | |
| s = c1 | |
| for b_ in bounds: # the upper layers with this route's slot | |
| self._run(up, ids, s, b_) | |
| nl = self.n_lo_g | |
| snaps[b_] = (lo_snap[b_][0], (self.conv[nl:].clone(), self.S[nl:].clone()), lo_snap[b_][2]) | |
| s = b_ | |
| out = self._run(up, ids, p, n, e["slots"], e["nopts"]) | |
| probs = out[: len(e["slots"])].cpu() | |
| self.last_route = route | |
| self._done(start, c1, n, x, p, bounds, snaps, route) | |
| return [probs[j, : e["nopts"][j]].tolist() for j in range(len(e["slots"]))] | |
| def _done(self, start, c0, n, x, p, bounds, snaps, route): | |
| """Bookkeeping after a row: LRU, statistics, the new checkpoints (one group sharing the K / V of [0, p)).""" | |
| self.stats["tokens_reused"] += c0 | |
| self.stats["tokens_computed"] += n - c0 | |
| if start is not None and id(start) in self.ckpts: | |
| self._touch(start) | |
| if bounds: | |
| g = _Group(x) | |
| g.K = self.K[:, :p].clone() | |
| g.V = self.V[:, :p].clone() | |
| g.nbytes = 2 * g.K.numel() * g.K.element_size() | |
| self.groups[id(g)] = g | |
| self.bytes += g.nbytes | |
| for b_ in bounds: | |
| lo, up, fs = snaps[b_] | |
| ck = _Ckpt(g, b_, route, lo, up, fs) | |
| g.ckpts.append(ck) | |
| self.ckpts[id(ck)] = ck | |
| self.bytes += ck.nbytes | |
| self.stats["ckpts_made"] += 1 | |
| self._evict() | |
| self._tick("save") | |
| class _NotCached(Exception): | |
| pass | |