fuel-tensor 0.19.1

// SPDX-License-Identifier: MIT OR Apache-2.0
//! 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.

use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, RwLock};

use fuel_dispatch::{
    optimize::OptimizedGraph,
    pipelined::{PipelinedExecutor, StorageCache},
};
use fuel_graph::{Graph, Node, NodeId, Op};
use fuel_ir::{DType, DeviceLocation, Error, Layout, Result, SType, Shape, SymEnv, SymId};
use fuel_memory::{BackendStorage, Storage};

use crate::Device;
use crate::decode_shape::KvAllocId;

// ===========================================================================
// KV cache primitive types
// ===========================================================================

/// 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.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct KvLayerId {
    pub layer_idx: usize,
    pub slot: KvSlot,
}

/// Which half of a KV pair a [`KvLayerId`] addresses.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum KvSlot {
    K,
    V,
}

/// 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.
///
/// The intended protocol (deferred until multi-GPU lands):
/// - **Host**: the host has the most recent version. Devices may
///   hold stale copies. The default at session start.
/// - **Device(loc)**: a device owns the most recent version. The
///   host's copy may be stale until an explicit flush.
#[derive(Debug, Clone)]
pub enum AuthorityState {
    Host,
    Device(DeviceLocation),
}

/// One layer's K + V slots in a backend-erased KV cache.
///
/// Storage Arcs are dispatch-erased: each Arc points at a
/// [`fuel_memory::BackendStorage`] enum variant (`Cpu`, `Cuda`,
/// `Vulkan`, `Metal`). The graph references these Arcs via the
/// persistent map in [`InferenceContext`]; readers see the bytes
/// through whichever Layout the graph carries for the consuming
/// NodeId.
pub struct KvLayer {
    pub k: Arc<RwLock<Storage>>,
    pub v: Arc<RwLock<Storage>>,
    /// View layout into `k` describing the live extent. For the grow-
    /// by-replacement pattern (today) the layout matches the storage's
    /// full shape. For pre-allocated buffers (Phase E.3.2) the layout
    /// describes the leading `[..., cached_len, ...]` slice.
    pub k_layout: Layout,
    pub v_layout: Layout,
    /// Monotonic write version for K. Bumps on every successful
    /// `Op::WriteSlice` targeting this slot. Used by the future
    /// multi-device coherence protocol to detect staleness.
    pub k_version: u64,
    pub v_version: u64,
    /// Coherence authority per slot. K and V can in principle diverge
    /// (a future optimization that decouples their update cadence
    /// could exploit this), though standard transformer decoding
    /// updates both together every step.
    pub k_authority: AuthorityState,
    pub v_authority: AuthorityState,
}

/// 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.
///
/// Pipeline-parallel friendly: when a model is split across devices,
/// each device's cache holds entries only for its layer range; other
/// indices are `None`. The model code queries `cache.layer(global_i)`
/// regardless of which device owns it; the storage Arc inside
/// [`KvLayer`] carries the device identity via its
/// [`fuel_memory::BackendStorage`] variant.
pub struct KvCache {
    pub layers: Vec<Option<KvLayer>>,
    pub cached_len: usize,
    pub n_kv_heads: usize,
    pub head_dim: usize,
    /// Identity of the *allocation* backing this cache — the thing a held
    /// [`DecodeSession`] is welded to. Private with an [`Self::alloc_id`]
    /// reader so it cannot be assigned by hand: every path that replaces
    /// storage ([`Self::clear`], [`Self::set_layer`]) re-mints it, and
    /// [`Self::truncate_to`] deliberately does not. See [`KvAllocId`].
    alloc_id: KvAllocId,
    /// Pre-allocated capacity along the sequence axis. `Some(n)` when
    /// constructed via [`KvCache::with_capacity`] — every layer's K + V
    /// storage holds `[1, n_kv_heads, n, head_dim]` of zeros and the
    /// forward path writes into it via `Op::WriteSlice`. `None` for the
    /// legacy [`KvCache::with_dims`] grow-by-replacement constructor.
    pub max_seq_len: Option<usize>,
    /// Dtype the pre-allocated buffers were allocated with. `Some(dt)`
    /// in the [`KvCache::with_capacity`] path, `None` in the legacy
    /// `with_dims` path (which leaves dtype to whatever the first
    /// inserted layer specifies).
    pub dtype: Option<DType>,
    /// Where this cache's storage lives — `Some(loc)` only while every
    /// buffer is one [`Self::with_capacity`] allocated on `loc`, `None`
    /// the moment that stops being provable.
    ///
    /// Exists for [`DecodeSession::rebind_kv`], which splices these `Arc`s
    /// into a graph whose placement stamps and residency `Op::Copy`s were
    /// baked against the *previous* cache. Shape it can prove from the
    /// storage itself; residency it cannot — `fuel_memory::Storage` exposes
    /// a backend *variant*, and `VulkanStorage` records no device ordinal at
    /// all, so a multi-GPU swap is invisible from the bytes. The cache is
    /// therefore the honest carrier: it is handed a `&Device` at the one
    /// moment residency is known for certain.
    ///
    /// `None` means "cannot prove", not "CPU" — and every consumer must
    /// treat it as a refusal. [`Self::with_dims`] never had a device;
    /// [`Self::set_layer`] installs caller-supplied storage from anywhere.
    /// Both therefore forfeit re-bind and fall back to a plan rebuild,
    /// which is what shipped in GAP-014 and is safe by construction.
    device_location: Option<DeviceLocation>,
}

impl KvCache {
    /// Legacy grow-by-replacement constructor. Layers start `None`; the
    /// caller calls [`Self::set_layer`] each forward step with a freshly
    /// allocated `KvLayer` whose storage is sized to the current cached
    /// length. Subsequent steps replace the layer entirely with a
    /// larger buffer.
    ///
    /// Replaced by [`Self::with_capacity`] when the caller knows the
    /// generation's max sequence length up-front (the common case for
    /// autoregressive decoding).
    pub fn with_dims(n_layers: usize, n_kv_heads: usize, head_dim: usize) -> Self {
        Self {
            layers: (0..n_layers).map(|_| None).collect(),
            cached_len: 0,
            n_kv_heads,
            head_dim,
            max_seq_len: None,
            dtype: None,
            alloc_id: KvAllocId::next(),
            // No device was named, and this constructor's whole shape is
            // grow-by-replacement — the storage a plan baked is replaced
            // every step anyway. Unprovable residency, and nothing lost.
            device_location: None,
        }
    }

    /// Pre-allocated KV cache. Every layer's K + V storage is allocated
    /// up-front as a `[1, n_kv_heads, max_seq_len, head_dim]` zero
    /// buffer on `device` with `dtype`. Subsequent forward steps write
    /// fresh K/V slabs into these buffers via `Op::WriteSlice` rather
    /// than growing-by-replacement.
    ///
    /// Memory footprint: `n_layers * 2 * n_kv_heads * max_seq_len *
    /// head_dim * dtype_size` bytes. For Llama-7B at max_seq_len=4096,
    /// bf16: 32 * 2 * 8 * 4096 * 128 * 2 ≈ 512 MiB.
    ///
    /// Returns `Err` if any per-layer allocation fails (e.g. CUDA OOM)
    /// or the requested device hasn't been wired up in
    /// [`crate::pipelined_bridge`] (Vulkan / Metal — D2H/H2D for those still
    /// goes through the legacy executor).
    pub fn with_capacity(
        n_layers: usize,
        n_kv_heads: usize,
        head_dim: usize,
        max_seq_len: usize,
        dtype: DType,
        device: &Device,
    ) -> Result<Self> {
        // Bridge-retirement Phase 3a (post-9c): allocate the 2N K/V
        // buffers via `Op::Alloc` graph emission, not direct backend
        // calls. The executor's `WorkItemKind::Alloc` arm dispatches
        // per-backend; the per-`DeviceLocation` match that used to
        // live in `alloc_zeroed_on` is now in the executor (the
        // architectural dispatch layer).
        let shape = Shape::from_dims(&[1, n_kv_heads, max_seq_len, head_dim]);
        let layout = Layout::contiguous(shape.clone());
        let target_loc = device.location();

        // Build the transient graph. For non-CPU targets the first
        // node is an `Op::Const` placeholder whose StorageCache entry
        // carries a small "device anchor" storage — the executor's
        // Alloc arm searches the cache for any storage on the target
        // backend to derive the device handle, so this entry makes
        // the first Op::Alloc's device lookup succeed. The anchor's
        // NodeId must exist in the graph so `realize_many`'s
        // `layout_cache` seeding (which calls `g.layout(id)` on every
        // cache key) doesn't panic.
        let graph = Arc::new(RwLock::new(Graph::new()));
        let mut cache = StorageCache::new();
        if let Some(seed) = crate::pipelined_bridge::device_seed_storage(device)? {
            let anchor_id = {
                let mut g = graph.write().map_err(|_| {
                    Error::Msg("graph lock poisoned during KvCache build".into()).bt()
                })?;
                g.push(Node {
                    op: Op::Const,
                    inputs: vec![],
                    shape: Shape::from_dims(&[4]),
                    dtype: DType::U8,
                })
            };
            cache.insert(anchor_id, Arc::new(RwLock::new(seed)));
        }
        // Emit pairs of (Op::Alloc, Op::ZeroFill) per K/V slot. The
        // ZeroFill consumes the Alloc destructively (in-place fill,
        // output adopts the Alloc's Storage Arc). Realizing the
        // ZeroFill IDs produces zero-initialized device storages.
        //
        // Phase 3a follow-up architecture: Op::Alloc gives uninit
        // memory; Op::ZeroFill explicitly fills with zero. On Vulkan
        // this replaces the host-staged-zeros path with device-side
        // `vkCmdFillBuffer` (~2× the bandwidth saved). On CUDA the
        // pair becomes uninit `cuMemAlloc` + async `cuMemsetD8Async`
        // — same total cost as the old `alloc_zeros` but exposes the
        // ZeroFill as a first-class graph node the optimizer can
        // see (and skip if a downstream op covers the full buffer).
        let zero_fill_ids: Vec<NodeId> = {
            let mut g = graph
                .write()
                .map_err(|_| Error::Msg("graph lock poisoned during KvCache build".into()).bt())?;
            (0..(2 * n_layers))
                .map(|_| {
                    let alloc_id = g.push(Node {
                        op: Op::Alloc { target: target_loc },
                        inputs: vec![],
                        shape: shape.clone(),
                        dtype,
                    });
                    g.push(Node {
                        op: Op::ZeroFill,
                        inputs: vec![alloc_id],
                        shape: shape.clone(),
                        dtype,
                    })
                })
                .collect()
        };

        // Realize all 2*n_layers Op::ZeroFill targets in one pass —
        // PipelinedExecutor::realize_many shares the compile/execute
        // pipeline so device-handle reuse is automatic, and the
        // executor's destructive_input cleanup evicts the Op::Alloc
        // intermediate NodeIds while the ZeroFill targets adopt the
        // same Arcs (post-fill bytes).
        let realized = PipelinedExecutor::realize_many(Arc::clone(&graph), &zero_fill_ids, cache)?;
        if realized.len() != 2 * n_layers {
            return Err(Error::Msg(format!(
                "KvCache::with_capacity: realize_many returned {} storages \
                 for {} Op::ZeroFill targets — internal bug",
                realized.len(),
                2 * n_layers,
            ))
            .bt());
        }

        // Never-panic: the length is checked above, but the production path
        // must not `.expect()` — surface a typed `Err` if the iterator runs
        // short. (Same pattern as `alloc_batched_kv` below.)
        let mut realized_iter = realized.into_iter();
        let short = |what: &str, l: usize| {
            Error::Msg(format!(
                "alloc_kv: realize_many yielded fewer storages than expected \
                 (ran out at layer {l} fetching {what})"
            ))
            .bt()
        };
        let mut layers: Vec<Option<KvLayer>> = Vec::with_capacity(n_layers);
        for l in 0..n_layers {
            let (k_arc, _) = realized_iter.next().ok_or_else(|| short("K", l))?;
            let (v_arc, _) = realized_iter.next().ok_or_else(|| short("V", l))?;
            layers.push(Some(KvLayer {
                k: k_arc,
                v: v_arc,
                k_layout: layout.clone(),
                v_layout: layout.clone(),
                k_version: 0,
                v_version: 0,
                k_authority: AuthorityState::Host,
                v_authority: AuthorityState::Host,
            }));
        }

        // Silence the unused-bind lint for elem_count when this
        // path no longer needs the legacy alloc helper.
        let _ = (n_kv_heads, max_seq_len, head_dim);

        Ok(Self {
            layers,
            cached_len: 0,
            n_kv_heads,
            head_dim,
            max_seq_len: Some(max_seq_len),
            dtype: Some(dtype),
            alloc_id: KvAllocId::next(),
            // Every buffer above was `Op::Alloc { target: target_loc }`, so
            // this is measured at the allocation, not inferred later.
            device_location: Some(target_loc),
        })
    }

