needle_core/
layers.rs

1//! Encoder and decoder layer stacks.
2//!
3//! Gated residual: each sub-layer has a learned scalar sigmoid gate g (init=0).
4//! output = x + sigmoid(g) * sublayer(norm(x))
5//!
6//! Encoder block: self-attn only (no FFN by default).
7//! Decoder block: self-attn → cross-attn → optional FFN.
8
9use crate::attn::{
10    cross_attn_incremental, self_attn_full, self_attn_incremental, AttnConfig, AttnWeights, KvCache,
11};
12use crate::config::FfnActivation;
13use crate::ffn::{ffn_forward, FfnWeights};
14use crate::norm::zc_rms_norm_vec;
15use crate::ops::sigmoid;
16use crate::rope::RopeCache;
17use alloc::vec::Vec;
18
19pub struct EncoderLayer {
20    pub self_attn: AttnWeights,
21    pub self_attn_gate: f32, // scalar sigmoid gate
22    pub norm: Vec<f32>,      // ZCRMSNorm scale [d_model]
23    pub ffn: Option<FfnWeights>,
24    pub ffn_gate: f32,
25    pub ffn_norm: Option<Vec<f32>>,
26    pub ffn_activation: Option<FfnActivation>,
27}
28
29pub struct DecoderLayer {
30    pub self_attn: AttnWeights,
31    pub self_attn_gate: f32,
32    pub self_attn_norm: Vec<f32>,
33    pub cross_attn: AttnWeights,
34    pub cross_attn_gate: f32,
35    pub cross_attn_norm: Vec<f32>,
36    pub ffn: Option<FfnWeights>,
37    pub ffn_gate: f32,
38    pub ffn_norm: Option<Vec<f32>>,
39    pub ffn_activation: Option<FfnActivation>,
40}
41
42/// Run a single encoder layer on the full sequence.
43/// `x`:      [seq_len, d_model] — modified in-place
44/// `normed`: scratch buffer [seq_len, d_model] — pre-allocated by caller, reused across layers
45#[allow(clippy::too_many_arguments)]
46pub fn encoder_layer_forward(
47    x: &mut [f32],
48    layer: &EncoderLayer,
49    cfg: &AttnConfig,
50    rope: &RopeCache,
51    seq_len: usize,
52    mask: Option<&[f32]>,
53    tmp: &mut Vec<f32>,
54    normed: &mut Vec<f32>,
55) {
56    let d = cfg.d_model;
57
58    // Pre-norm (copy x into pre-allocated scratch, normalize in-place)
59    normed.resize(seq_len * d, 0.0);
60    normed.copy_from_slice(&x[..seq_len * d]);
61    for t in 0..seq_len {
62        zc_rms_norm_vec(&mut normed[t * d..(t + 1) * d], &layer.norm);
63    }
64
65    // Self-attention
66    tmp.resize(seq_len * d, 0.0);
67    self_attn_full(normed, tmp, &layer.self_attn, cfg, rope, seq_len, mask);
68
69    // Gated residual
70    let gate = sigmoid(layer.self_attn_gate);
71    for i in 0..seq_len * d {
72        x[i] += gate * tmp[i];
73    }
74
75    // Optional FFN
76    if let (Some(ffn), Some(ffn_norm), Some(activation)) =
77        (&layer.ffn, &layer.ffn_norm, &layer.ffn_activation)
78    {
79        let d_ff = ffn.gate_proj.out_feat;
80        for t in 0..seq_len {
81            let row = &mut x[t * d..(t + 1) * d];
82            let mut normed_row = row.to_vec();
83            zc_rms_norm_vec(&mut normed_row, ffn_norm);
84            tmp.resize(d, 0.0);
85            ffn_forward(&normed_row, tmp, ffn, d_ff, activation);
86            let ffn_gate = sigmoid(layer.ffn_gate);
87            for (xi, &ti) in row.iter_mut().zip(tmp.iter()) {
88                *xi += ffn_gate * ti;
89            }
90        }
91    }
92}
93
94/// Run a single decoder layer incrementally (one token at a time).
95/// `x`:      current decoder token vector [d_model] — modified in-place with residual
96/// `enc_kv`: precomputed encoder KV (for cross-attn)
97/// `dec_kv_self`: decoder self-attn KV cache (appended to here)
98/// `normed`: scratch buffer [d_model] — pre-allocated by caller, reused for self-attn/cross-attn/ffn
99#[allow(clippy::too_many_arguments)]
100pub fn decoder_layer_forward(
101    x: &mut [f32],
102    layer: &DecoderLayer,
103    cfg: &AttnConfig,
104    rope: &RopeCache,
105    enc_kv: &KvCache,
106    dec_kv_self: &mut KvCache,
107    tmp: &mut Vec<f32>,
108    normed: &mut Vec<f32>,
109) {
110    let d = cfg.d_model;
111
112    // --- Self-attention ---
113    normed.resize(d, 0.0);
114    normed.copy_from_slice(x);
115    zc_rms_norm_vec(normed, &layer.self_attn_norm);
116
117    self_attn_incremental(normed, tmp, &layer.self_attn, cfg, rope, dec_kv_self, None);
118
119    let gate = sigmoid(layer.self_attn_gate);
120    for (xi, &ti) in x.iter_mut().zip(tmp.iter()) {
121        *xi += gate * ti;
122    }
123
124    // --- Cross-attention (reuse normed buffer) ---
125    normed.copy_from_slice(x);
126    zc_rms_norm_vec(normed, &layer.cross_attn_norm);
127
128    cross_attn_incremental(normed, tmp, &layer.cross_attn, cfg, enc_kv);
129
130    let cross_gate = sigmoid(layer.cross_attn_gate);
131    for (xi, &ti) in x.iter_mut().zip(tmp.iter()) {
132        *xi += cross_gate * ti;
133    }
134
135    // --- Optional FFN (reuse normed buffer) ---
136    if let (Some(ffn), Some(ffn_norm), Some(activation)) =
137        (&layer.ffn, &layer.ffn_norm, &layer.ffn_activation)
138    {
139        let d_ff = ffn.gate_proj.out_feat;
140        normed.copy_from_slice(x);
141        zc_rms_norm_vec(normed, ffn_norm);
142        tmp.resize(d, 0.0);
143        ffn_forward(normed, tmp, ffn, d_ff, activation);
144        let ffn_gate = sigmoid(layer.ffn_gate);
145        for (xi, &ti) in x.iter_mut().zip(tmp.iter()) {
146            *xi += ffn_gate * ti;
147        }
148    }
149}
needle_core/layers.rs

needle_core/
layers.rs
