ferrox_models/
hyper_connections.rs

1//! DeepSeek V4's "mHC" (multi-stream Hyper-Connection) residual mixing:
2//! instead of one residual stream, the model carries `hc` (`hc_mult`,
3//! real reference value 4) parallel streams per token, and each
4//! sub-layer (attention, FFN) is preceded by a learned, per-token gated
5//! merge of those streams into one input, followed by a Sinkhorn-
6//! normalized mix back into all `hc` streams.
7//!
8//! Transcribed directly from the real, merged reference implementation
9//! (llama.cpp PR #24162, `src/models/deepseek4.cpp`:
10//! `build_hc_pre`/`build_hc_post`/`build_hc_head`/`build_hc_sinkhorn`/
11//! `build_hc_weighted_sum`, read line-by-line), not derived by analogy
12//! -- this closes the "mHC's exact math was not read" gap from earlier
13//! research. The real implementation asserts `hc == 4` in `build_hc_pre`
14//! (the mix-tensor offset layout is hardcoded to that split), so this
15//! module does too rather than silently pretending to support other
16//! values.
17//!
18//! Not yet wired into a DeepSeek V4 decoder (no such decoder exists in
19//! ferrox yet) -- this is the residual-mixing primitive on its own,
20//! analogous to how `mla.rs`/`block_residual.rs` exist as standalone,
21//! tested modules before Kimi K3's decoder consumed them.
22
23use ferrox_core::weight_matrix::WeightMatrix;
24
25/// The real reference calls bare `ggml_rms_norm` here (no learned
26/// per-element scale, unlike a normal transformer RMSNorm layer) --
27/// this hyper-connection router computation has no weight tensor for
28/// it in `deepseek4.cpp`.
29fn rms_norm_no_weight(x: &[f32], eps: f32) -> Vec<f32> {
30    let mean_sq = x.iter().map(|v| v * v).sum::<f32>() / x.len() as f32;
31    let scale = 1.0 / (mean_sq + eps).sqrt();
32    x.iter().map(|v| v * scale).collect()
33}
34
35/// The real reference implementation's only supported hyper-connection
36/// multiplicity; `build_hc_pre`'s mix-tensor offsets are hardcoded to
37/// this split (`GGML_ASSERT(hc == 4)` in the real source).
38pub const HC_MULT: usize = 4;
39
40/// Weights for the pre-sub-layer merge (`build_hc_pre`): projects the
41/// flattened, RMS-normed `hc` streams to `(2 + hc) * hc` mix logits,
42/// split into `pre` (hc), `post` (hc), and `comb` (hc*hc).
43pub struct HyperConnectionPreWeights {
44    pub fn_proj: WeightMatrix, // [(2+hc)*hc, hc*n_embd]
45    /// `[scale_pre, scale_post, scale_comb]`, real tensor shape `{3}`.
46    pub scale: [f32; 3],
47    pub base_pre: [f32; HC_MULT],
48    pub base_post: [f32; HC_MULT],
49    pub base_comb: [f32; HC_MULT * HC_MULT],
50}
51
52/// Weights for the final output merge (`build_hc_head`): same
53/// structure as the `pre` half of [`HyperConnectionPreWeights`], but
54/// only ever produces the `hc`-wide merge gate (no `post`/`comb`,
55/// since there is no further sub-layer to re-inject into).
56pub struct HyperConnectionHeadWeights {
57    pub fn_proj: WeightMatrix, // [hc, hc*n_embd]
58    pub scale: f32,
59    pub base: [f32; HC_MULT],
60}
61
62fn sigmoid(x: f32) -> f32 {
63    1.0 / (1.0 + (-x).exp())
64}
65
66/// `build_hc_weighted_sum`: per-token weighted sum of the `hc` stream
67/// vectors (each `n_embd`-wide), `sum_h(x[h] * weights[h])`.
68pub fn weighted_sum(x: &[Vec<f32>; HC_MULT], weights: &[f32; HC_MULT]) -> Vec<f32> {
69    let n_embd = x[0].len();
70    let mut out = vec![0f32; n_embd];
71    for (xh, &w) in x.iter().zip(weights.iter()) {
72        for (o, v) in out.iter_mut().zip(xh.iter()) {
73            *o += v * w;
74        }
75    }
76    out
77}
78
79/// `build_hc_sinkhorn`: `comb[dst][src]`, real algorithm -- softmax
80/// over `dst` (per fixed `src`), `+eps`, one row-normalization (each
81/// `dst` row sums to 1 over `src`), then `sinkhorn_iters - 1` rounds of
82/// [column-normalize (each `src` column sums to 1 over `dst`),
83/// row-normalize].
84#[allow(clippy::needless_range_loop)]
85pub fn sinkhorn(comb: &mut [[f32; HC_MULT]; HC_MULT], iters: u32, eps: f32) {
86    for src in 0..HC_MULT {
87        let mut col = [0f32; HC_MULT];
88        for dst in 0..HC_MULT {
89            col[dst] = comb[dst][src];
90        }
91        let max = col.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
92        let mut sum = 0f32;
93        for v in col.iter_mut() {
94            *v = (*v - max).exp();
95            sum += *v;
96        }
97        for v in col.iter_mut() {
98            *v /= sum;
99        }
100        for dst in 0..HC_MULT {
101            comb[dst][src] = col[dst] + eps;
102        }
103    }
104
105    let norm_rows_over_src = |comb: &mut [[f32; HC_MULT]; HC_MULT]| {
106        for dst in 0..HC_MULT {
107            let sum: f32 = comb[dst].iter().sum::<f32>() + eps;
108            for src in 0..HC_MULT {
109                comb[dst][src] /= sum;
110            }
111        }
112    };
113    let norm_cols_over_dst = |comb: &mut [[f32; HC_MULT]; HC_MULT]| {
114        for src in 0..HC_MULT {
115            let sum: f32 = (0..HC_MULT).map(|dst| comb[dst][src]).sum::<f32>() + eps;
116            for dst in 0..HC_MULT {
117                comb[dst][src] /= sum;
118            }
119        }
120    };
121
122    norm_rows_over_src(comb);
123    for _ in 1..iters {
124        norm_cols_over_dst(comb);
125        norm_rows_over_src(comb);
126    }
127}
128
129/// `build_hc_pre`: merges `hc` residual streams into one sub-layer
130/// input, plus the `post` gate and Sinkhorn-normalized `comb` matrix
131/// needed by [`post`] afterward. Returns `(merged_input, post_gate,
132/// comb_matrix)`.
133#[allow(clippy::type_complexity, clippy::needless_range_loop)]
134pub fn pre(
135    weights: &HyperConnectionPreWeights,
136    x: &[Vec<f32>; HC_MULT],
137    rms_norm_eps: f32,
138    sinkhorn_iters: u32,
139    hc_eps: f32,
140) -> (Vec<f32>, [f32; HC_MULT], [[f32; HC_MULT]; HC_MULT]) {
141    let n_embd = x[0].len();
142    let mut flat = Vec::with_capacity(HC_MULT * n_embd);
143    for xh in x {
144        flat.extend_from_slice(xh);
145    }
146    let flat_norm = rms_norm_no_weight(&flat, rms_norm_eps);
147    let mixes = weights.fn_proj.apply(&flat_norm); // [(2+hc)*hc]
148
149    let mut pre_gate = [0f32; HC_MULT];
150    for h in 0..HC_MULT {
151        pre_gate[h] = sigmoid(mixes[h] * weights.scale[0] + weights.base_pre[h]) + hc_eps;
152    }
153
154    let mut post_gate = [0f32; HC_MULT];
155    for h in 0..HC_MULT {
156        post_gate[h] = sigmoid(mixes[HC_MULT + h] * weights.scale[1] + weights.base_post[h]) * 2.0;
157    }
158
159    let mut comb = [[0f32; HC_MULT]; HC_MULT];
160    for dst in 0..HC_MULT {
161        for src in 0..HC_MULT {
162            let idx = dst * HC_MULT + src;
163            comb[dst][src] = mixes[2 * HC_MULT + idx] * weights.scale[2] + weights.base_comb[idx];
164        }
165    }
166    sinkhorn(&mut comb, sinkhorn_iters, hc_eps);
167
168    let merged = weighted_sum(x, &pre_gate);
169    (merged, post_gate, comb)
170}
171
172/// `build_hc_post`: re-injects the sub-layer's single output back into
173/// all `hc` streams, each scaled by its `post` gate and mixed with the
174/// original (pre-merge) residual streams via the Sinkhorn `comb`
175/// matrix: `out[dst] = sub_layer_out * post[dst] + sum_src(residual[src]
176/// * comb[dst][src])`.
177pub fn post(
178    sub_layer_out: &[f32],
179    residual: &[Vec<f32>; HC_MULT],
180    post_gate: &[f32; HC_MULT],
181    comb: &[[f32; HC_MULT]; HC_MULT],
182) -> [Vec<f32>; HC_MULT] {
183    let n_embd = sub_layer_out.len();
184    std::array::from_fn(|dst| {
185        let mut out = vec![0f32; n_embd];
186        for (o, v) in out.iter_mut().zip(sub_layer_out.iter()) {
187            *o = v * post_gate[dst];
188        }
189        for src in 0..HC_MULT {
190            let w = comb[dst][src];
191            for (o, v) in out.iter_mut().zip(residual[src].iter()) {
192                *o += v * w;
193            }
194        }
195        out
196    })
197}
198
199/// `build_hc_head`: the final collapse of `hc` streams into one output
200/// vector before unembedding -- structurally the `pre`-gate half of
201/// [`pre`], with no `post`/`comb` since nothing follows it.
202pub fn head(
203    weights: &HyperConnectionHeadWeights,
204    x: &[Vec<f32>; HC_MULT],
205    rms_norm_eps: f32,
206    hc_eps: f32,
207) -> Vec<f32> {
208    let n_embd = x[0].len();
209    let mut flat = Vec::with_capacity(HC_MULT * n_embd);
210    for xh in x {
211        flat.extend_from_slice(xh);
212    }
213    let flat_norm = rms_norm_no_weight(&flat, rms_norm_eps);
214    let mixes = weights.fn_proj.apply(&flat_norm); // [hc]
215
216    let mut gate = [0f32; HC_MULT];
217    for h in 0..HC_MULT {
218        gate[h] = sigmoid(mixes[h] * weights.scale + weights.base[h]) + hc_eps;
219    }
220    weighted_sum(x, &gate)
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226    use ferrox_core::tensor::Tensor;
227
228    fn wm(data: &[f32], rows: usize, cols: usize) -> WeightMatrix {
229        assert_eq!(data.len(), rows * cols);
230        WeightMatrix::F32(Tensor::new(data.to_vec(), vec![rows, cols]))
231    }
232
233    #[test]
234    fn weighted_sum_with_one_hot_weights_selects_that_stream() {
235        let x: [Vec<f32>; HC_MULT] = [
236            vec![1.0, 2.0],
237            vec![10.0, 20.0],
238            vec![100.0, 200.0],
239            vec![1000.0, 2000.0],
240        ];
241        let weights = [0.0, 1.0, 0.0, 0.0];
242        let out = weighted_sum(&x, &weights);
243        assert_eq!(out, vec![10.0, 20.0]);
244    }
245
246    #[test]
247    fn sinkhorn_output_rows_and_columns_are_plausibly_normalized() {
248        let mut comb = [
249            [1.0, 0.5, 0.2, 0.1],
250            [0.3, 1.2, 0.4, 0.2],
251            [0.1, 0.2, 1.5, 0.3],
252            [0.2, 0.1, 0.3, 1.1],
253        ];
254        sinkhorn(&mut comb, 3, 1e-6);
255        for row in comb.iter() {
256            for &v in row.iter() {
257                assert!(v.is_finite() && v >= 0.0);
258            }
259        }
260    }
261
262    #[test]
263    fn pre_then_post_round_trip_preserves_finite_output_and_stream_count() {
264        let n_embd = 3;
265        let hc_mix_dim = (2 + HC_MULT) * HC_MULT;
266        let fn_proj = wm(
267            &vec![0.05; hc_mix_dim * (HC_MULT * n_embd)],
268            hc_mix_dim,
269            HC_MULT * n_embd,
270        );
271        let weights = HyperConnectionPreWeights {
272            fn_proj,
273            scale: [1.0, 1.0, 1.0],
274            base_pre: [0.0; HC_MULT],
275            base_post: [0.0; HC_MULT],
276            base_comb: [0.0; HC_MULT * HC_MULT],
277        };
278        let x: [Vec<f32>; HC_MULT] = std::array::from_fn(|h| vec![(h + 1) as f32; n_embd]);
279
280        let (merged, post_gate, comb) = pre(&weights, &x, 1e-5, 3, 1e-6);
281        assert_eq!(merged.len(), n_embd);
282        assert!(merged.iter().all(|v| v.is_finite()));
283
284        // Pretend the sub-layer is identity for this round-trip check.
285        let streams = post(&merged, &x, &post_gate, &comb);
286        assert_eq!(streams.len(), HC_MULT);
287        for s in streams.iter() {
288            assert_eq!(s.len(), n_embd);
289            assert!(s.iter().all(|v| v.is_finite()));
290        }
291    }
292
293    #[test]
294    fn head_collapses_streams_to_a_single_finite_vector() {
295        let n_embd = 3;
296        let fn_proj = wm(
297            &vec![0.05; HC_MULT * (HC_MULT * n_embd)],
298            HC_MULT,
299            HC_MULT * n_embd,
300        );
301        let weights = HyperConnectionHeadWeights {
302            fn_proj,
303            scale: 1.0,
304            base: [0.0; HC_MULT],
305        };
306        let x: [Vec<f32>; HC_MULT] = std::array::from_fn(|h| vec![(h + 1) as f32; n_embd]);
307        let out = head(&weights, &x, 1e-5, 1e-6);
308        assert_eq!(out.len(), n_embd);
309        assert!(out.iter().all(|v| v.is_finite()));
310    }
311}
ferrox_models/hyper_connections.rs

ferrox_models/
hyper_connections.rs