    /// Borrow the K-or-V storage Arc for a given layer + slot. Used by
    /// the forward path to bind cache storage to per-step Const nodes
    /// via [`InferenceContext::insert`]. Returns `None` if the layer
    /// isn't populated (with_dims path before first
    /// [`Self::set_layer`]) or `layer_idx` is out of range.
    pub fn slot_storage(&self, layer_idx: usize, slot: KvSlot) -> Option<Arc<RwLock<Storage>>> {
        let layer = self.layer(layer_idx)?;
        Some(match slot {
            KvSlot::K => Arc::clone(&layer.k),
            KvSlot::V => Arc::clone(&layer.v),
        })
    }

    /// Bump a slot's monotonic version counter. Called after every
    /// successful `Op::WriteSlice` targeting the slot — the forward
    /// path invokes this once per (layer, slot) per step.
    ///
    /// The version is consumed by the future multi-device coherence
    /// protocol (placeholder today — see [`AuthorityState`]).
    pub fn bump_version(&mut self, layer_idx: usize, slot: KvSlot) {
        if let Some(layer) = self.layer_mut(layer_idx) {
            match slot {
                KvSlot::K => layer.k_version += 1,
                KvSlot::V => layer.v_version += 1,
            }
        }
    }

    pub fn n_layers(&self) -> usize {
        self.layers.len()
    }

    pub fn layer(&self, li: usize) -> Option<&KvLayer> {
        self.layers.get(li).and_then(|o| o.as_ref())
    }

    pub fn layer_mut(&mut self, li: usize) -> Option<&mut KvLayer> {
        self.layers.get_mut(li).and_then(|o| o.as_mut())
    }

    /// Identity of the storage this cache currently owns — the value a held
    /// [`DecodeSession`]'s validity key carries so a cache **swap** cannot be
    /// mistaken for a continuation. See [`KvAllocId`] for why re-allocation
    /// re-mints it and rewinding does not.
    pub fn alloc_id(&self) -> KvAllocId {
        self.alloc_id
    }

    /// Where this cache's storage lives, or `None` when that is not provable.
    /// See the field's doc — `None` is a refusal, never a CPU default.
    pub fn device_location(&self) -> Option<DeviceLocation> {
        self.device_location
    }

    /// Install a layer's storage, replacing whatever was there.
    ///
    /// This **re-mints [`Self::alloc_id`]**: the storage `Arc` a held plan
    /// baked into its `base_cache` for this layer is no longer the storage
    /// this cache reads, so any plan welded to the old one must be rebuilt.
    /// It also clears [`Self::device_location`]: the storage arrives from the
    /// caller, so where it lives is no longer something this cache measured.
    /// That forfeits [`DecodeSession::rebind_kv`] for the rest of this cache's
    /// life — deliberately, because an unprovable residency spliced into a
    /// baked plan is the GAP-014 failure one layer down.
    pub fn set_layer(&mut self, li: usize, layer: KvLayer) {
        self.layers[li] = Some(layer);
        self.alloc_id = KvAllocId::next();
        self.device_location = None;
    }

    /// Drop every layer; reset `cached_len` to zero. Use between
    /// independent generations (the cache's K/V shapes are tied to
    /// a specific prompt prefix).
    ///
    /// **Re-mints [`Self::alloc_id`]** — every storage `Arc` is released, so
    /// this is an allocation change, not a rewind. Contrast
    /// [`Self::truncate_to`], which moves only `cached_len` and therefore must
    /// preserve the id.
    pub fn clear(&mut self) {
        for layer in &mut self.layers {
            *layer = None;
        }
        self.cached_len = 0;
        self.alloc_id = KvAllocId::next();
        // There is no storage left to have a residency. Repopulating goes
        // through `set_layer`, which cannot prove one either.
        self.device_location = None;
    }

    /// Shrink the cache to `new_len` cached positions. Speculative-
    /// decoding's reject path calls this after a draft batch is
    /// rejected by the target model.
    ///
    /// **Phase E.3.0 limitation**: this only updates `cached_len`;
    /// the underlying storages and their layouts are *not* shrunk.
    /// The grow-by-replacement pattern (today) means the next forward
    /// pass would still see the old layout's full extent. Phase E.3.2
    /// pairs this with pre-allocated buffers + `Op::WriteSlice` so
    /// truncate becomes a pure metadata update — the trailing rows
    /// of the pre-allocated buffer simply stop being read.
    ///
    /// **Deliberately does NOT touch [`Self::alloc_id`].** Rewinding the
    /// conversation is not re-allocating the buffers: a held plan is still
    /// welded to exactly the right storage, so invalidating here would forfeit
    /// plan reuse on every rejected draft batch and buy no correctness. The
    /// id names the allocation, never the conversation — see [`KvAllocId`].
    pub fn truncate_to(&mut self, new_len: usize) {
        if new_len < self.cached_len {
            self.cached_len = new_len;
        }
    }

    /// Write host-supplied K/V bytes into layer `layer_idx`'s `slot` buffer at
    /// position range `[start, start+n)` — the `KvCache`-side sibling of
    /// [`DeviceKvPool::write_block_bytes`](crate::kv_block_pool_device::DeviceKvPool::write_block_bytes).
    ///
    /// Builds a throwaway `Op::Const`(bytes) → `Op::WriteSlice` graph and
    /// realizes it ONCE into the SAME buffer `Arc` this cache already holds —
    /// it does **not** replace the Arc, so unlike [`Self::set_layer`] this
    /// does **not** touch [`Self::alloc_id`] or [`Self::device_location`]: a
    /// [`DecodeSession`] already welded to this cache's buffers (via
    /// `base_cache`) sees the write transparently on its next token, with no
    /// rebind and no rebuild — the write lands in the exact buffer the held
    /// plan's `Op::WriteSlice` nodes already target.
    ///
    /// `bytes` is `n * n_kv_heads * head_dim * dtype.size_in_bytes()` raw
    /// little-endian element bytes, row-major `[1, n_kv_heads, n, head_dim]`
    /// — **HEAD-MAJOR** (position is axis 2). This is NOT the same layout as
    /// [`DeviceKvPool`](crate::kv_block_pool_device::DeviceKvPool)'s
    /// **position-major** `[block_size, Hkv, D]` block layout: a caller
    /// moving bytes from the paged pool into a `KvCache` must transpose, a
    /// plain memcpy silently produces wrong attention.
    ///
    /// `Err` (never a panic): this cache has no fixed dtype/capacity (built
    /// via [`Self::with_dims`] rather than [`Self::with_capacity`] — this
    /// primitive is `with_capacity`-only, since the legacy grow-by-replacement
    /// path has no stable buffer to write into); `layer_idx` is out of range
    /// or unpopulated; `start + n` exceeds the allocated `max_seq_len`; or
    /// `bytes.len()` doesn't match the expected size for `n` positions.
    pub fn write_bytes(
        &self,
        layer_idx: usize,
        slot: KvSlot,
        start: usize,
        n: usize,
        bytes: &[u8],
        device: &Device,
    ) -> Result<()> {
        let dtype = self.dtype.ok_or_else(|| {
            Error::Msg(
                "KvCache::write_bytes: cache has no fixed dtype — not built via with_capacity"
                    .into(),
            )
            .bt()
        })?;
        let max_seq_len = self.max_seq_len.ok_or_else(|| {
            Error::Msg(
                "KvCache::write_bytes: cache has no fixed capacity — not built via with_capacity"
                    .into(),
            )
            .bt()
        })?;
        if start + n > max_seq_len {
            return Err(Error::Msg(format!(
                "KvCache::write_bytes: range [{start}, {}) exceeds max_seq_len {max_seq_len}",
                start + n,
            ))
            .bt());
        }
        let elem_bytes = kv_cache_dtype_elem_bytes(dtype)?;
        let want = n * self.n_kv_heads * self.head_dim * elem_bytes;
        if bytes.len() != want {
            return Err(Error::Msg(format!(
                "KvCache::write_bytes: {} bytes != expected {want} \
                 ([1, n_kv_heads {}, n {n}, head_dim {}] for {dtype:?})",
                bytes.len(),
                self.n_kv_heads,
                self.head_dim,
            ))
            .bt());
        }
        let layer = self.layer(layer_idx).ok_or_else(|| {
            Error::Msg(format!(
                "KvCache::write_bytes: layer {layer_idx} out of range or unpopulated"
            ))
            .bt()
        })?;
        let buf = match slot {
            KvSlot::K => Arc::clone(&layer.k),
            KvSlot::V => Arc::clone(&layer.v),
        };
        let full_shape = Shape::from_dims(&[1, self.n_kv_heads, max_seq_len, self.head_dim]);
        let src = kv_cache_const_from_bytes(
            dtype,
            bytes,
            Shape::from_dims(&[1, self.n_kv_heads, n, self.head_dim]),
            device,
        )?;
        let dest = src.const_placeholder_like(full_shape, dtype);
        let ranges = vec![
            (0, 1),
            (0, self.n_kv_heads),
            (start, start + n),
            (0, self.head_dim),
        ];
        let post = dest.write_slice(&src, ranges)?;
        let mut cache = StorageCache::new();
        cache.insert(dest.node_id(), buf);
        let _ = crate::pipelined_bridge::realize_one_bytes_with_initial(
            post.graph_handle(),
            post.node_id(),
            device,
            cache,
        )?;
        Ok(())
    }
}

/// Byte width of one `KvCache` element for byte-level movement — the
/// `KvCache`-side sibling of
/// [`DeviceKvPool::dtype_elem_bytes`](crate::kv_block_pool_device::DeviceKvPool).
/// F32/BF16 only, matching [`KvCache::write_bytes`]'s supported dtypes.
fn kv_cache_dtype_elem_bytes(dtype: DType) -> Result<usize> {
    match dtype {
        DType::F32 => Ok(4),
        DType::BF16 => Ok(2),
        other => Err(Error::Msg(format!(
            "KvCache: byte movement unsupported for dtype {other:?} (F32/BF16)",
        ))
        .bt()),
    }
}

/// Build a fresh-graph `Op::Const` of `dtype` from raw little-endian element
/// bytes — the `KvCache`-side sibling of
/// [`DeviceKvPool::const_from_bytes`](crate::kv_block_pool_device::DeviceKvPool).
/// Dispatches `dtype` to the matching typed graph constructor via
/// alignment-safe `try_cast_slice` (never panics).
fn kv_cache_const_from_bytes(
    dtype: DType,
    bytes: &[u8],
    shape: Shape,
    device: &Device,
) -> Result<crate::lazy::Tensor> {
    match dtype {
        DType::F32 => {
            let v: &[f32] = bytemuck::try_cast_slice(bytes).map_err(|e| {
                Error::Msg(format!("kv_cache_const_from_bytes: f32 cast: {e:?}")).bt()
            })?;
            crate::lazy::Tensor::from_f32(v.to_vec(), shape, device)
        }
        DType::BF16 => {
            let v: &[half::bf16] = bytemuck::try_cast_slice(bytes).map_err(|e| {
                Error::Msg(format!("kv_cache_const_from_bytes: bf16 cast: {e:?}")).bt()
            })?;
            let anchor =
                crate::lazy::Tensor::from_f32(vec![0.0f32], Shape::from_dims(&[1]), device)?;
            crate::lazy::Tensor::from_bf16_on(anchor.graph(), v.to_vec(), shape, device)
        }
        other => Err(Error::Msg(format!(
            "kv_cache_const_from_bytes: unsupported dtype {other:?} (F32/BF16)",
        ))
        .bt()),
    }
}

/// One layer's `(K, V)` cache storages. A KV cache is a `Vec` of these, one
/// per layer, and the two halves are always allocated and dropped together.
type KvStoragePair = (Arc<RwLock<Storage>>, Arc<RwLock<Storage>>);

/// 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.
///
/// Returns `Err` if any allocation fails (fail-on-OOM, spec #6).
pub fn alloc_batched_kv(
    k: usize,
    n_layers: usize,
    n_kv_heads: usize,
    head_dim: usize,
    max_seq_len: usize,
    dtype: DType,
    device: &Device,
) -> Result<Vec<KvStoragePair>> {
    let shape = Shape::from_dims(&[k, n_kv_heads, max_seq_len, head_dim]);
    let target_loc = device.location();

    let graph = Arc::new(RwLock::new(Graph::new()));
    let mut cache = StorageCache::new();
    if let Some(seed) = crate::pipelined_bridge::device_seed_storage(device)? {
        let anchor_id = {
            let mut g = graph.write().map_err(|_| {
                Error::Msg("graph lock poisoned during alloc_batched_kv".into()).bt()
            })?;
            g.push(Node {
                op: Op::Const,
                inputs: vec![],
                shape: Shape::from_dims(&[4]),
                dtype: DType::U8,
            })
        };
        cache.insert(anchor_id, Arc::new(RwLock::new(seed)));
    }

    let zero_fill_ids: Vec<NodeId> = {
        let mut g = graph
            .write()
            .map_err(|_| Error::Msg("graph lock poisoned during alloc_batched_kv".into()).bt())?;
        (0..(2 * n_layers))
            .map(|_| {
                let alloc_id = g.push(Node {
                    op: Op::Alloc { target: target_loc },
                    inputs: vec![],
                    shape: shape.clone(),
                    dtype,
                });
                g.push(Node {
                    op: Op::ZeroFill,
                    inputs: vec![alloc_id],
                    shape: shape.clone(),
                    dtype,
                })
            })
            .collect()
    };

    let realized = PipelinedExecutor::realize_many(Arc::clone(&graph), &zero_fill_ids, cache)?;
    if realized.len() != 2 * n_layers {
        return Err(Error::Msg(format!(
            "alloc_batched_kv: realize_many returned {} storages for {} Op::ZeroFill targets",
            realized.len(),
            2 * n_layers,
        ))
        .bt());
    }

    // Never-panic: the length is checked above, but the production path must
    // not `.expect()` — surface a typed `Err` if the iterator is short.
    let mut it = realized.into_iter();
    let mut pairs = Vec::with_capacity(n_layers);
    let short = || {
        Error::Msg(
            "alloc_batched_kv: realize_many yielded fewer storages than expected".to_string(),
        )
        .bt()
    };
    for _ in 0..n_layers {
        let (k_arc, _) = it.next().ok_or_else(short)?;
        let (v_arc, _) = it.next().ok_or_else(short)?;
        pairs.push((k_arc, v_arc));
    }
    Ok(pairs)
}

