metrale_model_arch/deepseek_v41_ref/
compress.rs

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

metrale_model_arch/deepseek_v41_ref/
compress.rs
