spark_model/layers/deepseek_v41_ref/
hc.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! DeepSeek-V4.1 **block scaffold**: hyper-connections (`hc_mixes` with the Sinkhorn split,
5//! `hc_pre`, `hc_post`) and RMSNorm, CPU reference against the golden.
6//!
7//! The residual stream is `hc_mult` parallel copies `[tokens, hc, dim]`. A sub-layer sits
8//! between `hc_pre` (collapse the copies to one input, weighted by a `pre` mix) and `hc_post`
9//! (expand its output back out, mixing the residual in through a doubly-stochastic `comb`).
10//! `hc_mixes` derives all three coefficient sets from the stream itself, one projection of the
11//! flattened `hc * dim` stream normalised by its RMS, then split by `hc_split_sinkhorn`:
12//! `pre = sigmoid(.) + eps`, `post = 2 sigmoid(.)`, `comb` = row softmax + eps, then alternating
13//! column / row normalisations (`sinkhorn_iters` in total).
14//!
15//! The mixes are DELAYED: a block's attention uses the `pre` its predecessor's FFN produced, and
16//! its FFN uses the `pre` its own attention produced (`Block.forward`). The tests reproduce that
17//! order exactly and check every intermediate the golden captured.
18//!
19//! Orientation of `comb`, the one place this is easy to get backwards: with `comb[j][k]`, the
20//! output copy is `k` and the residual copies `j` are summed,
21//! `y[k] = post[k] * x + sum_j comb[j][k] * residual[j]` (the reference's
22//! `sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2)`).
23
24use super::to_bf16_rne;
25
26fn sigmoid(x: f32) -> f32 {
27    1.0 / (1.0 + (-x).exp())
28}
29
30/// `ds41_ref_shims.hc_split_sinkhorn` for one token. `mixes`: `[(2 + hc) * hc]`. Returns
31/// (`pre[hc]`, `post[hc]`, `comb[hc * hc]` row-major `[j][k]`), all f32.
32pub fn split_sinkhorn(
33    mixes: &[f32],
34    scale: &[f32],
35    base: &[f32],
36    hc: usize,
37    iters: usize,
38    eps: f32,
39) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
40    let pre: Vec<f32> = (0..hc)
41        .map(|j| sigmoid(mixes[j] * scale[0] + base[j]) + eps)
42        .collect();
43    let post: Vec<f32> = (0..hc)
44        .map(|j| 2.0 * sigmoid(mixes[hc + j] * scale[1] + base[hc + j]))
45        .collect();
46    let mut comb: Vec<f32> = (0..hc * hc)
47        .map(|i| mixes[2 * hc + i] * scale[2] + base[2 * hc + i])
48        .collect();
49    // row softmax + eps
50    for j in 0..hc {
51        let row = &mut comb[j * hc..(j + 1) * hc];
52        let m = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
53        let mut sum = 0f32;
54        for v in row.iter_mut() {
55            *v = (*v - m).exp();
56            sum += *v;
57        }
58        for v in row.iter_mut() {
59            *v = *v / sum + eps;
60        }
61    }
62    let col_norm = |c: &mut [f32]| {
63        for k in 0..hc {
64            let s: f32 = (0..hc).map(|j| c[j * hc + k]).sum();
65            for j in 0..hc {
66                c[j * hc + k] /= s + eps;
67            }
68        }
69    };
70    let row_norm = |c: &mut [f32]| {
71        for j in 0..hc {
72            let s: f32 = c[j * hc..(j + 1) * hc].iter().sum();
73            for v in &mut c[j * hc..(j + 1) * hc] {
74                *v /= s + eps;
75            }
76        }
77    };
78    col_norm(&mut comb);
79    for _ in 1..iters {
80        row_norm(&mut comb);
81        col_norm(&mut comb);
82    }
83    (pre, post, comb)
84}
85
86/// `Block.hc_mixes`. `x`: `[tokens, hc, dim]` (bf16 values); `hc_fn`: `[(2+hc)*hc, hc*dim]` f32.
87/// Returns (`pre[tokens, hc]`, `post[tokens, hc]`, `comb[tokens, hc, hc]`).
88pub fn hc_mixes(
89    x: &[f32],
90    tokens: usize,
91    hc: usize,
92    dim: usize,
93    hc_fn: &[f32],
94    scale: &[f32],
95    base: &[f32],
96    iters: usize,
97    hc_eps: f32,
98    norm_eps: f32,
99) -> (Vec<f32>, Vec<f32>, Vec<f32>) {
100    let n = hc * dim;
101    let mix_hc = (2 + hc) * hc;
102    let (mut pre, mut post, mut comb) = (
103        Vec::with_capacity(tokens * hc),
104        Vec::with_capacity(tokens * hc),
105        Vec::with_capacity(tokens * hc * hc),
106    );
107    for t in 0..tokens {
108        let xf = &x[t * n..(t + 1) * n];
109        // one RMS over the whole flattened hc*dim stream, per token
110        let rsqrt = (xf.iter().map(|v| v * v).sum::<f32>() / n as f32 + norm_eps)
111            .sqrt()
112            .recip();
113        let mixes: Vec<f32> = (0..mix_hc)
114            .map(|m| {
115                hc_fn[m * n..(m + 1) * n]
116                    .iter()
117                    .zip(xf)
118                    .map(|(a, b)| a * b)
119                    .sum::<f32>()
120                    * rsqrt
121            })
122            .collect();
123        let (p, q, c) = split_sinkhorn(&mixes, scale, base, hc, iters, hc_eps);
124        pre.extend(p);
125        post.extend(q);
126        comb.extend(c);
127    }
128    (pre, post, comb)
129}
130
131/// `Block.hc_pre`: collapse the copies, `y[d] = sum_c pre[c] * x[c][d]`, back to bf16.
132pub fn hc_pre(x: &[f32], pre: &[f32], tokens: usize, hc: usize, dim: usize) -> Vec<f32> {
133    let mut y = vec![0f32; tokens * dim];
134    for t in 0..tokens {
135        for d in 0..dim {
136            let mut acc = 0f32;
137            for c in 0..hc {
138                acc += pre[t * hc + c] * x[(t * hc + c) * dim + d];
139            }
140            y[t * dim + d] = to_bf16_rne(acc);
141        }
142    }
143    y
144}
145
146/// `Block.hc_post`: `y[k][d] = post[k] * x[d] + sum_j comb[j][k] * residual[j][d]`, to bf16.
147pub fn hc_post(
148    x: &[f32],
149    residual: &[f32],
150    post: &[f32],
151    comb: &[f32],
152    tokens: usize,
153    hc: usize,
154    dim: usize,
155) -> Vec<f32> {
156    let mut y = vec![0f32; tokens * hc * dim];
157    for t in 0..tokens {
158        for k in 0..hc {
159            for d in 0..dim {
160                let mut acc = post[t * hc + k] * x[t * dim + d];
161                for j in 0..hc {
162                    acc += comb[(t * hc + j) * hc + k] * residual[(t * hc + j) * dim + d];
163                }
164                y[(t * hc + k) * dim + d] = to_bf16_rne(acc);
165            }
166        }
167    }
168    y
169}
170
171/// `RMSNorm.forward`: `w * (x * rsqrt(mean(x^2) + eps))` in f32, back to bf16.
172pub fn rms_norm(x: &[f32], w: &[f32], tokens: usize, dim: usize, eps: f32) -> Vec<f32> {
173    let mut y = vec![0f32; tokens * dim];
174    for t in 0..tokens {
175        let row = &x[t * dim..(t + 1) * dim];
176        let rsqrt = (row.iter().map(|v| v * v).sum::<f32>() / dim as f32 + eps)
177            .sqrt()
178            .recip();
179        for d in 0..dim {
180            y[t * dim + d] = to_bf16_rne(w[d] * (row[d] * rsqrt));
181        }
182    }
183    y
184}
185
186#[cfg(test)]
187#[path = "hc_tests.rs"]
188mod tests;
spark_model/layers/deepseek_v41_ref/hc.rs

spark_model/layers/deepseek_v41_ref/
hc.rs
