metrale_model_arch/deepseek_v41_ref/compress/
compressor.rs

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

metrale_model_arch/deepseek_v41_ref/compress/
compressor.rs
