metrale_model_arch/deepseek_v41_ref/
model.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: DeepSeek-V4.1 full tiny-graph model: the config, weights, state
5//! and trace types that `forward` (in `model/forward.rs`) fills. On the
6//! fixture: embed, expand to `hc_mult` copies, the one-hot initial pre-mix,
7//! engram before layers 1 and 4, six blocks with the delayed mixes,
8//! collapse, final RMSNorm, f32 logits. The per-layer caches and the shared
9//! slots carry across prefill and decode, so the golden's captures for every
10//! layer and the logits are checkable.
11//!
12//! Owner: model-arch, DeepSeek-V4.1 reference.
13//! Invariants: none beyond the types.
14
15use super::attn::freqs_cis;
16use super::compress::{
17    CompAttnCfg, CompAttnRun, IndexerCfg, LayerAttnState, SharedRuntime, yarn_freqs_cis,
18};
19use super::engram::{EngramTables, NgramHashState, regen_bf16_matrix, regen_bf16_qk, regen_table};
20use super::moe::MoeCfg;
21use super::{Golden, regen_param};
22
23mod forward;
24
25pub use forward::forward;
26
27pub struct ModelCfg {
28    pub dim: usize,
29    pub hc: usize,
30    pub iters: usize,
31    pub hc_eps: f32,
32    pub eps: f32,
33    pub n_layers: usize,
34    pub vocab: usize,
35    pub max_seq: usize,
36    pub n_heads: usize,
37    pub head_dim: usize,
38    pub rope_dim: usize,
39    pub q_rank: usize,
40    pub o_rank: usize,
41    pub groups: usize,
42    pub window: usize,
43    pub ratios: Vec<usize>,
44    pub kv_sources: Vec<usize>,
45    pub index_sources: Vec<usize>,
46    pub cand_src: i64,
47    pub cand_topk_blocks: usize,
48    pub cand_block: usize,
49    pub index_heads: usize,
50    pub index_hd: usize,
51    pub index_topk: usize,
52    pub rope_theta: f32,
53    pub compress_rope_theta: f32,
54    pub rope_factor: f32,
55    pub orig_seq: usize,
56    pub beta_fast: f32,
57    pub beta_slow: f32,
58    pub moe: MoeCfg,
59}
60
61impl ModelCfg {
62    pub fn from_golden(g: &Golden) -> Self {
63        let u = |k: &str| g.fixture_u64(k) as usize;
64        let f = |k: &str| g.fixture_f64(k) as f32;
65        let list = |k: &str| g.fixture_usize_list(k);
66        let dim = u("dim");
67        ModelCfg {
68            dim,
69            hc: u("hc_mult"),
70            iters: u("hc_sinkhorn_iters"),
71            hc_eps: f("hc_eps"),
72            eps: f("norm_eps"),
73            n_layers: u("n_layers"),
74            vocab: u("vocab_size"),
75            max_seq: u("max_seq_len"),
76            n_heads: u("n_heads"),
77            head_dim: u("head_dim"),
78            rope_dim: u("rope_head_dim"),
79            q_rank: u("q_lora_rank"),
80            o_rank: u("o_lora_rank"),
81            groups: u("o_groups"),
82            window: u("window_size"),
83            ratios: list("compress_ratios"),
84            kv_sources: list("kv_source_layers"),
85            index_sources: list("index_source_layers"),
86            cand_src: g.fixture_i64("candidate_source_layer"),
87            cand_topk_blocks: u("candidate_topk_blocks"),
88            cand_block: u("candidate_block_size"),
89            index_heads: u("index_n_heads"),
90            index_hd: u("index_head_dim"),
91            index_topk: u("index_topk"),
92            rope_theta: f("rope_theta"),
93            compress_rope_theta: f("compress_rope_theta"),
94            rope_factor: f("rope_factor"),
95            orig_seq: u("original_seq_len"),
96            beta_fast: f("beta_fast"),
97            beta_slow: f("beta_slow"),
98            moe: MoeCfg {
99                dim,
100                inter: u("moe_inter_dim"),
101                n_routed: u("n_routed_experts"),
102                topk: u("n_activated_experts"),
103                gate_temp: f("gate_temp"),
104                norm_topk_prob: true,
105                route_scale: f("route_scale"),
106                swiglu_limit: f("swiglu_limit"),
107            },
108        }
109    }
110
111    pub fn attn_cfg(&self, layer: usize) -> CompAttnCfg {
112        let ratio = self.ratios[layer];
113        CompAttnCfg {
114            dim: self.dim,
115            n_heads: self.n_heads,
116            head_dim: self.head_dim,
117            rope_dim: self.rope_dim,
118            q_rank: self.q_rank,
119            o_rank: self.o_rank,
120            groups: self.groups,
121            window: self.window,
122            eps: self.eps,
123            ratio,
124            is_kv_source: self.kv_sources.contains(&layer),
125            is_index_source: self.index_sources.contains(&layer),
126            is_candidate_source: self.cand_src == layer as i64,
127            uses_candidates: self.cand_src >= 0 && (self.cand_src as usize) < layer,
128        }
129    }
130
131    pub fn indexer_cfg(&self) -> IndexerCfg {
132        IndexerCfg {
133            n_heads: self.index_heads,
134            index_hd: self.index_hd,
135            rope_dim: self.rope_dim,
136            q_rank: self.q_rank,
137            dim: self.dim,
138            hd: self.head_dim,
139            index_topk: self.index_topk,
140            cand_topk_blocks: self.cand_topk_blocks,
141            cand_block: self.cand_block,
142            eps: self.eps,
143        }
144    }
145}
146
147pub struct EngramWeights {
148    pub table: Vec<f32>,
149    pub wkv: Vec<f32>,
150    pub q: Vec<f32>,
151    pub k: Vec<f32>,
152    pub hash_index: usize,
153}
154
155pub struct LayerWeights {
156    pub hc_attn_fn: Vec<f32>,
157    pub hc_attn_scale: Vec<f32>,
158    pub hc_attn_base: Vec<f32>,
159    pub hc_ffn_fn: Vec<f32>,
160    pub hc_ffn_scale: Vec<f32>,
161    pub hc_ffn_base: Vec<f32>,
162    pub attn_norm: Vec<f32>,
163    pub ffn_norm: Vec<f32>,
164    pub sink: Vec<f32>,
165    pub wq_a: Vec<f32>,
166    pub q_norm: Vec<f32>,
167    pub wq_b: Vec<f32>,
168    pub wkv: Vec<f32>,
169    pub kv_norm: Vec<f32>,
170    pub wo_a: Vec<f32>,
171    pub wo_b: Vec<f32>,
172    pub comp_wkv: Option<Vec<f32>>,
173    pub comp_wgate: Option<Vec<f32>>,
174    pub comp_norm: Option<Vec<f32>>,
175    pub idx_wq_b: Option<Vec<f32>>,
176    pub idx_weights_proj: Option<Vec<f32>>,
177    pub idx_wk: Option<Vec<f32>>,
178    pub idx_k_norm: Option<Vec<f32>>,
179    pub gate_w: Vec<f32>,
180    pub gate_bias: Vec<f32>,
181    pub experts: Vec<(Vec<f32>, Vec<f32>, Vec<f32>)>,
182    pub shared: (Vec<f32>, Vec<f32>, Vec<f32>),
183    pub engram: Option<EngramWeights>,
184}
185
186pub struct ModelWeights {
187    pub embed: Vec<f32>,
188    pub norm: Vec<f32>,
189    pub head: Vec<f32>,
190    pub layers: Vec<LayerWeights>,
191}
192
193impl ModelWeights {
194    pub fn from_golden(g: &Golden, c: &ModelCfg, t: &EngramTables) -> Self {
195        let layers = (0..c.n_layers)
196            .map(|l| {
197                let p = |n: &str| regen_param(g, &format!("layers.{l}.{n}"));
198                let opt = |n: &str, on: bool| if on { Some(p(n)) } else { None };
199                let ac = c.attn_cfg(l);
200                let has_comp = ac.is_kv_source;
201                let has_idx = ac.is_index_source;
202                let engram = t
203                    .layer_ids
204                    .iter()
205                    .position(|&x| x == l)
206                    .map(|hi| EngramWeights {
207                        table: regen_table(t, hi),
208                        wkv: regen_bf16_matrix(
209                            &format!("layers.{l}.engram.wkv.weight"),
210                            c.dim * (c.hc + 1),
211                            t.n_hash_cols() * t.head_dim,
212                        ),
213                        q: regen_bf16_qk(&format!("layers.{l}.engram.q_weight"), c.hc * c.dim),
214                        k: regen_bf16_qk(&format!("layers.{l}.engram.k_weight"), c.hc * c.dim),
215                        hash_index: hi,
216                    });
217                LayerWeights {
218                    hc_attn_fn: p("hc_attn_fn"),
219                    hc_attn_scale: p("hc_attn_scale"),
220                    hc_attn_base: p("hc_attn_base"),
221                    hc_ffn_fn: p("hc_ffn_fn"),
222                    hc_ffn_scale: p("hc_ffn_scale"),
223                    hc_ffn_base: p("hc_ffn_base"),
224                    attn_norm: p("attn_norm.weight"),
225                    ffn_norm: p("ffn_norm.weight"),
226                    sink: p("attn.attn_sink"),
227                    wq_a: p("attn.wq_a.weight"),
228                    q_norm: p("attn.q_norm.weight"),
229                    wq_b: p("attn.wq_b.weight"),
230                    wkv: p("attn.wkv.weight"),
231                    kv_norm: p("attn.kv_norm.weight"),
232                    wo_a: p("attn.wo_a.weight"),
233                    wo_b: p("attn.wo_b.weight"),
234                    comp_wkv: opt("attn.compressor.wkv.weight", has_comp),
235                    comp_wgate: opt("attn.compressor.wgate.weight", has_comp && ac.ratio > 1),
236                    comp_norm: opt("attn.compressor.norm.weight", has_comp),
237                    idx_wq_b: opt("attn.indexer.wq_b.weight", has_idx),
238                    idx_weights_proj: opt("attn.indexer.weights_proj.weight", has_idx),
239                    idx_wk: opt("attn.indexer.wk.weight", has_idx && has_comp),
240                    idx_k_norm: opt("attn.indexer.k_norm.weight", has_idx && has_comp),
241                    gate_w: p("ffn.gate.weight"),
242                    gate_bias: p("ffn.gate.bias"),
243                    experts: (0..c.moe.n_routed)
244                        .map(|e| {
245                            (
246                                p(&format!("ffn.experts.{e}.w1.weight")),
247                                p(&format!("ffn.experts.{e}.w2.weight")),
248                                p(&format!("ffn.experts.{e}.w3.weight")),
249                            )
250                        })
251                        .collect(),
252                    shared: (
253                        p("ffn.shared_experts.w1.weight"),
254                        p("ffn.shared_experts.w2.weight"),
255                        p("ffn.shared_experts.w3.weight"),
256                    ),
257                    engram,
258                }
259            })
260            .collect();
261        ModelWeights {
262            embed: regen_param(g, "embed.weight"),
263            norm: regen_param(g, "norm.weight"),
264            head: regen_param(g, "head.weight"),
265            layers,
266        }
267    }
268}
269
270pub struct ModelState<'a> {
271    pub hash: NgramHashState<'a>,
272    pub layers: Vec<LayerAttnState>,
273    pub shared: SharedRuntime,
274    pub fcs: Vec<Vec<(f32, f32)>>,
275}
276
277impl<'a> ModelState<'a> {
278    pub fn new(c: &ModelCfg, t: &'a EngramTables) -> Self {
279        ModelState {
280            hash: NgramHashState::new(t, 1, c.max_seq),
281            layers: (0..c.n_layers)
282                .map(|l| LayerAttnState::new(&c.attn_cfg(l), c.max_seq, c.index_hd))
283                .collect(),
284            shared: SharedRuntime::default(),
285            fcs: (0..c.n_layers)
286                .map(|l| {
287                    if c.ratios[l] == 0 {
288                        freqs_cis(c.rope_dim, c.max_seq, c.rope_theta)
289                    } else {
290                        yarn_freqs_cis(
291                            c.rope_dim,
292                            c.max_seq,
293                            c.orig_seq,
294                            c.compress_rope_theta,
295                            c.rope_factor,
296                            c.beta_fast,
297                            c.beta_slow,
298                        )
299                    }
300                })
301                .collect(),
302        }
303    }
304}
305
306pub struct LayerTrace {
307    pub engram_out: Option<Vec<f32>>,
308    pub h_in: Vec<f32>,
309    pub pre_mix_in: Vec<f32>,
310    pub attn_pre: Vec<f32>,
311    pub attn_post: Vec<f32>,
312    pub attn_comb: Vec<f32>,
313    pub attn_in: Vec<f32>,
314    pub attn: CompAttnRun,
315    pub ffn_pre: Vec<f32>,
316    pub ffn_post: Vec<f32>,
317    pub ffn_comb: Vec<f32>,
318    pub ffn_in: Vec<f32>,
319    pub moe_weights: Vec<f32>,
320    pub moe_indices: Vec<usize>,
321    pub ffn_out: Vec<f32>,
322    pub h_out: Vec<f32>,
323    // 2026-09-25: The `shared_*` fields snapshot the shared slots after this
324    // layer, for the source-layer captures.
325    pub shared_compress_kv: Vec<f32>,
326    pub shared_index_k: Vec<f32>,
327    pub shared_topk: Vec<i32>,
328    pub shared_topk_w: usize,
329    pub shared_cand: Vec<bool>,
330    pub shared_cand_w: usize,
331}
332
333pub struct StepTrace {
334    pub tokens: usize,
335    pub embed: Vec<f32>,
336    pub layers: Vec<LayerTrace>,
337    pub h_final: Vec<f32>,
338    pub head_in: Vec<f32>,
339    pub logits: Vec<f32>,
340}
341
342#[cfg(test)]
343#[path = "model_tests.rs"]
344mod tests;
metrale_model_arch/deepseek_v41_ref/model.rs

metrale_model_arch/deepseek_v41_ref/
model.rs
