metrale_model_arch/deepseek_v41_ref/
hc.rs

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

metrale_model_arch/deepseek_v41_ref/
hc.rs
