perf(prefill): export only final MLX snapshot

Eliminate quadratic full-cache serialization by computing every token normally but encoding and compressing only the final chained-prefix snapshot.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 2 people

docs/ops/distributed-prefill-kv-network.md

@@ -361,6 +361,10 @@ prevents adaptive shrink from consuming active reservations, then atomically
 publishes and leases the final snapshot before adding optional intermediate
 boundaries. The 16GB allens deployment uses a 1 GiB cache floor and 0.5 GiB
 memory reserve. Capacity failures are rejected before Prefill starts.
+The MLX worker exports and compresses only the final chained-prefix snapshot.
+The final hash commits every preceding token block, so exporting a growing full
+snapshot at every 64-token boundary is redundant and creates quadratic
+serialization/compression work for long Critic contexts.
 Interactive prompt templates are deterministic and contain no per-run nonce, so
 repeating the same task can reuse allens cold-tier and Primary hot-tier KV.
 

inference_engine/backends/mlx/prefill_worker.py

@@ -45,17 +45,7 @@ def compute_prefill(
                 raise InterruptedError("prefill job cancelled")
             first_end = min(size, len(tokens))
             self.verifier.prefill(tokens[:first_end])
-            blocks: list[CacheBlock] = [
-                self._snapshot(
-                    token_count=first_end,
-                    block_hash=block_hashes[0],
-                    compression=compression,
-                ),
-            ]
-            for block_index, start in enumerate(
-                range(first_end, len(tokens), size),
-                start=1,
-            ):
+            for start in range(first_end, len(tokens), size):
                 if cancelled.is_set():
                     raise InterruptedError("prefill job cancelled")
                 block = tokens[start:start + size]
@@ -65,12 +55,17 @@ def compute_prefill(
                     accepted=len(block),
                 )
                 self.verifier.next_token_logits = logits[-1].clone()
-                blocks.append(self._snapshot(
-                    token_count=min(start + size, len(tokens)),
-                    block_hash=block_hashes[block_index],
+            # Export exactly once. Intermediate full snapshots make encoding
+            # and compression quadratic in prompt length and are not required
+            # for correctness because the final chained hash commits every
+            # preceding token block.
+            return (
+                self._snapshot(
+                    token_count=len(tokens),
+                    block_hash=block_hashes[-1],
                     compression=compression,
-                ))
-            return blocks
+                ),
+            )
 
     def _snapshot(
         self,

inference_engine/distributed/prefill_cache.py

@@ -297,8 +297,11 @@ def publish_and_lease(
         lease_seconds: float = DEFAULT_LEASE_SECONDS,
     ) -> PrefixLease:
         """Atomically publish and pin the final computed snapshot."""
-        if not blocks or len(blocks) != len(block_hashes):
-            raise ValueError("one computed snapshot is required per block hash")
+        if not blocks or len(blocks) not in (1, len(block_hashes)):
+            raise ValueError(
+                "computed snapshots must contain only the final snapshot "
+                "or one snapshot per block hash",
+            )
         if lease_seconds <= 0:
             raise ValueError("lease_seconds must be > 0")
         final = blocks[-1]

inference_engine/distributed/prefill_worker.py

@@ -282,9 +282,10 @@ def _run(self, job_id: str) -> None:
             ))
             if job.cancelled.is_set():
                 raise InterruptedError("prefill job cancelled")
