GroveMinting · Sep 22, 2026
diff --git a/serve.py b/serve.py
@@ -232,11 +232,25 @@ def build_prompt(messages, budget, prime=""):
     return prompt[-budget:] if len(prompt) > budget else prompt
 
 
+def _prefill_cache(model, tokens, chunk=512):
+    """Encode one self-contained window and return its cache and last logits."""
+    if tokens.shape[1] > model.cfg.block:
+        raise ValueError("prefill exceeds the model context")
+    caches = model.empty_caches()
+    logits = None
+    offset = 0
+    chunk = max(1, int(chunk))
+    for i in range(0, tokens.shape[1], chunk):
+        part = tokens[:, i:i + chunk]
+        logits = model(part, caches=caches, pos_offset=offset)[0]
+        offset += part.shape[1]
+    return caches, logits, offset
+
+
 @torch.no_grad()
 def stream(prompt, max_new):
     from minagi.config import get as _g, load as _lc
     from minagi.decode import pick_next
-    from minagi.stream import trim_caches
 
     c = _lc()
     strength = _g(c, "decoding.adapt_strength", 2.5)
@@ -246,22 +260,6 @@ def stream(prompt, max_new):
     # it is a far shorter span than pool.segment_chars, which is for reading.
     reselect = _g(c, "pool.reselect_chars", 64)
 
-    def where(caches):
-        """
-        The position the next character sits at: however much history the
-        cache still holds after trimming.
-
-        NOT a running count. Rotary tables are built for positions 0 to
-        block-1, so a counter that saturates at `block` asks for position
-        `block` on the very next character and the model refuses. Reading it
-        back off the cache cannot drift, because the cache is the thing the
-        positions have to agree with.
-        """
-        for c in caches:
-            if c.get("k") is not None:
-                return c["k"].shape[-2]
-        return 0
-
     model, tok = STATE["model"], STATE["tok"]
     device = next(model.parameters()).device
     ids = tok.encode(prompt).ids[-model.cfg.block:]
@@ -275,18 +273,11 @@ def where(caches):
     # keeps a working set steady while reading a stream, because a prompt is a
     # deliberate change of subject.
     model.choose_for(out, free=True)
-    caches = model.empty_caches()
-
     # Prefill in chunks, the way training reads a corpus. Feeding a long
     # prompt in one pass materialises activations for every position across
     # every block application at once, which is what puts a long context out
     # of reach; the cache carries the reach instead.
-    CHUNK = 512
-    logits = None
-    for i in range(0, out.shape[1], CHUNK):
-        part = out[:, i:i + CHUNK]
-        trim_caches(caches, model.cfg.block - part.shape[1])
-        logits = model(part, caches=caches, pos_offset=where(caches))[0]
+    caches, logits, offset = _prefill_cache(model, out)
 
     cur = out[:, -1:]
     produced = []
@@ -304,8 +295,20 @@ def where(caches):
                 yield {"swap": {"at": i, "moved": int(moved),
                                 "pool": resident_experts()}}
         if logits is None:
