spark_model/layers/deepseek_v41_ref/compress/
compressor.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! The `Compressor` of DeepSeek-V4.1 compressed attention: `Compressor.forward` (softmax pooling
5//! of `ratio` tokens into one latent, or a plain projection at ratio 1) with its weights and the
6//! decode-time fill state. Split out of `compress.rs` for the file-size cap; the quantisers and
7//! `linear_f32` it needs stay there.
8
9use super::super::hc::rms_norm;
10use super::super::moe::linear_bf16;
11use super::super::to_bf16_rne;
12use super::linear_f32;
13
14pub struct CompressorWeights<'a> {
15    /// f32 at ratio > 1, bf16 values at ratio 1
16    pub wkv: &'a [f32],
17    /// ratio > 1 only
18    pub wgate: Option<&'a [f32]>,
19    pub norm: &'a [f32],
20}
21
22pub struct CompressorState {
23    pub kv_state: Vec<f32>,
24    pub score_state: Vec<f32>,
25}
26
27impl CompressorState {
28    pub fn new(ratio: usize, hd: usize) -> Self {
29        CompressorState {
30            kv_state: vec![0f32; ratio * hd],
31            score_state: vec![f32::NEG_INFINITY; ratio * hd],
32        }
33    }
34}
35
36/// `Compressor.forward`: the pre-RoPE latent `[groups][hd]` (bf16 values), or `None` while a
37/// group is still filling up in decode.
38pub fn compressor(
39    x: &[f32],
40    seqlen: usize,
41    start_pos: usize,
42    ratio: usize,
43    dim: usize,
44    hd: usize,
45    w: &CompressorWeights,
46    st: &mut CompressorState,
47    eps: f32,
48) -> Option<Vec<f32>> {
49    if ratio == 1 {
50        return Some(rms_norm(
51            &linear_bf16(x, w.wkv, seqlen, dim, hd),
52            w.norm,
53            seqlen,
54            hd,
55            eps,
56        ));
57    }
58    let kv = linear_f32(x, w.wkv, seqlen, dim, hd);
59    let score = linear_f32(x, w.wgate.expect("wgate at ratio > 1"), seqlen, dim, hd);
60    let pool = |kvg: &[f32], scg: &[f32]| -> Vec<f32> {
61        // softmax over the `ratio` members per dimension, then the weighted sum
62        let mut out = vec![0f32; hd];
63        for d in 0..hd {
64            let m = (0..ratio)
65                .map(|r| scg[r * hd + d])
66                .fold(f32::NEG_INFINITY, f32::max);
67            let ex: Vec<f32> = (0..ratio).map(|r| (scg[r * hd + d] - m).exp()).collect();
68            let sum: f32 = ex.iter().sum();
69            out[d] = (0..ratio).map(|r| kvg[r * hd + d] * (ex[r] / sum)).sum();
70        }
71        out
72    };
73    let pooled: Vec<f32> = if start_pos == 0 {
74        let remainder = seqlen % ratio;
75        let cutoff = seqlen - remainder;
76        if remainder > 0 {
77            st.kv_state[..remainder * hd].copy_from_slice(&kv[cutoff * hd..]);
78            st.score_state[..remainder * hd].copy_from_slice(&score[cutoff * hd..]);
79        }
80        if seqlen < ratio {
81            return None;
82        }
83        (0..cutoff / ratio)
84            .flat_map(|g| {
85                pool(
86                    &kv[g * ratio * hd..(g + 1) * ratio * hd],
87                    &score[g * ratio * hd..(g + 1) * ratio * hd],
88                )
89            })
90            .collect()
91    } else {
92        let slot = start_pos % ratio;
93        st.kv_state[slot * hd..(slot + 1) * hd].copy_from_slice(&kv);
94        st.score_state[slot * hd..(slot + 1) * hd].copy_from_slice(&score);
95        if !(start_pos + 1).is_multiple_of(ratio) {
96            return None;
97        }
98        pool(&st.kv_state, &st.score_state)
99    };
100    let groups = pooled.len() / hd;
101    let as_bf16: Vec<f32> = pooled.into_iter().map(to_bf16_rne).collect();
102    Some(rms_norm(&as_bf16, w.norm, groups, hd, eps))
103}
spark_model/layers/deepseek_v41_ref/compress/compressor.rs

spark_model/layers/deepseek_v41_ref/compress/
compressor.rs
