spark_model/layers/deepseek_v41_ref/
compress.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! DeepSeek-V4.1 **compressed attention** (`compress_ratio > 0`): the `Compressor` (softmax
5//! pooling of `ratio` tokens into one latent, or a plain projection at ratio 1), the `Indexer`
6//! (fp4 query heads against one shared fp4 key per compressed position, rectified scores combined
7//! by `weights_proj`), `select_candidate_blocks` (level one of the two-level top-k), the
8//! `SharedAttentionRuntime` slots layers hand down the stack, YaRN RoPE with
9//! `compress_rope_theta`, the two fp4 quantisers, and the full `Attention.forward` for any ratio.
10//! CPU reference against the golden; the ratio-0 pieces are shared with `attn.rs`.
11//!
12//! Precision follows the reference: the ratio-2 compressor is f32 end to end until its RMSNorm
13//! casts to bf16; the indexer's q and k go through the fp4 e2m1 round trip (e8m0 scales, block
14//! 32) and its scores are bf16 tensors; the compressed KV latent goes through e2m1 with e4m3
15//! scales (block 16). Index selections are compared exactly, tensors at bf16 tolerance.
16
17use super::attn::{act_quant_inplace, apply_rotary, pow2_ceil, sparse_attn, window_topk_idxs};
18use super::engram::{e4m3_to_f32, f32_to_e4m3_rne};
19use super::hc::rms_norm;
20use super::moe::linear_bf16;
21use super::to_bf16_rne;
22
23mod compressor;
24mod indexer;
25
26pub use compressor::{CompressorState, CompressorWeights, compressor};
27pub use indexer::{
28    IndexerCfg, IndexerWeights, SharedRuntime, indexer, select_candidate_blocks, torch_cpu_topk_set,
29};
30
31/// `fp4_block_size` (model.py:28).
32pub const FP4_BLOCK: usize = 32;
33/// Block size the compressed KV latent is quantised with (`fp4_act_quant(latent, 16, ...)`).
34pub const LATENT_BLOCK: usize = 16;
35const FP4_MAX: f32 = 6.0;
36const E2M1: [f32; 8] = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0];
37
38/// Round onto the e2m1 grid with ties to the even code, keeping the sign (`_to_e2m1_rne`).
39pub fn to_e2m1_rne(y: f32) -> f32 {
40    let a = y.abs();
41    let mut lo = 0usize;
42    for (i, &v) in E2M1.iter().enumerate() {
43        if v <= a {
44            lo = i;
45        }
46    }
47    let hi = (lo + 1).min(7);
48    let (dlo, dhi) = (a - E2M1[lo], E2M1[hi] - a);
49    let pick_hi = dhi < dlo || (dhi == dlo && hi.is_multiple_of(2) && lo % 2 == 1);
50    let v = if pick_hi { E2M1[hi] } else { E2M1[lo] };
51    v.copysign(y)
52}
53
54/// `fp4_act_quant(x, block, inplace=True)` with e8m0 scales: scale = pow2 ceiling of
55/// `amax / 6` (amax floored at `6 * 2^-126`), values to e2m1, dequantised back to bf16.
56pub fn fp4_quant_e8m0_inplace(x: &mut [f32], block: usize) {
57    for blk in x.chunks_mut(block) {
58        let amax = blk
59            .iter()
60            .fold(0f32, |a, v| a.max(v.abs()))
61            .max(6.0 * 2f32.powi(-126));
62        let s = pow2_ceil(amax / FP4_MAX);
63        for v in blk.iter_mut() {
64            *v = to_bf16_rne(to_e2m1_rne((*v / s).clamp(-FP4_MAX, FP4_MAX)) * s);
65        }
66    }
67}
68
69/// `fp4_act_quant(x, 16, inplace=True, scale_dtype=e4m3)`: amax floored at `6 * 2^-9`, scale =
70/// `amax / 6` rounded through e4m3, values to e2m1, dequantised back to bf16.
71pub fn fp4_quant_e4m3_inplace(x: &mut [f32], block: usize) {
72    for blk in x.chunks_mut(block) {
73        let amax = blk
74            .iter()
75            .fold(0f32, |a, v| a.max(v.abs()))
76            .max(6.0 * 2f32.powi(-9));
77        let s = e4m3_to_f32(f32_to_e4m3_rne(amax / FP4_MAX));
78        for v in blk.iter_mut() {
79            *v = to_bf16_rne(to_e2m1_rne((*v / s).clamp(-FP4_MAX, FP4_MAX)) * s);
80        }
81    }
82}
83
84/// `precompute_freqs_cis` WITH YaRN (`original_seq_len > 0`), as the compress layers build it
85/// from `compress_rope_theta`. `[seqlen][dim/2]` of (cos, sin).
86pub fn yarn_freqs_cis(
87    dim: usize,
88    seqlen: usize,
89    original_seq_len: usize,
90    base: f32,
91    factor: f32,
92    beta_fast: f32,
93    beta_slow: f32,
94) -> Vec<(f32, f32)> {
95    let half = dim / 2;
96    let mut freqs: Vec<f32> = (0..half)
97        .map(|k| 1.0 / base.powf((2 * k) as f32 / dim as f32))
98        .collect();
99    let corrected = |rot: f64| -> f64 {
100        dim as f64 * (original_seq_len as f64 / (rot * 2.0 * std::f64::consts::PI)).ln()
101            / (2.0 * (base as f64).ln())
102    };
103    let low = corrected(beta_fast as f64).floor().max(0.0) as f32;
104    let high = (corrected(beta_slow as f64).ceil()).min((dim - 1) as f64) as f32;
105    for (i, f) in freqs.iter_mut().enumerate() {
106        let ramp = ((i as f32 - low) / (high - low).max(1e-3)).clamp(0.0, 1.0);
107        let smooth = 1.0 - ramp;
108        *f = *f / factor * (1.0 - smooth) + *f * smooth;
109    }
110    let mut out = Vec::with_capacity(seqlen * half);
111    for p in 0..seqlen {
112        for &f in &freqs {
113            let a = p as f32 * f;
114            out.push((a.cos(), a.sin()));
115        }
116    }
117    out
118}
119
120/// f32 `F.linear` (the ratio-2 compressor's `wkv` / `wgate` are fp32 weights on an fp32 input).
121pub fn linear_f32(x: &[f32], w: &[f32], rows: usize, in_dim: usize, out_dim: usize) -> Vec<f32> {
122    let mut y = vec![0f32; rows * out_dim];
123    for i in 0..rows {
124        for o in 0..out_dim {
125            y[i * out_dim + o] = x[i * in_dim..(i + 1) * in_dim]
126                .iter()
127                .zip(&w[o * in_dim..(o + 1) * in_dim])
128                .map(|(a, b)| a * b)
129                .sum();
130        }
131    }
132    y
133}
134
135pub struct CompAttnCfg {
136    pub dim: usize,
137    pub n_heads: usize,
138    pub head_dim: usize,
139    pub rope_dim: usize,
140    pub q_rank: usize,
141    pub o_rank: usize,
142    pub groups: usize,
143    pub window: usize,
144    pub eps: f32,
145    pub ratio: usize,
146    pub is_kv_source: bool,
147    pub is_index_source: bool,
148    pub is_candidate_source: bool,
149    pub uses_candidates: bool,
150}
151
152pub struct LayerAttnState {
153    pub window: Vec<f32>,
154    pub compressor: Option<CompressorState>,
155    pub compress_kv_cache: Option<Vec<f32>>,
156    pub k_cache: Option<Vec<f32>>,
157}
158
159impl LayerAttnState {
160    pub fn new(c: &CompAttnCfg, max_seq: usize, index_hd: usize) -> Self {
161        let groups = if c.ratio > 0 { max_seq / c.ratio } else { 0 };
162        LayerAttnState {
163            window: vec![0f32; c.window * c.head_dim],
164            compressor: if c.is_kv_source {
165                Some(CompressorState::new(c.ratio.max(1), c.head_dim))
166            } else {
167                None
168            },
169            compress_kv_cache: if c.is_kv_source {
170                Some(vec![0f32; groups * c.head_dim])
171            } else {
172                None
173            },
174            k_cache: if c.is_kv_source {
175                Some(vec![0f32; groups * index_hd])
176            } else {
177                None
178            },
179        }
180    }
181}
182
183pub struct CompAttnRun {
184    pub q: Vec<f32>,
185    pub kv_rows: Vec<f32>,
186    pub idx: Vec<i32>,
187    pub topk: usize,
188    pub o: Vec<f32>,
189    pub out: Vec<f32>,
190}
191
192/// `Attention.forward` for any `compress_ratio`, mutating the layer state and the shared slots.
193#[allow(clippy::too_many_arguments)]
194pub fn attention_any(
195    x: &[f32],
196    seqlen: usize,
197    start_pos: usize,
198    w: &super::attn::AttnWeights,
199    comp: Option<&CompressorWeights>,
200    idxw: Option<&IndexerWeights>,
201    icfg: &IndexerCfg,
202    c: &CompAttnCfg,
203    fc: &[(f32, f32)],
204    st: &mut LayerAttnState,
205    shared: &mut SharedRuntime,
206) -> CompAttnRun {
207    let (hd, rd, nh) = (c.head_dim, c.rope_dim, c.n_heads);
208    let pos: Vec<usize> = (0..seqlen).map(|t| start_pos + t).collect();
209    let head_pos: Vec<usize> = pos
210        .iter()
211        .flat_map(|&p| std::iter::repeat_n(p, nh))
212        .collect();
213
214    let qr = rms_norm(
215        &linear_bf16(x, w.wq_a, seqlen, c.dim, c.q_rank),
216        w.q_norm,
217        seqlen,
218        c.q_rank,
219        c.eps,
220    );
221    let mut q = linear_bf16(&qr, w.wq_b, seqlen, c.q_rank, nh * hd);
222    apply_rotary(&mut q, hd, rd, &head_pos, fc, false);
223
224    let mut kv = rms_norm(
225        &linear_bf16(x, w.wkv, seqlen, c.dim, hd),
226        w.kv_norm,
227        seqlen,
228        hd,
229        c.eps,
230    );
231    apply_rotary(&mut kv, hd, rd, &pos, fc, false);
232    act_quant_inplace(&mut kv);
233    let win = c.window;
234    let mut kv_rows = if start_pos == 0 {
235        if seqlen <= win {
236            st.window[..seqlen * hd].copy_from_slice(&kv);
237        } else {
238            let cutoff = seqlen % win;
239            let tail = &kv[(seqlen - win) * hd..];
240            st.window[cutoff * hd..win * hd].copy_from_slice(&tail[..(win - cutoff) * hd]);
241            st.window[..cutoff * hd].copy_from_slice(&tail[(win - cutoff) * hd..]);
242        }
243        kv.clone()
244    } else {
245        let slot = start_pos % win;
246        st.window[slot * hd..(slot + 1) * hd].copy_from_slice(&kv);
247        st.window.clone()
248    };
249    let (mut idx, mut topk) = window_topk_idxs(win, seqlen, start_pos);
250
251    if c.ratio > 0 {
252        let ratio = c.ratio;
253        let offset = kv_rows.len() / hd;
254        let compress_len = (start_pos + seqlen) / ratio;
255        let latent = if c.is_kv_source {
256            let r = compressor(
257                x,
258                seqlen,
259                start_pos,
260                ratio,
261                c.dim,
262                hd,
263                comp.expect("compressor weights"),
264                st.compressor.as_mut().expect("compressor state"),
265                c.eps,
266            );
267            shared.compress_kv = st.compress_kv_cache.as_ref().expect("cache").clone();
268            r
269        } else {
270            None
271        };
272        // the indexer needs the latent before RoPE, so it runs before the cache is written
273        let (cidx, ctopk) = if !c.is_index_source {
274            (shared.topk_idxs.clone(), shared.topk)
275        } else if compress_len == 0 {
276            (Vec::new(), 0)
277        } else {
278            let r = indexer(
279                x,
280                &qr,
281                latent.as_deref(),
282                seqlen,
283                start_pos,
284                offset,
285                ratio,
286                icfg,
287                idxw.expect("indexer weights"),
288                fc,
289                st.k_cache.as_mut(),
290                shared,
291                c.is_candidate_source,
292                c.uses_candidates,
293            );
294            shared.topk_idxs = r.0.clone();
295            shared.topk = r.1;
296            r
297        };
298        if let Some(mut lat) = latent {
299            let groups = lat.len() / hd;
300            let lpos: Vec<usize> = if start_pos == 0 {
301                (0..groups).map(|g| g * ratio).collect()
302            } else {
303                vec![start_pos + 1 - ratio]
304            };
305            apply_rotary(&mut lat, hd, rd, &lpos, fc, false);
306            fp4_quant_e4m3_inplace(&mut lat, LATENT_BLOCK);
307            let cache = st.compress_kv_cache.as_mut().expect("cache");
308            let at = start_pos / ratio;
309            cache[at * hd..(at + groups) * hd].copy_from_slice(&lat);
310            shared.compress_kv = cache.clone();
311        }
312        // read after the write
313        kv_rows.extend_from_slice(&shared.compress_kv[..compress_len * hd]);
314        let mut merged = Vec::with_capacity(seqlen * (topk + ctopk));
315        for t in 0..seqlen {
316            merged.extend_from_slice(&idx[t * topk..(t + 1) * topk]);
317            merged.extend_from_slice(&cidx[t * ctopk..(t + 1) * ctopk]);
318        }
319        idx = merged;
320        topk += ctopk;
321    }
322
323    let scale = (hd as f32).powf(-0.5);
324    let o_sa = sparse_attn(&q, &kv_rows, w.sink, &idx, seqlen, nh, hd, topk, scale);
325    let mut o = o_sa.clone();
326    apply_rotary(&mut o, hd, rd, &head_pos, fc, true);
327    let gw = nh * hd / c.groups;
328    let mut og = vec![0f32; seqlen * c.groups * c.o_rank];
329    for t in 0..seqlen {
330        for g in 0..c.groups {
331            let ov = &o[t * nh * hd + g * gw..t * nh * hd + (g + 1) * gw];
332            for r in 0..c.o_rank {
333                let wr = &w.wo_a[(g * c.o_rank + r) * gw..(g * c.o_rank + r + 1) * gw];
334                og[(t * c.groups + g) * c.o_rank + r] =
335                    to_bf16_rne(ov.iter().zip(wr).map(|(a, b)| a * b).sum());
336            }
337        }
338    }
339    let out = linear_bf16(&og, w.wo_b, seqlen, c.groups * c.o_rank, c.dim);
340    CompAttnRun {
341        q,
342        kv_rows,
343        idx,
344        topk,
345        o: o_sa,
346        out,
347    }
348}
349
350#[cfg(test)]
351#[path = "compress_tests.rs"]
352mod tests;
spark_model/layers/deepseek_v41_ref/compress.rs

spark_model/layers/deepseek_v41_ref/
compress.rs
