Struct DecodeSession 

pub struct DecodeSession { /* private fields */ }

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.

The ~1.8×/token win lives here: the decode-step graph is built + prepared (D2H Op::Copy spliced at the root) + optimize_graph’d (placement stamps + residency Op::Copy + layout Op::Contiguize baked in place) ONCE, then re-realized via the D2a prebuilt seam (InferenceContext::realize_prebuilt_as_with_env) which SKIPS both prepare and optimize_graph.

Held by the generate-loop owner (the caller of forward_with_kv_context_persistent), NOT by the immutable LlamaModel (the model is read-only weights; a session is per- generation state). Constructed lazily as Option<DecodeSession> and dropped-then-rebuilt on any validity-key mismatch / TopologyChanged / a non-seq==1 step (see forward_with_kv_context_persistent).

Why the graph is HELD here

D1 rebuilds a fresh Tensor graph every token and drops it after realize. D2 keeps the Arc<RwLock<Graph>> alive on the session so the already-optimized structure (the stamps + inserted copies) survives. The cached OptimizedGraph holds only {roots, generation} — it bakes NO Const data / storage / SymEnv, so reusing it is sound iff the graph structure + topology generation are unchanged (D1’s input-independent-graph guarantee).

Implementations

impl DecodeSession

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, mask_node: NodeId, kv_nodes: Vec<(NodeId, NodeId)>, offset_node: Option<NodeId>, cached_len_sym: SymId, attended_len_sym: SymId, base_cache: StorageCache, seq: usize, max_seq_len: usize, n_layers: usize, cache_dtype: DType, shape_key: u64, kv: &dyn KvRebindSource, ) -> Self

Assemble a held session from the artifacts a first-realize prebuild produced. graph is the (already prepared + optimized) held decode graph; effective_target / optimized come from InferenceContext::prebuild_optimized_as_with_env; the node ids are the STABLE data-Const NodeIds the builder minted.

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

Realize the held graph for one decode token via the D2a prebuilt seam (SKIP prepare + optimize_graph). data supplies the per-token re-bound data-Const Arcs (token-ids / RoPE cos+sin / mask); they overwrite the held base_cache’s entries for this token (the weight + KV Arcs stay from base_cache). sym_env binds cached_len. TopologyChanged surfaces typed (caller invalidates the session).

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

Whether this held session is valid for the given decode step. Rebuild (drop + build fresh) on any mismatch. seq must be 1 (the held graph is the seq==1 decode graph); a change in max_seq_len / n_layers / cache_dtype means a different model/cache → the held graph’s shapes are stale. May this held plan be reused for a step with these parameters?

shape_key is the caller’s crate::decode_shape key. Geometry alone is NOT sufficient and was the whole reason that module exists: a held graph bakes its op structure and its weight Consts, so two models at identical dims — a different architecture, or the same architecture with different weights — would otherwise be judged interchangeable and one would silently execute the other’s baked plan.

alloc_id is the live cache’s KvCache::alloc_id. It is a SEPARATE question from shape_key and neither subsumes the other: shape_key names the model whose weights are baked in, alloc_id names the KV buffers the plan writes into. One model serving many concurrent requests has one shape_key and one allocation per request, and it is exactly that case — the slot-pooled server — where geometry-only keying hands request B a plan welded to request A’s KV. Nothing errors: B just decodes over A’s context at full speed.

Both directions matter. Returning true too readily is a wrong answer at full speed; returning false too readily forfeits plan reuse (measured 223× on CUDA) while every correctness test stays green. That second direction is why alloc_id tracks the ALLOCATION and not the conversation — see KvCache::truncate_to, which must not invalidate.

pub fn validity_for( &self, seq: usize, max_seq_len: usize, n_layers: usize, cache_dtype: DType, shape_key: u64, alloc_id: KvAllocId, ) -> PlanValidity

The same comparisons as Self::is_valid_for, reporting which kind of mismatch occurred rather than collapsing to a bool.

The split exists so GAP-028 does not introduce a second predicate. A caller that wants to try Self::rebind_kv needs to know “everything matches except the allocation”, and answering that with an is_valid_ignoring_alloc twin would leave two comparison lists to keep in sync — the exact drift that made lazy::invalidate_decode_pair_if_stale a shared free function instead of a per-model copy. There is one list, here.

Note the asymmetry between the two non-Valid answers. AllocationChanged is a conditional verdict: it says a re-bind is worth attempting, NOT that one is safe. Everything it reports on is a property of the held plan; whether the incoming storage is actually interchangeable is Self::rebind_kv’s question, and it is answered against the storage itself.

pub fn rebind_kv( &mut self, kv: &dyn KvRebindSource, ) -> Result<(), RebindRefusal>

Re-point the held plan’s baked KV Consts at a replacement allocation, or refuse — the GAP-028 alternative to throwing the plan away.

What this trades

GAP-014 shipped invalidate-and-rebuild, which is safe by construction: a plan welded to storage it no longer serves is simply destroyed, and no guard can be wrong because there is no guard. This method replaces that with a guarantee that holds by proof. If the checks below ever accept an incompatible allocation, the failure mode is GAP-014 verbatim — a plan executing one request’s tokens over another’s KV, at full speed, with a plausible distribution and nothing to report. Every check here is therefore load-bearing for correctness, not for performance, and the cost of refusing wrongly is only a rebuild.

