Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions docs/adr/0016-distributed-prefill-kv-cache-network.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion docs/ops/distributed-prefill-kv-network.md
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
27 changes: 16 additions & 11 deletions inference_engine/distributed/prefill_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
11 changes: 11 additions & 0 deletions tests/inference_engine/distributed/test_prefill_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
Loading