Struct NeedleModel 

pub struct NeedleModel {
    pub cfg: TransformerConfig,
    pub embedding: Vec<f32>,
    pub encoder_layers: Vec<EncoderLayer>,
    pub decoder_layers: Vec<DecoderLayer>,
    pub encoder_final_norm: Vec<f32>,
    pub decoder_final_norm: Vec<f32>,
    /* private fields */
}

Fields

cfg: TransformerConfig

embedding: Vec<f32>

Token embedding matrix [vocab_size, d_model] (also used as LM head transposed)

encoder_layers: Vec<EncoderLayer>

decoder_layers: Vec<DecoderLayer>

encoder_final_norm: Vec<f32>

Encoder final norm scale [d_model]

decoder_final_norm: Vec<f32>

Decoder final norm scale [d_model]

Implementations

impl NeedleModel

pub fn new( cfg: TransformerConfig, embedding: Vec<f32>, encoder_layers: Vec<EncoderLayer>, decoder_layers: Vec<DecoderLayer>, encoder_final_norm: Vec<f32>, decoder_final_norm: Vec<f32>, ) -> Self

pub fn encode( &self, input_ids: &[u32], enc_kv_caches: &mut [KvCache], ) -> Vec<f32>

Run the encoder over the full input sequence. Returns encoder hidden states [enc_len, d_model]. Also fills enc_kv_caches (one per decoder layer) with cross-attn K/V projections.

pub fn decode_step( &self, token_id: u32, enc_kv_caches: &[KvCache], dec_kv_caches: &mut [KvCache], logits: &mut [f32], )

Autoregressive decode: generates one token given decoder input token and KV caches. Returns logits over the full vocabulary.

pub fn make_enc_kv_caches(&self, enc_len: usize) -> Vec<KvCache>

Allocate fresh KV caches for cross-attention (one per decoder layer, sized for enc input).

pub fn make_dec_kv_caches(&self) -> Vec<KvCache>

Allocate fresh KV caches for decoder self-attention.