spark_model/layers/deepseek_v41_ref/
attn.rs

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

spark_model/layers/deepseek_v41_ref/
attn.rs
