From fda044e7df77232c3c0aff36cc59ce1b112d7f93 Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Thu, 16 Jul 2026 20:58:15 +0800 Subject: [PATCH] fix(prefill): resolve promoted final snapshots directly Use the longest available chained boundary so Primary hot promotion stores one complete snapshot without copying every intermediate checkpoint, while preserving full-prefix integrity. Co-authored-by: Cursor --- ...16-distributed-prefill-kv-cache-network.md | 6 +++-- docs/ops/distributed-prefill-kv-network.md | 2 +- inference_engine/distributed/prefill_cache.py | 27 +++++++++++-------- .../distributed/test_prefill_cache.py | 11 ++++++++ 4 files changed, 32 insertions(+), 14 deletions(-) diff --git a/docs/adr/0016-distributed-prefill-kv-cache-network.md b/docs/adr/0016-distributed-prefill-kv-cache-network.md index 7f208df6..d3e32ac1 100644 --- a/docs/adr/0016-distributed-prefill-kv-cache-network.md +++ b/docs/adr/0016-distributed-prefill-kv-cache-network.md @@ -29,7 +29,7 @@ Remote cache access happens only before decode: 1. Tokenize the request and compute chained fixed-size block hashes. 2. Query the local cache and compatible live peers concurrently. -3. Select the longest contiguous prefix whose transfer/import cost is lower +3. Select the longest available chained-prefix snapshot whose transfer/import cost is lower than local prefill recomputation. 4. Transfer one immutable snapshot at the selected prefix boundary. 5. Import it, compute the missing suffix locally, then decode entirely locally. @@ -162,7 +162,9 @@ allowing suffix-only prefill. ### Arbitrary block/hole reuse Rejected. Later K/V depends on the complete preceding token sequence and -positions. Only the longest contiguous prefix is safe. +positions. A stored snapshot must cover that complete sequence. Intermediate +checkpoint entries may be absent, however, because the chained hash and final +snapshot already commit and contain the full prefix through that boundary. ### SMB/NFS snapshot files diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index f2e9c0e8..8d098254 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -206,7 +206,7 @@ Expected invariants: - both node cards appear within two gossip intervals; - stale nodes disappear after TTL; -- remote lookup returns only the longest contiguous prefix; +- remote lookup returns the longest available chained-prefix snapshot; - imported snapshot checksum and compatibility fingerprint match; - remote failure falls back to local prefill; - no remote RPC occurs in autoregressive decode; diff --git a/inference_engine/distributed/prefill_cache.py b/inference_engine/distributed/prefill_cache.py index 638bd151..fa97f6d3 100644 --- a/inference_engine/distributed/prefill_cache.py +++ b/inference_engine/distributed/prefill_cache.py @@ -3,7 +3,7 @@ The cache stores an opaque restorable snapshot at selected token-block boundaries. Model-specific adapters own serialization/import; this module owns deterministic prefix hashing, exact compatibility matching, -longest-contiguous-prefix lookup, leases, accounting, and memory-pressure +longest chained-prefix lookup, leases, accounting, and memory-pressure eviction. A hit transfers only the snapshot at the longest matched boundary. Decode never reads this store. A requester imports a hit once, computes the @@ -208,30 +208,35 @@ def lookup( lease_seconds: float = DEFAULT_LEASE_SECONDS, now: float | None = None, ) -> PrefixLease: - """Lease the longest contiguous prefix held by this store.""" + """Lease the longest available chained-prefix snapshot. + + A chained hash at boundary N commits all preceding token blocks, so a + promoted final snapshot remains valid even when intermediate boundary + snapshots are absent or have been evicted. + """ if lease_seconds <= 0: raise ValueError("lease_seconds must be > 0") now = time.time() if now is None else now with self._lock: self._expire_leases(now) - matched: list[CacheBlock] = [] - for raw_hash in block_hashes: + snapshot = None + hit_block_count = 0 + for index, raw_hash in enumerate(block_hashes): block_hash = bytes(raw_hash) block = self._blocks.get(block_hash) - if block is None: - break - matched.append(block) - self._blocks.move_to_end(block_hash) - if not matched: + if block is not None: + snapshot = block + hit_block_count = index + 1 + if snapshot is None: self._lookup_misses += 1 return PrefixLease("", (), 0, 0, 0, self._epoch, now, bytes(32)) + self._blocks.move_to_end(snapshot.block_hash) self._lookup_hits += 1 lease_id = secrets.token_urlsafe(18) - snapshot = matched[-1] lease = PrefixLease( lease_id=lease_id, block_hashes=(snapshot.block_hash,), - hit_block_count=len(matched), + hit_block_count=hit_block_count, hit_token_count=snapshot.token_count, transfer_bytes=snapshot.nbytes, cache_epoch=self._epoch, diff --git a/tests/inference_engine/distributed/test_prefill_cache.py b/tests/inference_engine/distributed/test_prefill_cache.py index f6fcb9f9..0895eca3 100644 --- a/tests/inference_engine/distributed/test_prefill_cache.py +++ b/tests/inference_engine/distributed/test_prefill_cache.py @@ -72,6 +72,17 @@ def test_store_miss_expiry_collision_and_lru(): store.fetch(lease.lease_id, now=22.0) +def test_lookup_uses_sparse_promoted_final_snapshot(): + store = PrefixCacheStore(_compat(), max_bytes=100, node_id="head") + hashes = chained_block_hashes([1, 2, 3, 4], _compat()) + final = CacheBlock.create(hashes[1], 4, b"full-snapshot") + store.put(final) + lease = store.lookup(hashes) + assert lease.hit_block_count == 2 + assert lease.hit_token_count == 4 + assert store.fetch(lease.lease_id) == (final,) + + def test_validation_and_stats(): with pytest.raises(ValueError, match="max_bytes"): PrefixCacheStore(_compat(), max_bytes=0, node_id="x")