Struct V2Model 

pub struct V2Model {
    pub cfg: V2Config,
    pub embedding: CqWeight,
    pub layers: Vec<V2Layer>,
    pub mhc: V2Mhc,
    pub engrams: Vec<V2Engram>,
    pub final_norm: Vec<f32>,
    /* private fields */
}

A loaded v2 model. Immutable during inference; per-sequence state lives in V2State.

Fields

cfg: V2Config

embedding: CqWeight

Tied embedding, [vocab_size, d_model] — also the LM head, transposed.

layers: Vec<V2Layer>

mhc: V2Mhc

engrams: Vec<V2Engram>

final_norm: Vec<f32>

Implementations

impl V2Model

pub fn prefill_batch( &self, tokens: &[u32], state: &mut V2State, batch: &mut V2Batch, logits: Option<&mut [f32]>, ) -> Result<(), &'static str>

Feed tokens at the current position, batching the projections.

Chunks internally at V2Batch::cap. When logits is Some, it receives the logits for the final token only — the intermediate positions skip the tied LM head, which is what makes prefill cheaper than the same number of step calls even before batching.

Leaves state exactly as the equivalent run of step_prefill would, so decoding can continue with step.

impl V2Model

pub const ENGRAM_GATE_SOURCE: &'static str = "decode.py (_forward_cached), gate on u"

Where this port follows the decode reference rather than the training graph.

architecture.Block gates the Engram contribution on the lane-stacked x (einsum "btd,sbtd->sbt", scaled by site_flags). decode._forward_cached gates it on u, the lane-mixed vector the block actually consumes, with an explicit per-layer branch. This port follows decode.py: it is the incremental path, and the one the parity fixtures come from. Do not “fix” this toward architecture.py.

pub fn new( cfg: V2Config, embedding: CqWeight, layers: Vec<V2Layer>, mhc: V2Mhc, engrams: Vec<V2Engram>, final_norm: Vec<f32>, ) -> Result<Self, &'static str>

pub fn make_state(&self) -> V2State

Allocate per-sequence state sized for the model’s own attention window.

This is what generation wants. The KV cache is a ring of kv_window positions rather than max_seq_len, which is the difference between 14 MB and 113 MB at the shipped geometry.

Use make_state_full_causal for the probe heads, which attend over the whole sequence and therefore need the whole cache.

pub fn make_state_full_causal(&self) -> V2State

Allocate per-sequence state for full causal attention.

Sizes the cache to max_seq_len and sets the window to 0. Needed by architecture.encode_contrastive and forward_confidence, which default to full causal even though the LM path is windowed.

pub fn step( &self, token: u32, state: &mut V2State, logits: &mut [f32], ) -> Result<(), &'static str>

Process one token at the next position and write logits over the vocabulary.

Feed a prompt by calling this once per token in order; the logits from the last call are the ones to sample from.

pub fn step_prefill( &self, token: u32, state: &mut V2State, ) -> Result<(), &'static str>

Process one token without evaluating the LM head.

Prefill needs logits only at the final prompt position, and the tied LM head is the single widest matmul in a step (vocab_size x d_model — 8192 rows against the 512 of any projection). Skipping it on the preceding tokens removes that cost from every one of them.

last_hidden is still valid afterwards.

pub fn step_traced<F: FnMut(&str, usize, &[f32])>( &self, token: u32, state: &mut V2State, logits: &mut [f32], trace: &mut F, ) -> Result<(), &'static str>

step, with a callback at each of the capture points tools/gen_v2_forward_parity.py records — the ladder used to locate a divergence instead of only detecting one.

trace(name, layer, values); layer is usize::MAX for whole-forward points. Generic rather than a trait object, so the no-op closure step passes inlines away entirely.

pub fn last_hidden(state: &V2State) -> &[f32]

The hidden state the probe heads pool over, for the position just stepped.

Valid only immediately after a successful step; it is the post-final-norm vector the LM head consumed.