qualia_core_db/inference/
topk.rs

1//! STELLAR §A A1a — GPU top-K reduction: CPU oracle, host merge, and the WGSL kernel.
2//!
3//! Decode is memory-bandwidth-bound; the sentinel/sampler only needs the high-probability
4//! mass, not the 49 k near-zero tail. Instead of reading back the full logit vector
5//! (~196 KB/token) and doing a CPU argmax, the GPU reduces each block of the vocabulary to
6//! its top-K candidates; the host merges those into the global top-K. The CPU oracle here is
7//! the byte-for-byte reference the on-device kernel is verified against (see `topk_gpu.rs`),
8//! exactly as `ternary.rs` anchors `ternary_gpu.rs`.
9//!
10//! Contract (must match `shaders/topk_reduction.wgsl`): NaN → −∞ (never selected); ties broken
11//! toward the LOWER token id (deterministic); K=1 == argmax.
12
13/// One top-K entry: a token id and its raw logit.
14#[derive(Debug, Clone, Copy, PartialEq)]
15pub struct TopKItem {
16    pub token_id: u32,
17    pub logit: f32,
18}
19
20/// The block-reduction kernel (auto bind-group layout; entry `topk_block`).
21pub const TOPK_REDUCTION_WGSL: &str = include_str!("../shaders/topk_reduction.wgsl");
22
23/// Block size each workgroup reduces — must equal `MAX_BLOCK` in the WGSL (`var<workgroup>` cap).
24pub const TOPK_BLOCK_SIZE: usize = 1024;
25
26/// Largest K the host paths support (kept generous; the kernel itself is K-agnostic per round).
27pub const TOPK_MAX_K: usize = 64;
28
29/// 16-byte `Params` uniform: `n, k, block_size, cand_base` (cand_base=0 for single-chunk).
30pub fn topk_params_bytes(n: u32, k: u32, block_size: u32) -> [u8; 16] {
31    topk_params_bytes_with_base(n, k, block_size, 0)
32}
33
34/// Like [`topk_params_bytes`] with a non-zero candidate write base for multi-chunk mega-pass.
35pub fn topk_params_bytes_with_base(n: u32, k: u32, block_size: u32, cand_base: u32) -> [u8; 16] {
36    let mut b = [0u8; 16];
37    b[0..4].copy_from_slice(&n.to_le_bytes());
38    b[4..8].copy_from_slice(&k.to_le_bytes());
39    b[8..12].copy_from_slice(&block_size.to_le_bytes());
40    b[12..16].copy_from_slice(&cand_base.to_le_bytes());
41    b
42}
43
44/// Normalize: NaN → −∞ so it can never win a comparison.
45#[inline]
46fn clean(x: f32) -> f32 {
47    if x.is_nan() {
48        f32::NEG_INFINITY
49    } else {
50        x
51    }
52}
53
54/// Order two `(id, logit)` candidates: higher logit first, lower id on ties.
55#[inline]
56fn cmp_desc(a: &(u32, f32), b: &(u32, f32)) -> std::cmp::Ordering {
57    b.1.partial_cmp(&a.1)
58        .unwrap_or(std::cmp::Ordering::Equal)
59        .then(a.0.cmp(&b.0))
60}
61
62/// CPU reference top-K over a full logit vector. Drops −∞ (e.g. masked) entries.
63pub fn topk_cpu(logits: &[f32], k: usize) -> Vec<TopKItem> {
64    let mut v: Vec<(u32, f32)> = logits
65        .iter()
66        .enumerate()
67        .map(|(i, &x)| (i as u32, clean(x)))
68        .collect();
69    v.sort_by(cmp_desc);
70    v.into_iter()
71        .filter(|(_, val)| *val > f32::NEG_INFINITY)
72        .take(k)
73        .map(|(token_id, logit)| TopKItem { token_id, logit })
74        .collect()
75}
76
77/// Merge per-block GPU candidates (`num_blocks × k` pairs) into the global top-K.
78/// Blocks cover disjoint index ranges, so candidate ids are unique. Drops −∞ entries
79/// and any id in `masked` (the governance/sieve veto — "a masked token never returned").
80pub fn merge_block_candidates(
81    cand_val: &[f32],
82    cand_idx: &[u32],
83    k: usize,
84    masked: Option<&dyn Fn(u32) -> bool>,
85) -> Vec<TopKItem> {
86    let mut v: Vec<(u32, f32)> = cand_val
87        .iter()
88        .zip(cand_idx.iter())
89        .map(|(val, idx)| (*idx, clean(*val)))
90        .filter(|(idx, val)| *val > f32::NEG_INFINITY && masked.map(|m| !m(*idx)).unwrap_or(true))
91        .collect();
92    v.sort_by(cmp_desc);
93    v.truncate(k);
94    v.into_iter()
95        .map(|(token_id, logit)| TopKItem { token_id, logit })
96        .collect()
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn cpu_topk_orders_and_breaks_ties() {
105        // Two 5.0s at ids 1 and 4 → lower id (1) wins the tie.
106        let logits = [1.0, 5.0, 3.0, 2.0, 5.0, f32::NAN, -1.0];
107        let top = topk_cpu(&logits, 3);
108        assert_eq!(
109            top[0],
110            TopKItem {
111                token_id: 1,
112                logit: 5.0
113            }
114        );
115        assert_eq!(
116            top[1],
117            TopKItem {
118                token_id: 4,
119                logit: 5.0
120            }
121        );
122        assert_eq!(
123            top[2],
124            TopKItem {
125                token_id: 2,
126                logit: 3.0
127            }
128        );
129    }
130
131    #[test]
132    fn cpu_topk_k1_is_argmax() {
133        let logits = [0.1, -2.0, 9.9, 9.9, 1.0];
134        let top = topk_cpu(&logits, 1);
135        assert_eq!(top.len(), 1);
136        assert_eq!(top[0].token_id, 2); // first of the tied maxima
137    }
138
139    #[test]
140    fn merge_drops_masked_and_neg_inf() {
141        // Simulate 2 blocks × k=2 candidates.
142        let cand_val = [9.0, 7.0, 8.0, f32::NEG_INFINITY];
143        let cand_idx = [3u32, 10, 42, 99];
144        // Mask id 3 (a governance veto) → next-best survives.
145        let masked = |id: u32| id == 3;
146        let merged = merge_block_candidates(&cand_val, &cand_idx, 2, Some(&masked));
147        assert_eq!(
148            merged[0],
149            TopKItem {
150                token_id: 42,
151                logit: 8.0
152            }
153        );
154        assert_eq!(
155            merged[1],
156            TopKItem {
157                token_id: 10,
158                logit: 7.0
159            }
160        );
161    }
162}
qualia_core_db/inference/topk.rs

qualia_core_db/inference/
topk.rs