// ===========================================================================
// LatentKvCache — N-slot generalization of KvCache (MLA increment 4)
// ===========================================================================

/// 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.
pub struct LatentSlot {
    pub storage: Arc<RwLock<Storage>>,
    /// View layout into `storage`. For the pre-allocated (capacity)
    /// buffers [`LatentKvCache::with_capacity`] allocates, the layout
    /// matches the storage's full `[max_seq_len, ...trailing]` shape.
    pub layout: Layout,
    /// Monotonic write version. Bumps on every successful
    /// `Op::WriteSlice` targeting this slot. Placeholder for the future
    /// multi-device coherence protocol — mirrors [`KvLayer::k_version`] /
    /// `v_version`.
    pub version: u64,
    /// Coherence authority. Placeholder — see [`AuthorityState`].
    pub authority: AuthorityState,
}

/// Cross-graph, cross-forward-pass persistent **N-slot** decode cache —
/// the [`KvCache`] generalization for latent-caching attention
/// architectures.
///
/// [`KvCache`] hardwires a symmetric K/V pair (two buffers per layer,
/// same shape). That's wrong for compression architectures:
///
///   - **Multi-head Latent Attention (DeepSeek-V2 MLA)**: per layer, a
///     low-rank compressed latent `[kv_lora_rank]` (slot 0) **and** a
///     single-head RoPE key `[qk_rope_head_dim]` (slot 1) — two slots of
///     *different* trailing shape, neither a `[n_kv_heads, head_dim]` K/V.
///   - **Two-projection attention / QKV pruning**: a single retained
///     projection per layer — one slot.
///
/// `LatentKvCache` is this module's counterpart to
/// [`crate::lazy_latent_cache::LatentCache`] — that type makes the
/// exact same **per-layer, ordered list of latent buffers with
/// independent trailing shapes** generalization for the **per-forward-
/// pass** (single [`fuel_graph::Graph`]-anchored, functional
/// append-and-thread) lifecycle. `LatentKvCache` is the **persistent**
/// sibling: device-resident `Arc<RwLock<Storage>>` buffers that survive
/// across graphs/forward calls and are mutated in place via
/// `Op::WriteSlice`, exactly the relationship [`KvCache`] has to a plain
/// per-pass K/V cache. Read both modules' docs together for the two
/// lifecycles side by side.
///
/// # Shape contract
///
/// Slot `s` of every layer is a buffer `[max_seq_len, ...slot_trailing[s]]`
/// — the capacity axis is dim **0**. This differs from [`KvCache`]'s `[1,
/// n_kv_heads, max_seq_len, head_dim]` convention (capacity axis 2):
/// a latent is a per-token *vector*, not a per-head *plane*, so there is
/// no head axis to lead with. Matches [`crate::lazy_latent_cache::
/// LatentCache`]'s dim-0 convention exactly — this type's per-slot
/// buffer is byte-for-byte what that type's `slot_buffer_full` realizes,
/// modulo device residency.
pub struct LatentKvCache {
    /// `layers[l]` is `None` until populated; `Some(slots)` holds one
    /// [`LatentSlot`] per slot index (pipeline-parallel friendly — mirrors
    /// [`KvCache::layers`]).
    pub layers: Vec<Option<Vec<LatentSlot>>>,
    pub cached_len: usize,
    /// Per-slot trailing shape (past the leading capacity axis); its
    /// length is the per-layer slot count.
    pub slot_trailing: Vec<Vec<usize>>,
    pub max_seq_len: usize,
    pub dtype: DType,
    /// Identity of the allocation backing this cache. Same role and same
    /// re-mint rules as [`KvCache::alloc_id`] — MLA reuses [`DecodeSession`]
    /// verbatim (its 2 latent slots ride `kv_nodes` as `(latent, kpe)` pairs),
    /// so it inherits the same welding and needs the same guard.
    alloc_id: KvAllocId,
    /// Where the slot buffers live. Unconditional here, unlike
    /// [`KvCache::device_location`]: this type has exactly one constructor
    /// and it is handed a `&Device`, so residency is never unprovable.
    device_location: DeviceLocation,
}

impl LatentKvCache {
    /// Pre-allocated N-slot latent cache. Every layer's slot buffers are
    /// allocated up-front as zero buffers on `device` with `dtype`,
    /// shaped `[max_seq_len, ...slot_trailing[s]]`. Mirrors [`KvCache::
    /// with_capacity`]'s `Op::Alloc` → `Op::ZeroFill` graph-emission
    /// pattern exactly, generalized from a fixed 2-buffer (K, V) layer to
    /// `slot_trailing.len()` buffers of independent shape.
    ///
    /// Returns `Err` for degenerate geometry (`n_layers == 0`,
    /// `max_seq_len == 0`, or no slots) or if any per-layer allocation
    /// fails (e.g. CUDA OOM / an unwired device).
    pub fn with_capacity(
        n_layers: usize,
        max_seq_len: usize,
        slot_trailing: Vec<Vec<usize>>,
        dtype: DType,
        device: &Device,
    ) -> Result<Self> {
        if n_layers == 0 {
            return Err(
                Error::Msg("LatentKvCache::with_capacity: n_layers must be >= 1".into()).bt(),
            );
        }
        if max_seq_len == 0 {
            return Err(Error::Msg(
                "LatentKvCache::with_capacity: max_seq_len must be >= 1".into(),
            )
            .bt());
        }
        if slot_trailing.is_empty() {
            return Err(Error::Msg(
                "LatentKvCache::with_capacity: need at least one slot (slot_trailing empty)".into(),
            )
            .bt());
        }
        let n_slots = slot_trailing.len();
        let shapes: Vec<Shape> = slot_trailing
            .iter()
            .map(|trailing| {
                let mut dims = Vec::with_capacity(1 + trailing.len());
                dims.push(max_seq_len);
                dims.extend_from_slice(trailing);
                Shape::from_dims(&dims)
            })
            .collect();
        let layouts: Vec<Layout> = shapes
            .iter()
            .map(|s| Layout::contiguous(s.clone()))
            .collect();
        let target_loc = device.location();

        // Transient graph, same device-anchor trick + Alloc/ZeroFill
        // emission pattern as KvCache::with_capacity — see that
        // constructor's doc for the full rationale.
        let graph = Arc::new(RwLock::new(Graph::new()));
        let mut cache = StorageCache::new();
        if let Some(seed) = crate::pipelined_bridge::device_seed_storage(device)? {
            let anchor_id = {
                let mut g = graph.write().map_err(|_| {
                    Error::Msg("graph lock poisoned during LatentKvCache build".into()).bt()
                })?;
                g.push(Node {
                    op: Op::Const,
                    inputs: vec![],
                    shape: Shape::from_dims(&[4]),
                    dtype: DType::U8,
                })
            };
            cache.insert(anchor_id, Arc::new(RwLock::new(seed)));
        }

        // Emit (Op::Alloc, Op::ZeroFill) pairs for every (layer, slot),
        // each shaped per its own slot's shape.
        let mut zero_fill_ids: Vec<NodeId> = Vec::with_capacity(n_layers * n_slots);
        {
            let mut g = graph.write().map_err(|_| {
                Error::Msg("graph lock poisoned during LatentKvCache build".into()).bt()
            })?;
            for _ in 0..n_layers {
                for shape in &shapes {
                    let alloc_id = g.push(Node {
                        op: Op::Alloc { target: target_loc },
                        inputs: vec![],
                        shape: shape.clone(),
                        dtype,
                    });
                    let zf_id = g.push(Node {
                        op: Op::ZeroFill,
                        inputs: vec![alloc_id],
                        shape: shape.clone(),
                        dtype,
                    });
                    zero_fill_ids.push(zf_id);
                }
            }
        }

        // Realize all n_layers * n_slots Op::ZeroFill targets in one pass
        // — see KvCache::with_capacity's doc for why a single realize_many
        // call is used instead of one realize per buffer.
        let realized = PipelinedExecutor::realize_many(Arc::clone(&graph), &zero_fill_ids, cache)?;
        if realized.len() != n_layers * n_slots {
            return Err(Error::Msg(format!(
                "LatentKvCache::with_capacity: realize_many returned {} storages \
                 for {} Op::ZeroFill targets — internal bug",
                realized.len(),
                n_layers * n_slots,
            ))
            .bt());
        }

        // Never-panic: length checked above; still surface a typed `Err`
        // rather than `.expect()` on a production path.
        let mut realized_iter = realized.into_iter();
        let short = |l: usize| {
            Error::Msg(format!(
                "alloc_latent: realize_many yielded fewer storages than \
                 expected (ran out at layer {l})"
            ))
            .bt()
        };
        let mut layers: Vec<Option<Vec<LatentSlot>>> = Vec::with_capacity(n_layers);
        for l in 0..n_layers {
            let mut slots = Vec::with_capacity(n_slots);
            for layout in &layouts {
                let (arc, _) = realized_iter.next().ok_or_else(|| short(l))?;
                slots.push(LatentSlot {
                    storage: arc,
                    layout: layout.clone(),
                    version: 0,
                    authority: AuthorityState::Host,
                });
            }
            layers.push(Some(slots));
        }

        Ok(Self {
            layers,
            cached_len: 0,
            slot_trailing,
            max_seq_len,
            dtype,
            alloc_id: KvAllocId::next(),
            device_location: target_loc,
        })
    }

    /// Borrow the slot `s` storage Arc for layer `layer_idx`. `None` if
    /// the layer is unpopulated or either index is out of range. Used by
    /// the forward path to bind cache storage to per-step Const nodes via
    /// [`InferenceContext::insert`]. Mirrors [`KvCache::slot_storage`].
    pub fn slot_storage(&self, layer_idx: usize, slot: usize) -> Option<Arc<RwLock<Storage>>> {
        let layer = self.layers.get(layer_idx)?.as_ref()?;
        layer.get(slot).map(|s| Arc::clone(&s.storage))
    }

    /// Bump slot `s`'s monotonic version counter for `layer_idx`. No-op
    /// (never panics) if either index is out of range. Mirrors
    /// [`KvCache::bump_version`].
    pub fn bump_version(&mut self, layer_idx: usize, slot: usize) {
        if let Some(Some(layer)) = self.layers.get_mut(layer_idx)
            && let Some(s) = layer.get_mut(slot)
        {
            s.version += 1;
        }
    }

    pub fn n_layers(&self) -> usize {
        self.layers.len()
    }

    /// Number of slots per layer.
    pub fn n_slots(&self) -> usize {
        self.slot_trailing.len()
    }

    /// Trailing shape of slot `s` (past the leading capacity axis).
    pub fn slot_trailing(&self, s: usize) -> &[usize] {
        &self.slot_trailing[s]
    }

    /// Identity of the storage this cache currently owns. Mirrors
    /// [`KvCache::alloc_id`]; see [`KvAllocId`].
    pub fn alloc_id(&self) -> KvAllocId {
        self.alloc_id
    }

    /// Drop every layer; reset `cached_len` to zero. Mirrors [`KvCache::clear`],
    /// including the [`Self::alloc_id`] re-mint — every storage `Arc` is
    /// released, so a plan welded to them must be rebuilt.
    pub fn clear(&mut self) {
        for layer in &mut self.layers {
            *layer = None;
        }
        self.cached_len = 0;
        self.alloc_id = KvAllocId::next();
    }

    /// Shrink the cache to `new_len` cached positions — metadata-only,
    /// same Phase E.3.0 limitation and rationale as [`KvCache::
    /// truncate_to`]: the underlying storages are not shrunk; the
    /// pre-allocated buffers' trailing rows simply stop being read until
    /// overwritten by a later `Op::WriteSlice` at the same positions.
    ///
    /// Deliberately preserves [`Self::alloc_id`] — a rewind is not a
    /// re-allocation. Same rule as [`KvCache::truncate_to`].
    pub fn truncate_to(&mut self, new_len: usize) {
        if new_len < self.cached_len {
            self.cached_len = new_len;
        }
    }
}

