cortiq-engine 0.5.62

//! `gemm_nt_f32`'s w-side cache against a REUSED buffer.
//!
//! Every batched attention in this engine allocates one `k`/`v` scratch
//! pair per call and refills it per head, so the device-side cache sees
//! the same address holding a different matrix over and over. Keyed on
//! the address alone it returns head 0's keys for every head — silently,
//! and only on the GPU, which is exactly the shape of bug that survives
//! a CPU test suite. Run with the `gpu` feature on a machine that has a
//! device; without one `gemm_nt_f32` declines and the test skips.

#![cfg(feature = "gpu")]

use cortiq_engine::gpu_wgpu::gemm_nt_f32;

/// y = x · wᵀ, row-major, the definition.
fn reference(x: &[f32], w: &[f32], n: usize, k: usize, m: usize) -> Vec<f32> {
    let mut y = vec![0f32; n * m];
    for i in 0..n {
        for j in 0..m {
            let mut acc = 0f64;
            for l in 0..k {
                acc += (x[i * k + l] as f64) * (w[j * k + l] as f64);
            }
            y[i * m + j] = acc as f32;
        }
    }
    y
}

fn fill(buf: &mut [f32], seed: u64) {
    let mut s = seed | 1;
    for v in buf.iter_mut() {
        s ^= s << 13;
        s ^= s >> 7;
        s ^= s << 17;
        *v = ((s >> 40) as f32 / 8_388_608.0) - 1.0;
    }
}

#[test]
fn refilling_the_w_buffer_is_not_a_cache_hit() {
    // Past gemm_nt_f32's own size floor, or it declines and runs on the
    // host — which would pass whatever the cache did.
    let (n, k, m) = (256usize, 128usize, 256usize);
    let mut x = vec![0f32; n * k];
    let mut w = vec![0f32; m * k];
    fill(&mut x, 0x1234_5678);
    fill(&mut w, 0xabcd_ef01);

    let mut y = vec![0f32; n * m];
    if !gemm_nt_f32(&x, &w, &mut y, n, k, m) {
        eprintln!("no GPU (or CMF_BAKE_GPU=0) — skipping");
        return;
    }
    let want = reference(&x, &w, n, k, m);
    let worst = |a: &[f32], b: &[f32]| {
        a.iter()
            .zip(b)
            .fold(0f32, |acc, (&p, &q)| acc.max((p - q).abs()))
    };
    let d0 = worst(&y, &want);
    assert!(d0 < 2e-3, "first call already wrong: max {d0:.3e}");

    // The same allocation, a different matrix — one head later.
    fill(&mut w, 0x5555_aaaa);
    let mut y2 = vec![0f32; n * m];
    assert!(gemm_nt_f32(&x, &w, &mut y2, n, k, m));
    let want2 = reference(&x, &w, n, k, m);
    let d1 = worst(&y2, &want2);
    println!("first {d0:.3e}, after refill {d1:.3e}");
    assert!(
        d1 < 2e-3,
        "stale w served from the cache: max {d1:.3e} (a hit on the address, not the contents)"
    );

    // And the original contents again: a real hit, still correct.
    fill(&mut w, 0xabcd_ef01);
    let mut y3 = vec![0f32; n * m];
    assert!(gemm_nt_f32(&x, &w, &mut y3, n, k, m));
    let d2 = worst(&y3, &want);
    assert!(d2 < 2e-3, "cache hit returned the wrong matrix: max {d2:.3e}");
}