Module inference_context 

Persistent inference state for the pipelined executor (Phase 7.6 step 9c, sub-phase E.3.0).

What this module owns

• InferenceContext: a host-side per-session struct holding long-lived Arc<RwLock<Storage>> references that survive across realize calls. Replaces the legacy executor’s const_pool for the pipelined dispatch world. Today’s backing is a simple in-memory HashMap<NodeId, Arc<...>>; future work may swap this for an mmap-backed coherent store per project_unified_durable_tensor_store.md.

• KvCache, KvLayer, KvLayerId, KvSlot, AuthorityState: the backend-erased KV cache primitive that replaced lazy::KVCache<B> and lazy::LlamaKVCache (both retired — E.3.3.D and Unification Session 4 E.3.4). Single concrete type with no generic-over-B parameter; each layer holds Arc<RwLock<Storage>> for K and V plus side-table metadata (layout, monotonic version, coherence authority).

What this module deliberately does NOT do (yet)

• Pre-allocated buffers (added Phase E.3.3.A, 2026-05-20): KvCache::with_capacity eagerly allocates [1, n_kv_heads, max_seq_len, head_dim] zero-initialized K + V buffers on a target device. Subsequent forward steps write into these via Op::WriteSlice rather than growing-by-replacement. The legacy KvCache::with_dims constructor stays available for callers not yet on the pre-allocated path (grow-by-replace works without changes).
• Multi-device coherence protocol: the authority and version fields exist as placeholders. No protocol consults them yet — every layer starts and stays AuthorityState::Host for the lifetime of the session in single-device usage. Phase J (multi-GPU) activates the protocol.
• Weight persistence: weights stay in the graph’s storage_map per the design discussion; the persistent map handles KV layers and transient cross-step state only.

Structs

DecodeSession

Plan-once persistent decode state for one LlamaModel + one KvCache capacity/dtype. Built on the first seq == 1 decode token; reused (graph + optimize view HELD, only the per-token data Consts + SymEnv re-bound) for every subsequent token.

DecodeTokenData

Per-token re-bound data-Const Arcs for a persistent decode step (Phase D · D2b). Each is a device-resident fuel_memory::Storage built from freshly recomputed host bytes (token-ids = the new token, RoPE cos+sin at position = cached_len, mask with the shifted -inf boundary). DecodeSession::realize_token overwrites the held base_cache’s stable data-Const entries with these for the token.

InferenceContext

Per-session host-side context holding long-lived storage Arcs that survive across realize calls.

KvCache

Backend-erased KV cache. Replaces lazy::KVCache<B> and lazy::LlamaKVCache. Indexed by global layer index; layers[i] is None if layer i hasn’t been populated yet.

KvLayer

One layer’s K + V slots in a backend-erased KV cache.

KvLayerId

Identifies a single key-or-value tensor within a KvCache. Stable across realize calls — the graph-level NodeId may churn per step (each Op::WriteSlice produces a fresh NodeId pointing at the same Storage Arc) but the logical KV slot identity stays.

LatentKvCache

Cross-graph, cross-forward-pass persistent N-slot decode cache — the KvCache generalization for latent-caching attention architectures.

LatentSlot

One slot’s persistent latent buffer within a LatentKvCache layer. The per-slot analogue of KvLayer’s k/v pair, generalized to an arbitrary slot count and independent per-slot trailing shape.

PagedDecodeSession

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.

PagedDecodeTokenData

Per-token re-bound data Consts for one PAGED decode step — the paged twin of DecodeTokenData. The shapes are stable across tokens (that is the whole point of the plan-once path); only the bytes change, so each is an Arc<RwLock<Storage>> overwriting the held PagedDecodeSession’s base-cache entry for its stable NodeId.

Enums

AuthorityState

Coherence authority for one KV slot’s storage. Placeholder for the multi-device protocol; today every slot starts and stays AuthorityState::Host for the lifetime of the session unless future code explicitly transitions it.

KvSlot

Which half of a KV pair a KvLayerId addresses.

PagedDecodePlan

Runtime policy for the paged persistent decode path (LlamaModel::forward_paged_step_persistent): whether to build the plan ONCE and reuse it across tokens, or re-plan every token (the pre-plan-once behavior). The paged driver holds this per config and passes it each step; the driver default is PlanOnce (set by PagedSessionScheduler::new), with Replan available as an explicit opt-out.

PlanValidity

Whether a held plan may serve a step, and if not, why not — one set of comparisons, so the “is it valid” and “is only the allocation different” questions cannot drift apart. DecodeSession::is_valid_for is a thin view over this.

RebindRefusal

Why a held plan could not be re-bound onto a replacement KV allocation.

Traits

KvRebindSource

A KV allocation a held DecodeSession can be re-bound onto — the one abstraction over KvCache (K/V pairs) and LatentKvCache (MLA’s (latent, kpe) pairs), which already share DecodeSession verbatim and therefore share its welding.

Functions

alloc_batched_kv

Allocate n_layers pairs of zero-initialized [k, n_kv_heads, max_seq_len, head_dim] K/V buffers on device with dtype — the shared batch-slot KV buffer for the multi-session batched decode arm (serving Increment 1, C3). Same Op::Alloc→Op::ZeroFill emission + PipelinedExecutor::realize_many pattern as KvCache::with_capacity, with the leading dim k (the batch slot) instead of 1. Returns the per-layer (K_arc, V_arc) pairs; the caller binds them into a batched decode graph.