ferrox_core/gdn.rs
1//! The gated delta-net step (Qwen3-Next / Qwen3.5 linear attention), as
2//! llama.cpp's autoregressive graph computes it
3//! (`delta-net-base.cpp:289-365`, `build_delta_net_autoregressive`).
4//!
5//! Per token, per V head `h` reading K head `kh`:
6//!
7//! ```text
8//! q' = q / sqrt(S) :319-321
9//! S = S * exp(g_h) :339-340 g_h <= 0, one scalar per head
10//! pred[j] = sum_i S[j][i] * k[i] :343-345 ("sk", the state's guess at v)
11//! d[j] = (v[j] - pred[j]) * beta_h :348-350
12//! S[j][i] += k[i] * d[j] :357-361
13//! o[j] = sum_i S[j][i] * q'[i] :362-363
14//! ```
15//!
16//! with the state `[n_v_heads][S][S]` as `S[j][i]`, `i` the key index
17//! innermost (ggml's `{S_v, S_v, H_v}` with `ne0` the key dim: `:344`
18//! multiplies `k` along `ne0`, and the output at `:362` multiplies `q`
19//! along it too). The chunked prefill kernel (`:16-287`) is the same
20//! recurrence in a different summation order.
21//!
22//! Which K head a V head reads is the caller's: `Qwen3.5` tiles
23//! (`h % n_k_heads`, `llama-model.cpp:524-526`, the converter having
24//! reordered V heads for `ggml_repeat`), `Qwen3-Next` groups
25//! (`h / (n_v / n_k)`). [`HeadMap`] names the two so the wrong one
26//! cannot be assumed.
27//!
28//! This file is the arithmetic only; `ferrox_models::gdn` owns the
29//! projections, the conv, the norms and the gates around it.
30
31/// How V head `h` finds its K head.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum HeadMap {
34 /// `h % n_k_heads` (Qwen3.5: `llama-model.cpp:526`, "k0_v0, k1_v1,
35 /// k0_v2, k1_v3").
36 Tiled,
37 /// `h / (n_v_heads / n_k_heads)` (Qwen3-Next: `:525`, "k0_v0, k0_v1,
38 /// k1_v2, k1_v3").
39 Grouped,
40}
41
42impl HeadMap {
43 pub fn k_head(self, v_head: usize, n_k_heads: usize, n_v_heads: usize) -> usize {
44 match self {
45 HeadMap::Tiled => v_head % n_k_heads,
46 HeadMap::Grouped => v_head / (n_v_heads / n_k_heads),
47 }
48 }
49}
50
51/// The geometry one step needs. `head_dim` is both the key and the
52/// value width: the autoregressive graph asserts `S_k == S_v` (`:307`).
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub struct DeltaDims {
55 pub n_k_heads: usize,
56 pub n_v_heads: usize,
57 pub head_dim: usize,
58 pub map: HeadMap,
59}
60
61impl DeltaDims {
62 /// Floats in one sequence's state.
63 pub fn state_len(self) -> usize {
64 self.n_v_heads * self.head_dim * self.head_dim
65 }
66}
67
68/// `x / max(||x||, eps)`, ggml's `ggml_l2_norm` (ops.cpp:4185-4210): the
69/// sum of squares in double, the divisor clamped by `eps` from below.
70pub fn l2_normalize(x: &mut [f32], eps: f32) {
71 let sum: f64 = x.iter().map(|v| (*v as f64) * (*v as f64)).sum();
72 let scale = 1.0 / (sum as f32).sqrt().max(eps);
73 for v in x.iter_mut() {
74 *v *= scale;
75 }
76}
77
78/// One token of the delta rule, in place on `state`
79/// (`[n_v_heads][head_dim][head_dim]`, value index outer, key index
80/// inner).
81///
82/// `q` and `k` are `[n_k_heads][head_dim]` (already l2-normed; `q` is
83/// scaled by `1/sqrt(head_dim)` HERE, `:319`), `v` is
84/// `[n_v_heads][head_dim]`, `g` is `[n_v_heads]` (the log decay, the
85/// `exp` taken here, `:339`), `beta` is `[n_v_heads]` (after the
86/// sigmoid), `out` is `[n_v_heads][head_dim]`.
87#[allow(clippy::too_many_arguments)] // the six operands ggml_gated_delta_net takes, plus the dims and the output
88pub fn delta_step(
89 dims: DeltaDims,
90 state: &mut [f32],
91 q: &[f32],
92 k: &[f32],
93 v: &[f32],
94 g: &[f32],
95 beta: &[f32],
96 out: &mut [f32],
97) {
98 let DeltaDims {
99 n_k_heads,
100 n_v_heads,
101 head_dim: s,
102 map,
103 } = dims;
104 assert_eq!(state.len(), dims.state_len());
105 assert_eq!(q.len(), n_k_heads * s);
106 assert_eq!(k.len(), n_k_heads * s);
107 assert_eq!(v.len(), n_v_heads * s);
108 assert_eq!(g.len(), n_v_heads);
109 assert_eq!(beta.len(), n_v_heads);
110 assert_eq!(out.len(), n_v_heads * s);
111 assert!(n_k_heads > 0 && n_v_heads.is_multiple_of(n_k_heads), ":308");
112 let scale = 1.0 / (s as f32).sqrt();
113 // Heads are independent (each owns an `S x S` state and `S`
114 // outputs), so they run as one parallel region over the head axis;
115 // a serial version of this loop was 15% of a Bonsai-2-27B decode
116 // step (48 layers x 48 heads x 128 x 128 state floats per token).
117 // The two inner reductions are written over four accumulators so
118 // the compiler can vectorise them (a single-accumulator float sum
119 // cannot be reordered without fast-math).
120 crate::par::chunks_mut2_by(state, out, s * s, s, 1, |h, st, oh| {
121 let kh = map.k_head(h, n_k_heads, n_v_heads);
122 let (qh, kk) = (&q[kh * s..(kh + 1) * s], &k[kh * s..(kh + 1) * s]);
123 let vh = &v[h * s..(h + 1) * s];
124 let decay = g[h].exp();
125 let mut d = vec![0.0f32; s];
126 // :340 then :343-350: decay, the state's prediction, the error
127 // scaled by beta.
128 for j in 0..s {
129 let row = &mut st[j * s..(j + 1) * s];
130 let pred = decay_and_dot(row, kk, decay);
131 d[j] = (vh[j] - pred) * beta[h];
132 }
133 // :357-363: the rank-one update, then the read-out with the
134 // scaled query.
135 for j in 0..s {
136 let row = &mut st[j * s..(j + 1) * s];
137 oh[j] = update_and_dot(row, kk, d[j], qh) * scale;
138 }
139 });
140}
141
142/// `row *= decay`, then `row . k`, over four lanes of accumulation.
143#[inline]
144fn decay_and_dot(row: &mut [f32], k: &[f32], decay: f32) -> f32 {
145 let mut acc = [0.0f32; 4];
146 let (rb, rt) = row.as_chunks_mut::<4>();
147 let (kb, kt) = k.as_chunks::<4>();
148 for (r, kk) in rb.iter_mut().zip(kb) {
149 for l in 0..4 {
150 r[l] *= decay;
151 acc[l] += r[l] * kk[l];
152 }
153 }
154 let mut tail = 0.0f32;
155 for (r, kk) in rt.iter_mut().zip(kt) {
156 *r *= decay;
157 tail += *r * *kk;
158 }
159 acc[0] + acc[1] + acc[2] + acc[3] + tail
160}
161
162/// `row += k * d`, then `row . q`, over four lanes of accumulation.
163#[inline]
164fn update_and_dot(row: &mut [f32], k: &[f32], d: f32, q: &[f32]) -> f32 {
165 let mut acc = [0.0f32; 4];
166 let (rb, rt) = row.as_chunks_mut::<4>();
167 let (kb, kt) = k.as_chunks::<4>();
168 let (qb, qt) = q.as_chunks::<4>();
169 for ((r, kk), qq) in rb.iter_mut().zip(kb).zip(qb) {
170 for l in 0..4 {
171 r[l] += kk[l] * d;
172 acc[l] += r[l] * qq[l];
173 }
174 }
175 let mut tail = 0.0f32;
176 for ((r, kk), qq) in rt.iter_mut().zip(kt).zip(qt) {
177 *r += *kk * d;
178 tail += *r * *qq;
179 }
180 acc[0] + acc[1] + acc[2] + acc[3] + tail
181}
182
183#[cfg(test)]
184mod tests {
185 use super::*;
186
187 #[test]
188 fn l2_normalize_clamps_the_divisor() {
189 let mut v = [3.0f32, 4.0];
190 l2_normalize(&mut v, 1e-6);
191 assert!((v[0] - 0.6).abs() < 1e-6 && (v[1] - 0.8).abs() < 1e-6);
192 let mut z = [0.0f32, 0.0];
193 l2_normalize(&mut z, 1e-6);
194 assert_eq!(z, [0.0, 0.0]);
195 }
196
197 /// One head, `S = 1`: the scalar recurrence by hand.
198 #[test]
199 fn the_scalar_recurrence() {
200 let dims = DeltaDims {
201 n_k_heads: 1,
202 n_v_heads: 1,
203 head_dim: 1,
204 map: HeadMap::Tiled,
205 };
206 let mut st = vec![2.0f32];
207 let mut out = [0.0f32];
208 // decay exp(0) = 1; pred = 2 * k(1) = 2; d = (v(5) - 2) * beta(0.5) = 1.5;
209 // S = 2 + 1 * 1.5 = 3.5; o = 3.5 * q(1) * 1.
210 delta_step(
211 dims,
212 &mut st,
213 &[1.0],
214 &[1.0],
215 &[5.0],
216 &[0.0],
217 &[0.5],
218 &mut out,
219 );
220 assert!((st[0] - 3.5).abs() < 1e-6 && (out[0] - 3.5).abs() < 1e-6);
221 }
222
223 /// Two K heads, four V heads: tiled reads `h % 2`, grouped `h / 2`.
224 #[test]
225 fn the_two_head_maps_differ_and_are_the_documented_ones() {
226 assert_eq!(HeadMap::Tiled.k_head(3, 2, 4), 1);
227 assert_eq!(HeadMap::Grouped.k_head(3, 2, 4), 1);
228 assert_eq!(HeadMap::Tiled.k_head(1, 2, 4), 1);
229 assert_eq!(HeadMap::Grouped.k_head(1, 2, 4), 0);
230 let mk = |map| DeltaDims {
231 n_k_heads: 2,
232 n_v_heads: 4,
233 head_dim: 1,
234 map,
235 };
236 let (q, k) = ([1.0f32, 1.0], [1.0f32, 10.0]);
237 let v = [1.0f32; 4];
238 let mut out_t = [0.0f32; 4];
239 let mut out_g = [0.0f32; 4];
240 delta_step(
241 mk(HeadMap::Tiled),
242 &mut [0.0; 4],
243 &q,
244 &k,
245 &v,
246 &[0.0; 4],
247 &[1.0; 4],
248 &mut out_t,
249 );
250 delta_step(
251 mk(HeadMap::Grouped),
252 &mut [0.0; 4],
253 &q,
254 &k,
255 &v,
256 &[0.0; 4],
257 &[1.0; 4],
258 &mut out_g,
259 );
260 // S = k * v after one step (from zero): o = k * q.
261 assert_eq!(out_t, [1.0, 10.0, 1.0, 10.0]);
262 assert_eq!(out_g, [1.0, 1.0, 10.0, 10.0]);
263 }
264}