ferrox_core/
recurrent_state.rs

1//! What a layer with no KV history carries between tokens instead.
2//!
3//! A Mamba layer (and every other recurrent block llama.cpp keeps in
4//! `llama_memory_recurrent`) has no per-position rows to attend over;
5//! it has a fixed-size state that the next token reads and overwrites.
6//! LFM2's short convolution is the exception that proves the rule: its
7//! state IS the last `l_cache - 1` inputs, so `ferrox_models::shortconv`
8//! keeps it as the layer's KV history and needs nothing here. A Mamba
9//! state is a reduction over the whole prefix, not a window of it, and
10//! that is the one property every consumer of a per-layer cache has to
11//! know about:
12//!
13//! - it CLONES with the cache (a prefix-cache fork is a fork of the
14//!   state), and CLEARS with it;
15//! - it cannot be TRUNCATED to a middle position. llama.cpp's
16//!   `llama_memory_recurrent::seq_rm` refuses a `p0 > 0` for the same
17//!   reason and its server re-prefills. So [`KvCache::truncate`] on a
18//!   cache that holds one refuses anything but "to zero" or "to where
19//!   it is", and the callers that roll back -- the prefix cache,
20//!   speculative verification, the draft model, the whole-response
21//!   cache's back-off -- ask [`KvCache::can_truncate_to`] first or are
22//!   fenced off the model.
23//!
24//! The buffers are flat and the LAYER owns their geometry (its weights
25//! say what `d_conv`, the conv width and the scan dims are), so this
26//! type cannot disagree with the block about a shape: it is created by
27//! the block, on first use, at the size the block asks for.
28//!
29//! [`KvCache::truncate`]: crate::cache::KvCache::truncate
30//! [`KvCache::can_truncate_to`]: crate::cache::KvCache::can_truncate_to
31
32/// One sequence's state for one recurrent layer.
33#[derive(Debug, Clone, PartialEq)]
34pub struct RecurrentState {
35    /// The conv window, `[d_conv - 1][width]`, oldest row first
36    /// (`llama_hparams::n_embd_r`).
37    pub conv: Vec<f32>,
38    /// The SSM state, `[n_head][head_dim][d_state]`
39    /// (`llama_hparams::n_embd_s`).
40    pub ssm: Vec<f32>,
41}
42
43impl RecurrentState {
44    /// A fresh sequence's state: zeros, as `build_rs` zeroes a new
45    /// sequence's (`llama-graph.cpp`, `llm_graph_input_rs`).
46    pub fn zeros(conv_len: usize, ssm_len: usize) -> Self {
47        Self {
48            conv: vec![0.0; conv_len],
49            ssm: vec![0.0; ssm_len],
50        }
51    }
52
53    /// Bytes this state holds.
54    pub fn bytes(&self) -> usize {
55        (self.conv.len() + self.ssm.len()) * std::mem::size_of::<f32>()
56    }
57}
ferrox_core/recurrent_state.rs

ferrox_core/
recurrent_state.rs
