Struct InferenceContext 

pub struct InferenceContext { /* private fields */ }

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

The persistent map is the seam the unified-storage / mmap-coherence backplane work later replaces (see project_unified_durable_tensor_store.md); today it’s a simple in-memory HashMap. Each realize call clones the Arcs into the executor’s input cache; persistent entries reuse those Arcs across calls instead of re-uploading.

Lifecycle

1. Construct with a target device: InferenceContext::new(device).
2. Insert long-lived storages (KV layer storages, anything else that should survive across realize calls): ctx.insert(node_id, arc).
3. Realize: ctx.realize_one_as::<f32>(&graph, target_node). The persistent map is seeded into the executor’s input cache; new Op::Const nodes the graph references that aren’t in the persistent map get uploaded fresh from the graph’s storage map.
4. Subsequent realize calls reuse the persistent Arcs — no re-upload, no D2H/H2D round-trip.

Implementations

impl InferenceContext

pub fn new(device: Device) -> Self

pub fn device(&self) -> &Device

pub fn take_decode_session(&mut self) -> Option<DecodeSession>

Take the held decode plan out of the context, leaving None.

Paired with Self::put_decode_session. The take/put dance exists for the borrow checker: the persistent forward needs &mut InferenceContext AND &mut Option<DecodeSession> at once, which it cannot have if the session lives inside the context. Callers outside forward_decode_step should not need either method.

pub fn put_decode_session(&mut self, session: Option<DecodeSession>)

Put a decode plan back into the context. See Self::take_decode_session.

pub fn has_decode_session(&self) -> bool

Whether a decode plan is currently held. Observability hook — a caller that has decoded at least one token through LlamaModel::forward_decode_step holds one.

pub fn insert(&mut self, node_id: NodeId, storage: Arc<RwLock<Storage>>)

Insert a long-lived storage at node_id. The Arc survives across subsequent realize calls and is treated as pre-realized for any graph that references node_id as a Const-shaped leaf.

Replaces any prior entry at node_id.

pub fn remove(&mut self, node_id: NodeId) -> Option<Arc<RwLock<Storage>>>

Remove the persistent entry at node_id. Returns the Arc if one was present. After this, the next realize that touches node_id will re-fetch from the graph’s storage_map (and fail if the slot isn’t populated there either).

pub fn contains(&self, node_id: NodeId) -> bool

Whether the persistent map has an entry at node_id.

pub fn get(&self, node_id: NodeId) -> Option<&Arc<RwLock<Storage>>>

Borrow the persistent entry at node_id if present.

pub fn len(&self) -> usize

Number of persistent slots.

pub fn is_empty(&self) -> bool

pub fn clear_persistent(&mut self)

Drop every persistent entry. The Arcs are released; storages whose only remaining reference was the context’s are freed (host memory for CPU, deferred-on-stream for async backends).

pub fn realize_one_as<T: Pod + WithDType>( &self, graph: &Arc<RwLock<Graph>>, target: NodeId, ) -> Result<Vec<T>>

Realize a single target. The persistent map is seeded into the executor’s input cache via Arc::clone (the persistent entries survive the call). Op::Const NodeIds in the graph that aren’t already in the persistent map get uploaded fresh from graph.storage_for(id) per the existing pipelined-bridge pattern.

pub fn realize_one_as_with_env<T: Pod + WithDType>( &self, graph: &Arc<RwLock<Graph>>, target: NodeId, env: &SymEnv, ) -> Result<Vec<T>>

Env-carrying counterpart of Self::realize_one_as: supplies a per-pass SymEnv binding the runtime values of any DynScalar op params (Phase D symbolic extents — e.g. the decode KV-cache write offset cached_len). The env is per-pass (re-supplied every forward step) while the persistent map is per-session; an empty env is byte-identical to Self::realize_one_as.

pub fn prebuild_optimized_as_with_env<T: Pod + WithDType>( &self, graph: &Arc<RwLock<Graph>>, target: NodeId, env: &SymEnv, ) -> Result<(NodeId, OptimizedGraph, Vec<T>)>

Phase D · D2a — first-realize that ALSO captures the reusable optimize artifacts for a later prebuilt (plan-once) re-realize.

Runs the full prepare + optimize_graph + dispatch path ONCE (like Self::realize_one_as_with_env) and returns (effective_target, OptimizedGraph, result). The caller (D2b’s DecodeSession) holds the effective_target (the spliced D2H Op::Copy root) + the cached OptimizedGraph and feeds them to Self::realize_prebuilt_as_with_env on later tokens to SKIP the re-plan. See crate::pipelined_bridge::prebuild_optimized_env.

pub fn prebuild_optimized_capturing_as_with_env<T: Pod + WithDType>( &self, graph: &Arc<RwLock<Graph>>, target: NodeId, env: &SymEnv, ) -> Result<(NodeId, OptimizedGraph, StorageCache, Vec<T>)>

Phase D · D2b — first-realize that captures the reusable optimize artifacts AND the full realized StorageCache (all weight Consts uploaded by build_const_cache, merged over this context’s persistent Arcs). The DecodeSession holds the cache so later prebuilt realizes (which SKIP the const-cache walk) still resolve every weight Const. Returns (effective_target, OptimizedGraph, full_cache, result).

pub fn realize_prebuilt_as_with_env<T: Pod + WithDType>( &self, graph: &Arc<RwLock<Graph>>, effective_target: NodeId, optimized: &OptimizedGraph, env: &SymEnv, ) -> Result<Vec<T>>

Phase D · D2a — plan-once re-realize over a graph already prepared + optimized by a prior Self::prebuild_optimized_as_with_env. Skips BOTH prepare (no D2H re-splice / const-cache walk) AND optimize_graph (no re-plan / double-insert). Re-binds only the per-call StorageCache (this context’s persistent Arcs, incl. the re-bound per-token data Consts) + env, then dispatches straight to the executor. effective_target + optimized come from the prebuild.

A TopologyChanged error surfaces to the caller (typed, not retried) — the cached view is stale; invalidate + rebuild the session.

pub fn realize_many_as<T: Pod + WithDType>( &self, graph: &Arc<RwLock<Graph>>, targets: &[NodeId], ) -> Result<Vec<Vec<T>>>

Multi-target counterpart of Self::realize_one_as.

pub fn realize_many_as_with_env<T: Pod + WithDType>( &self, graph: &Arc<RwLock<Graph>>, targets: &[NodeId], env: &SymEnv, ) -> Result<Vec<Vec<T>>>

Env-carrying counterpart of Self::realize_many_as (Phase D symbolic extents). An empty env is byte-identical to Self::realize_many_as.

pub fn realize_split_as<T: Pod + WithDType>( &self, graph: &Arc<RwLock<Graph>>, targets: &[NodeId], n_host: usize, ) -> Result<SplitRealizeOutput<T>>

Realize-split counterpart of Self::realize_many_as: the first n_host targets are downloaded to host Vec<T>s, the rest come back as device-resident (storage, layout) pairs — no D2H for results that feed the next step’s graph. See crate::pipelined_bridge::realize_split_as_with_initial.