What is checked, and against what

Not against the cache’s own description of itself — against the storage already baked into base_cache, which is ground truth and cannot disagree with the graph the way a KvCache field can. Per slot: dtype, encoding, backend variant, bundling, and byte extent. Byte extent is the load-bearing one: it collapses batch × n_kv_heads × max_seq_len × head_dim into a single comparison that needs no reconstruction of the baked shape.

Residency is the exception and is why KvRebindSource::rebind_location exists. fuel_memory::Storage exposes a backend variant but no device ordinal — VulkanStorage records none at all — so a same-geometry cache on a different GPU of the same backend is invisible from the bytes. It would pass every check above and be spliced into a graph whose placement stamps and residency Op::Copys were baked for the old device. An unprovable residency is therefore a refusal, never a default.

What the caller still owes

A re-bind invalidates any capture. A CapturedDecodeSession records a CUDA graph over FIXED device addresses drawn from base_cache; after this returns Ok, some of those addresses name the old allocation. Callers must clear the capture alongside, on the same reader-before-owner rule that governs lazy::invalidate_decode_pair_if_stale.

Structural validity is NOT rechecked here — call this only when Self::validity_for answered PlanValidity::AllocationChanged.

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

The held graph handle (the caller re-binds data Consts + realizes through it).

pub fn optimized(&self) -> &OptimizedGraph

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

pub fn effective_target(&self) -> NodeId

The D2H Op::Copy root the executor is asked for.

pub fn logits_node(&self) -> NodeId

The logits node (pre-D2H-splice) — capture_decode hard-rejects any Op::Copy/Move in the captured graph (cross-device, not single-device-CUDA-capturable), so [fuel_dispatch::pipelined:: CapturedDecodeSession::capture] must target THIS node, not Self::effective_target (the spliced D2H Op::Copy root). The D2H happens OUTSIDE the capture, after replay_token returns the device-resident output Arc.

pub fn token_ids_node(&self) -> NodeId

pub fn rope_cos_node(&self) -> NodeId

pub fn rope_sin_node(&self) -> NodeId

pub fn mask_node(&self) -> NodeId

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

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

The device-resident KV-write offset Const NodeId, Some only on the device-offset decode path (CUDA/CPU); None on Vulkan’s SymEnv path. Used by the CapturedRun replay wiring (Phase 3) and the per-token offset re-bind.

pub fn cached_len_sym(&self) -> SymId

pub fn attended_len_sym(&self) -> SymId

pub fn max_seq_len(&self) -> usize

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

Build the per-token SymEnv for one decode step: bind cached_len_sym = cached_len (the KV-write offset) AND attended_len_sym = cached_len + seq (the flash-arm k_len — the live attended prefix including this token). Both are bound each token; the attended-length binding is unreferenced on today’s f32 decode graph (no flash arm) and becomes load-bearing the moment a bf16/f16 CUDA decode offers the arm. Write-once per pass (a conflicting rebind surfaces a typed error, never a panic).

DO NOT DELETE THE cached_len_sym BINDING BECAUSE IT LOOKS DEAD

LlamaModel and PhiModel take DIFFERENT decode paths on CPU/F32, measured 2026-08-12 (GAP-029 2b):

llama: offset_node.is_some() == true   → device-offset path; SymEnv INERT
phi:   offset_node.is_some() == false  → SymEnv path;        SymEnv LIVE

On the device-offset path the KV write offset rides a device-resident buffer (Op::WriteSliceDoff reads the start from DecodeTokenData::offset at launch), so no symbol here drives Llama’s write and both bindings are inert for it. On the SymEnv path the offset rides cached_len_sym, and it is load-bearing for Phi — proved by negative control: binding it wrong moves Phi’s logits (phi_attended_len_sym_is_unreferenced_negative_control).

The trap this comment exists to prevent: someone profiling or tracing the Llama path sees a binding nothing consumes and removes it as dead code. Their evidence would be about Llama; the requirement comes from Phi, and the breakage would be silent — wrong KV write offsets, not a compile error. Whichever model you were looking at, you were looking at half the callers.

attended_len_sym is separately inert on both paths today (F32, CPU, no flash arm) and is bound anyway: it becomes load-bearing the moment a bf16/f16 CUDA decode offers the flash arm, and binding it unconditionally is what makes that switch a no-op rather than a hunt. That claim is scoped — re-run the controls above before relying on it under any other dtype/graph.

pub fn graph_node_count(&self) -> usize

The held graph’s node count. The D2b 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 / a builder sneaking a cached_len-dependent shape back in).

pub fn base_cache(&self) -> &StorageCache

The full realized StorageCache from the first realize (every weight Const + the KV Arcs + the initial per-token data Consts). The CapturedRun capture-building step clones this cheaply (Arc-clones only) and overwrites the per-token entries with fresh FIXED-address Arcs before handing the merged cache to [fuel_dispatch::pipelined::CapturedDecodeSession::capture] as inputs.