// ===========================================================================
// Guarded KV re-bind (GAP-028)
// ===========================================================================

/// 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.
///
/// Deliberately narrow: identity, residency, and the storage behind a baked
/// `kv_nodes` slot. Geometry is **not** here, because the guard does not ask
/// this trait for it — it compares the incoming storage against the storage
/// already baked into the plan, which is ground truth and cannot drift from
/// whatever a cache believes about itself.
pub trait KvRebindSource {
    /// Identity of this allocation. See [`KvAllocId`].
    fn alloc_id(&self) -> KvAllocId;
    /// Where the storage lives, or `None` when that is not provable — which
    /// [`DecodeSession::rebind_kv`] treats as a refusal, never as a default.
    fn rebind_location(&self) -> Option<DeviceLocation>;
    /// Storage for slot `which` (0 or 1) of `layer` — the two halves of the
    /// plan's `kv_nodes[layer]` pair, in the order the builder bound them.
    fn rebind_slot(&self, layer: usize, which: usize) -> Option<Arc<RwLock<Storage>>>;
}

impl KvRebindSource for KvCache {
    fn alloc_id(&self) -> KvAllocId {
        KvCache::alloc_id(self)
    }
    fn rebind_location(&self) -> Option<DeviceLocation> {
        self.device_location
    }
    fn rebind_slot(&self, layer: usize, which: usize) -> Option<Arc<RwLock<Storage>>> {
        // The builder pushes `(k_id, v_id)` — this order is the contract.
        let slot = match which {
            0 => KvSlot::K,
            1 => KvSlot::V,
            _ => return None,
        };
        self.slot_storage(layer, slot)
    }
}

impl KvRebindSource for LatentKvCache {
    fn alloc_id(&self) -> KvAllocId {
        LatentKvCache::alloc_id(self)
    }
    fn rebind_location(&self) -> Option<DeviceLocation> {
        Some(self.device_location)
    }
    fn rebind_slot(&self, layer: usize, which: usize) -> Option<Arc<RwLock<Storage>>> {
        // MLA binds `(latent, kpe)` as slots 0 and 1 — same positional
        // contract as K/V above, which is why one `DecodeSession` serves both.
        self.slot_storage(layer, which)
    }
}

/// Why a held plan could not be re-bound onto a replacement KV allocation.
///
/// A refusal is a normal outcome, not an error: the caller falls back to the
/// GAP-014 rebuild, which is correct by construction. The variants exist so
/// the guard's reasons are *nameable in a test* — the point of GAP-028 is that
/// this guard is load-bearing for correctness, so "it refused" is not enough
/// to know it refused for the right reason.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RebindRefusal {
    /// The replacement cannot say where its storage lives, so residency
    /// cannot be proven equal to the plan's. See [`KvCache::device_location`].
    ResidencyUnknown,
    /// The plan's placement stamps and residency copies were baked for a
    /// different device.
    ResidencyDiffers {
        baked: DeviceLocation,
        incoming: DeviceLocation,
    },
    /// The replacement has no storage for a slot the plan writes.
    MissingSlot { layer: usize, which: usize },
    /// A baked KV `Const` has no `base_cache` entry — an internal
    /// inconsistency in the held plan, not a property of the replacement.
    UnbakedSlot { layer: usize, which: usize },
    /// The replacement storage is not byte-interchangeable with the baked
    /// storage: different dtype, encoding, backend variant, or extent.
    StorageDiffers { layer: usize, which: usize },
    /// A storage lock was poisoned while comparing. Refuse rather than guess.
    LockPoisoned,
}

/// 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.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PlanValidity {
    /// Reuse the plan as-is.
    Valid,
    /// Everything the plan bakes still matches **except which KV allocation
    /// it is welded to.** A guarded [`DecodeSession::rebind_kv`] may rescue
    /// it; without one, the plan MUST NOT be reused — that is GAP-014.
    AllocationChanged,
    /// Something structural differs (shape, layer count, dtype, or the model
    /// itself). Rebuild.
    Stale,
}

/// Everything about a storage that a baked graph depends on and that is
/// readable from the storage itself. Deliberately excludes residency, which
/// is NOT readable — see [`DecodeSession::rebind_kv`].
#[derive(PartialEq, Eq)]
struct StorageFingerprint {
    dtype: DType,
    stype: SType,
    bundled: bool,
    backend: std::mem::Discriminant<BackendStorage>,
    len_bytes: usize,
}

fn fingerprint(s: &Arc<RwLock<Storage>>) -> Option<StorageFingerprint> {
    let g = s.read().ok()?;
    Some(StorageFingerprint {
        dtype: g.dtype,
        stype: g.stype.clone(),
        bundled: g.bundle.is_some(),
        backend: std::mem::discriminant(&g.inner),
        len_bytes: g.inner.len_bytes(),
    })
}

/// May `incoming` stand in for `baked` inside an already-optimized graph?
/// `None` means a lock was poisoned and the question could not be answered —
/// which callers must treat as "no", not as "probably fine".
///
/// The two locks are taken one at a time and released before the next, so this
/// cannot deadlock even if some future caller passes overlapping storage.
fn storage_is_interchangeable(
    baked: &Arc<RwLock<Storage>>,
    incoming: &Arc<RwLock<Storage>>,
) -> Option<bool> {
    if Arc::ptr_eq(baked, incoming) {
        return Some(true);
    }
    let a = fingerprint(baked)?;
    let b = fingerprint(incoming)?;
    Some(a == b)
}

// ===========================================================================
// DecodeSession (Phase D · D2b — plan-once persistent decode)
// ===========================================================================

/// 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.
pub struct DecodeTokenData {
    pub token_ids: Arc<RwLock<Storage>>,
    pub rope_cos: Arc<RwLock<Storage>>,
    pub rope_sin: Arc<RwLock<Storage>>,
    pub mask: Arc<RwLock<Storage>>,
    /// The device-resident KV-write offset (`cached_len` as a rank-0
    /// `I64`), present ONLY on the device-offset decode path (CUDA/CPU
    /// targets, where `Op::WriteSliceDoff` reads the start from this
    /// buffer at kernel launch — capture-ready). `None` on the SymEnv
    /// path (Vulkan), where the KV write is `Op::WriteSlice` with a
    /// `DynScalar::Sym(cached_len)` offset resolved host-side per token.
    pub offset: Option<Arc<RwLock<Storage>>>,
}

/// 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 +
/// `prepare`d (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).
pub struct DecodeSession {
    /// The held decode-step graph, ALREADY optimized in place (stamps +
    /// residency/layout copies + D2H root splice baked in). Structure is
    /// stable across tokens (D1 guarantee).
    graph: Arc<RwLock<Graph>>,
    /// The cached optimize view from the first realize. Holds only
    /// `{roots, generation}`; valid while `graph` structure + topology
    /// generation are unchanged.
    optimized: OptimizedGraph,
    /// The realize root the executor was asked for — the D2H `Op::Copy`
    /// NodeId `prepare` spliced (NOT the logits node itself). Stable
    /// across tokens.
    effective_target: NodeId,
    /// The logits node (pre-D2H-splice) — `effective_target`'s input.
    /// Retained for D3 attribution / debugging; the executor is asked
    /// for `effective_target`. Also the CapturedRun replay TARGET
    /// (`CapturedDecodeSession::capture` must be given the pre-splice
    /// node — see [`Self::logits_node`]'s doc for why).
    logits_node: NodeId,
    /// Stable re-bindable token-ids Const (`[seq]` U32).
    token_ids_node: NodeId,
    /// Stable re-bindable RoPE cos table Const (`[seq, head_dim]` F32).
    rope_cos_node: NodeId,
    /// Stable re-bindable RoPE sin table Const (`[seq, head_dim]` F32).
    rope_sin_node: NodeId,
    /// Stable re-bindable causal mask Const (`[1, 1, seq, max_seq_len]`
    /// F32) — hoisted to ONE shared Const (was per-layer in D1).
    mask_node: NodeId,
    /// Per-layer `(k_const, v_const)` stable KV placeholder NodeIds. The
    /// KV Arcs are re-bound once at build time and mutated in place by
    /// `Op::WriteSlice`/`Op::WriteSliceDoff` each token (never re-inserted
    /// per token).
    kv_nodes: Vec<(NodeId, NodeId)>,
    /// The stable rank-0 `I64` KV-write offset Const, present ONLY on the
    /// **device-offset** decode path (CUDA/CPU — `Op::WriteSliceDoff`
    /// reads `cached_len` from this buffer device-side at launch, making
    /// the decode step CUDA-graph-capturable). `None` on the SymEnv path
    /// (Vulkan — the offset is host-resolved via `cached_len_sym`). When
    /// `Some`, `realize_token` re-binds the per-token offset Arc into the
    /// base-cache clone alongside token-ids/RoPE/mask.
    offset_node: Option<NodeId>,
    /// The symbol the per-pass `SymEnv` binds to `cached_len` each token.
    cached_len_sym: SymId,
    /// The symbol the per-pass `SymEnv` binds to the live **attended
    /// prefix length** (`cached_len + seq`) each token. This is the
    /// `k_len` the optimizer-emitted CUDA flash-decode arm resolves
    /// against (`decode_flash::DecodeFlashSpec::k_len`); distinct from
    /// `cached_len_sym` because the flash kernel attends `[0, k_len)`
    /// (the whole prefix including this token), while the KV-write lands
    /// at `[cached_len, cached_len + seq)`. Unreferenced on today's f32
    /// decode graph (no flash arm offered) — a harmless extra binding.
    attended_len_sym: SymId,
    /// The full realized [`StorageCache`] from the first realize — every
    /// reachable `Op::Const` (weights + the KV Arcs + the initial data
    /// Consts) that `build_const_cache` uploaded. Held because the
    /// prebuilt-realize path SKIPS the const-cache walk: each token we
    /// clone this and overwrite ONLY the per-token data-Const entries
    /// (token-ids / RoPE / mask). The weight Arcs + KV Arcs are stable
    /// (KV mutates in place via `Op::WriteSlice`).
    base_cache: StorageCache,
    /// The `seq` the held graph is shape-keyed to (always 1 for D2 decode).
    seq: usize,
    // ---- validity keys — rebuild if any change vs. the live cache/model ----
    max_seq_len: usize,
    n_layers: usize,
    cache_dtype: DType,
    /// The [`crate::decode_shape`] key this plan was baked against — model
    /// structure + weight identity. See [`Self::is_valid_for`].
    shape_key: u64,
    /// The KV **allocation** whose storage `Arc`s are baked into `base_cache`.
    /// `rebind_and_realize_prebuilt` never re-binds `kv_nodes`, so this plan
    /// reads and writes THESE buffers regardless of which cache the caller
    /// hands in — unless [`Self::rebind_kv`] re-points them — making the id,
    /// not the geometry, the thing that decides reuse. See [`KvAllocId`] and
    /// [`Self::is_valid_for`].
    alloc_id: KvAllocId,
    /// Where the KV `Arc`s baked into `base_cache` lived when this graph was
    /// optimized. The placement stamps and residency `Op::Copy`s in `graph`
    /// were chosen against THIS location, so [`Self::rebind_kv`] must refuse a
    /// replacement that lives anywhere else. `None` when the source could not
    /// prove it, which is likewise a refusal.
    kv_location: Option<DeviceLocation>,
    /// `order_for` fix track (`docs/design/incremental-consumer-index.md`
    /// §5): the memoized dispatch order for this session's
    /// `(graph, effective_target, optimized)` triple, computed once on the
    /// first `realize_token` call and reused on every later one — measured
    /// as the single largest named per-token sub-cost prior to this cache
    /// (`extract_runs_multi` + `non_chosen_arm_nodes`, ~56ms/step on an RTX
    /// 4070). Empty at construction; `realize_token` never resets it, since
    /// a `DecodeSession` holds one shape-keyed graph for its whole lifetime
    /// (see `is_valid_for` — any mismatch rebuilds a fresh session with a
    /// fresh, empty cache, never reuses this one against a different graph).
    order_cache: std::sync::OnceLock<Vec<NodeId>>,
}

impl DecodeSession {
    /// 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.
    #[allow(clippy::too_many_arguments)]
    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 {
        Self {
            shape_key,
            // Identity and residency are read from the SAME source in the same
            // breath, so they cannot describe different allocations — which a
            // pair of separate `alloc_id` / `location` parameters could.
            alloc_id: kv.alloc_id(),
            kv_location: kv.rebind_location(),
            graph,
            optimized,
            effective_target,
            logits_node,
            token_ids_node,
            rope_cos_node,
            rope_sin_node,
            mask_node,
            kv_nodes,
            offset_node,
            cached_len_sym,
            attended_len_sym,
            base_cache,
            seq,
            max_seq_len,
            n_layers,
            cache_dtype,
            order_cache: std::sync::OnceLock::new(),
        }
    }

