Skip to content

Proposal: additive KV cache connector seam at paged-attention block pool #2308

Description

@Travis-Gilbert

Proposal

Add an optional, no-op-by-default KV cache connector seam around the existing paged-attention prefix-cache block flow.

The current code already computes content-addressed block hashes in mistralrs-core/src/paged_attention/block_hash.rs, resolves prefix hits through KVCacheManager::get_computed_blocks, assigns physical blocks through allocate_slots, stores newly full blocks through cache_blocks, and evicts cached block hashes when BlockPool::get_new_blocks reallocates an LRU free block.

A small connector trait at that seam would let integrations observe or implement external KV tiers without changing default behavior:

  • look up a BlockHash plus KV cache group IDs before or alongside local prefix-cache lookup
  • observe the selected physical block IDs after a hit is allocated
  • observe newly cacheable full blocks when cache_blocks records them
  • observe evicted block hashes when a cached block is reallocated

Motivation

This is narrower than disaggregated prefill/network transfer. The immediate goal is an additive hook for external KV stores such as RAM/disk/blob registries, while keeping mistral.rs behavior unchanged unless a connector is installed.

The default implementation would return misses and only record no-op observations, so existing tests and users should see no behavior change.

Related prior broader issue: #2000.

Sketch

pub trait KvCacheConnector: Send + Sync {
    fn lookup_blocks(&self, _block_hashes: &[BlockHash], _group_ids: &[u32], _max_blocks: usize) -> Option<Vec<usize>> { None }
    fn observe_hit(&self, _block_hashes: &[BlockHash], _block_ids: &[usize], _num_computed_tokens: usize) {}
    fn observe_store(&self, _block_hashes: &[BlockHash], _block_ids: &[usize], _num_full_blocks: usize) {}
    fn observe_evict(&self, _block_hashes: &[BlockHashWithGroupId], _block_id: usize) {}
}

I am happy to adjust the exact signatures to match maintainer preference; the key property is that the seam is optional and defaults to current local-only behavior.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions