Struct DeviceKvPool 

pub struct DeviceKvPool { /* private fields */ }

Device-backed KV block pool: the pure KvBlockPool core plus the physical K/V pool buffers it hands out slots into. Owns the core (so device-aware evict/restore can move bytes and update metadata atomically); the core is reachable via core / core_mut for the pure capacity/lifecycle verbs.

Implementations

impl DeviceKvPool

pub fn new(geom: KvGeometry, dtype: DType, device: &Device) -> Result<Self>

Build a device-backed pool: allocate n_layers × 2 zero-initialized [num_blocks, block_size, Hkv, D] buffers of dtype on device, and wrap a pure KvBlockPool over the same geometry.

geom.elem_size must equal dtype.size_in_bytes() — the geometry’s byte accounting (kv_bytes_resident) and the physical buffers must agree, so a mismatch is a build-time Err, never a silently wrong budget signal.

pub fn alloc_id(&self) -> KvAllocId

Identity of the physical pool buffers a held paged plan is welded to. See KvAllocId.

There is no re-mint path: DeviceKvPool allocates once in Self::new and never replaces a layer buffer. Block allocation churn lives in core (a KvBlockPool index) and reaches the graph through the per-token block_table rebind, so it is correctly NOT part of this identity — over-keying it would rebuild the plan on every new block.

pub fn core(&self) -> &KvBlockPool

The pure metadata core (capacity/lifecycle verbs live here).

pub fn core_mut(&mut self) -> &mut KvBlockPool

Mutable core — open/append/evict/splice/… route through this.

pub fn geometry(&self) -> KvGeometry

Pool geometry (delegates to the core).

pub fn n_layers(&self) -> usize

Number of transformer layers (= number of K buffers = number of V buffers).

pub fn pool_shape(&self) -> &Shape

[num_blocks, block_size, Hkv, D] — every layer buffer’s shape.

pub fn pool_shape_flat(&self) -> Shape

[num_blocks · block_size, Hkv, D] — the pool buffer viewed with its leading (block, slot) axes collapsed into ONE flat slot axis. A physical (phys, slot) then addresses a single row phys · block_size + slot, so the paged decode’s new-token K/V write is one dynamic start on axis 0 (Op::WriteSliceDyn / Op::WriteSliceDoff) instead of a concrete two-axis [(phys,phys+1),(slot,slot+1),…] slab whose ranges change every step. Self::build_decode_attn_off binds a placeholder of THIS shape to the same pool buffer — a pure metadata reshape of Self::pool_shape (same element count, row-major order preserved), so the write still mutates the persistent buffer in place.

pub fn dtype(&self) -> DType

Pool buffer element dtype.

pub fn device(&self) -> &Device

The device the pool buffers live on — a paged forward builds + realizes on this device (so a CUDA pool runs on CUDA, not CPU).

pub fn k_pool(&self, layer: usize) -> Option<&Arc<RwLock<Storage>>>

Layer l’s K pool buffer, if l < n_layers.

pub fn v_pool(&self, layer: usize) -> Option<&Arc<RwLock<Storage>>>

Layer l’s V pool buffer, if l < n_layers.

pub fn block_write_ranges(&self, phys: PhysBlockId) -> Vec<(usize, usize)>

The Op::WriteSlice rectangular ranges that select physical block phys in a [num_blocks, block_size, Hkv, D] pool buffer: [(phys, phys+1), (0, block_size), (0, Hkv), (0, D)]. Public so a consumer building its own decode graph can write a token’s K/V slab straight into a block without re-deriving the offset.

pub fn write_block( &self, layer: usize, kind: BlockKind, phys: PhysBlockId, data: &[f32], ) -> Result<()>

Write data (block_size · Hkv · D f32, row-major [block_size, Hkv, D]) into physical block phys of layer layer’s K-or-V buffer, in place, through the executor (Op::WriteSlice) — so it is device-general (the same path a CUDA/Vulkan pool uses), not a host poke.

f32-only for now (the CPU-gated increment); the pool’s real dtype is carried in Self::dtype and a byte/dtype-generic form is the follow-up tracked for the CUDA bf16 pool.

pub fn read_block( &self, layer: usize, kind: BlockKind, phys: PhysBlockId, ) -> Result<Vec<f32>>

Read physical block phys of layer layer’s K-or-V buffer back to host as block_size · Hkv · D f32 (row-major [block_size, Hkv, D]). Reads through the executor (Slice of the cache-bound buffer, realized) so it is device-general — a CUDA/Vulkan pool D2H-copies the block the same way. f32-only for now (see Self::write_block).

