metrale_model_arch/deepseek_v41_ref/compress/
indexer.rs

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

metrale_model_arch/deepseek_v41_ref/compress/
indexer.rs
