From a68042476b14fdb221f1d40c06b8bb4a0141aa71 Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Fri, 17 Jul 2026 16:11:35 +0800 Subject: [PATCH] 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 --- docs/ops/distributed-prefill-kv-network.md | 4 + .../backends/mlx/prefill_worker.py | 27 +++---- inference_engine/distributed/prefill_cache.py | 7 +- .../distributed/prefill_worker.py | 5 +- .../distributed/test_prefill_cache.py | 17 +++- .../distributed/test_prefill_worker.py | 80 +++++++++++++++++++ 6 files changed, 119 insertions(+), 21 deletions(-) diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 11d9fdaf..a351d13f 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/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. diff --git a/inference_engine/backends/mlx/prefill_worker.py b/inference_engine/backends/mlx/prefill_worker.py index 6b8f1942..594237c9 100644 --- a/inference_engine/backends/mlx/prefill_worker.py +++ b/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, diff --git a/inference_engine/distributed/prefill_cache.py b/inference_engine/distributed/prefill_cache.py index 0b8029d3..99025180 100644 --- a/inference_engine/distributed/prefill_cache.py +++ b/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] diff --git a/inference_engine/distributed/prefill_worker.py b/inference_engine/distributed/prefill_worker.py index 809e2fd1..dec13fca 100644 --- a/inference_engine/distributed/prefill_worker.py +++ b/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, diff --git a/tests/inference_engine/distributed/test_prefill_cache.py b/tests/inference_engine/distributed/test_prefill_cache.py index c5ad2d62..447a3c9a 100644 --- a/tests/inference_engine/distributed/test_prefill_cache.py +++ b/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( diff --git a/tests/inference_engine/distributed/test_prefill_worker.py b/tests/inference_engine/distributed/test_prefill_worker.py index 357b46ff..3760be9f 100644 --- a/tests/inference_engine/distributed/test_prefill_worker.py +++ b/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 = []