metrale_storage/
projection.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The high-speed-swap predictor's random Gaussian projection `P`, built
4//! on the host from a seed; `Predictor` uploads it once.
5//!
6//! Owner: storage, high-speed swap.
7//! Invariants:
8//! - `build_projection` is a pure function of its shape and seed.
9
10use half::bf16;
11use rand::SeedableRng;
12use rand::distr::Distribution;
13use rand_chacha::ChaCha8Rng;
14use rand_distr::StandardNormal;
15
16#[derive(Clone, Copy, Debug, PartialEq, Eq)]
17pub struct PredictorShape {
18    pub head_dim: usize,
19    pub r: usize,
20}
21
22impl PredictorShape {
23    pub fn new(head_dim: usize, r: usize) -> Self {
24        assert!(head_dim > 0 && r > 0, "head_dim/r must be positive");
25        assert!(head_dim <= 256, "MAX_HEAD_DIM=256 in kv_lowrank_project.cu");
26        assert!(r <= 128, "predictor_score block dim caps r at 128");
27        Self { head_dim, r }
28    }
29}
30
31/// 2026-09-25: Random Gaussian projection `P`, `[head_dim, r]` row-major BF16, from a
32/// ChaCha8 stream seeded with `seed`. Entries have variance `1 / head_dim`, so for a
33/// key `k` a projected coordinate has variance `|k|^2 / head_dim`.
34pub fn build_projection(shape: PredictorShape, seed: u64) -> Vec<bf16> {
35    let mut rng = ChaCha8Rng::seed_from_u64(seed);
36    let inv_sqrt_d = 1.0_f32 / (shape.head_dim as f32).sqrt();
37    let dist = StandardNormal;
38    let n = shape.head_dim * shape.r;
39    let mut out = Vec::with_capacity(n);
40    for _ in 0..n {
41        let v: f32 = dist.sample(&mut rng);
42        out.push(bf16::from_f32(v * inv_sqrt_d));
43    }
44    out
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50
51    #[test]
52    fn determinism() {
53        let s = PredictorShape::new(128, 32);
54        let a = build_projection(s, 0xCAFE_F00D);
55        let b = build_projection(s, 0xCAFE_F00D);
56        assert_eq!(a, b);
57        let c = build_projection(s, 0xDEAD_BEEF);
58        assert_ne!(a, c);
59    }
60
61    #[test]
62    fn shape() {
63        let s = PredictorShape::new(128, 32);
64        let p = build_projection(s, 1);
65        assert_eq!(p.len(), 128 * 32);
66    }
67
68    /// 2026-09-30: The stream itself, pinned: the default seed's `P` bit for bit, as rand 0.8 /
69    /// rand_chacha 0.3 / rand_distr 0.4 produced it and rand 0.9 / rand_chacha 0.9 /
70    /// rand_distr 0.5 still do. A dependency bump that changes the predictor's projection
71    /// fails here instead of passing `determinism`, which compares two runs of one build.
72    #[test]
73    fn the_default_seed_gives_the_pinned_projection() {
74        let p = build_projection(PredictorShape::new(128, 32), 0xCAFE_F00D);
75        let bits: Vec<u16> = p.iter().map(|x| x.to_bits()).collect();
76        assert_eq!(
77            bits[..8],
78            [48010, 48438, 15723, 15775, 48398, 48642, 48035, 48715]
79        );
80        let fnv = bits.iter().fold(0xcbf2_9ce4_8422_2325_u64, |h, b| {
81            (h ^ u64::from(*b)).wrapping_mul(0x0100_0000_01b3)
82        });
83        assert_eq!(fnv, 0x5498_6948_5a9b_0c37);
84    }
85}
metrale_storage/projection.rs

metrale_storage/
projection.rs