    /// 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 realize_token(
        &self,
        device: &Device,
        data: DecodeTokenData,
        sym_env: &SymEnv,
    ) -> Result<Vec<f32>> {
        let mut cache = self.base_cache.clone();
        cache.insert(self.token_ids_node, data.token_ids);
        cache.insert(self.rope_cos_node, data.rope_cos);
        cache.insert(self.rope_sin_node, data.rope_sin);
        cache.insert(self.mask_node, data.mask);
        // Device-offset path (CUDA/CPU): re-bind the per-token KV-write
        // offset (`cached_len` as I64) into the base-cache clone so the
        // held `Op::WriteSliceDoff` nodes read the live position from a
        // fixed device buffer. On the SymEnv path both are `None` (the
        // offset rides `cached_len_sym` in `sym_env` instead).
        if let (Some(offset_node), Some(offset)) = (self.offset_node, data.offset) {
            cache.insert(offset_node, offset);
        }
        crate::pipelined_bridge::realize_one_prebuilt_cached_env::<f32>(
            &self.graph,
            self.effective_target,
            &self.optimized,
            &self.order_cache,
            device,
            cache,
            sym_env,
        )
    }

    /// 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 `Const`s, 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 is_valid_for(
        &self,
        seq: usize,
        max_seq_len: usize,
        n_layers: usize,
        cache_dtype: DType,
        shape_key: u64,
        alloc_id: KvAllocId,
    ) -> bool {
        matches!(
            self.validity_for(seq, max_seq_len, n_layers, cache_dtype, shape_key, alloc_id),
            PlanValidity::Valid,
        )
    }

    /// 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 validity_for(
        &self,
        seq: usize,
        max_seq_len: usize,
        n_layers: usize,
        cache_dtype: DType,
        shape_key: u64,
        alloc_id: KvAllocId,
    ) -> PlanValidity {
        let structural = self.seq == seq
            && self.max_seq_len == max_seq_len
            && self.n_layers == n_layers
            && self.cache_dtype == cache_dtype
            && self.shape_key == shape_key;
        if !structural {
            return PlanValidity::Stale;
        }
        if self.alloc_id == alloc_id {
            PlanValidity::Valid
        } else {
            PlanValidity::AllocationChanged
        }
    }

    /// Re-point the held plan's baked KV `Const`s 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::Copy`s 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 rebind_kv(
        &mut self,
        kv: &dyn KvRebindSource,
    ) -> core::result::Result<(), RebindRefusal> {
        let incoming_loc = kv
            .rebind_location()
            .ok_or(RebindRefusal::ResidencyUnknown)?;
        let baked_loc = self.kv_location.ok_or(RebindRefusal::ResidencyUnknown)?;
        if baked_loc != incoming_loc {
            return Err(RebindRefusal::ResidencyDiffers {
                baked: baked_loc,
                incoming: incoming_loc,
            });
        }

        // Collect first, mutate second. A partial re-bind — some slots on the
        // new allocation, the rest still on the old — is the one state strictly
        // worse than either endpoint, and the only way to avoid it is to refuse
        // before touching `base_cache`.
        let mut staged: Vec<(NodeId, Arc<RwLock<Storage>>)> =
            Vec::with_capacity(self.kv_nodes.len() * 2);
        for (layer, &(k_node, v_node)) in self.kv_nodes.iter().enumerate() {
            for (which, node) in [(0usize, k_node), (1usize, v_node)] {
                let incoming = kv
                    .rebind_slot(layer, which)
                    .ok_or(RebindRefusal::MissingSlot { layer, which })?;
                let baked = self
                    .base_cache
                    .get(&node)
                    .ok_or(RebindRefusal::UnbakedSlot { layer, which })?;
                if !storage_is_interchangeable(baked, &incoming)
                    .ok_or(RebindRefusal::LockPoisoned)?
                {
                    return Err(RebindRefusal::StorageDiffers { layer, which });
                }
                staged.push((node, incoming));
            }
        }

        for (node, arc) in staged {
            self.base_cache.insert(node, arc);
        }
        self.alloc_id = kv.alloc_id();
        Ok(())
    }

    /// The held graph handle (the caller re-binds data Consts + realizes
    /// through it).
    pub fn graph(&self) -> &Arc<RwLock<Graph>> {
        &self.graph
    }

    /// The cached optimize view (fed to the prebuilt-realize seam).
    pub fn optimized(&self) -> &OptimizedGraph {
        &self.optimized
    }

    /// The D2H `Op::Copy` root the executor is asked for.
    pub fn effective_target(&self) -> NodeId {
        self.effective_target
    }

    /// 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 logits_node(&self) -> NodeId {
        self.logits_node
    }
    pub fn token_ids_node(&self) -> NodeId {
        self.token_ids_node
    }
    pub fn rope_cos_node(&self) -> NodeId {
        self.rope_cos_node
    }
    pub fn rope_sin_node(&self) -> NodeId {
        self.rope_sin_node
    }
    pub fn mask_node(&self) -> NodeId {
        self.mask_node
    }
    pub fn kv_nodes(&self) -> &[(NodeId, NodeId)] {
        &self.kv_nodes
    }
    /// 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 offset_node(&self) -> Option<NodeId> {
        self.offset_node
    }
    pub fn cached_len_sym(&self) -> SymId {
        self.cached_len_sym
    }
    pub fn attended_len_sym(&self) -> SymId {
        self.attended_len_sym
    }
    pub fn max_seq_len(&self) -> usize {
        self.max_seq_len
    }

    /// 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):**
    ///
    /// ```text
    /// 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 per_token_sym_env(&self, cached_len: usize) -> Result<SymEnv> {
        let mut env = SymEnv::new();
        env.bind(self.cached_len_sym, cached_len)?;
        env.bind(self.attended_len_sym, cached_len + self.seq)?;
        Ok(env)
    }

    /// 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 graph_node_count(&self) -> usize {
        self.graph.read().map(|g| g.len()).unwrap_or(0)
    }

    /// 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`.
    pub fn base_cache(&self) -> &StorageCache {
        &self.base_cache
    }
}

/// 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`](Self::PlanOnce)** (set by
/// `PagedSessionScheduler::new`), with [`Replan`](Self::Replan) available as an
/// explicit opt-out.
///
/// The A/B rollout this flag was built for is DONE, and the default moved with
/// the result. Measured (nsys, Lightbulb, 2026-08-01; 16 tokens, same binary and
/// model, only the plan mode differing): `Replan` → `PlanOnce` is **29.7×** on
/// ms/token (8,192.0 → 275.8) and cuts host-ward DtoH 61% (101,691 → 40,088 MB)
/// at IDENTICAL kernel-launch count (17,802 in both arms — same compute, less
/// data movement). Shipping `Replan` as the default meant every consumer who did
/// not know to opt in paid that 29.7×.
///
/// It remains a real runtime flag — not a test hook — because toggling it in ONE
/// process is what lets the correctness gate run both arms off a single setup,
/// and because `Replan` is the parity reference the plan-once path is checked
/// against. On [`Replan`](Self::Replan) the persistent forward drops any held
/// session (so nothing stale lingers) and delegates to the re-planning
/// `LlamaModel::forward_paged_step`;
/// on [`PlanOnce`](Self::PlanOnce) it
/// builds-once / rebinds.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PagedDecodePlan {
    /// Re-plan (build + optimize the graph) every token — the pre-plan-once
    /// path. Drops any held [`PagedDecodeSession`] and routes to
    /// `LlamaModel::forward_paged_step`.
    Replan,
    /// Build + optimize the graph ONCE (into a [`PagedDecodeSession`]) and reuse
    /// it (rebinding only the per-token data) for every subsequent token.
    PlanOnce,
}

/// 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.
pub struct PagedDecodeTokenData {
    /// The single decode token id (`[1]` U32).
    pub token_ids: Arc<RwLock<Storage>>,
    /// RoPE cos table at THIS token's absolute position (`[1, head_dim]` F32).
    pub rope_cos: Arc<RwLock<Storage>>,
    /// RoPE sin table at THIS token's absolute position (`[1, head_dim]` F32).
    pub rope_sin: Arc<RwLock<Storage>>,
    /// Capacity-padded block table (`[1, max_blocks_cap]` U32) — the shape is
    /// FIXED by `max_blocks_cap`; the content grows a physical block whenever
    /// this token crosses a block boundary (padded tail stays 0, masked out).
    pub block_table: Arc<RwLock<Storage>>,
    /// Attended prefix length for this row (`[1]` U32) — the `Op::PagedAttn`
    /// mask cutoff (`Ge(pos, context_len) → −inf`) that neutralizes both the
    /// future positions AND the padded block_table tail.
    pub context_lens: Arc<RwLock<Storage>>,
    /// The flattened KV-write offset (`linear = phys·block_size + slot` as a
    /// rank-0 I64), present ONLY on the device-offset path (CPU/CUDA, where
    /// `Op::WriteSliceDoff` reads the start device-side). `None` on the SymEnv
    /// path (Vulkan), where the offset rides `write_sym` in the per-pass env.
    pub offset: Option<Arc<RwLock<Storage>>>,
}

/// 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.
pub struct PagedDecodeSession {
    /// The held paged decode-step graph, ALREADY optimized in place. Structure
    /// is stable across tokens (variabilities #1–#3 removed).
    graph: Arc<RwLock<Graph>>,
    /// The cached optimize view from the first realize. Holds only
    /// `{roots, generation}`; valid while `graph` structure + topology
    /// generation are unchanged.
    optimized: OptimizedGraph,
    /// The realize root the executor is asked for (the D2H `Op::Copy` splice).
    effective_target: NodeId,
    /// The logits node (pre-D2H-splice) — retained for debugging/attribution.
    logits_node: NodeId,
    /// Stable re-bindable single-token-id Const (`[1]` U32).
    token_ids_node: NodeId,
    /// Stable re-bindable RoPE cos table Const (`[1, head_dim]` F32).
    rope_cos_node: NodeId,
    /// Stable re-bindable RoPE sin table Const (`[1, head_dim]` F32).
    rope_sin_node: NodeId,
    /// Stable re-bindable padded block_table Const (`[1, max_blocks_cap]` U32).
    block_table_node: NodeId,
    /// Stable re-bindable context_lens Const (`[1]` U32).
    context_lens_node: NodeId,
    /// Per-layer `(k_pool_ph, v_pool_ph)` stable KV placeholder NodeIds, bound
    /// ONCE to the pool buffers (viewed at `pool_shape_flat`) and mutated in
    /// place by the flattened-offset write each token.
    kv_nodes: Vec<(NodeId, NodeId)>,
    /// The stable rank-0 `I64` KV-write offset Const, present ONLY on the
    /// device-offset path (CPU/CUDA — `Op::WriteSliceDoff`). `None` on the
    /// SymEnv path (Vulkan — the offset rides `write_sym`).
    offset_node: Option<NodeId>,
    /// The symbol the per-pass `SymEnv` binds to the flattened write offset
    /// (`linear`) each token. Load-bearing on the SymEnv path; a harmless
    /// (unreferenced) extra binding on the device-offset path.
    write_sym: SymId,
    /// The full realized [`StorageCache`] from the first realize (weights +
    /// the pool KV Arcs + the initial per-token data Consts). Each token we
    /// clone this and overwrite ONLY the per-token data-Const entries; the
    /// weight + pool Arcs are stable (the pool mutates in place).
    base_cache: StorageCache,
    // ---- validity keys — rebuild if any change vs. the live pool ----
    max_blocks_cap: usize,
    n_layers: usize,
    block_size: usize,
    cache_dtype: DType,
    /// The model's `decode_shape_key()` at build time — family + structure-
    /// affecting config + WEIGHT IDENTITY (a never-recycled `ModelInstanceId`).
    ///
    /// Geometry alone cannot separate two models: `PagedDecodeSession` bakes the
    /// weight `Const`s, so two same-shaped models produce identical
    /// `(max_blocks_cap, n_layers, block_size, cache_dtype)` and a plan built for
    /// one would be judged valid for the other — silently computing the wrong
    /// architecture, or the right architecture with the wrong weights, at full
    /// speed. The contiguous twin closed this first; this is the paged half of
    /// the same hole.
    shape_key: u64,
    /// The [`DeviceKvPool`](crate::kv_block_pool_device::DeviceKvPool) whose
    /// per-layer block buffers are baked into `base_cache`. The paged half of
    /// the same welding: `rebind_and_realize_paged_prebuilt` re-binds the
    /// per-token data (token id, RoPE, block_table, context_lens, offset) and
    /// never the pool, so this plan reads THESE buffers whichever pool the
    /// caller passes. Two same-geometry pools are otherwise indistinguishable
    /// to the key. See [`KvAllocId`].
    alloc_id: KvAllocId,
    /// Count of REBIND realizes this session has served (one per
    /// [`Self::realize_token`]). The build itself realizes through the prebuild
    /// seam, NOT `realize_token`, so it stays 0 right after construction and
    /// increments only on each plan-reuse token. A DEVICE-INDEPENDENT plan-HIT
    /// signal (the `optimize_calls_thread_local` counter is CPU-scoped — GPU
    /// per-token uploads bump it), so the correctness gate can assert "this
    /// token reused the plan" on any backend. Interior-mutable (`realize_token`
    /// takes `&self`); `Relaxed` — it is telemetry, ordered by the caller.
    realize_count: AtomicUsize,
    /// `order_for` fix track (`docs/design/incremental-consumer-index.md`
    /// §5) — see [`DecodeSession::order_cache`]'s doc for the full
    /// argument; identical role here for the paged path.
    order_cache: std::sync::OnceLock<Vec<NodeId>>,
}

