qualia_core_db/inference/lab/
micro.rs1use 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
62pub 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 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 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 let notes = if ok {
93 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}