pub fn splice_prefix_shifted( &mut self, prefix: PrefixId, dst: SessionHandle, rope_base: f64, ) -> Result<usize>

rung-2 product — materialize a registered prefix at a block-aligned, NON-ZERO position offset in dst via a uniform θ·M RoPE delta-rotation of its cached keys. Unlike the pure-core zero-copy splice_prefix_from (rung-1, prefix at position 0), this COPIES: the prefix’s keys were rotated for positions 0..N and are numerically wrong at M..M+N, so fresh blocks are allocated and written with rotated K + copied V. rope_base is the model’s θ — a plain param, so a pool-driving consumer supplies its own (the §15 seam). Returns the shared token count N. M = filled_tokens(dst) must be block-aligned; validation + fresh-block bookkeeping is KvBlockPool::alloc_shifted_prefix_slots, which is transactional (a refusal allocates nothing). f32 pool.

pub fn write_block_bytes( &self, layer: usize, kind: BlockKind, phys: PhysBlockId, bytes: &[u8], ) -> Result<()>

Dtype-agnostic byte-level block WRITE — the movement core that the f32 Self::write_block wrapper and the evict/restore/CoW paths build on. bytes is the block’s raw content in the pool’s dtype (block_elems · dtype.size_in_bytes() bytes, row-major [block_size, Hkv, D]). Backend-neutral: builds a dtype-tagged Op::Const from the bytes and Op::WriteSlices it into the persistent pool buffer through the executor — the same path a CUDA/Vulkan (or future) pool uses; no host poke, no per-backend code.

pub fn read_block_bytes( &self, layer: usize, kind: BlockKind, phys: PhysBlockId, ) -> Result<Vec<u8>>

Dtype-agnostic byte-level block READ — sibling of Self::write_block_bytes. Returns the block’s raw bytes in the pool’s dtype via the byte-view entry realize_one_bytes (a byte reinterpret of the realized Slice, correct for ANY dtype by design — there is no per-dtype read path; GAP-327 keeps this UNGUARDED).

pub fn materialize_block_table( &self, sessions: &[SessionHandle], ) -> Result<PageTableHost, KvAllocError>

Project the core’s per-session block tables into the Op::PagedAttn [B, max_blocks] u32 block_table + [B] u32 context_lens, one row per session in sessions order (that order defines the op’s batch axis).

Fails (typed, never a panic) if any session is unknown or has an externalized slot — a session must be resident to back live attention, so a mis-sequenced call surfaces as KvAllocError rather than routing through a reclaimed block. Shorter rows are zero-padded (see PageTableHost::block_table).

pub fn materialize_block_table_padded( &self, sessions: &[SessionHandle], max_blocks_cap: usize, ) -> Result<PageTableHost, KvAllocError>

Capacity-padded sibling of Self::materialize_block_table: every row is right-padded with 0 to a FIXED width max_blocks_cap, so the block_table shape [B, max_blocks_cap] is constant across decode steps — the paged plan-once prerequisite. Padded entries (0) are in-bounds and neutralized by Op::PagedAttn’s context_len mask (a padded key position is always >= context_len). Errors if a session already occupies more than max_blocks_cap blocks.

pub fn build_decode_attn( &self, k_pool_ph: &Tensor, v_pool_ph: &Tensor, q: &Tensor, k_new: &Tensor, v_new: &Tensor, block_table: &Tensor, context_lens: &Tensor, phys_block: PhysBlockId, slot: usize, scale: f32, ) -> Result<Tensor>

Build (does NOT realize) one decode step’s paged attention for one layer: write the new token’s [Hkv, D] K/V into its physical pool slot, then run Op::PagedAttn over the (post-write) pool via the session’s block table. The graph-builder half of paged decode — the model’s forward calls this per layer, binds k_pool(layer)/v_pool(layer) into the realize cache (the caller owns the binding + the single realize for the whole forward), and threads the result into the o-projection.

All tensors are on q’s graph: k_pool_ph/v_pool_ph are const_placeholder_like placeholders the caller binds to this layer’s pool buffers; block_table/context_lens are the materialize_block_table tensors; q is [1, Hq, 1, D]; k_new/v_new are [1, Hkv, 1, D] (the new token, post-RoPE for K). (phys_block, slot) is where the new token lands — phys_block = resident_block(session, tok_pos / block_size), slot = tok_pos % block_size, after the session has been appended to cover tok_pos. Returns the attention output [1, Hq, 1, D].

The write mutates the bound pool buffer in place (so it persists for the next step) AND feeds this step’s Op::PagedAttn — one graph, one realize, exactly like the contiguous write_slice-then-attend path it replaces.