impl PagedDecodeSession {
    /// 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.
    #[allow(clippy::too_many_arguments)]
    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 {
        Self {
            alloc_id,
            graph,
            optimized,
            effective_target,
            logits_node,
            token_ids_node,
            rope_cos_node,
            rope_sin_node,
            block_table_node,
            context_lens_node,
            kv_nodes,
            offset_node,
            write_sym,
            base_cache,
            max_blocks_cap,
            n_layers,
            block_size,
            cache_dtype,
            shape_key,
            realize_count: AtomicUsize::new(0),
            order_cache: std::sync::OnceLock::new(),
        }
    }

    /// 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 realize_token(
        &self,
        device: &Device,
        data: PagedDecodeTokenData,
        sym_env: &SymEnv,
    ) -> Result<Vec<f32>> {
        let mut cache = self.base_cache.clone();
        cache.insert(self.token_ids_node, data.token_ids);
        cache.insert(self.rope_cos_node, data.rope_cos);
        cache.insert(self.rope_sin_node, data.rope_sin);
        cache.insert(self.block_table_node, data.block_table);
        cache.insert(self.context_lens_node, data.context_lens);
        if let (Some(offset_node), Some(offset)) = (self.offset_node, data.offset) {
            cache.insert(offset_node, offset);
        }
        // Device-independent plan-HIT tick: this token reuses the held plan
        // (the prebuild seam below SKIPs optimize). Bumped BEFORE the realize
        // so a realize error still records the attempted rebind.
        self.realize_count.fetch_add(1, Ordering::Relaxed);
        crate::pipelined_bridge::realize_one_prebuilt_cached_env::<f32>(
            &self.graph,
            self.effective_target,
            &self.optimized,
            &self.order_cache,
            device,
            cache,
            sym_env,
        )
    }

    /// 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`](crate::kv_block_pool_device::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 is_valid_for(
        &self,
        max_blocks_cap: usize,
        n_layers: usize,
        block_size: usize,
        cache_dtype: DType,
        shape_key: u64,
        alloc_id: KvAllocId,
    ) -> bool {
        self.max_blocks_cap == max_blocks_cap
            && self.n_layers == n_layers
            && self.block_size == block_size
            && self.cache_dtype == cache_dtype
            && self.shape_key == shape_key
            && self.alloc_id == alloc_id
    }

    /// 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 per_token_sym_env(&self, linear: usize) -> Result<SymEnv> {
        let mut env = SymEnv::new();
        env.bind(self.write_sym, linear)?;
        Ok(env)
    }

    /// The held graph handle.
    pub fn graph(&self) -> &Arc<RwLock<Graph>> {
        &self.graph
    }
    /// The cached optimize view (fed to the prebuilt-realize seam).
    pub fn optimized(&self) -> &OptimizedGraph {
        &self.optimized
    }
    pub fn effective_target(&self) -> NodeId {
        self.effective_target
    }
    pub fn logits_node(&self) -> NodeId {
        self.logits_node
    }
    pub fn token_ids_node(&self) -> NodeId {
        self.token_ids_node
    }
    pub fn rope_cos_node(&self) -> NodeId {
        self.rope_cos_node
    }
    pub fn rope_sin_node(&self) -> NodeId {
        self.rope_sin_node
    }
    pub fn block_table_node(&self) -> NodeId {
        self.block_table_node
    }
    pub fn context_lens_node(&self) -> NodeId {
        self.context_lens_node
    }
    pub fn kv_nodes(&self) -> &[(NodeId, NodeId)] {
        &self.kv_nodes
    }
    pub fn offset_node(&self) -> Option<NodeId> {
        self.offset_node
    }
    pub fn write_sym(&self) -> SymId {
        self.write_sym
    }
    pub fn max_blocks_cap(&self) -> usize {
        self.max_blocks_cap
    }

    /// 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 realize_count(&self) -> usize {
        self.realize_count.load(Ordering::Relaxed)
    }

    /// 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 graph_node_count(&self) -> usize {
        self.graph.read().map(|g| g.len()).unwrap_or(0)
    }

    /// The full realized [`StorageCache`] from the first realize.
    pub fn base_cache(&self) -> &StorageCache {
        &self.base_cache
    }
}

// `alloc_zeroed_on` retired 2026-05-22 (bridge-retirement Phase 3a).
// Zero-init device allocation now flows through `Op::Alloc` graph
// emission + the executor's `WorkItemKind::Alloc` arm. The per-
// `DeviceLocation` match this function used to carry lives in
// `fuel-storage::pipelined::execute_work_item`'s Alloc arm; the
// residual "0-byte device anchor" helper lives in
// `crate::pipelined_bridge::device_seed_storage`.

// ===========================================================================
// InferenceContext
// ===========================================================================

/// 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.
pub struct InferenceContext {
    device: Device,
    persistent: HashMap<NodeId, Arc<RwLock<Storage>>>,
    /// Held plan-once decode plan for `LlamaModel::forward_decode_step`
    /// — the carrier that lets a caller-driven decode loop get plan reuse
    /// WITHOUT threading an `Option<DecodeSession>` of its own.
    ///
    /// This exists because the ergonomic call shape and the fast call shape had
    /// drifted apart. `forward_with_kv_context(tokens, cache, ctx)` is the raw
    /// rebuild primitive; its persistent sibling takes a fourth argument, so a
    /// consumer hand-rolling a decode loop reaches the slow one and never learns
    /// the fast one exists. Not hypothetical: that gap cost a measured consumer
    /// 5,901 → 26.47 ms/token before anyone noticed (nsys, 2026-08-01).
    ///
    /// **One context per model.** Already this codebase's invariant — speculative
    /// decoding builds a separate `InferenceContext` for the draft and the target
    /// — and the held plan self-heals regardless: `DecodeSession::is_valid_for`
    /// drops it on any cache-geometry or dtype change.
    decode_session: Option<DecodeSession>,
}

impl InferenceContext {
    pub fn new(device: Device) -> Self {
        Self {
            device,
            persistent: HashMap::new(),
            decode_session: None,
        }
    }

    pub fn device(&self) -> &Device {
        &self.device
    }

    /// 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 take_decode_session(&mut self) -> Option<DecodeSession> {
        self.decode_session.take()
    }

    /// Put a decode plan back into the context. See [`Self::take_decode_session`].
    pub fn put_decode_session(&mut self, session: Option<DecodeSession>) {
        self.decode_session = session;
    }

    /// 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 has_decode_session(&self) -> bool {
        self.decode_session.is_some()
    }

    /// 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 insert(&mut self, node_id: NodeId, storage: Arc<RwLock<Storage>>) {
        self.persistent.insert(node_id, 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 remove(&mut self, node_id: NodeId) -> Option<Arc<RwLock<Storage>>> {
        self.persistent.remove(&node_id)
    }

    /// Whether the persistent map has an entry at `node_id`.
    pub fn contains(&self, node_id: NodeId) -> bool {
        self.persistent.contains_key(&node_id)
    }

    /// Borrow the persistent entry at `node_id` if present.
    pub fn get(&self, node_id: NodeId) -> Option<&Arc<RwLock<Storage>>> {
        self.persistent.get(&node_id)
    }

    /// Number of persistent slots.
    pub fn len(&self) -> usize {
        self.persistent.len()
    }

    pub fn is_empty(&self) -> bool {
        self.persistent.is_empty()
    }

    /// 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 clear_persistent(&mut self) {
        self.persistent.clear();
    }

    /// 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<T: bytemuck::Pod + fuel_ir::WithDType>(
        &self,
        graph: &Arc<RwLock<Graph>>,
        target: NodeId,
    ) -> Result<Vec<T>> {
        self.realize_one_as_with_env::<T>(graph, target, &SymEnv::default())
    }

    /// 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 realize_one_as_with_env<T: bytemuck::Pod + fuel_ir::WithDType>(
        &self,
        graph: &Arc<RwLock<Graph>>,
        target: NodeId,
        env: &SymEnv,
    ) -> Result<Vec<T>> {
        crate::pipelined_bridge::realize_one_as_with_initial_env::<T>(
            graph,
            target,
            &self.device,
            self.cloned_persistent(),
            env,
        )
    }

    /// 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_as_with_env<T: bytemuck::Pod + fuel_ir::WithDType>(
        &self,
        graph: &Arc<RwLock<Graph>>,
        target: NodeId,
        env: &SymEnv,
    ) -> Result<(NodeId, fuel_dispatch::optimize::OptimizedGraph, Vec<T>)> {
        crate::pipelined_bridge::prebuild_optimized_env::<T>(
            graph,
            target,
            &self.device,
            self.cloned_persistent(),
            env,
        )
    }

    /// 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 prebuild_optimized_capturing_as_with_env<T: bytemuck::Pod + fuel_ir::WithDType>(
        &self,
        graph: &Arc<RwLock<Graph>>,
        target: NodeId,
        env: &SymEnv,
    ) -> Result<(
        NodeId,
        fuel_dispatch::optimize::OptimizedGraph,
        StorageCache,
        Vec<T>,
    )> {
        crate::pipelined_bridge::prebuild_optimized_env_capturing_cache::<T>(
            graph,
            target,
            &self.device,
            self.cloned_persistent(),
            env,
        )
    }

    /// 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_prebuilt_as_with_env<T: bytemuck::Pod + fuel_ir::WithDType>(
        &self,
        graph: &Arc<RwLock<Graph>>,
        effective_target: NodeId,
        optimized: &fuel_dispatch::optimize::OptimizedGraph,
        env: &SymEnv,
    ) -> Result<Vec<T>> {
        crate::pipelined_bridge::realize_one_prebuilt_env::<T>(
            graph,
            effective_target,
            optimized,
            &self.device,
            self.cloned_persistent(),
            env,
        )
    }

    /// Multi-target counterpart of [`Self::realize_one_as`].
    pub fn realize_many_as<T: bytemuck::Pod + fuel_ir::WithDType>(
        &self,
        graph: &Arc<RwLock<Graph>>,
        targets: &[NodeId],
    ) -> Result<Vec<Vec<T>>> {
        self.realize_many_as_with_env::<T>(graph, targets, &SymEnv::default())
    }

    /// 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_many_as_with_env<T: bytemuck::Pod + fuel_ir::WithDType>(
        &self,
        graph: &Arc<RwLock<Graph>>,
        targets: &[NodeId],
        env: &SymEnv,
    ) -> Result<Vec<Vec<T>>> {
        crate::pipelined_bridge::realize_many_as_with_initial_env::<T>(
            graph,
            targets,
            &self.device,
            self.cloned_persistent(),
            env,
        )
    }

    /// 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`].
    pub fn realize_split_as<T: bytemuck::Pod + fuel_ir::WithDType>(
        &self,
        graph: &Arc<RwLock<Graph>>,
        targets: &[NodeId],
        n_host: usize,
    ) -> Result<crate::pipelined_bridge::SplitRealizeOutput<T>> {
        crate::pipelined_bridge::realize_split_as_with_initial::<T>(
            graph,
            targets,
            n_host,
            &self.device,
            self.cloned_persistent(),
        )
    }

    /// Build a [`StorageCache`] containing Arc-clones of every
    /// persistent entry. The clone is cheap (Arc refcount bumps); the
    /// returned `StorageCache` is consumed by the realize call but
    /// the original Arcs in `self.persistent` survive.
    fn cloned_persistent(&self) -> StorageCache {
        let mut out = StorageCache::with_capacity(self.persistent.len());
        for (id, arc) in &self.persistent {
            out.insert(*id, Arc::clone(arc));
        }
        out
    }
}

