metrale_storage/
predictor_ref.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Host reference for the predictor kernels, with the kernels' tensor
4//! layouts; `tests/predictor_parity.rs` compares the GPU results against it.
5//!
6//! Owner: storage, high-speed swap.
7//! Invariants: none beyond the types.
8//!
9//! Layouts:
10//! - `q`: `[num_q_heads, head_dim]`; `k_block`: `[block_size, num_kv_heads, head_dim]`;
11//! - `p`: `[head_dim, r]`; `q_proj`: `[num_q_heads, r]`;
12//! - one block's low-rank keys: `[num_kv_heads, block_size, r]`, and a sequence of
13//!   them `[num_blocks, num_kv_heads, block_size, r]`;
14//! - scores: `[num_blocks]`, the max over query heads and tokens of the dot product.
15
16use half::bf16;
17
18#[inline]
19fn b2f(x: bf16) -> f32 {
20    x.to_f32()
21}
22#[inline]
23fn f2b(x: f32) -> bf16 {
24    bf16::from_f32(x)
25}
26
27pub fn project_q_ref(
28    q: &[bf16],
29    p: &[bf16],
30    num_q_heads: usize,
31    head_dim: usize,
32    r: usize,
33) -> Vec<bf16> {
34    let mut out = vec![bf16::from_f32(0.0); num_q_heads * r];
35    for h in 0..num_q_heads {
36        for o in 0..r {
37            let mut acc = 0.0_f32;
38            for i in 0..head_dim {
39                acc += b2f(q[h * head_dim + i]) * b2f(p[i * r + o]);
40            }
41            out[h * r + o] = f2b(acc);
42        }
43    }
44    out
45}
46
47/// 2026-09-25: Project every token of one block; output `[num_kv_heads, block_size, r]`.
48pub fn project_kv_block_ref(
49    k_block: &[bf16],
50    p: &[bf16],
51    block_size: usize,
52    num_kv_heads: usize,
53    head_dim: usize,
54    r: usize,
55) -> Vec<bf16> {
56    let mut out = vec![bf16::from_f32(0.0); num_kv_heads * block_size * r];
57    for kh in 0..num_kv_heads {
58        for tok in 0..block_size {
59            let k_base = (tok * num_kv_heads + kh) * head_dim;
60            for o in 0..r {
61                let mut acc = 0.0_f32;
62                for i in 0..head_dim {
63                    acc += b2f(k_block[k_base + i]) * b2f(p[i * r + o]);
64                }
65                out[(kh * block_size + tok) * r + o] = f2b(acc);
66            }
67        }
68    }
69    out
70}
71
72pub fn predictor_score_ref(
73    q_proj: &[bf16],
74    k_lr_seq: &[bf16],
75    num_q_heads: usize,
76    num_kv_heads: usize,
77    block_size: usize,
78    r: usize,
79    num_active_blocks: usize,
80) -> Vec<f32> {
81    assert!(num_q_heads.is_multiple_of(num_kv_heads));
82    let gqa = num_q_heads / num_kv_heads;
83    let per_block = num_kv_heads * block_size * r;
84    let mut scores = vec![f32::NEG_INFINITY; num_active_blocks];
85    for blk in 0..num_active_blocks {
86        let mut best = f32::NEG_INFINITY;
87        for qh in 0..num_q_heads {
88            let kh = qh / gqa;
89            for tok in 0..block_size {
90                let mut dot = 0.0_f32;
91                for i in 0..r {
92                    let q = b2f(q_proj[qh * r + i]);
93                    let k = b2f(k_lr_seq[blk * per_block + (kh * block_size + tok) * r + i]);
94                    dot += q * k;
95                }
96                if dot > best {
97                    best = dot;
98                }
99            }
100        }
101        scores[blk] = best;
102    }
103    scores
104}
105
106/// 2026-09-25: Attention weight of each block: for every query head, the softmax
107/// over all tokens summed over the block's tokens; the max over query heads.
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    #[allow(clippy::needless_range_loop)]
123    for qh in 0..num_q_heads {
124        let kh = qh / gqa;
125        let mut logits = vec![0.0_f32; total_tokens];
126        for t in 0..total_tokens {
127            let blk = t / block_size;
128            let off = t % block_size;
129            let k_base = (off * num_kv_heads + kh) * head_dim;
130            let mut dot = 0.0_f32;
131            for i in 0..head_dim {
132                dot += b2f(q[qh * head_dim + i])
133                    * b2f(k[blk * block_size * num_kv_heads * head_dim + k_base + i]);
134            }
135            logits[t] = dot * scale;
136        }
137        let lmax = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
138        let sum: f32 = logits.iter().map(|l| (l - lmax).exp()).sum();
139        for blk in 0..num_blocks {
140            let mut bw = 0.0_f32;
141            for off in 0..block_size {
142                let t = blk * block_size + off;
143                bw += (logits[t] - lmax).exp() / sum;
144            }
145            if bw > block_weights[blk] {
146                block_weights[blk] = bw;
147            }
148        }
149    }
150    block_weights
151}
metrale_storage/predictor_ref.rs

metrale_storage/
predictor_ref.rs
