spark_model/layers/deepseek_v41_ref/model/
forward.rs

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

spark_model/layers/deepseek_v41_ref/model/
forward.rs
