metrale_model_arch/deepseek_v41_ref/
attn.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: DeepSeek-V4.1 attention for ratio-0 layers (the reference's
5//! `Attention.forward` at `compress_ratio == 0`), CPU reference against the
6//! golden: low-rank Q with RMSNorm, one latent KV row per token (`wkv` +
7//! `kv_norm`), RoPE on the last `rope_dim` of both, the KV row rounded
8//! through fp8 in place (block 32, power-of-two scale), a sliding-window ring
9//! cache, `sparse_attn` with the attention sink, the inverse rotation on the
10//! output, then the grouped low-rank output projection (`wo_a` per group,
11//! then `wo_b`). Also the RoPE, window and fp8 helpers the compressed layers
12//! and `attn_v41` share.
13//!
14//! Precision: every projection is `linear_bf16` (f32 accumulate, bf16
15//! result), RMSNorm is f32 back to bf16, RoPE multiplies in f32 and writes
16//! bf16, the fp8 round trip is bf16-exact (e4m3 times a power of two), and
17//! `sparse_attn` runs in f32 and returns bf16. The tests compare tensors at
18//! bf16 tolerance and the window indices exactly.
19//!
20//! The ring cache needs a sequential run to check: prefill of the fixture's
21//! 12 tokens into 8 slots seeds `slots[4..8] = tokens 4..8` and `slots[0..4] =
22//! tokens 8..12`; decode step 12 writes slot 4 and reads the whole ring
23//! oldest-first; step 13 writes slot 5. The tests run the three regimes in
24//! that order through one cache and check `sa_kv` (the ring) at each decode step.
25//!
26//! Owner: model-arch, DeepSeek-V4.1 reference.
27//! Invariants: none beyond the types.
28
29use super::engram::{e4m3_to_f32, f32_to_e4m3_rne};
30use super::hc::rms_norm;
31use super::moe::linear_bf16;
32use super::to_bf16_rne;
33
34const FP8_MAX: f32 = 448.0;
35/// 2026-09-25: Values per fp8 scale block.
36pub const FP8_BLOCK: usize = 32;
37
38/// 2026-09-25: `2 ** ceil(log2(v))` for positive normal f32 via the IEEE fields,
39/// as `_pow2_ceil_log2` in `ds41_ref_shims.py`.
40pub fn pow2_ceil(v: f32) -> f32 {
41    let bits = v.to_bits();
42    let e = ((bits >> 23) & 0xFF) as i32 - 127;
43    let man = bits & 0x7F_FFFF;
44    let n = e + if man != 0 { 1 } else { 0 };
45    f32::from_bits(((n + 127) as u32) << 23)
46}
47
48/// 2026-09-25: The shims' `act_quant` with a power-of-two scale, in place: per
49/// block of `FP8_BLOCK`, scale = the power-of-two ceiling of `amax / 448`
50/// (amax floored at 1e-4), values to e4m3 (RNE, clamped), then dequantised
51/// and rounded to bf16.
52pub fn act_quant_inplace(x: &mut [f32]) {
53    for blk in x.chunks_mut(FP8_BLOCK) {
54        let amax = blk.iter().fold(0f32, |a, v| a.max(v.abs())).max(1e-4);
55        let s = pow2_ceil(amax / FP8_MAX);
56        for v in blk.iter_mut() {
57            let q = f32_to_e4m3_rne((*v / s).clamp(-FP8_MAX, FP8_MAX));
58            *v = to_bf16_rne(e4m3_to_f32(q) * s);
59        }
60    }
61}
62
63/// 2026-09-25: The RoPE table without YaRN, which the ratio-0 layers use:
64/// `[seqlen][dim/2]` of (cos, sin).
65pub fn freqs_cis(dim: usize, seqlen: usize, base: f32) -> Vec<(f32, f32)> {
66    let half = dim / 2;
67    let freqs: Vec<f32> = (0..half)
68        .map(|k| 1.0 / base.powf((2 * k) as f32 / dim as f32))
69        .collect();
70    let mut out = Vec::with_capacity(seqlen * half);
71    for p in 0..seqlen {
72        for &f in &freqs {
73            let a = p as f32 * f;
74            out.push((a.cos(), a.sin()));
75        }
76    }
77    out
78}
79
80/// 2026-09-25: Rotate the last `rope_dim` elements of each `row_len`-wide row of
81/// `x`, row `r` at position `pos[r]` (negated angle when `inverse`); adjacent
82/// pairs are complex numbers, multiplied in f32, written back bf16.
83pub fn apply_rotary(
84    x: &mut [f32],
85    row_len: usize,
86    rope_dim: usize,
87    pos: &[usize],
88    fc: &[(f32, f32)],
89    inverse: bool,
90) {
91    let half = rope_dim / 2;
92    for (r, &p) in pos.iter().enumerate() {
93        let row = &mut x[r * row_len + row_len - rope_dim..(r + 1) * row_len];
94        for k in 0..half {
95            let (c, s) = fc[p * half + k];
96            let s = if inverse { -s } else { s };
97            let (a, b) = (row[2 * k], row[2 * k + 1]);
98            row[2 * k] = to_bf16_rne(a * c - b * s);
99            row[2 * k + 1] = to_bf16_rne(a * s + b * c);
100        }
101    }
102}
103
104/// 2026-09-25: Window selection `([queries][cols], cols)` of slot indices, -1 for
105/// a slot holding nothing. At `start_pos == 0`: one causal window per query
106/// into the chunk itself. Otherwise: the whole ring, oldest first, for one query.
107pub fn window_topk_idxs(win: usize, seqlen: usize, start_pos: usize) -> (Vec<i32>, usize) {
108    if start_pos == 0 {
109        let cols = seqlen.min(win);
110        let mut idx = Vec::with_capacity(seqlen * cols);
111        for t in 0..seqlen {
112            let base = t.saturating_sub(win - 1);
113            for j in 0..cols {
114                let i = base + j;
115                idx.push(if i > t { -1 } else { i as i32 });
116            }
117        }
118        (idx, cols)
119    } else {
120        let oldest = start_pos % win + 1;
121        let order = (oldest..win).chain(0..oldest);
122        (
123            order
124                .map(|i| if i > start_pos { -1 } else { i as i32 })
125                .collect(),
126            win,
127        )
128    }
129}
130
131/// 2026-09-25: `sparse_attn` as in `ds41_ref_shims.py`: per query and head,
132/// softmax over the gathered rows plus the sink (denominator only), all in
133/// f32; output bf16. `q`: `[s][h][d]`, `kv`: `[rows][d]`, `idx`: `[s][topk]`
134/// with -1 skipped.
135pub fn sparse_attn(
136    q: &[f32],
137    kv: &[f32],
138    sink: &[f32],
139    idx: &[i32],
140    s: usize,
141    h: usize,
142    d: usize,
143    topk: usize,
144    scale: f32,
145) -> Vec<f32> {
146    let mut o = vec![0f32; s * h * d];
147    for t in 0..s {
148        let ids = &idx[t * topk..(t + 1) * topk];
149        for hh in 0..h {
150            let qv = &q[(t * h + hh) * d..(t * h + hh + 1) * d];
151            let mut scores: Vec<f32> = ids
152                .iter()
153                .map(|&i| {
154                    if i < 0 {
155                        f32::NEG_INFINITY
156                    } else {
157                        let kr = &kv[i as usize * d..(i as usize + 1) * d];
158                        qv.iter().zip(kr).map(|(a, b)| a * b).sum::<f32>() * scale
159                    }
160                })
161                .collect();
162            let m = scores.iter().cloned().fold(sink[hh], f32::max);
163            let mut denom = (sink[hh] - m).exp();
164            for v in scores.iter_mut() {
165                *v = (*v - m).exp();
166                denom += *v;
167            }
168            let ov = &mut o[(t * h + hh) * d..(t * h + hh + 1) * d];
169            for (j, &i) in ids.iter().enumerate() {
170                if i < 0 {
171                    continue;
172                }
173                let p = scores[j] / denom;
174                let kr = &kv[i as usize * d..(i as usize + 1) * d];
175                for dd in 0..d {
176                    ov[dd] += p * kr[dd];
177                }
178            }
179            for v in ov.iter_mut() {
180                *v = to_bf16_rne(*v);
181            }
182        }
183    }
184    o
185}
186
187/// 2026-09-25: The layer's parameters, all bf16 values except the f32 sink;
188/// projections are `[out, in]`.
189pub struct AttnWeights<'a> {
190    pub sink: &'a [f32],
191    pub wq_a: &'a [f32],
192    pub q_norm: &'a [f32],
193    pub wq_b: &'a [f32],
194    pub wkv: &'a [f32],
195    pub kv_norm: &'a [f32],
196    pub wo_a: &'a [f32],
197    pub wo_b: &'a [f32],
198}
199
200pub struct AttnCfg {
201    pub dim: usize,
202    pub n_heads: usize,
203    pub head_dim: usize,
204    pub rope_dim: usize,
205    pub q_rank: usize,
206    pub o_rank: usize,
207    pub groups: usize,
208    pub window: usize,
209    pub eps: f32,
210}
211
212/// 2026-09-25: The sliding-window ring cache, `[window][head_dim]`.
213pub struct WindowCache {
214    pub slots: Vec<f32>,
215}
216
217impl WindowCache {
218    pub fn new(c: &AttnCfg) -> Self {
219        WindowCache {
220            slots: vec![0f32; c.window * c.head_dim],
221        }
222    }
223}
224
225/// 2026-09-25: The intermediates the tests compare, plus the output. `o` is
226/// `sparse_attn`'s result before the inverse rotation, the golden's `sa_o`.
227pub struct AttnRun {
228    pub q: Vec<f32>,
229    pub kv_rows: Vec<f32>,
230    pub idx: Vec<i32>,
231    pub topk: usize,
232    pub o: Vec<f32>,
233    pub out: Vec<f32>,
234}
235
236/// 2026-09-25: The ratio-0 attention forward. `x`: `[seqlen][dim]` bf16 values;
237/// `fc` from [`freqs_cis`] over the maximum sequence length. Writes the new kv
238/// rows into `cache`'s ring.
239pub fn attention(
240    x: &[f32],
241    seqlen: usize,
242    start_pos: usize,
243    w: &AttnWeights,
244    c: &AttnCfg,
245    fc: &[(f32, f32)],
246    cache: &mut WindowCache,
247) -> AttnRun {
248    let (hd, rd, nh) = (c.head_dim, c.rope_dim, c.n_heads);
249    let pos: Vec<usize> = (0..seqlen).map(|t| start_pos + t).collect();
250
251    let qr = rms_norm(
252        &linear_bf16(x, w.wq_a, seqlen, c.dim, c.q_rank),
253        w.q_norm,
254        seqlen,
255        c.q_rank,
256        c.eps,
257    );
258    let mut q = linear_bf16(&qr, w.wq_b, seqlen, c.q_rank, nh * hd);
259    let head_pos: Vec<usize> = pos
260        .iter()
261        .flat_map(|&p| std::iter::repeat_n(p, nh))
262        .collect();
263    apply_rotary(&mut q, hd, rd, &head_pos, fc, false);
264
265    let mut kv = rms_norm(
266        &linear_bf16(x, w.wkv, seqlen, c.dim, hd),
267        w.kv_norm,
268        seqlen,
269        hd,
270        c.eps,
271    );
272    apply_rotary(&mut kv, hd, rd, &pos, fc, false);
273    act_quant_inplace(&mut kv);
274
275    let win = c.window;
276    let kv_rows = if start_pos == 0 {
277        if seqlen <= win {
278            cache.slots[..seqlen * hd].copy_from_slice(&kv);
279        } else {
280            let cutoff = seqlen % win;
281            let tail = &kv[(seqlen - win) * hd..];
282            cache.slots[cutoff * hd..win * hd].copy_from_slice(&tail[..(win - cutoff) * hd]);
283            cache.slots[..cutoff * hd].copy_from_slice(&tail[(win - cutoff) * hd..]);
284        }
285        kv.clone()
286    } else {
287        let slot = start_pos % win;
288        cache.slots[slot * hd..(slot + 1) * hd].copy_from_slice(&kv);
289        cache.slots.clone()
290    };
291    let (idx, topk) = window_topk_idxs(win, seqlen, start_pos);
292
293    let scale = (hd as f32).powf(-0.5);
294    // 2026-09-25: Keep the pre-rotation `sparse_attn` output (`sa_o`) for the
295    // comparison; the inverse rotation runs on a copy.
296    let o_sa = sparse_attn(&q, &kv_rows, w.sink, &idx, seqlen, nh, hd, topk, scale);
297    let mut o = o_sa.clone();
298    apply_rotary(&mut o, hd, rd, &head_pos, fc, true);
299
300    // 2026-09-25: wo_a is [groups][o_rank][heads_per_group * hd], then wo_b.
301    let gw = nh * hd / c.groups;
302    let mut og = vec![0f32; seqlen * c.groups * c.o_rank];
303    for t in 0..seqlen {
304        for g in 0..c.groups {
305            let ov = &o[t * nh * hd + g * gw..t * nh * hd + (g + 1) * gw];
306            for r in 0..c.o_rank {
307                let wr = &w.wo_a[(g * c.o_rank + r) * gw..(g * c.o_rank + r + 1) * gw];
308                og[(t * c.groups + g) * c.o_rank + r] =
309                    to_bf16_rne(ov.iter().zip(wr).map(|(a, b)| a * b).sum());
310            }
311        }
312    }
313    let out = linear_bf16(&og, w.wo_b, seqlen, c.groups * c.o_rank, c.dim);
314    AttnRun {
315        q,
316        kv_rows,
317        idx,
318        topk,
319        o: o_sa,
320        out,
321    }
322}
323
324#[cfg(test)]
325#[path = "attn_tests.rs"]
326mod tests;
metrale_model_arch/deepseek_v41_ref/attn.rs

metrale_model_arch/deepseek_v41_ref/
attn.rs
