spark_storage/
predictor_ref.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// Pure-Rust reference implementation of the predictor math. Used to validate
4// the CUDA kernels (BF16 numerics + reduction order parity) and to compute
5// ground-truth scores for the recall@K test.
6//
7// Tensor layout matches the GPU kernels exactly:
8//   Q       : [num_q_heads, head_dim]
9//   K_block : [block_size,  num_kv_heads, head_dim]
10//   P       : [head_dim,    r]
11//   q_proj  : [num_q_heads, r]
12//   A_g     : [num_blocks,  num_kv_heads, r]
13//   scores  : [num_blocks]   (max over q_heads of dot product, lossless mode)
14
15use half::bf16;
16
17#[inline]
18fn b2f(x: bf16) -> f32 {
19    x.to_f32()
20}
21#[inline]
22fn f2b(x: f32) -> bf16 {
23    bf16::from_f32(x)
24}
25
26pub fn project_q_ref(
27    q: &[bf16],
28    p: &[bf16],
29    num_q_heads: usize,
30    head_dim: usize,
31    r: usize,
32) -> Vec<bf16> {
33    let mut out = vec![bf16::from_f32(0.0); num_q_heads * r];
34    for h in 0..num_q_heads {
35        for o in 0..r {
36            let mut acc = 0.0_f32;
37            for i in 0..head_dim {
38                acc += b2f(q[h * head_dim + i]) * b2f(p[i * r + o]);
39            }
40            out[h * r + o] = f2b(acc);
41        }
42    }
43    out
44}
45
46/// Per-token projection. Output layout `[num_kv_heads, block_size, r]`.
47pub fn project_kv_block_ref(
48    k_block: &[bf16],
49    p: &[bf16],
50    block_size: usize,
51    num_kv_heads: usize,
52    head_dim: usize,
53    r: usize,
54) -> Vec<bf16> {
55    let mut out = vec![bf16::from_f32(0.0); num_kv_heads * block_size * r];
56    for kh in 0..num_kv_heads {
57        for tok in 0..block_size {
58            let k_base = (tok * num_kv_heads + kh) * head_dim;
59            for o in 0..r {
60                let mut acc = 0.0_f32;
61                for i in 0..head_dim {
62                    acc += b2f(k_block[k_base + i]) * b2f(p[i * r + o]);
63                }
64                out[(kh * block_size + tok) * r + o] = f2b(acc);
65            }
66        }
67    }
68    out
69}
70
71pub fn predictor_score_ref(
72    q_proj: &[bf16],
73    k_lr_seq: &[bf16],
74    num_q_heads: usize,
75    num_kv_heads: usize,
76    block_size: usize,
77    r: usize,
78    num_active_blocks: usize,
79) -> Vec<f32> {
80    assert!(num_q_heads.is_multiple_of(num_kv_heads));
81    let gqa = num_q_heads / num_kv_heads;
82    let per_block = num_kv_heads * block_size * r;
83    let mut scores = vec![f32::NEG_INFINITY; num_active_blocks];
84    for blk in 0..num_active_blocks {
85        let mut best = f32::NEG_INFINITY;
86        for qh in 0..num_q_heads {
87            let kh = qh / gqa;
88            for tok in 0..block_size {
89                let mut dot = 0.0_f32;
90                for i in 0..r {
91                    let q = b2f(q_proj[qh * r + i]);
92                    let k = b2f(k_lr_seq[blk * per_block + (kh * block_size + tok) * r + i]);
93                    dot += q * k;
94                }
95                if dot > best {
96                    best = dot;
97                }
98            }
99        }
100        scores[blk] = best;
101    }
102    scores
103}
104
105/// Ground-truth attention-weight per block: max over q_heads of the
106/// softmax-normalized attention to the block's tokens. Used by the recall
107/// test as the "oracle" the predictor's scores are compared against.
108pub fn ground_truth_block_weights(
109    q: &[bf16],
110    k: &[bf16],
111    num_q_heads: usize,
112    num_kv_heads: usize,
113    head_dim: usize,
114    block_size: usize,
115    num_blocks: usize,
116) -> Vec<f32> {
117    let gqa = num_q_heads / num_kv_heads;
118    let scale = 1.0_f32 / (head_dim as f32).sqrt();
119    let total_tokens = num_blocks * block_size;
120    let mut block_weights = vec![0.0_f32; num_blocks];
121
122    // Per-query-head softmax over all tokens, then per-block sum, then max.
123    #[allow(clippy::needless_range_loop)]
124    for qh in 0..num_q_heads {
125        let kh = qh / gqa;
126        // Logits = q_h · k_t  for each t.
127        let mut logits = vec![0.0_f32; total_tokens];
128        for t in 0..total_tokens {
129            let blk = t / block_size;
130            let off = t % block_size;
131            let k_base = (off * num_kv_heads + kh) * head_dim;
132            let mut dot = 0.0_f32;
133            for i in 0..head_dim {
134                dot += b2f(q[qh * head_dim + i])
135                    * b2f(k[blk * block_size * num_kv_heads * head_dim + k_base + i]);
136            }
137            logits[t] = dot * scale;
138        }
139        let lmax = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
140        let sum: f32 = logits.iter().map(|l| (l - lmax).exp()).sum();
141        for blk in 0..num_blocks {
142            let mut bw = 0.0_f32;
143            for off in 0..block_size {
144                let t = blk * block_size + off;
145                bw += (logits[t] - lmax).exp() / sum;
146            }
147            if bw > block_weights[blk] {
148                block_weights[blk] = bw;
149            }
150        }
151    }
152    block_weights
153}
spark_storage/predictor_ref.rs

spark_storage/
predictor_ref.rs