-            if len(blocks) != len(job.block_hashes):
+            if len(blocks) not in (1, len(job.block_hashes)):
                 raise RuntimeError(
-                    "prefill engine must return one snapshot per block hash",
+                    "prefill engine must return only the final snapshot "
+                    "or one snapshot per block hash",
                 )
             lease = self.cache_store.publish_and_lease(
                 blocks,

tests/inference_engine/distributed/test_prefill_cache.py

@@ -189,6 +189,21 @@ def test_publish_and_lease_atomically_pins_final_snapshot():
     assert not store.resize(1)
 
 
+def test_publish_and_lease_accepts_final_snapshot_only():
+    store = PrefixCacheStore(_compat(), max_bytes=10, node_id="x")
+    hashes = chained_block_hashes([1, 2, 3, 4], _compat())
+    final = CacheBlock.create(hashes[-1], 4, b"final!")
+    store.reserve("job", 6)
+    lease = store.publish_and_lease(
+        [final],
+        hashes,
+        reservation_id="job",
+    )
+    assert lease.hit_block_count == 2
+    assert lease.hit_token_count == 4
+    assert store.fetch(lease.lease_id) == (final,)
+
+
 def test_reservation_and_atomic_publish_validation_guards():
     store = PrefixCacheStore(_compat(), max_bytes=10, node_id="x")
     with pytest.raises(ValueError, match="reservation id"):
@@ -197,7 +212,7 @@ def test_reservation_and_atomic_publish_validation_guards():
     with pytest.raises(ValueError, match="duplicate"):
         store.reserve("job", 1)
     block = CacheBlock.create(bytes(32), 1, b"abc")
-    with pytest.raises(ValueError, match="one computed snapshot"):
+    with pytest.raises(ValueError, match="computed snapshots"):
         store.publish_and_lease([], [], reservation_id="job")
     with pytest.raises(ValueError, match="lease_seconds"):
         store.publish_and_lease(

tests/inference_engine/distributed/test_prefill_worker.py

@@ -8,6 +8,7 @@
 import pytest
 import pytest_asyncio
 
+from inference_engine.backends.mlx.prefill_worker import MLXPrefillComputeEngine
 from inference_engine.distributed.capability import (
     CacheCompatibility,
     CompressionCodec,
@@ -55,6 +56,55 @@ def compute_prefill(self, token_ids, block_hashes, *, compression, cancelled):
         ]
 
 
+def test_mlx_worker_exports_only_final_snapshot(monkeypatch):
+    class Logit:
+        def clone(self):
+            return self
+
+    class Verifier:
+        def __init__(self):
+            self.next_token_logits = None
+            self.forwarded = []
+
+        def prefill(self, tokens):
+            self.forwarded.extend(tokens)
+
+        def forward_block(self, tokens):
+            self.forwarded.extend(tokens)
+            return [Logit() for _ in tokens]
+
+        def commit_or_truncate(self, *, forwarded, accepted):
+            assert forwarded == accepted
+
+    verifier = Verifier()
+    engine = MLXPrefillComputeEngine(verifier, COMPAT)
+    snapshots = []
+
+    def snapshot(**kwargs):
+        snapshots.append(kwargs)
+        return CacheBlock.create(
+            kwargs["block_hash"],
+            kwargs["token_count"],
+            b"final",
+        )
+
+    monkeypatch.setattr(engine, "_snapshot", snapshot)
+    hashes = [b"a" * 32, b"b" * 32, b"c" * 32]
+    blocks = engine.compute_prefill(
+        [1, 2, 3, 4, 5],
+        hashes,
+        compression=CompressionCodec.NONE,
+        cancelled=threading.Event(),
+    )
+    assert verifier.forwarded == [1, 2, 3, 4, 5]
+    assert len(blocks) == 1
+    assert snapshots == [{
+        "token_count": 5,
+        "block_hash": hashes[-1],
+        "compression": CompressionCodec.NONE,
+    }]
+
+
 @pytest_asyncio.fixture
 async def worker():
     engine = _Engine()
@@ -282,6 +332,36 @@ def test_job_reservation_preserves_unrelated_restore_snapshot():
         jobs.close()
 
 
+def test_job_accepts_final_snapshot_only():
+    class FinalOnlyEngine:
+        def compute_prefill(self, token_ids, block_hashes, **_kwargs):
+            return [
+                CacheBlock.create(
+                    block_hashes[-1],
+                    len(token_ids),
+                    b"final-only",
+                ),
+            ]
+
+    cache = PrefixCacheStore(COMPAT, max_bytes=1024, node_id="final-only")
+    jobs = PrefillJobStore(FinalOnlyEngine(), cache)
+    try:
+        job = jobs.submit(
+            request_id="final-only",
+            tenant_id="tenant",
+            token_ids=[1, 2, 3, 4],
+            block_hashes=[b"a" * 32, b"b" * 32],
+            compatibility=COMPAT,
+            compression=CompressionCodec.NONE,
+        )
+        job.future.result(timeout=1)
+        assert job.state == PrefillJobState.COMPLETED
+        assert cache.block_hashes() == (b"b" * 32,)
+        assert cache.fetch(job.lease_id)[0].payload == b"final-only"
+    finally:
+        jobs.close()
+
+
 def test_factory_engine_is_warmed_and_used_on_same_compute_thread():
     cache = PrefixCacheStore(COMPAT, max_bytes=1024, node_id="w")
     created_on = []