metrale_model_arch/deepseek_v41_ref/model/
forward.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: The forward for one chunk of the DeepSeek-V4.1 tiny graph: embed,
5//! expand to `hc_mult` copies, the one-hot initial pre-mix, engram, the
6//! blocks with the delayed mixes, collapse, final RMSNorm, f32 logits.
7//!
8//! Owner: model-arch, DeepSeek-V4.1 reference.
9//! Invariants: none beyond the types.
10
11use super::super::attn::AttnWeights;
12use super::super::compress::{CompressorWeights, IndexerWeights, attention_any};
13use super::super::engram::{EngramTables, embed_rows, gate_and_add};
14use super::super::hc::{hc_mixes, hc_post, hc_pre, rms_norm};
15use super::super::moe::{MoeWeights, moe};
16use super::{LayerTrace, ModelCfg, ModelState, ModelWeights, StepTrace};
17
18/// 2026-09-25: The forward for one chunk (`[seqlen]` ids at `start_pos`), text only.
19pub fn forward(
20    ids: &[i64],
21    start_pos: usize,
22    c: &ModelCfg,
23    w: &ModelWeights,
24    st: &mut ModelState,
25    t: &EngramTables,
26) -> StepTrace {
27    let (dim, hc, s) = (c.dim, c.hc, ids.len());
28    let hashes = st.hash.forward(ids, s, start_pos);
29    let (nl, cols) = (t.layer_ids.len(), t.n_hash_cols());
30
31    let embed: Vec<f32> = ids
32        .iter()
33        .flat_map(|&id| {
34            w.embed[id as usize * dim..(id as usize + 1) * dim]
35                .iter()
36                .copied()
37        })
38        .collect();
39    let mut h: Vec<f32> = (0..s)
40        .flat_map(|tk| std::iter::repeat_n(embed[tk * dim..(tk + 1) * dim].to_vec(), hc).flatten())
41        .collect();
42    let mut pre_mix: Vec<f32> = (0..s)
43        .flat_map(|_| (0..hc).map(|stream| if stream == 0 { 1.0 } else { 0.0 }))
44        .collect();
45    let icfg = c.indexer_cfg();
46    let mut traces = Vec::with_capacity(c.n_layers);
47
48    for l in 0..c.n_layers {
49        let lw = &w.layers[l];
50        let hs = &hashes;
51        let engram_out = lw.engram.as_ref().map(|e| {
52            let hi = e.hash_index;
53            let ids_l: Vec<i64> = (0..s)
54                .flat_map(|tk| (0..cols).map(move |cc| hs[(tk * nl + hi) * cols + cc]))
55                .collect();
56            let emb = embed_rows(&e.table, t.head_dim, &ids_l);
57            let in_f = cols * t.head_dim;
58            let kv = super::super::engram::linear_bf16(&emb, s, in_f, &e.wkv, dim * (hc + 1));
59            let out = gate_and_add(&h, &kv, &e.q, &e.k, s, hc, dim, c.eps);
60            h = out.clone();
61            out
62        });
63        let h_in = h.clone();
64        let pre_mix_in = pre_mix.clone();
65
66        let (attn_pre, attn_post, attn_comb) = hc_mixes(
67            &h,
68            s,
69            hc,
70            dim,
71            &lw.hc_attn_fn,
72            &lw.hc_attn_scale,
73            &lw.hc_attn_base,
74            c.iters,
75            c.hc_eps,
76            c.eps,
77        );
78        let attn_in = rms_norm(
79            &hc_pre(&h, &pre_mix, s, hc, dim),
80            &lw.attn_norm,
81            s,
82            dim,
83            c.eps,
84        );
85        let aw = AttnWeights {
86            sink: &lw.sink,
87            wq_a: &lw.wq_a,
88            q_norm: &lw.q_norm,
89            wq_b: &lw.wq_b,
90            wkv: &lw.wkv,
91            kv_norm: &lw.kv_norm,
92            wo_a: &lw.wo_a,
93            wo_b: &lw.wo_b,
94        };
95        let comp = lw.comp_wkv.as_ref().map(|wkv| CompressorWeights {
96            wkv,
97            wgate: lw.comp_wgate.as_deref(),
98            norm: lw.comp_norm.as_deref().expect("comp norm"),
99        });
100        let idxw = lw.idx_wq_b.as_ref().map(|wq_b| IndexerWeights {
101            wq_b,
102            weights_proj: lw.idx_weights_proj.as_deref().expect("weights_proj"),
103            wk: lw.idx_wk.as_deref(),
104            k_norm: lw.idx_k_norm.as_deref(),
105        });
106        let ac = c.attn_cfg(l);
107        let attn = attention_any(
108            &attn_in,
109            s,
110            start_pos,
111            &aw,
112            comp.as_ref(),
113            idxw.as_ref(),
114            &icfg,
115            &ac,
116            &st.fcs[l],
117            &mut st.layers[l],
118            &mut st.shared,
119        );
120        let h_mid = hc_post(&attn.out, &h, &attn_post, &attn_comb, s, hc, dim);
121
122        let (ffn_pre, ffn_post, ffn_comb) = hc_mixes(
123            &h_mid,
124            s,
125            hc,
126            dim,
127            &lw.hc_ffn_fn,
128            &lw.hc_ffn_scale,
129            &lw.hc_ffn_base,
130            c.iters,
131            c.hc_eps,
132            c.eps,
133        );
134        let ffn_in = rms_norm(
135            &hc_pre(&h_mid, &attn_pre, s, hc, dim),
136            &lw.ffn_norm,
137            s,
138            dim,
139            c.eps,
140        );
141        let mw = MoeWeights {
142            gate_w: &lw.gate_w,
143            gate_bias: &lw.gate_bias,
144            experts: lw
145                .experts
146                .iter()
147                .map(|(a, b, cc)| (a.as_slice(), b.as_slice(), cc.as_slice()))
148                .collect(),
149            shared: (&lw.shared.0, &lw.shared.1, &lw.shared.2),
150        };
151        let (ffn_out, moe_weights, moe_indices) = moe(&ffn_in, s, &mw, &c.moe);
152        let h_out = hc_post(&ffn_out, &h_mid, &ffn_post, &ffn_comb, s, hc, dim);
153
154        h = h_out.clone();
155        pre_mix = ffn_pre.clone();
156        traces.push(LayerTrace {
157            engram_out,
158            h_in,
159            pre_mix_in,
160            attn_pre,
161            attn_post,
162            attn_comb,
163            attn_in,
164            attn,
165            ffn_pre,
166            ffn_post,
167            ffn_comb,
168            ffn_in,
169            moe_weights,
170            moe_indices,
171            ffn_out,
172            h_out,
173            shared_compress_kv: st.shared.compress_kv.clone(),
174            shared_index_k: st.shared.index_k.clone(),
175            shared_topk: st.shared.topk_idxs.clone(),
176            shared_topk_w: st.shared.topk,
177            shared_cand: st.shared.candidates.clone(),
178            shared_cand_w: st.shared.cand_width,
179        });
180    }
181
182    let h_final = hc_pre(&h, &pre_mix, s, hc, dim);
183    let head_in = rms_norm(&h_final, &w.norm, s, dim, c.eps);
184    let mut logits = vec![0f32; s * c.vocab];
185    for tk in 0..s {
186        for v in 0..c.vocab {
187            logits[tk * c.vocab + v] = head_in[tk * dim..(tk + 1) * dim]
188                .iter()
189                .zip(&w.head[v * dim..(v + 1) * dim])
190                .map(|(a, b)| a * b)
191                .sum();
192        }
193    }
194    StepTrace {
195        tokens: s,
196        embed,
197        layers: traces,
198        h_final,
199        head_in,
200        logits,
201    }
202}
metrale_model_arch/deepseek_v41_ref/model/forward.rs

metrale_model_arch/deepseek_v41_ref/model/
forward.rs
