Struct PagedDecodeSession 

pub struct PagedDecodeSession { /* private fields */ }

Plan-once persistent decode state for one PAGED (block-pool) decode stream — the paged twin of DecodeSession. Built on the first seq == 1 paged decode token; reused (graph + optimize view HELD, only the per-token data Consts + SymEnv re-bound) for every subsequent token.

Paged decode’s structural variabilities that a re-plan would otherwise pay EVERY token are removed here: (1) the fresh Tensor graph root becomes a held graph with stable re-bindable Const placeholders; (2) the L-varying block_table shape is pinned to [1, max_blocks_cap] (Task 1 padded materialize); (3) the per-step KV-write range becomes ONE flattened dynamic offset (Task 2 build_decode_attn_off). What remains per token is shape-stable data re-bound into a clone of base_cache + one bound scalar.

The pool K/V buffer Arcs are bound ONCE into base_cache and mutate IN PLACE via Op::WriteSliceDoff/Op::WriteSliceDyn each token (never re-inserted) — exactly the contiguous KV relationship, so the write this step lands in the pool and is visible next step.

Held by the paged generate-loop owner alongside the DeviceKvPool + SessionHandle; the caller must NOT swap the pool/model out from under a live session (the held base_cache references THIS pool’s buffers). A validity-key mismatch (is_valid_for) or TopologyChanged drops + rebuilds it (see LlamaModel::forward_paged_step_persistent). Single-session (B = 1) today; the ragged batched persistent path is a follow-on.

Implementations

impl PagedDecodeSession

pub fn new( graph: Arc<RwLock<Graph>>, optimized: OptimizedGraph, effective_target: NodeId, logits_node: NodeId, token_ids_node: NodeId, rope_cos_node: NodeId, rope_sin_node: NodeId, block_table_node: NodeId, context_lens_node: NodeId, kv_nodes: Vec<(NodeId, NodeId)>, offset_node: Option<NodeId>, write_sym: SymId, base_cache: StorageCache, max_blocks_cap: usize, n_layers: usize, block_size: usize, cache_dtype: DType, shape_key: u64, alloc_id: KvAllocId, ) -> Self

Assemble a held paged session from the artifacts a first-realize prebuild produced. The node ids are the STABLE data-Const NodeIds the paged builder minted.

pub fn realize_token( &self, device: &Device, data: PagedDecodeTokenData, sym_env: &SymEnv, ) -> Result<Vec<f32>>

Realize the held paged graph for one decode token via the prebuilt seam (SKIP prepare + optimize_graph). data supplies the per-token re-bound data-Const Arcs; they overwrite the held base_cache’s entries for this token (the weight + pool KV Arcs stay from base_cache). sym_env binds write_sym (load-bearing on the SymEnv path, harmless on the device-offset path). TopologyChanged surfaces typed (caller invalidates the session).

pub fn is_valid_for( &self, max_blocks_cap: usize, n_layers: usize, block_size: usize, cache_dtype: DType, shape_key: u64, alloc_id: KvAllocId, ) -> bool

Whether this held session is valid for the given paged decode step. Rebuild (drop + build fresh) on any mismatch: a change in max_blocks_cap (the block_table shape key), n_layers, block_size, or cache_dtype means the held graph’s shapes / offsets are stale.

alloc_id is the live pool’s DeviceKvPool::alloc_id — which POOL’s buffers, as distinct from shape_key’s which MODEL’s weights. Note it is deliberately not sensitive to block allocation: which physical blocks a session owns rides the per-token block_table rebind, so a generation that grows into new blocks keeps its plan.

pub fn per_token_sym_env(&self, linear: usize) -> Result<SymEnv>

Build the per-token SymEnv for one paged decode step: bind write_sym = linear (the flattened KV-write offset). Load-bearing on the SymEnv path; a harmless unreferenced binding on the device-offset path (a conflicting rebind surfaces a typed error, never a panic).

pub fn graph(&self) -> &Arc<RwLock<Graph>>

The held graph handle.

pub fn optimized(&self) -> &OptimizedGraph

The cached optimize view (fed to the prebuilt-realize seam).

pub fn effective_target(&self) -> NodeId

pub fn logits_node(&self) -> NodeId

pub fn token_ids_node(&self) -> NodeId

pub fn rope_cos_node(&self) -> NodeId

pub fn rope_sin_node(&self) -> NodeId

pub fn block_table_node(&self) -> NodeId

pub fn context_lens_node(&self) -> NodeId

pub fn kv_nodes(&self) -> &[(NodeId, NodeId)]

pub fn offset_node(&self) -> Option<NodeId>

pub fn write_sym(&self) -> SymId

pub fn max_blocks_cap(&self) -> usize

pub fn realize_count(&self) -> usize

How many REBIND (plan-reuse) tokens this session has served — a device-independent plan-HIT signal (see Self::realize_count). Zero right after the build; n after n reuse tokens. The correctness gate asserts this advances by exactly one per plan-once decode step.

pub fn graph_node_count(&self) -> usize

The held graph’s node count. The Task 3 born-red test asserts this is stable from token 2 onward (no per-token node growth — the guard against an accidental re-splice / re-insert).

pub fn base_cache(&self) -> &StorageCache

The full realized StorageCache from the first realize.