// ===========================================================================
// Tests
// ===========================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use fuel_graph::{Node, Op};
    use fuel_ir::{DType, Shape};

    /// `InferenceContext::insert` + immediate retrieval round-trips.
    #[test]
    fn context_insert_retrieve() {
        let mut ctx = InferenceContext::new(Device::cpu());
        let storage = fuel_memory::from_slice_cpu(&[1.0_f32, 2.0, 3.0]);
        let arc = Arc::new(RwLock::new(storage));
        let id = NodeId(42);

        assert!(!ctx.contains(id));
        assert_eq!(ctx.len(), 0);

        ctx.insert(id, Arc::clone(&arc));

        assert!(ctx.contains(id));
        assert_eq!(ctx.len(), 1);
        // The returned Arc is the same as what we inserted.
        let retrieved = ctx.get(id).expect("just inserted");
        assert!(Arc::ptr_eq(retrieved, &arc));
    }

    /// Inserted Arcs survive across realize calls. Build a minimal
    /// 3-node graph (Const + Const + Add). Both Const slots come
    /// from the context's persistent map (no graph.storage_map
    /// involvement). Realize twice; verify the persistent Arcs are
    /// still held by the context after each call.
    #[test]
    fn context_persistent_arc_survives_realize() {
        let lhs_arc = Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
            1.0_f32, 2.0, 3.0,
        ])));
        let rhs_arc = Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
            10.0_f32, 20.0, 30.0,
        ])));

        let graph = Arc::new(RwLock::new(Graph::new()));
        let (lhs_id, rhs_id, add_id) = {
            let mut g = graph.write().unwrap();
            let lhs = g.push(Node {
                op: Op::Const,
                inputs: vec![],
                shape: Shape::from_dims(&[3]),
                dtype: DType::F32,
            });
            let rhs = g.push(Node {
                op: Op::Const,
                inputs: vec![],
                shape: Shape::from_dims(&[3]),
                dtype: DType::F32,
            });
            let add = g.push(Node {
                op: Op::Add,
                inputs: vec![lhs, rhs],
                shape: Shape::from_dims(&[3]),
                dtype: DType::F32,
            });
            (lhs, rhs, add)
        };

        let mut ctx = InferenceContext::new(Device::cpu());
        ctx.insert(lhs_id, Arc::clone(&lhs_arc));
        ctx.insert(rhs_id, Arc::clone(&rhs_arc));

        // Before realize: each Arc held by the test (1) + context (1) = 2.
        assert_eq!(Arc::strong_count(&lhs_arc), 2);
        assert_eq!(Arc::strong_count(&rhs_arc), 2);

        let out1 = ctx
            .realize_one_as::<f32>(&graph, add_id)
            .expect("first realize");
        assert_eq!(out1, vec![11.0, 22.0, 33.0]);

        // After first realize: same count. The realize call's
        // cloned_persistent() bumped the counts to 3 during the call
        // and dropped them back to 2 when the call returned.
        assert_eq!(Arc::strong_count(&lhs_arc), 2);
        assert_eq!(Arc::strong_count(&rhs_arc), 2);

        // Realize again — same Arcs get reused without re-upload.
        // This is the autoregressive-decoding pattern: weights stay
        // resident across forward passes.
        let out2 = ctx
            .realize_one_as::<f32>(&graph, add_id)
            .expect("second realize");
        assert_eq!(out2, vec![11.0, 22.0, 33.0]);
        assert_eq!(Arc::strong_count(&lhs_arc), 2);
        assert_eq!(Arc::strong_count(&rhs_arc), 2);
    }

    /// Read back layer `layer_idx`'s `slot` buffer in full, as a flat
    /// row-major `[n_kv_heads * max_seq_len * head_dim]` f32 vec — test-only
    /// readback counterpart to [`KvCache::write_bytes`], deliberately NOT
    /// sliced, so a wrong write RANGE shows up as corruption anywhere in the
    /// buffer rather than being hidden by only reading the written window.
    fn read_kv_layer_f32(
        cache: &KvCache,
        layer_idx: usize,
        slot: KvSlot,
        device: &Device,
    ) -> Vec<f32> {
        let max_seq_len = cache.max_seq_len.expect("with_capacity cache");
        let full_shape = Shape::from_dims(&[1, cache.n_kv_heads, max_seq_len, cache.head_dim]);
        let buf = match slot {
            KvSlot::K => Arc::clone(&cache.layer(layer_idx).unwrap().k),
            KvSlot::V => Arc::clone(&cache.layer(layer_idx).unwrap().v),
        };
        let anchor =
            crate::lazy::Tensor::from_f32(vec![0.0f32], Shape::from_dims(&[1]), device).unwrap();
        let dest = anchor.const_placeholder_like(full_shape, DType::F32);
        let mut sc = StorageCache::new();
        sc.insert(dest.node_id(), buf);
        let bytes = crate::pipelined_bridge::realize_one_bytes_with_initial(
            dest.graph_handle(),
            dest.node_id(),
            device,
            sc,
        )
        .expect("readback realize");
        bytemuck::try_cast_slice::<u8, f32>(&bytes)
            .expect("f32 cast")
            .to_vec()
    }

    /// Board item 86 (OverMind self-fork / Lightbulb KV-cache handoff design):
    /// [`KvCache::write_bytes`] must write EXACTLY the `[start, start+n)`
    /// position range on the right axis (2, head-major) and leave every other
    /// position untouched — a wrong range (e.g. accidentally position-major,
    /// axis 0) would either corrupt unrelated positions or silently land on
    /// the wrong head/position entirely, and reading back the FULL buffer
    /// (not just the written window) is what catches either failure.
    #[test]
    fn write_bytes_lands_at_the_correct_position_and_nowhere_else() {
        let device = Device::cpu();
        let (n_layers, n_kv_heads, head_dim, max_seq_len) = (1, 2, 3, 5);
        let cache = KvCache::with_capacity(
            n_layers,
            n_kv_heads,
            head_dim,
            max_seq_len,
            DType::F32,
            &device,
        )
        .expect("with_capacity");

        let alloc_before = cache.alloc_id();
        let loc_before = cache.device_location();

        // Write positions [1, 3) (n=2) with a distinctive, non-zero pattern —
        // zero would be indistinguishable from the buffer's own zero-init.
        let (start, n) = (1usize, 2usize);
        let payload: Vec<f32> = (0..(n * n_kv_heads * head_dim))
            .map(|i| 1000.0 + i as f32)
            .collect();
        cache
            .write_bytes(
                0,
                KvSlot::K,
                start,
                n,
                bytemuck::cast_slice(&payload),
                &device,
            )
            .expect("write_bytes");

        // THE CORE CLAIM: no Arc replacement happened, so neither identity
        // marker moved — this is what makes a held DecodeSession safe to
        // keep using without a rebind or rebuild.
        assert_eq!(
            cache.alloc_id(),
            alloc_before,
            "write_bytes must not re-mint alloc_id (that would mean the Arc was replaced)",
        );
        assert_eq!(
            cache.device_location(),
            loc_before,
            "write_bytes must not clear device_location (that would mean residency became unprovable)",
        );

        let full = read_kv_layer_f32(&cache, 0, KvSlot::K, &device);
        for h in 0..n_kv_heads {
            for p in 0..max_seq_len {
                for d in 0..head_dim {
                    let got = full[h * max_seq_len * head_dim + p * head_dim + d];
                    if (start..start + n).contains(&p) {
                        let local_p = p - start;
                        let want = 1000.0 + (h * n * head_dim + local_p * head_dim + d) as f32;
                        assert_eq!(
                            got, want,
                            "head {h} position {p} dim {d}: written value must land exactly here",
                        );
                    } else {
                        assert_eq!(
                            got,
                            0.0,
                            "head {h} position {p} dim {d}: outside [{start},{}) must stay \
                             untouched (zero-init) — a wrong write range would corrupt this",
                            start + n,
                        );
                    }
                }
            }
        }
    }

    /// `write_bytes` is `with_capacity`-only (needs a fixed dtype + a stable
    /// buffer to write into) — a `with_dims` (grow-by-replacement) cache has
    /// neither, and must refuse with a typed error, never panic.
    #[test]
    fn write_bytes_refuses_a_with_dims_cache() {
        let device = Device::cpu();
        let cache = KvCache::with_dims(1, 2, 3);
        let err = cache
            .write_bytes(0, KvSlot::K, 0, 1, &[0u8; 24], &device)
            .expect_err("with_dims cache has no fixed dtype/capacity");
        assert!(
            err.to_string().contains("with_capacity"),
            "refusal must name the real constraint, got: {err}",
        );
    }

    /// A range exceeding `max_seq_len` must refuse before touching anything —
    /// never silently clamp (which would write at the wrong position) and
    /// never panic.
    #[test]
    fn write_bytes_refuses_an_out_of_range_position() {
        let device = Device::cpu();
        let cache = KvCache::with_capacity(1, 2, 3, 5, DType::F32, &device).unwrap();
        let payload = vec![0.0f32; 2 * 2 * 3]; // n=2 positions, valid byte length
        let err = cache
            .write_bytes(
                0,
                KvSlot::K,
                4, // start=4, n=2 -> [4,6) exceeds max_seq_len=5
                2,
                bytemuck::cast_slice(&payload),
                &device,
            )
            .expect_err("range must be refused, not clamped");
        assert!(
            err.to_string().contains("max_seq_len"),
            "refusal must name the real constraint, got: {err}",
        );
    }

    /// A `bytes` length that doesn't match `n * n_kv_heads * head_dim *
    /// dtype_size` must refuse — a mismatched length is the shape of exactly
    /// the bug this check exists to catch (a caller passing position-major
    /// pool bytes straight through without transposing to head-major).
    #[test]
    fn write_bytes_refuses_a_mismatched_byte_length() {
        let device = Device::cpu();
        let cache = KvCache::with_capacity(1, 2, 3, 5, DType::F32, &device).unwrap();
        let wrong_payload = vec![0.0f32; 5]; // not n_kv_heads*n*head_dim for any valid n
        let err = cache
            .write_bytes(
                0,
                KvSlot::K,
                0,
                1,
                bytemuck::cast_slice(&wrong_payload),
                &device,
            )
            .expect_err("mismatched byte length must be refused");
        assert!(
            err.to_string().contains("bytes"),
            "refusal must name the byte-length mismatch, got: {err}",
        );
    }

    /// Phase D symbolic extents — the decode-shaped path: a persistent
    /// fixed-capacity buffer + an `Op::WriteSlice` whose start offset is
    /// a per-pass `SymEnv` binding (the KV-cache append at `cached_len`).
    /// `realize_one_as_with_env` threads the env from the session through
    /// the bridge to the executor; the width-2 slab lands at the bound
    /// offset (3), not the static placeholder (0). Re-realizing the SAME
    /// graph with a different binding lands the slab at the new offset —
    /// the input-independent-graph property persistent decode relies on.
    #[test]
    fn realize_one_as_with_env_resolves_write_slice_offset() {
        use fuel_ir::{DynScalar, SymId};
        let dest_arc = Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[0.0_f32; 6])));
        let src_arc = Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[7.0_f32, 8.0])));
        let graph = Arc::new(RwLock::new(Graph::new()));
        let sym = SymId(0);
        let (dest_id, src_id, ws_id) = {
            let mut g = graph.write().unwrap();
            let dest = g.push(Node {
                op: Op::Const,
                inputs: vec![],
                shape: Shape::from_dims(&[6]),
                dtype: DType::F32,
            });
            let src = g.push(Node {
                op: Op::Const,
                inputs: vec![],
                shape: Shape::from_dims(&[2]),
                dtype: DType::F32,
            });
            let ws = g.push(Node {
                op: Op::WriteSlice {
                    ranges: vec![(0, 2)],
                    dyn_offset: Some((0, DynScalar::Sym(sym))),
                },
                inputs: vec![dest, src],
                shape: Shape::from_dims(&[6]),
                dtype: DType::F32,
            });
            (dest, src, ws)
        };
        let mut ctx = InferenceContext::new(Device::cpu());
        ctx.insert(dest_id, Arc::clone(&dest_arc));
        ctx.insert(src_id, Arc::clone(&src_arc));

        // Bind cached_len = 3: the slab must append at indices [3, 4].
        let mut env = SymEnv::new();
        env.bind(sym, 3).unwrap();
        let out = ctx
            .realize_one_as_with_env::<f32>(&graph, ws_id, &env)
            .expect("realize_one_as_with_env");
        assert_eq!(
            out,
            vec![0.0, 0.0, 0.0, 7.0, 8.0, 0.0],
            "KV slab must land at the SymEnv-bound cached_len=3, not the placeholder 0",
        );

        // An empty env on a graph with an unbound dyn_offset surfaces a
        // typed error (never a panic) — the write-once contract's
        // "presence ⇒ produced" read at realize.
        let err = ctx.realize_one_as_with_env::<f32>(&graph, ws_id, &SymEnv::new());
        assert!(
            err.is_err(),
            "unbound cached_len must surface a typed error"
        );
    }

    /// `KvCache::with_dims` produces a fresh cache of the right
    /// shape with all layers `None` and `cached_len = 0`.
    #[test]
    fn kv_cache_with_dims_constructs_fresh() {
        let cache = KvCache::with_dims(4, 8, 64);
        assert_eq!(cache.n_layers(), 4);
        assert_eq!(cache.cached_len, 0);
        assert_eq!(cache.n_kv_heads, 8);
        assert_eq!(cache.head_dim, 64);
        for li in 0..4 {
            assert!(cache.layer(li).is_none());
        }
    }

    /// `KvCache::set_layer` + `layer` round-trip.
    #[test]
    fn kv_cache_set_and_get_layer() {
        let mut cache = KvCache::with_dims(2, 4, 8);
        let k_arc = Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[0.0_f32; 32])));
        let v_arc = Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[0.0_f32; 32])));
        let layout = Layout::contiguous(Shape::from_dims(&[1, 4, 1, 8]));
        cache.set_layer(
            0,
            KvLayer {
                k: Arc::clone(&k_arc),
                v: Arc::clone(&v_arc),
                k_layout: layout.clone(),
                v_layout: layout.clone(),
                k_version: 0,
                v_version: 0,
                k_authority: AuthorityState::Host,
                v_authority: AuthorityState::Host,
            },
        );
        let layer = cache.layer(0).expect("just set");
        assert!(Arc::ptr_eq(&layer.k, &k_arc));
        assert!(Arc::ptr_eq(&layer.v, &v_arc));
        assert_eq!(layer.k_version, 0);
        assert!(matches!(layer.k_authority, AuthorityState::Host));
        assert!(cache.layer(1).is_none());
    }

    /// `KvCache::clear` drops all layers and resets `cached_len`.
    #[test]
    fn kv_cache_clear_drops_layers() {
        let mut cache = KvCache::with_dims(2, 4, 8);
        let k = Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[0.0_f32; 32])));
        let v = Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[0.0_f32; 32])));
        let layout = Layout::contiguous(Shape::from_dims(&[1, 4, 1, 8]));
        cache.set_layer(
            0,
            KvLayer {
                k: Arc::clone(&k),
                v: Arc::clone(&v),
                k_layout: layout.clone(),
                v_layout: layout,
                k_version: 1,
                v_version: 1,
                k_authority: AuthorityState::Host,
                v_authority: AuthorityState::Host,
            },
        );
        cache.cached_len = 7;
        // Before clear: cache holds an Arc to k → strong_count 2.
        assert_eq!(Arc::strong_count(&k), 2);

        cache.clear();

        assert!(cache.layer(0).is_none());
        assert_eq!(cache.cached_len, 0);
        // After clear: cache's Arc dropped → strong_count 1.
        assert_eq!(Arc::strong_count(&k), 1);
    }

    /// `KvCache::truncate_to` updates `cached_len` only (Phase E.3.0
    /// limitation; pre-allocated buffer + WriteSlice handles the
    /// in-place truncate in Phase E.3.2).
    #[test]
    fn kv_cache_truncate_updates_cached_len() {
        let mut cache = KvCache::with_dims(2, 4, 8);
        cache.cached_len = 16;
        cache.truncate_to(10);
        assert_eq!(cache.cached_len, 10);
        // Truncating to a value >= current is a no-op.
        cache.truncate_to(15);
        assert_eq!(cache.cached_len, 10);
        // Truncating to zero clears the live extent.
        cache.truncate_to(0);
        assert_eq!(cache.cached_len, 0);
    }

    /// Authority + version fields are placeholders today; just
    /// confirm they default sensibly via `KvLayer` construction.
    #[test]
    fn kv_layer_default_authority_is_host() {
        let arc = Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[0.0_f32; 4])));
        let layout = Layout::contiguous(Shape::from_dims(&[4]));
        let layer = KvLayer {
            k: Arc::clone(&arc),
            v: Arc::clone(&arc),
            k_layout: layout.clone(),
            v_layout: layout,
            k_version: 0,
            v_version: 0,
            k_authority: AuthorityState::Host,
            v_authority: AuthorityState::Host,
        };
        assert!(matches!(layer.k_authority, AuthorityState::Host));
        assert!(matches!(layer.v_authority, AuthorityState::Host));
        assert_eq!(layer.k_version, 0);
        assert_eq!(layer.v_version, 0);
    }

    // ---- KvCache::with_capacity + accessors (Phase E.3.3.A) -----------------

    /// `with_capacity` allocates n_layers fresh K + V buffers on the
    /// CPU device. Each buffer is `n_kv_heads * max_seq_len * head_dim`
    /// elements of the requested dtype, zero-initialized, with the
    /// `[1, n_kv_heads, max_seq_len, head_dim]` layout pre-populated.
    #[test]
    // `BackendStorage::{Cuda,Vulkan,Metal}` are cfg-gated, so at default
    // features `Cpu` is the only variant and this `if let` is irrefutable;
    // under --features cuda/vulkan/metal the pattern is refutable, so a plain
    // `let` would be an E0005 there. Allow the default-features lint.
    #[allow(irrefutable_let_patterns)]
    fn kv_cache_with_capacity_allocates_all_layers_on_cpu() {
        let device = Device::cpu();
        let cache = KvCache::with_capacity(
            /* n_layers     */ 3,
            /* n_kv_heads   */ 4,
            /* head_dim     */ 16,
            /* max_seq_len  */ 32,
            DType::F32,
            &device,
        )
        .expect("with_capacity");

        assert_eq!(cache.n_layers(), 3);
        assert_eq!(cache.cached_len, 0);
        assert_eq!(cache.n_kv_heads, 4);
        assert_eq!(cache.head_dim, 16);
        assert_eq!(cache.max_seq_len, Some(32));
        assert_eq!(cache.dtype, Some(DType::F32));

        // Every layer is populated and has the expected layout.
        for li in 0..3 {
            let layer = cache.layer(li).expect("layer populated");
            assert_eq!(layer.k_layout.shape().dims(), &[1, 4, 32, 16]);
            assert_eq!(layer.v_layout.shape().dims(), &[1, 4, 32, 16]);
            // Bytes are zero-initialized: 4 (n_kv) * 32 (seq) * 16 (head)
            // * 4 (f32) = 8192 bytes per slot.
            let k_guard = layer.k.read().unwrap();
            assert_eq!(k_guard.inner.len_bytes(), 8192);
            assert_eq!(k_guard.dtype, DType::F32);
            // Spot-check zero init.
            if let BackendStorage::Cpu(c) = &k_guard.inner {
                let typed: &[f32] = c.as_slice().unwrap();
                assert!(typed.iter().all(|&x| x == 0.0));
            } else {
                panic!("expected CPU storage");
            }
        }
    }

    /// `slot_storage(li, K)` returns the same Arc as `layer(li).k`.
    /// Used by the forward path to bind cache storage to per-step
    /// Const nodes.
    #[test]
    fn kv_cache_slot_storage_returns_layer_arc() {
        let device = Device::cpu();
        let cache =
            KvCache::with_capacity(2, 4, 8, 16, DType::F32, &device).expect("with_capacity");

        let k_via_layer = Arc::clone(&cache.layer(0).unwrap().k);
        let k_via_slot = cache.slot_storage(0, KvSlot::K).expect("layer 0 K");
        assert!(
            Arc::ptr_eq(&k_via_layer, &k_via_slot),
            "slot_storage should return the same Arc as layer().k",
        );

        let v_via_layer = Arc::clone(&cache.layer(1).unwrap().v);
        let v_via_slot = cache.slot_storage(1, KvSlot::V).expect("layer 1 V");
        assert!(Arc::ptr_eq(&v_via_layer, &v_via_slot));
    }

    /// `slot_storage` returns `None` for an unpopulated layer (the
    /// with_dims path before the first set_layer) and for an out-of-
    /// range layer index.
    #[test]
    fn kv_cache_slot_storage_returns_none_for_unpopulated() {
        let cache = KvCache::with_dims(2, 4, 8);
        assert!(cache.slot_storage(0, KvSlot::K).is_none());
        assert!(cache.slot_storage(1, KvSlot::V).is_none());
        // Out of range.
        assert!(cache.slot_storage(5, KvSlot::K).is_none());
    }

    /// `bump_version` advances the per-slot version counter
    /// independently for K and V.
    #[test]
    fn kv_cache_bump_version_advances_per_slot() {
        let device = Device::cpu();
        let mut cache =
            KvCache::with_capacity(1, 4, 8, 16, DType::F32, &device).expect("with_capacity");
        assert_eq!(cache.layer(0).unwrap().k_version, 0);
        assert_eq!(cache.layer(0).unwrap().v_version, 0);

        cache.bump_version(0, KvSlot::K);
        cache.bump_version(0, KvSlot::K);
        cache.bump_version(0, KvSlot::V);

        assert_eq!(cache.layer(0).unwrap().k_version, 2);
        assert_eq!(cache.layer(0).unwrap().v_version, 1);
    }

    /// Bumping a version on an unpopulated layer is a no-op (not a
    /// panic). The `with_dims` path leaves layers `None`; the
    /// forward code shouldn't blow up if it calls bump_version
    /// before the layer is wired.
    #[test]
    fn kv_cache_bump_version_unpopulated_is_noop() {
        let mut cache = KvCache::with_dims(2, 4, 8);
        cache.bump_version(0, KvSlot::K);
        cache.bump_version(99, KvSlot::V); // out of range, also no-op
        assert!(cache.layer(0).is_none());
    }

    /// Sanity: bf16 capacity allocation produces 2-byte elements.
    #[test]
    fn kv_cache_with_capacity_bf16_byte_count() {
        let device = Device::cpu();
        let cache =
            KvCache::with_capacity(1, 2, 4, 8, DType::BF16, &device).expect("with_capacity bf16");
        let layer = cache.layer(0).unwrap();
        // 2 (n_kv) * 8 (seq) * 4 (head) * 2 (bf16) = 128 bytes per slot.
        assert_eq!(layer.k.read().unwrap().inner.len_bytes(), 128);
        assert_eq!(cache.dtype, Some(DType::BF16));
    }

    // ---- LatentKvCache (MLA increment 4 — persistent N-slot decode cache) --

    /// `LatentKvCache::with_capacity` allocates `n_layers * n_slots` fresh
    /// zero-initialized buffers on the CPU device, one per (layer, slot),
    /// shaped `[max_seq_len, ...slot_trailing[s]]`.
    #[test]
    // Irrefutable only at default features: `BackendStorage::{Cuda,Vulkan,
    // Metal}` are cfg-gated, so under --features cuda/vulkan/metal this pattern
    // is refutable and a plain `let` would be an E0005. Allow the default lint.
    #[allow(irrefutable_let_patterns)]
    fn latent_kv_cache_with_capacity_allocates_all_layers_on_cpu() {
        let device = Device::cpu();
        let cache = LatentKvCache::with_capacity(
            /* n_layers    */ 3,
            /* max_seq_len */ 8,
            // MLA-shaped: slot 0 latent trailing [5], slot 1 k_pe trailing [2].
            vec![vec![5], vec![2]],
            DType::F32,
            &device,
        )
        .expect("with_capacity");

        assert_eq!(cache.n_layers(), 3);
        assert_eq!(cache.n_slots(), 2);
        assert_eq!(cache.cached_len, 0);
        assert_eq!(cache.max_seq_len, 8);
        assert_eq!(cache.dtype, DType::F32);
        assert_eq!(cache.slot_trailing(0), &[5]);
        assert_eq!(cache.slot_trailing(1), &[2]);

        for li in 0..3 {
            let latent = cache.slot_storage(li, 0).expect("latent slot populated");
            let kpe = cache.slot_storage(li, 1).expect("k_pe slot populated");
            // 8 (max_seq) * 5 (trailing) * 4 (f32) = 160 bytes.
            assert_eq!(latent.read().unwrap().inner.len_bytes(), 160);
            // 8 (max_seq) * 2 (trailing) * 4 (f32) = 64 bytes.
            assert_eq!(kpe.read().unwrap().inner.len_bytes(), 64);
            let guard = latent.read().unwrap();
            if let BackendStorage::Cpu(c) = &guard.inner {
                let typed: &[f32] = c.as_slice().unwrap();
                assert!(typed.iter().all(|&x| x == 0.0));
            } else {
                panic!("expected CPU storage");
            }
        }
    }

    /// `slot_storage` returns `None` for an out-of-range layer or slot
    /// index (never panics).
    #[test]
    fn latent_kv_cache_slot_storage_returns_none_for_oob() {
        let device = Device::cpu();
        let cache = LatentKvCache::with_capacity(2, 4, vec![vec![3]], DType::F32, &device)
            .expect("with_capacity");
        assert!(cache.slot_storage(5, 0).is_none()); // layer OOB
        assert!(cache.slot_storage(0, 5).is_none()); // slot OOB
    }

    /// `bump_version` advances the per-slot version counter independently
    /// per slot; out-of-range (layer, slot) is a no-op, never a panic.
    #[test]
    fn latent_kv_cache_bump_version_advances_per_slot_and_oob_is_noop() {
        let device = Device::cpu();
        let mut cache =
            LatentKvCache::with_capacity(1, 4, vec![vec![2], vec![1]], DType::F32, &device)
                .expect("with_capacity");

        cache.bump_version(0, 0);
        cache.bump_version(0, 0);
        cache.bump_version(0, 1);
        cache.bump_version(99, 0); // layer OOB, no-op
        cache.bump_version(0, 99); // slot OOB, no-op

        let layer0 = cache.layers[0].as_ref().expect("layer 0 populated");
        assert_eq!(layer0[0].version, 2);
        assert_eq!(layer0[1].version, 1);
    }

    /// Degenerate geometry (`n_layers == 0`, `max_seq_len == 0`, or no
    /// slots) surfaces as a typed `Err`, never a panic.
    #[test]
    fn latent_kv_cache_with_capacity_rejects_bad_geometry() {
        let device = Device::cpu();
        assert!(LatentKvCache::with_capacity(0, 4, vec![vec![2]], DType::F32, &device).is_err());
        assert!(LatentKvCache::with_capacity(1, 0, vec![vec![2]], DType::F32, &device).is_err());
        assert!(LatentKvCache::with_capacity(1, 4, vec![], DType::F32, &device).is_err());
    }
}