metrale_model_arch/deepseek_v41_ref/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: DeepSeek-V4.1 Flash tiny-graph reference: the V4.1 graph on the
5//! CPU, checked against `ds41_tiny_golden.json`.
6//!
7//! Nothing here runs on a GPU. The GPU attention (`attn_v41`) uses its RoPE
8//! tables, window indices and top-k selections, and `model::forward` is the
9//! oracle the GPU tests compare against.
10//!
11//! # Where the truth comes from
12//!
13//! `ds41_tiny_golden.json` is produced by `gen_ds41_tiny_golden.py` (beside
14//! it), which runs DeepSeek's `inference/model.py` (the revision the golden's
15//! `fixture.reference` records) on a tiny synthetic model: hc_mult=4
16//! hyper-connections with delayed mixes, engram on two layers with the
17//! reference's hash (`engram.py`), shared compressed attention with two kv
18//! sources, three index sources and a candidate prefilter whose consumers
19//! share its ratio, sqrt-softplus routing with a correction bias,
20//! route_scale, the swiglu clamp and a shared expert. The tilelang kernels
21//! the reference imports are replaced by `ds41_ref_shims.py` (pure torch).
22//!
23//! # Determinism, and why only outputs are committed
24//!
25//! Weights and inputs are RNG-free and order-free: `fixed_value(name, i,
26//! scale, offset)` here reproduces the generator's filler (fnv1a64 salt,
27//! splitmix64, 24-bit uniform, f64 arithmetic, f32 result, bf16 RNE where the
28//! parameter is stored bf16). So the golden carries only outputs: each
29//! capture as `{shape, n, stride, ck, data}` where `data` is a prime-strided
30//! sample and `ck` is the fp64 index-weighted checksum over the whole tensor.
31//! `tests.rs` checks the regeneration against the golden.
32//!
33//! # The engram hash parameters come from the GGUF
34//!
35//! Measured 2026-09-15 with `gen_ds41_engram_hash_check.py`: the token map,
36//! bucket primes, slot offsets and hash multipliers that `engram.py`
37//! computes from the V4.1 tokenizer equal the `deepseek41.engram.*` metadata
38//! of the Q2_K GGUF, so the loader builds the engram tables from that
39//! metadata alone (`EngramHashTables::from_flat`).
40//!
41//! Owner: model-arch, DeepSeek-V4.1 reference.
42//! Invariants: none beyond the types.
43
44use serde_json::Value;
45
46/// 2026-09-25: The tracked golden, produced by `gen_ds41_tiny_golden.py`.
47pub const GOLDEN_JSON: &str = include_str!("ds41_tiny_golden.json");
48
49const GOLD: u64 = 0x9E37_79B9_7F4A_7C15;
50
51/// 2026-09-25: The splitmix64 finaliser of `gen_ds41_tiny_golden.py::splitmix64`.
52pub fn splitmix64(z: u64) -> u64 {
53    let mut z = z.wrapping_add(GOLD);
54    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
55    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
56    z ^ (z >> 31)
57}
58
59/// 2026-09-25: FNV-1a 64 over the UTF-8 bytes of `s`; the per-tensor salt.
60pub fn fnv1a64(s: &str) -> u64 {
61    let mut h: u64 = 0xCBF2_9CE4_8422_2325;
62    for b in s.as_bytes() {
63        h ^= *b as u64;
64        h = h.wrapping_mul(0x0000_0100_0000_01B3);
65    }
66    h
67}
68
69/// 2026-09-25: Raw stream value for element `i` of the tensor named `name`.
70pub fn fixed_raw(name: &str, i: u64) -> u64 {
71    splitmix64(fnv1a64(name) ^ i.wrapping_mul(GOLD))
72}
73
74/// 2026-09-25: The generator's filler: `u = (z >> 40) / 2^24`,
75/// `f32((2u - 1) * scale + offset)` in f64.
76pub fn fixed_value(name: &str, i: u64, scale: f64, offset: f64) -> f32 {
77    let u = (fixed_raw(name, i) >> 40) as f64 / (1u64 << 24) as f64;
78    ((2.0 * u - 1.0) * scale + offset) as f32
79}
80
81/// 2026-09-25: The generator's integer filler (`fixed_ints`): `fixed_raw % modulus`.
82pub fn fixed_int(name: &str, i: u64, modulus: u64) -> u64 {
83    fixed_raw(name, i) % modulus
84}
85
86/// 2026-09-25: Round an f32 to bf16 precision, round-to-nearest-even; NaN passes through.
87pub fn to_bf16_rne(x: f32) -> f32 {
88    if x.is_nan() {
89        return x;
90    }
91    let b = x.to_bits();
92    let lsb = (b >> 16) & 1;
93    f32::from_bits(b.wrapping_add(0x7FFF + lsb) & 0xFFFF_0000)
94}
95
96/// 2026-09-25: fp64 index-weighted checksum `sum(v[i] * (i + 1))`, the generator's
97/// `ck`. Summed sequentially here and by a torch reduction there, so callers
98/// compare with a relative tolerance.
99pub fn checksum<I: IntoIterator<Item = f64>>(v: I) -> f64 {
100    v.into_iter()
101        .enumerate()
102        .map(|(i, x)| x * (i as f64 + 1.0))
103        .sum()
104}
105
106/// 2026-09-25: One committed capture.
107pub struct GoldenTensor {
108    pub shape: Vec<usize>,
109    pub n: usize,
110    pub stride: usize,
111    pub ck: f64,
112    pub data: Vec<f64>,
113}
114
115/// 2026-09-25: One parameter's regeneration rule, from `weights_meta`.
116pub struct WeightMeta {
117    pub name: String,
118    pub n: usize,
119    pub scale: f64,
120    pub offset: f64,
121    pub kind: String,
122    pub dtype: String,
123    pub ck: f64,
124}
125
126pub struct Golden(Value);
127
128impl Golden {
129    pub fn load() -> Self {
130        Golden(serde_json::from_str(GOLDEN_JSON).expect("ds41_tiny_golden.json parses"))
131    }
132
133    pub fn regimes(&self) -> Vec<String> {
134        self.0["fixture"]["regimes"]
135            .as_array()
136            .expect("fixture.regimes")
137            .iter()
138            .map(|v| v.as_str().expect("regime name").to_string())
139            .collect()
140    }
141
142    pub fn fixture_u64(&self, key: &str) -> u64 {
143        self.0["fixture"][key]
144            .as_u64()
145            .unwrap_or_else(|| panic!("fixture.{key} is not an integer"))
146    }
147
148    pub fn fixture_f64(&self, key: &str) -> f64 {
149        self.0["fixture"][key]
150            .as_f64()
151            .unwrap_or_else(|| panic!("fixture.{key} is not numeric"))
152    }
153
154    pub fn fixture_i64(&self, key: &str) -> i64 {
155        self.0["fixture"][key]
156            .as_i64()
157            .unwrap_or_else(|| panic!("fixture.{key} is not an integer"))
158    }
159
160    pub fn fixture_usize_list(&self, key: &str) -> Vec<usize> {
161        self.0["fixture"][key]
162            .as_array()
163            .unwrap_or_else(|| panic!("fixture.{key} is not a list"))
164            .iter()
165            .map(|v| v.as_u64().expect("usize") as usize)
166            .collect()
167    }
168
169    pub fn lcg_probe(&self) -> Vec<u64> {
170        self.0["fixture"]["lcg_probe"]
171            .as_array()
172            .expect("fixture.lcg_probe")
173            .iter()
174            .map(|v| v.as_u64().expect("u64 probe"))
175            .collect()
176    }
177
178    pub fn capture_names(&self, regime: &str) -> Vec<String> {
179        self.0[regime]
180            .as_object()
181            .unwrap_or_else(|| panic!("regime {regime} missing"))
182            .keys()
183            .cloned()
184            .collect()
185    }
186
187    pub fn tensor(&self, regime: &str, name: &str) -> GoldenTensor {
188        let t = &self.0[regime][name];
189        assert!(!t.is_null(), "missing capture {regime}.{name}");
190        let as_usize = |k: &str| {
191            t[k].as_u64()
192                .unwrap_or_else(|| panic!("{regime}.{name}.{k}")) as usize
193        };
194        GoldenTensor {
195            shape: t["shape"]
196                .as_array()
197                .expect("shape")
198                .iter()
199                .map(|v| v.as_u64().expect("dim") as usize)
200                .collect(),
201            n: as_usize("n"),
202            stride: as_usize("stride"),
203            ck: t["ck"].as_f64().expect("ck"),
204            data: t["data"]
205                .as_array()
206                .expect("data")
207                .iter()
208                .map(|v| v.as_f64().expect("numeric"))
209                .collect(),
210        }
211    }
212
213    /// 2026-09-25: One parameter's regeneration rule by name; panics when
214    /// `weights_meta` has no such entry.
215    pub fn weight_meta(&self, name: &str) -> WeightMeta {
216        let m = &self.0["weights_meta"][name];
217        assert!(!m.is_null(), "weights_meta has no '{name}'");
218        WeightMeta {
219            name: name.to_string(),
220            n: m["n"].as_u64().expect("n") as usize,
221            scale: m["scale"].as_f64().expect("scale"),
222            offset: m["offset"].as_f64().expect("offset"),
223            kind: m["kind"].as_str().expect("kind").to_string(),
224            dtype: m["dtype"].as_str().expect("dtype").to_string(),
225            ck: m["ck"].as_f64().expect("ck"),
226        }
227    }
228
229    pub fn weights_meta(&self) -> Vec<WeightMeta> {
230        self.0["weights_meta"]
231            .as_object()
232            .expect("weights_meta")
233            .iter()
234            .map(|(name, m)| WeightMeta {
235                name: name.clone(),
236                n: m["n"].as_u64().expect("n") as usize,
237                scale: m["scale"].as_f64().expect("scale"),
238                offset: m["offset"].as_f64().expect("offset"),
239                kind: m["kind"].as_str().expect("kind").to_string(),
240                dtype: m["dtype"].as_str().expect("dtype").to_string(),
241                ck: m["ck"].as_f64().expect("ck"),
242            })
243            .collect()
244    }
245}
246
247/// 2026-09-25: Regenerate an f32-kind parameter (bf16- or f32-stored) as the
248/// generator initialised it: the rule (scale, offset) and the storage dtype
249/// come from `weights_meta`. Panics on another kind; the fp8 tables are
250/// regenerated by `engram::regen_table`.
251pub fn regen_param(g: &Golden, name: &str) -> Vec<f32> {
252    let m = g.weight_meta(name);
253    assert_eq!(
254        m.kind, "f32",
255        "{name}: regen_param handles f32-kind parameters only"
256    );
257    let bf16 = m.dtype == "bfloat16";
258    (0..m.n as u64)
259        .map(|i| {
260            let v = fixed_value(name, i, m.scale, m.offset);
261            if bf16 { to_bf16_rne(v) } else { v }
262        })
263        .collect()
264}
265
266/// 2026-09-25: Elementwise max-abs comparison that names the worst index and both values.
267#[track_caller]
268pub fn assert_close(what: &str, got: &[f64], want: &[f64], tol: f64) {
269    assert_eq!(
270        got.len(),
271        want.len(),
272        "{what}: length {} vs {}",
273        got.len(),
274        want.len()
275    );
276    let mut worst = 0.0f64;
277    let mut at = 0usize;
278    for (i, (g, w)) in got.iter().zip(want).enumerate() {
279        let d = (g - w).abs();
280        if d > worst {
281            worst = d;
282            at = i;
283        }
284    }
285    assert!(
286        worst <= tol,
287        "{what}: max abs diff {worst:e} > {tol:e} at index {at} (got {}, want {})",
288        got[at],
289        want[at]
290    );
291}
292
293pub mod attn;
294pub mod compress;
295pub mod engram;
296pub mod hc;
297pub mod model;
298pub mod moe;
299
300/// 2026-09-25: Shared comparison bars for the component tests.
301#[cfg(test)]
302pub(crate) mod testutil {
303    use super::{GoldenTensor, assert_close, checksum};
304
305    // 2026-09-25: `EXACT` allows only f64 summation-order noise; `BF16_CK_REL`
306    // is 2^-7, the bf16 relative step, for outputs the reference rounds to
307    // bf16; `F32_TOL` is for f32 chains summed in another order (the hc mixes).
308    pub const EXACT: f64 = 1e-12;
309    pub const BF16_CK_REL: f64 = 0.0078125;
310    pub const F32_TOL: f64 = 1e-4;
311
312    /// 2026-09-25: Check the sample elementwise at `tol`, and the whole-tensor
313    /// checksum against `ck_rel` times the index-weighted magnitude, so
314    /// last-bit accumulation-order differences pass.
315    #[track_caller]
316    pub fn check_capture(what: &str, got: &[f64], g: &GoldenTensor, tol: f64, ck_rel: f64) {
317        assert_eq!(got.len(), g.n, "{what}: numel");
318        let sample: Vec<f64> = got.iter().step_by(g.stride).copied().collect();
319        assert_close(what, &sample, &g.data, tol);
320        let ck = checksum(got.iter().copied());
321        let mag: f64 = got
322            .iter()
323            .enumerate()
324            .map(|(i, v)| v.abs() * (i as f64 + 1.0))
325            .sum();
326        let bound = ck_rel * mag.max(1.0);
327        assert!(
328            (ck - g.ck).abs() <= bound,
329            "{what}: checksum {ck} vs golden {} (|diff| {:e} > bound {:e})",
330            g.ck,
331            (ck - g.ck).abs(),
332            bound
333        );
334    }
335
336    pub fn bf16_tol(g: &GoldenTensor) -> f64 {
337        let m = g.data.iter().fold(0f64, |a, v| a.max(v.abs()));
338        2.0 * 2f64.powi(-8) * m + 1e-6
339    }
340
341    /// 2026-09-25: A capture stored at full resolution (stride 1), as f32.
342    pub fn full_f32(g: &GoldenTensor, what: &str) -> Vec<f32> {
343        assert_eq!(g.stride, 1, "{what} must be in FULL_CAPTURES (stride 1)");
344        g.data.iter().map(|&v| v as f32).collect()
345    }
346
347    pub fn as_f64(v: &[f32]) -> Vec<f64> {
348        v.iter().map(|&x| x as f64).collect()
349    }
350}
351
352#[cfg(test)]
353mod tests;
metrale_model_arch/deepseek_v41_ref/mod.rs

metrale_model_arch/deepseek_v41_ref/
mod.rs
