diff --git a/serve.py b/serve.py index 5558057..b8e3ce2 100644 --- 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 new file mode 100644 index 0000000..a7e5174 --- /dev/null +++ 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()