spark_model/layers/deepseek_v41_ref/compress/
indexer.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! The `Indexer` of DeepSeek-V4.1 compressed attention: `Indexer.forward` (fp4 query heads
5//! against one shared fp4 key per compressed position, rectified scores combined by
6//! `weights_proj`), `select_candidate_blocks` (level one of the two-level top-k), the
7//! `SharedAttentionRuntime` slots, and the torch CPU `topk` tie-order emulation. Split out of
8//! `compress.rs` for the file-size cap; the fp4 quantiser it needs stays there.
9
10use super::super::attn::apply_rotary;
11use super::super::hc::rms_norm;
12use super::super::moe::linear_bf16;
13use super::super::to_bf16_rne;
14use super::{FP4_BLOCK, fp4_quant_e8m0_inplace};
15
16pub struct IndexerWeights<'a> {
17    pub wq_b: &'a [f32],
18    pub weights_proj: &'a [f32],
19    /// index-key owners (kv sources) only
20    pub wk: Option<&'a [f32]>,
21    pub k_norm: Option<&'a [f32]>,
22}
23
24/// `SharedAttentionRuntime`: what source layers publish and later layers read.
25#[derive(Default)]
26pub struct SharedRuntime {
27    /// the source's full cache `[max_groups][hd]`
28    pub compress_kv: Vec<f32>,
29    /// the source's full index-key cache `[max_groups][index_hd]`
30    pub index_k: Vec<f32>,
31    /// `[queries][topk]`, already offset for the window rows
32    pub topk_idxs: Vec<i32>,
33    pub topk: usize,
34    /// `[queries][width]`
35    pub candidates: Vec<bool>,
36    pub cand_width: usize,
37}
38
39pub struct IndexerCfg {
40    pub n_heads: usize,
41    pub index_hd: usize,
42    pub rope_dim: usize,
43    pub q_rank: usize,
44    pub dim: usize,
45    pub hd: usize,
46    pub index_topk: usize,
47    pub cand_topk_blocks: usize,
48    pub cand_block: usize,
49    pub eps: f32,
50}
51
52/// The index SET torch's CPU `topk(k, largest=True)` returns for `vals`, ties included. For the
53/// row lengths here (`k * 64 > n`) torch runs `std::nth_element` on (value, index) pairs with a
54/// descending comparator and returns the first `k` pairs, so among equal values (the indexer's
55/// many -inf entries) the picks are whatever libstdc++'s introselect leaves in front. That order
56/// is deterministic and this is a line-for-line emulation of it (`__introselect`,
57/// `__move_median_to_first`, `__unguarded_partition`, `__insertion_sort`). Pinned by
58/// `topk_tie_order_matches_torch_cpu` against this box's torch.
59pub fn torch_cpu_topk_set(vals: &[f32], k: usize) -> Vec<usize> {
60    let n = vals.len();
61    assert!(k <= n && k > 0);
62    assert!(
63        k * 64 > n,
64        "torch takes the partial_sort branch for k*64 <= n; not emulated"
65    );
66    let mut q: Vec<(f32, usize)> = vals
67        .iter()
68        .copied()
69        .enumerate()
70        .map(|(i, v)| (v, i))
71        .collect();
72    // comp(x, y): x before y when x is NaN and y is not, or x.value > y.value
73    let comp = |x: &(f32, usize), y: &(f32, usize)| (x.0.is_nan() && !y.0.is_nan()) || x.0 > y.0;
74    let nth = k - 1;
75    let (mut first, mut last) = (0usize, n);
76    let mut depth = 2 * (usize::BITS - n.leading_zeros()) as usize;
77    while last - first > 3 {
78        if depth == 0 {
79            // heap_select fallback; never reached for these sizes
80            unreachable!("introselect depth limit reached");
81        }
82        depth -= 1;
83        let mid = first + (last - first) / 2;
84        // __move_median_to_first(first, first+1, mid, last-1)
85        let (a, b, c) = (first + 1, mid, last - 1);
86        let r = if comp(&q[a], &q[b]) {
87            if comp(&q[b], &q[c]) {
88                b
89            } else if comp(&q[a], &q[c]) {
90                c
91            } else {
92                a
93            }
94        } else if comp(&q[a], &q[c]) {
95            a
96        } else if comp(&q[b], &q[c]) {
97            c
98        } else {
99            b
100        };
101        q.swap(first, r);
102        // __unguarded_partition(first+1, last, pivot=first)
103        let pivot = q[first];
104        let (mut lo, mut hi) = (first + 1, last);
105        let cut = loop {
106            while comp(&q[lo], &pivot) {
107                lo += 1;
108            }
109            hi -= 1;
110            while comp(&pivot, &q[hi]) {
111                hi -= 1;
112            }
113            if lo >= hi {
114                break lo;
115            }
116            q.swap(lo, hi);
117            lo += 1;
118        };
119        if cut <= nth {
120            first = cut;
121        } else {
122            last = cut;
123        }
124    }
125    // __insertion_sort(first, last)
126    for i in (first + 1)..last {
127        let val = q[i];
128        if comp(&val, &q[first]) {
129            for j in (first + 1..=i).rev() {
130                q[j] = q[j - 1];
131            }
132            q[first] = val;
133        } else {
134            let mut j = i;
135            while comp(&val, &q[j - 1]) {
136                q[j] = q[j - 1];
137                j -= 1;
138            }
139            q[j] = val;
140        }
141    }
142    q[..k].iter().map(|&(_, i)| i).collect()
143}
144
145/// `select_candidate_blocks`: `logits[queries][width]` with unreachable positions at -inf;
146/// `compress_lens[q]`. Returns the keep mask `[queries][width]`.
147pub fn select_candidate_blocks(
148    logits: &[f32],
149    queries: usize,
150    width: usize,
151    compress_lens: &[usize],
152    topk_blocks: usize,
153    block: usize,
154) -> Vec<bool> {
155    let nb = width.div_ceil(block);
156    let mut keep = vec![false; queries * width];
157    for q in 0..queries {
158        let row = &logits[q * width..(q + 1) * width];
159        let mut scores: Vec<f32> = (0..nb)
160            .map(|b| {
161                (b * block..((b + 1) * block).min(width))
162                    .map(|i| row[i])
163                    .fold(f32::NEG_INFINITY, f32::max)
164            })
165            .collect();
166        // the block holding this query's newest position is pinned in
167        let last = (compress_lens[q] as i64 - 1).div_euclid(block as i64);
168        if last >= 0 && (last as usize) < nb {
169            scores[last as usize] = f32::INFINITY;
170        }
171        for b in torch_cpu_topk_set(&scores, topk_blocks.min(nb)) {
172            if scores[b] > f32::NEG_INFINITY {
173                for i in b * block..((b + 1) * block).min(width) {
174                    keep[q * width + i] = true;
175                }
176            }
177        }
178    }
179    keep
180}
181
182/// `Indexer.forward`. Returns `[queries][topk]` compressed-position indices, offset by
183/// `offset`, -1 where unreachable; also publishes index keys when this layer owns them.
184#[allow(clippy::too_many_arguments)]
185pub fn indexer(
186    x: &[f32],
187    qr: &[f32],
188    latent: Option<&[f32]>,
189    seqlen: usize,
190    start_pos: usize,
191    offset: usize,
192    ratio: usize,
193    c: &IndexerCfg,
194    w: &IndexerWeights,
195    fc: &[(f32, f32)],
196    k_cache: Option<&mut Vec<f32>>,
197    shared: &mut SharedRuntime,
198    is_candidate_source: bool,
199    uses_candidates: bool,
200) -> (Vec<i32>, usize) {
201    let (nh, ihd, rd) = (c.n_heads, c.index_hd, c.rope_dim);
202    let end_pos = start_pos + seqlen;
203
204    if let (Some(lat), Some(cache)) = (latent, k_cache) {
205        let groups = lat.len() / c.hd;
206        let mut k = rms_norm(
207            &linear_bf16(lat, w.wk.expect("wk"), groups, c.hd, ihd),
208            w.k_norm.expect("k_norm"),
209            groups,
210            ihd,
211            c.eps,
212        );
213        let pos: Vec<usize> = if start_pos == 0 {
214            (0..groups).map(|g| g * ratio).collect()
215        } else {
216            vec![start_pos + 1 - ratio]
217        };
218        apply_rotary(&mut k, ihd, rd, &pos, fc, false);
219        fp4_quant_e8m0_inplace(&mut k, FP4_BLOCK);
220        let at = start_pos / ratio;
221        cache[at * ihd..(at + groups) * ihd].copy_from_slice(&k);
222        shared.index_k = cache.clone();
223    }
224
225    let mut q = linear_bf16(qr, w.wq_b, seqlen, c.q_rank, nh * ihd);
226    let head_pos: Vec<usize> = (0..seqlen)
227        .flat_map(|t| std::iter::repeat_n(start_pos + t, nh))
228        .collect();
229    apply_rotary(&mut q, ihd, rd, &head_pos, fc, false);
230    fp4_quant_e8m0_inplace(&mut q, FP4_BLOCK);
231
232    let width = end_pos / ratio;
233    let index_k = &shared.index_k[..width * ihd];
234    let wscale = (ihd as f32).powf(-0.5) * (nh as f32).powf(-0.5);
235    let weights: Vec<f32> = linear_bf16(x, w.weights_proj, seqlen, c.dim, nh)
236        .into_iter()
237        .map(|v| to_bf16_rne(v * wscale))
238        .collect();
239
240    // index_score[q][t] = sum_h bf16(relu(bf16(q_h . k_t)) * weights[h]), each stage bf16
241    let mut score = vec![0f32; seqlen * width];
242    for t in 0..seqlen {
243        for p in 0..width {
244            let kr = &index_k[p * ihd..(p + 1) * ihd];
245            let mut acc = 0f32;
246            for h in 0..nh {
247                let qv = &q[(t * nh + h) * ihd..(t * nh + h + 1) * ihd];
248                let dot = to_bf16_rne(qv.iter().zip(kr).map(|(a, b)| a * b).sum::<f32>());
249                acc += to_bf16_rne(dot.max(0.0) * weights[t * nh + h]);
250            }
251            score[t * width + p] = to_bf16_rne(acc);
252        }
253    }
254    let compress_lens: Vec<usize> = if start_pos == 0 {
255        (0..seqlen).map(|t| (t + 1) / ratio).collect()
256    } else {
257        vec![end_pos / ratio; seqlen]
258    };
259    if start_pos == 0 {
260        for t in 0..seqlen {
261            for p in compress_lens[t]..width {
262                score[t * width + p] = f32::NEG_INFINITY;
263            }
264        }
265    }
266    if is_candidate_source {
267        shared.candidates = select_candidate_blocks(
268            &score,
269            seqlen,
270            width,
271            &compress_lens,
272            c.cand_topk_blocks,
273            c.cand_block,
274        );
275        shared.cand_width = width;
276    } else if uses_candidates {
277        for i in 0..seqlen * width {
278            if !shared.candidates[i] {
279                score[i] = f32::NEG_INFINITY;
280            }
281        }
282    }
283    let topk = c.index_topk.min(end_pos / ratio);
284    let mut out = Vec::with_capacity(seqlen * topk);
285    for t in 0..seqlen {
286        let row = &score[t * width..(t + 1) * width];
287        let mut picked = torch_cpu_topk_set(row, topk);
288        picked.sort_unstable();
289        for i in picked {
290            out.push(if i < compress_lens[t] {
291                (i + offset) as i32
292            } else {
293                -1
294            });
295        }
296    }
297    (out, topk)
298}
spark_model/layers/deepseek_v41_ref/compress/indexer.rs

spark_model/layers/deepseek_v41_ref/compress/
indexer.rs
