Struct SharedPagedKv 

pub struct SharedPagedKv { /* private fields */ }

Per-layer PagedKvStores that many concurrent requests share.

Why a lock per layer, and why two phases

ferrox-server runs generation on spawn_blocking with, in its own words, “no I/O and no shared lock”. KvBlockPool survives that because it only bounds a count: each KvCache owns a private Vec, and the pool mutex is taken briefly at acquire and release, never during a forward. A PagedKvStore is the opposite – it IS the backing memory – so sharing one across concurrent requests needs an answer to “who may touch these bytes when”.

The answer the API already implies: attention takes &PagedKvStore and only push takes &mut. So the accesses split cleanly into many concurrent readers and one short exclusive write per position, which is exactly an RwLock – and one per LAYER rather than one for the whole model, so two requests contend only when both are writing the same layer at the same instant.

A caller must therefore take the write guard for the push alone and drop it before attending under a read guard. Holding the write guard across attention would serialise the expensive half and give back a global lock with extra steps. Nothing breaks in the gap: a sequence’s block table and length are its own, and another request’s push in between only touches blocks it exclusively holds.

Deadlock

Self::write_all is the one place several layers are held at once, and it takes them in ascending layer order. Every caller getting the same order is what makes that safe; there is no other multi-layer acquisition in the codebase, and a new one must follow the same rule.

Poisoning

A panic while holding a store leaves the KV mid-write, which is not recoverable state, but it is also not unsound – the bytes are plain f32. Poison is stepped over with into_inner, matching how ferrox-server already treats its pool mutex: a poisoned lock should not turn one request’s panic into a permanently dead server.

Implementations

impl SharedPagedKv

pub fn new( n_layers: usize, block_size: usize, blocks_per_layer: usize, n_kv_heads: usize, head_dim: usize, ) -> Self

One store per layer, each with blocks_per_layer blocks.

pub fn from_stores(stores: Vec<PagedKvStore>) -> Self

Wraps stores the caller built, for tests and for callers that size layers differently.

pub fn layer_count(&self) -> usize

pub fn read(&self, layer: usize) -> RwLockReadGuard<'_, PagedKvStore>

Shared access to one layer, for attention.

pub fn write(&self, layer: usize) -> RwLockWriteGuard<'_, PagedKvStore>

Exclusive access to one layer, for a push. Hold it for the push and nothing else – see the type docs.

pub fn write_all(&self) -> Vec<RwLockWriteGuard<'_, PagedKvStore>>

Every layer at once, in ascending order, so a multi-layer append is atomic against other requests.

This is what makes “all layers advance or none do” hold under concurrency rather than only single-threaded: checking free space and then appending are separate steps, and without the guards spanning both, another request can take the blocks in between and leave this one half-written.

Ascending order is the deadlock rule; see the type docs.

pub fn free_blocks(&self, layer: usize) -> usize

Free blocks in one layer, for admission control. A snapshot: by the time a caller acts on it another request may have taken them, which is why the append itself re-checks under the guard.

pub fn acquire_group(&self) -> Option<PageGroup>

Takes one block from EVERY layer as a single group, refcount 1.

All layers or none: a group that existed in some layers and not others could not answer “which block holds position p in layer l”, which is the only question it exists to answer.

pub fn retain_group(&self, group: PageGroup)

One more holder of group.

Called when a second sequence adopts a cached prefix. Without it, the first sequence to finish frees pages the second is still attending over – a use-after-free that shows up as another conversation’s tokens rather than as a crash.

pub fn release_group(&self, group: PageGroup) -> bool

One fewer holder. At zero the blocks go back to their layers.

Returns whether this was the last holder, so a caller can assert on it rather than guess.

pub fn group_blocks(&self, group: PageGroup) -> Vec<usize>

Which block in each layer this group owns, indexed by layer.

pub fn group_refs(&self, group: PageGroup) -> u32

How many holders group has. Zero means it does not exist.

pub fn free_groups(&self) -> usize

Groups that could still be allocated, bounded by the layer with the fewest free blocks: a group needs one from each.