From 50a261cee36576989ffa05953eb4599bf45aa31b Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 4 Aug 2026 17:05:25 +0000 Subject: [PATCH] llm: simplify recurrent cache fix --- tinygrad/llm/kernels/__init__.py | 2 +- tinygrad/llm/model.py | 18 ++++++++---------- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/tinygrad/llm/kernels/__init__.py b/tinygrad/llm/kernels/__init__.py index aeff29b77d..191d58e8f6 100644 --- a/tinygrad/llm/kernels/__init__.py +++ b/tinygrad/llm/kernels/__init__.py @@ -63,8 +63,8 @@ def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, core, kq = Tensor.empty_like(v), (q*k).sum(-1).contiguous() srcs = (core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq) if start_pos is None: return Tensor.custom_kernel(*srcs, fxn=kernel)[0] - if start_pos.uop.op is not Ops.BIND: return Tensor.custom_kernel(*srcs, fxn=functools.partial(kernel, start_pos=start_pos.uop))[0] contig = tuple(x.uop if x.uop.op is Ops.AFTER else x.uop.contiguous() for x in srcs) params = tuple(UOp.placeholder_like(x, slot=i) for i,x in enumerate(contig)) + assert start_pos.uop.op is Ops.BIND call = kernel(*params, start_pos.uop.src[0]).call(*contig, start_pos.uop) return Tensor(contig[0].after(call)) diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index 80a6d96e18..fc6e47850a 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -413,22 +413,20 @@ class Transformer: return model, kv def get_start_pos(self, tokens:list[int]) -> int: - prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens))) if self.has_recurrent_block: - if prefix_len == len(self._cached_tokens): return prefix_len - checkpoint_prefix = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._checkpoint_tokens))) - return len(self._checkpoint_tokens) if checkpoint_prefix == len(self._checkpoint_tokens) else 0 + for cached in (self._cached_tokens, self._checkpoint_tokens): + if cached and len(cached) < len(tokens) and tokens[:len(cached)] == cached: return len(cached) + return 0 + prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens))) return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk) def _checkpoint_state(self, tokens:list[int]|None=None): states = [getattr(block, name) for block in self.blk for name in ("conv_state", "recurrent_state") if hasattr(block, name)] if not states: return - if not self._state_checkpoint: - self._state_checkpoint = [state.clone().realize() for state in states] - if tokens is not None: self._checkpoint_tokens = list(tokens) - return - dst, src = (self._state_checkpoint, states) if tokens is not None else (states, self._state_checkpoint) - Tensor.realize(*(d.assign(s) for d,s in zip(dst, src))) + if not self._state_checkpoint: self._state_checkpoint = [state.clone().realize() for state in states] + else: + dst, src = (self._state_checkpoint, states) if tokens is not None else (states, self._state_checkpoint) + Tensor.realize(*(d.assign(s) for d,s in zip(dst, src))) if tokens is not None: self._checkpoint_tokens = list(tokens) def warmup(self, chunk_size:int=256):