Trait PrefixCache 

pub trait PrefixCache: Send + Sync {
    // Required methods
    fn lookup(
        &self,
        tokens: &[u32],
        block_size: usize,
        session_hash: u64,
        adapter_id: u64,
    ) -> PrefixMatch;
    fn insert(
        &self,
        tokens: &[u32],
        block_table: &[u32],
        disk_block_ids: &[u32],
        block_size: usize,
        matched_tokens: usize,
        adapter_id: u64,
    ) -> InsertAcquired;
    fn insert_with_snapshot(
        &self,
        tokens: &[u32],
        block_table: &[u32],
        disk_block_ids: &[u32],
        block_size: usize,
        snapshot_id: usize,
        session_hash: u64,
        matched_tokens: usize,
        adapter_id: u64,
    ) -> (Option<usize>, InsertAcquired);
    fn insert_intermediate_snapshot(
        &self,
        tokens: &[u32],
        block_table: &[u32],
        disk_block_ids: &[u32],
        block_size: usize,
        snapshot_id: usize,
        session_hash: u64,
        matched_tokens: usize,
        adapter_id: u64,
    ) -> Option<usize>;
    fn insert_tail_snapshot(
        &self,
        tokens: &[u32],
        snapshot_id: usize,
        session_hash: u64,
        adapter_id: u64,
    ) -> Vec<usize>;
    fn insert_tail_sibling_snapshot(
        &self,
        tokens: &[u32],
        snapshot_id: usize,
        session_hash: u64,
        adapter_id: u64,
    ) -> Option<usize>;
    fn release(&self, tokens: &[u32], block_size: usize, adapter_id: u64);
    fn release_matched(
        &self,
        tokens: &[u32],
        block_size: usize,
        matched_tokens: usize,
        adapter_id: u64,
    );
    fn evict(&self, num_blocks: usize) -> EvictedBlocks;
    fn evict_snapshot_lru(&self) -> Option<usize>;
    fn snapshot_count(&self) -> usize;
    fn stats(&self) -> (usize, usize);

    // Provided methods
    fn is_active(&self) -> bool { ... }
    fn peek_matched_tokens(
        &self,
        _tokens: &[u32],
        _block_size: usize,
        _adapter_id: u64,
    ) -> usize { ... }
    fn evict_snapshot_to_tier(&self, min_tokens: usize) -> Option<TierEvict> { ... }
    fn promote_snapshot(&self, key: u64, new_slot: usize) -> bool { ... }
    fn forget_snapshot_tier_key(&self, key: u64) -> bool { ... }
}

Trait for prefix caching strategies.

All methods take &self — implementations use interior mutability (e.g., Mutex) for thread safety. This allows the prefix cache to be shared between the model (prefill) and scheduler (free_sequence) without requiring &mut self.

Required Methods

fn lookup( &self, tokens: &[u32], block_size: usize, session_hash: u64, adapter_id: u64, ) -> PrefixMatch

Look up a token sequence and return cached KV blocks for the longest matching prefix (block-aligned).

Increments ref_count on matched nodes so they survive eviction while the sequence is active. session_hash is used for SSM snapshot isolation (0 = legacy/no session tracking).

Task #24: adapter_id keys the KV/prefix + SSM-snapshot cache so a request reuses ONLY blocks computed under the same adapter. 0 = base / no adapter, which keys byte-identically to the pre-LoRA token-only cache.

fn insert( &self, tokens: &[u32], block_table: &[u32], disk_block_ids: &[u32], block_size: usize, matched_tokens: usize, adapter_id: u64, ) -> InsertAcquired

Insert a completed prefill’s blocks into the cache.

block_table[i] is the physical block for tokens [i*block_size .. (i+1)*block_size].

disk_block_ids parallels block_table for --high-speed-swap (Phase 6.1.e). Empty when HSS is not in use; same length as block_table when populated. The cache stores these alongside the physical block IDs and returns them in EvictedBlocks so the caller can dec_disk_ref the orchestrator’s per-block refcount.

