Skip to main content

qualia_core_db/inference/lab/
micro.rs

1//! Isolated Q4_K SoA GEMV microbench (CPU oracle + optional CUDA).
2
3use std::time::Instant;
4
5use crate::ggml_quants::{
6    dequantize_row_into, q4k_block_to_soa, BLOCK_Q4K_SOA_BYTES, GGML_TYPE_Q4_K_SOA,
7};
8use crate::inference_kernel_parity::{max_abs_err, max_ulp_diff};
9use crate::inference_modes::{set_inference_mode, InferenceMode};
10
11#[derive(Debug, Clone)]
12pub struct MicrobenchResult {
13    pub n_in: usize,
14    pub n_out: usize,
15    pub cpu_ms: f64,
16    pub cpu_gflops: f64,
17    pub cuda_ms: Option<f64>,
18    pub cuda_ok: bool,
19    pub max_abs_err: f32,
20    pub max_ulp: u64,
21    pub notes: String,
22}
23
24fn synth_soa_weight(n_in: usize, n_out: usize) -> (Vec<u8>, Vec<f32>) {
25    let n_in = n_in.max(256) & !255;
26    let mut stock = [0u8; 144];
27    stock[0] = 0x00;
28    stock[1] = 0x3c;
29    stock[2] = 0x00;
30    stock[3] = 0x38;
31    for i in 4..144 {
32        stock[i] = (i as u8).wrapping_mul(17);
33    }
34    let mut soa = [0u8; BLOCK_Q4K_SOA_BYTES];
35    q4k_block_to_soa(&stock, &mut soa).expect("soa");
36    let n_blocks = n_in / 256;
37    let mut weight = Vec::with_capacity(n_out * n_blocks * BLOCK_Q4K_SOA_BYTES);
38    for _ in 0..n_out {
39        for _ in 0..n_blocks {
40            weight.extend_from_slice(&soa);
41        }
42    }
43    let x: Vec<f32> = (0..n_in).map(|i| (i as f32) * 0.01).collect();
44    (weight, x)
45}
46
47fn cpu_gemv(n_in: usize, n_out: usize, weight: &[u8], x: &[f32], out: &mut [f32]) {
48    let row_bytes = (n_in / 256) * BLOCK_Q4K_SOA_BYTES;
49    let mut row = vec![0.0f32; n_in];
50    for r in 0..n_out {
51        dequantize_row_into(
52            &weight[r * row_bytes..(r + 1) * row_bytes],
53            GGML_TYPE_Q4_K_SOA,
54            n_in,
55            &mut row,
56        )
57        .unwrap();
58        out[r] = row.iter().zip(x.iter()).map(|(a, b)| a * b).sum();
59    }
60}
61
62/// Run CPU Q4 SoA GEMV; if CUDA mode available, differential vs GPU.
63pub fn run_q4k_soa_microbench(n_in: usize, n_out: usize) -> MicrobenchResult {
64    let n_in = n_in.max(256) & !255;
65    let n_out = n_out.max(1).min(512);
66    let (weight, x) = synth_soa_weight(n_in, n_out);
67    let mut cpu_out = vec![0.0f32; n_out];
68
69    // Warm (untimed) — no Instant here (avoids unused `t0` if warm is dropped).
70    cpu_gemv(n_in, n_out, &weight, &x, &mut cpu_out);
71    let cpu_start = Instant::now();
72    cpu_gemv(n_in, n_out, &weight, &x, &mut cpu_out);
73    let cpu_ms = cpu_start.elapsed().as_secs_f64() * 1e3;
74    let flops = (n_in as f64) * (n_out as f64) * 2.0;
75    let cpu_gflops = (flops / (cpu_ms / 1e3).max(1e-9)) / 1e9;
76
77    let mut cuda_ms = None;
78    let mut cuda_ok = false;
79    let mut max_abs = 0.0f32;
80    let mut max_ulp = 0u64;
81
82    // CUDA differential when available.
83    let prev = std::env::var("QUALIA_INFERENCE_MODE").ok();
84    set_inference_mode(InferenceMode::CudaTc);
85    std::env::set_var("QUALIA_INFERENCE_MODE", "cuda");
86    let mut gpu_out = vec![0.0f32; n_out];
87    #[cfg(feature = "cuda")]
88    let ok = crate::try_q4k_soa_gemv(n_in, n_out, &x, &weight, &mut gpu_out);
89    #[cfg(not(feature = "cuda"))]
90    let ok = false;
91    // Single assignment (not `mut notes` + overwrite) so unused_assignments stays quiet.
92    let notes = if ok {
93        // Warm GPU path untimed; second call times sticky weights.
94        let cuda_start = Instant::now();
95        #[cfg(feature = "cuda")]
96        let ok2 = crate::try_q4k_soa_gemv(n_in, n_out, &x, &weight, &mut gpu_out);
97        #[cfg(not(feature = "cuda"))]
98        let ok2 = false;
99        cuda_ms = Some(cuda_start.elapsed().as_secs_f64() * 1e3);
100        cuda_ok = ok2;
101        max_abs = max_abs_err(&cpu_out, &gpu_out);
102        max_ulp = max_ulp_diff(&cpu_out, &gpu_out);
103        #[cfg(feature = "cuda")]
104        let wcount = crate::q4k_device_weight_count();
105        #[cfg(not(feature = "cuda"))]
106        let wcount = 0usize;
107        format!("cuda differential max_abs={max_abs:.6e} max_ulp={max_ulp} weights={wcount}")
108    } else {
109        "cuda unavailable or prefer_tensor_core_gemm false - cpu only".into()
110    };
111    match prev {
112        Some(v) => std::env::set_var("QUALIA_INFERENCE_MODE", v),
113        None => std::env::remove_var("QUALIA_INFERENCE_MODE"),
114    }
115    set_inference_mode(InferenceMode::Portable);
116
117    MicrobenchResult {
118        n_in,
119        n_out,
120        cpu_ms,
121        cpu_gflops,
122        cuda_ms,
123        cuda_ok,
124        max_abs_err: max_abs,
125        max_ulp,
126        notes,
127    }
128}
129
130impl MicrobenchResult {
131    pub fn format_report(&self) -> String {
132        format!(
133            "Q4_K SoA GEMV microbench\n  shape:     {} x {}\n  cpu_ms:    {:.4}  ({:.2} GFLOP/s)\n  cuda_ms:   {}\n  cuda_ok:   {}\n  max_abs:   {:.6e}\n  max_ulp:   {}\n  {}\n",
134            self.n_in,
135            self.n_out,
136            self.cpu_ms,
137            self.cpu_gflops,
138            self.cuda_ms
139                .map(|m| format!("{m:.4}"))
140                .unwrap_or_else(|| "—".into()),
141            self.cuda_ok,
142            self.max_abs_err,
143            self.max_ulp,
144            self.notes
145        )
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152
153    #[test]
154    fn micro_cpu_runs() {
155        let r = run_q4k_soa_microbench(256, 4);
156        assert!(r.cpu_ms > 0.0);
157        assert!(r.cpu_gflops > 0.0);
158    }
159}