-            trim_caches(caches, model.cfg.block - cur.shape[1])
-            logits = model(cur, caches=caches, pos_offset=where(caches))[0]
+            if offset + cur.shape[1] > model.cfg.block:
+                # Cached keys have already been rotated at their old positions.
+                # Trimming them and deriving a new offset from the shorter cache
+                # reuses positions and corrupts every relative distance. Start
+                # a self-contained recent window instead, as RecurCoder.generate
+                # does, leaving half the context free before the next rebuild.
+                keep = max(1, model.cfg.block // 2)
+                recent = out[:, -keep:]
+                del caches                 # let the replacement reuse its VRAM
+                caches, logits, offset = _prefill_cache(
+                    model, recent)
+            else:
+                logits = model(cur, caches=caches, pos_offset=offset)[0]
+                offset += cur.shape[1]
         nxt = pick_next(logits[:, -1, :].float(), out, temperature=0.0,
                         adapt_strength=strength, adapt_decay=decay)
         out = torch.cat([out, nxt], dim=1)

diff --git a/tests/test_serve_context.py b/tests/test_serve_context.py
@@ -0,0 +1,102 @@
+import types
+import unittest
+import weakref
+
+import torch
+import torch.nn as nn
+
+import serve
+from minagi.recur import RecurCoder, RecurConfig
+from minagi.tokenizer import ByteTokenizer
+
+
+class RecordingModel(nn.Module):
+    def __init__(self, block=8):
+        super().__init__()
+        self.cfg = types.SimpleNamespace(block=block)
+        self.anchor = nn.Parameter(torch.zeros(()))
+        self.calls = []
+        self.tokens = []
+        self.full_cache_ref = None
+        self.released_before_reset = None
+
+    def empty_caches(self):
+        return [{"k": None, "v": None}]
+
+    def choose_for(self, tokens, free=False):
+        return 0
+
+    def forward(self, tokens, caches=None, pos_offset=0):
+        length = tokens.shape[1]
+        if pos_offset + length > self.cfg.block:
+            raise AssertionError("forward exceeded rotary table")
+        cache = caches[0]
+        previous = 0 if cache["k"] is None else cache["k"].shape[-2]
+        if pos_offset == 0 and previous == 0 and self.full_cache_ref is not None:
+            self.released_before_reset = self.full_cache_ref() is None
+        self.calls.append((pos_offset, length, previous))
+        self.tokens.append(tokens.detach().clone())
+        total = previous + length
+        cache["k"] = torch.zeros(1, 1, total, 1)
+        cache["v"] = torch.zeros(1, 1, total, 1)
+        if total == self.cfg.block:
+            self.full_cache_ref = weakref.ref(cache["k"])
+        logits = torch.full((tokens.shape[0], length, 265), -100.0)
+        nxt = (tokens + 1).clamp_max(255).unsqueeze(-1)
+        logits.scatter_(-1, nxt, 100.0)
+        return logits, None
+
+
+class ServeContextTest(unittest.TestCase):
+    def test_prefill_chunks_use_monotonic_positions(self):
+        model = RecordingModel(block=8)
+        tokens = torch.arange(8).unsqueeze(0)
+
+        caches, logits, offset = serve._prefill_cache(model, tokens, chunk=3)
+
+        self.assertEqual(model.calls, [(0, 3, 0), (3, 3, 3), (6, 2, 6)])
+        self.assertEqual(caches[0]["k"].shape[-2], 8)
+        self.assertEqual(logits.shape, (1, 2, 265))
+        self.assertEqual(offset, 8)
+
+    def test_generation_reprefills_recent_window_at_context_limit(self):
+        model = RecordingModel(block=8)
+        old_state = dict(serve.STATE)
+        serve.STATE.update(model=model, tok=ByteTokenizer(), learner=None)
+        try:
+            events = list(serve.stream("abcdefgh", max_new=10))
+        finally:
+            serve.STATE.clear()
+            serve.STATE.update(old_state)
+
+        self.assertEqual(model.calls[0], (0, 8, 0))
+        self.assertEqual(model.calls[1], (0, 4, 0))
+        self.assertGreaterEqual(model.calls.count((0, 4, 0)), 2)
+        self.assertEqual(model.tokens[1].tolist(), [[102, 103, 104, 105]])
+        reset = [tokens.tolist() for call, tokens in zip(model.calls, model.tokens)
+                 if call == (0, 4, 0)]
+        self.assertIn([[107, 108, 109, 110]], reset)
+        self.assertTrue(model.released_before_reset)
+        self.assertTrue(all(pos + length <= model.cfg.block
+                            for pos, length, _ in model.calls))
+        self.assertEqual(sum("t" in event for event in events), 10)
+        self.assertEqual("".join(event["t"] for event in events if "t" in event),
+                         "ijklmnopqr")
+
+    def test_chunked_prefill_matches_full_forward(self):
+        torch.manual_seed(0)
+        cfg = RecurConfig(vocab_size=265, d_model=8, n_head=1, d_ff=16,
+                          block=8, n_prelude=1, n_recur=1, n_coda=0,
+                          max_steps=1)
+        model = RecurCoder(cfg).eval()
+        tokens = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]])
+
+        full = model(tokens)[0][:, -1]
+        _, chunked, offset = serve._prefill_cache(model, tokens, chunk=3)
+
+        torch.testing.assert_close(chunked[:, -1], full)
+        self.assertEqual(offset, tokens.shape[1])
+
+
+if __name__ == "__main__":
+    unittest.main()