spark_storage/
projection.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// Random Gaussian projection matrix `P` for the Atlas high-speed-swap
4// predictor. Generated once at predictor init from a fixed seed
5// (Johnson–Lindenstrauss embedding); never re-derived per call. Stored on the
6// host so we can hand it to the GPU as BF16 without an extra dtype dance.
7
8use half::bf16;
9use rand::SeedableRng;
10use rand::distributions::Distribution;
11use rand_chacha::ChaCha8Rng;
12use rand_distr::StandardNormal;
13
14#[derive(Clone, Copy, Debug, PartialEq, Eq)]
15pub struct PredictorShape {
16    pub head_dim: usize,
17    pub r: usize,
18}
19
20impl PredictorShape {
21    pub fn new(head_dim: usize, r: usize) -> Self {
22        assert!(head_dim > 0 && r > 0, "head_dim/r must be positive");
23        assert!(head_dim <= 256, "MAX_HEAD_DIM=256 in kv_lowrank_project.cu");
24        assert!(r <= 128, "predictor_score block dim caps r at 128");
25        Self { head_dim, r }
26    }
27}
28
29/// Random Gaussian projection matrix `P` of shape `[head_dim, r]`, BF16, in
30/// row-major layout. Variance 1/head_dim so that ⟨k, p⟩ has unit variance
31/// for unit-norm `k`. Standard JL setting (KVSwap §2.2).
32pub fn build_projection(shape: PredictorShape, seed: u64) -> Vec<bf16> {
33    let mut rng = ChaCha8Rng::seed_from_u64(seed);
34    let inv_sqrt_d = 1.0_f32 / (shape.head_dim as f32).sqrt();
35    let dist = StandardNormal;
36    let n = shape.head_dim * shape.r;
37    let mut out = Vec::with_capacity(n);
38    for _ in 0..n {
39        let v: f32 = dist.sample(&mut rng);
40        out.push(bf16::from_f32(v * inv_sqrt_d));
41    }
42    out
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48
49    #[test]
50    fn determinism() {
51        let s = PredictorShape::new(128, 32);
52        let a = build_projection(s, 0xCAFE_F00D);
53        let b = build_projection(s, 0xCAFE_F00D);
54        assert_eq!(a, b);
55        let c = build_projection(s, 0xDEAD_BEEF);
56        assert_ne!(a, c);
57    }
58
59    #[test]
60    fn shape() {
61        let s = PredictorShape::new(128, 32);
62        let p = build_projection(s, 1);
63        assert_eq!(p.len(), 128 * 32);
64    }
65}
spark_storage/projection.rs

spark_storage/
projection.rs
