"""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) @torch.no_grad() 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