spark_model/layers/deepseek_v41_ref/
mod.rs

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

spark_model/layers/deepseek_v41_ref/
mod.rs