Disk-ref obligation (Issue #17 fix): the returned vec lists every disk_block_id on which this insert call newly took an ownership ref (a node was created OR an existing node had its disk_block_id populated for the first time). The caller MUST inc_disk_ref each returned ID so the swap allocator’s refcount matches the cache’s reachability. Already-cached portions (matched-prefix entries, or blocks a prior intermediate insert already covered) are NOT in the returned vec — re-incing them would leak the cache’s refcount.

matched_tokens is the number of tokens the inserting sequence already acquired via lookup()’s inc_refs (0 for a cache-miss request). Tokens past this offset are “seq-owned” — the inserting sequence’s eventual release() will decrement them — so insert must bump their ref_count to keep the cache’s own reference alive after the release. See the release/lookup dance at the top of radix_tree.rs.

fn insert_with_snapshot( &self, tokens: &[u32], block_table: &[u32], disk_block_ids: &[u32], block_size: usize, snapshot_id: usize, session_hash: u64, matched_tokens: usize, adapter_id: u64, ) -> (Option<usize>, InsertAcquired)

Insert blocks with an SSM state snapshot registered in the snapshot index.

The snapshot ID references a slot in an external SsmSnapshotPool. On future lookups matching this prefix, the snapshot ID is returned in PrefixMatch::ssm_snapshot so the caller can restore SSM state. session_hash tags the snapshot for session-scoped isolation. matched_tokens has the same semantics as in insert. Returns (displaced_snapshot_id, newly_acquired_disk_ids). The disk-ref obligation matches insert: caller inc_disk_refs each returned ID.

fn insert_intermediate_snapshot( &self, tokens: &[u32], block_table: &[u32], disk_block_ids: &[u32], block_size: usize, snapshot_id: usize, session_hash: u64, matched_tokens: usize, adapter_id: u64, ) -> Option<usize>

Insert an SSM snapshot at an intermediate token boundary.

tokens is the token sequence up to and including the snapshot point. block_table contains the physical block indices for those tokens. session_hash tags the snapshot for session-scoped isolation. matched_tokens has the same semantics as in insert. Returns the displaced snapshot ID if an existing entry was overwritten.

fn insert_tail_snapshot( &self, tokens: &[u32], snapshot_id: usize, session_hash: u64, adapter_id: u64, ) -> Vec<usize>

Register the per-session TAIL snapshot in the index WITHOUT touching the radix tree (the final chunk’s insert covers those blocks). Supersedes this session’s previous tail; returns displaced snapshot ids to free.

fn insert_tail_sibling_snapshot( &self, tokens: &[u32], snapshot_id: usize, session_hash: u64, adapter_id: u64, ) -> Option<usize>

Register the tail’s EARLY sibling (tb - bs) in the index. Must be called after insert_tail_snapshot in the same finalize (the tail insert sweeps the session’s previous tail + sibling). Returns a displaced snapshot id to free, if the prefix was already registered.

fn release(&self, tokens: &[u32], block_size: usize, adapter_id: u64)

Release ref_counts on blocks that were acquired via lookup.

Called when a sequence finishes. Decrements ref_count on cache nodes matching the token prefix, making them eligible for eviction. Task #24: adapter_id must match the one used at lookup/insert.

fn release_matched( &self, tokens: &[u32], block_size: usize, matched_tokens: usize, adapter_id: u64, )

Release exactly the block-aligned prefix acquired by one lookup.

Batched-prefill admission may acquire several candidates, then reject the batch before any sequence owns their blocks. Releasing the full token slice in that case would be racy: another request could insert a longer matching suffix between acquisition and rollback. Implementors must therefore decrement no more than matched_tokens.

fn evict(&self, num_blocks: usize) -> EvictedBlocks

Evict up to num_blocks cached blocks, returning their physical indices and parallel disk-block IDs (Phase 6.1.e).

Picks LRU zero-ref leaf nodes. Returns fewer than requested if not enough evictable blocks exist. The caller is responsible for dec_disk_ref-ing every entry in disk_block_ids (releasing the cache’s HSS-side refcount). When HSS isn’t in use the disk_block_ids vec is empty.

fn evict_snapshot_lru(&self) -> Option<usize>

Evict the least-recently-used SSM snapshot from the snapshot index. Returns the snapshot ID so the caller can free it in SsmSnapshotPool.

fn snapshot_count(&self) -> usize

Number of SSM snapshots currently stored in the snapshot index.

fn stats(&self) -> (usize, usize)

(entries, cached_blocks) for logging.

Provided Methods

fn is_active(&self) -> bool

Whether this implementation is active (i.e., a real cache that actually inserts/holds refs). NoPrefixCaching returns false; RadixTree returns true. Callers use this to skip ref-bookkeeping that’s only meaningful when the cache holds refs (e.g., the manual kv_cache.inc_ref in cache_sequence that pairs with eviction’s return_evicted_block).

fn peek_matched_tokens( &self, _tokens: &[u32], _block_size: usize, _adapter_id: u64, ) -> usize

Read-only longest-prefix probe: number of tokens (block-aligned) lookup would match, WITHOUT taking refs, touching LRU state, or counting a hit/miss. Used by the prefill tail-checkpoint split to detect conversation reuse before deciding to pay the extra pass. Task #24: keyed by adapter_id so a cross-adapter peek reports a miss.

fn evict_snapshot_to_tier(&self, min_tokens: usize) -> Option<TierEvict>

Phase 1b spill tier: pick a spill victim (same policy as evict_snapshot_lru, HBM-resident only) and decide whether it is worth spilling — see TierEvict. A victim shallower than min_tokens cannot repay the spill’s fixed cost, so its entry is dropped outright rather than left findable-but-empty. min_tokens == 0 disables the gate. None when nothing resident remains; default None (no tier).

fn promote_snapshot(&self, key: u64, new_slot: usize) -> bool

Phase 1b spill tier: after the caller faulted a spilled snapshot’s bytes into new_slot, re-home its index entry to HBM. Returns false if the key is unknown. Default: false.

fn forget_snapshot_tier_key(&self, key: u64) -> bool

Phase 1b spill tier: the FAILED-fault-in twin of Self::promote_snapshot. The caller’s store.get(key) MISSED, so this entry is findable by lookup_tiered with no bytes behind it. Left in place, every warm turn on this prefix repeats the whole doomed cycle — spill a LIVE 66 MB victim D2H to free a slot, fault in, miss, free the slot — and then recomputes anyway; under ATLAS_SSM_TIER_DISK_GB that doomed spill evicts one MORE tier record, so the cap’s own pressure re-amplifies itself. Dropping the entry degrades the prefix to a plain recompute ONCE.

Only removes an entry that is still tiered: a resident entry’s snapshot_id is a LIVE pool slot that only its owner may free, so a by-key remove of one would leak it. Returns whether an entry was dropped. Default: false (no tier).

Implementors

impl PrefixCache for RadixTree

impl PrefixCache for NoPrefixCaching