pub fn build_decode_attn_batched( &self, k_pool_ph: &Tensor, v_pool_ph: &Tensor, q: &Tensor, k_new: &Tensor, v_new: &Tensor, block_table: &Tensor, context_lens: &Tensor, writes: &[(PhysBlockId, usize)], scale: f32, ) -> Result<Tensor>

Batched (B = K) sibling of Self::build_decode_attn — one decode step over K sessions sharing this layer’s pool buffers (paged-storage PS4a). Each session i contributes one new token: its [Hkv, D] K/V (row i of k_new/v_new, shaped [K, Hkv, 1, D]) is written into its own physical slot writes[i] = (phys, slot); then a single Op::PagedAttn at batch K reads all K sessions via the [K, max_blocks] block_table + [K] context_lens. q is [K, Hq, 1, D]; returns [K, Hq, 1, D].

The K per-session writes chain on the same bound pool buffer (each write_slice returns the post-write handle) so all land before the op reads — one graph, one realize, exactly like the single-session path. writes must be non-empty and in block_table/q batch-row order.

pub fn build_decode_attn_off( &self, k_pool_ph: &Tensor, v_pool_ph: &Tensor, q: &Tensor, k_new: &Tensor, v_new: &Tensor, block_table: &Tensor, context_lens: &Tensor, write_off: Option<&Tensor>, write_sym: SymId, scale: f32, ) -> Result<Tensor>

Plan-once sibling of Self::build_decode_attn — the new token’s K/V is written at a FLATTENED, runtime-resolved offset instead of a concrete two-axis slab, so the write op is structurally identical across decode steps (only the bound offset changes). The pool buffer is viewed as Self::pool_shape_flat [num_blocks·block_size, Hkv, D]; the new token [1, Hkv, D] lands at row linear = phys·block_size + slot via a single dynamic start on axis 0. The post-write flat buffer is reshaped back to [num_blocks, block_size, Hkv, D] for Op::PagedAttn.

k_pool_ph/v_pool_ph must be const_placeholder_like(pool_shape_flat()) placeholders the caller binds to this layer’s pool buffers (NOTE: the flat shape, not Self::pool_shape). The offset is supplied one of two ways, mirroring the contiguous KV-write split in apply_layer_with_kv_writes:

• write_off = Some(off): a rank-0 I64 tensor read device-side at kernel launch (Op::WriteSliceDoff) — the CUDA-graph-capturable arm.
• write_off = None: Op::WriteSliceDyn with DynScalar::Sym(write_sym), the caller binding write_sym → linear in the per-pass SymEnv — the backend-generic (CPU/Vulkan) arm.

Both land the same slab at the same offset — bit-identical to the concrete Self::build_decode_attn. Returns the attention output [1, Hq, 1, D].

pub fn ensure_writable_block( &mut self, session: SessionHandle, logical_block: usize, ) -> Result<PhysBlockId>

Copy-on-write guard for the paged decode write path. Ensure the physical block backing logical slot logical_block of session is exclusively owned before a new token is written into it. If the block is shared (spliced, refcount > 1), break the share (cow_break → a fresh block) and COPY the shared block’s bytes (all layers, K + V) into it, so the session keeps the shared-prefix content while its write no longer mutates a block another session still references. Returns the (now private) physical block to write into; an exclusive block is returned unchanged. Must be called after append and before writing the token — otherwise decoding a spliced session silently corrupts its co-sharers.

pub fn evict(&mut self, s: SessionHandle) -> Result<DeviceEvicted>

Evict a session’s entire resident state: capture the bytes of its exclusive blocks (D2H, all layers, K and V) into a DeviceEvicted handle, then hand the physical blocks back to the pool via the core. restore re-materializes it byte-for-byte. Shared (spliced) blocks are never detached and carry no bytes here (they stay resident).

pub fn evict_blocks( &mut self, s: SessionHandle, indices: &[usize], ) -> Result<DeviceEvicted>

Evict a SET of a session’s logical blocks (sub-session tiering). Captures the bytes of exactly the exclusive blocks in indices (a shared block stays resident — its sharer still needs it), then calls the core’s refcount-aware KvBlockPool::evict_blocks. The returned handle carries the detached bytes; the rest of the session keeps decoding.

pub fn restore(&mut self, s: SessionHandle, handle: DeviceEvicted) -> Result<()>

Re-materialize an evicted session: the core allocates fresh physical blocks for the externalized slots, then the saved bytes are written back (H2D) into whichever blocks the pool handed out — so the restored content is byte-identical regardless of physical reuse. Fails (typed) if the pool can’t fit the session’s blocks again (OutOfBlocks).