qualia_core_db/inference/cuda_lane/
weight_cache.rs

1//! Host-side dense weight cache for CUDA tensor-core batch GEMM.
2//!
3//! Caches densified f32 weight matrices by FNV-1a content fingerprint so
4//! repeated layers don't re-upload the same matrix. Allocations are 2 MiB-aligned
5//! and backed by transparent huge pages on Linux (`MADV_HUGEPAGE`) to reduce TLB
6//! pressure during H2D uploads (Phase 7C).
7
8use std::alloc::Layout;
9use std::collections::HashMap;
10use std::sync::{Mutex, OnceLock};
11
12use crate::wgsl_forge::dispatch::{ensure_cuda_runtime_path, gemm_f32_tc_reduced};
13
14/// Max cached weight matrices (each can be tens of MB).
15const MAX_WEIGHT_ENTRIES: usize = 24;
16/// Max f32 elements per densified matrix (~192 MiB) — covers 3B FFN (~25M) with headroom.
17pub const MAX_DENSE_ELEMS: usize = 48 * 1024 * 1024;
18
19/// Alignment for huge-page-backed host weight buffers (2 MiB).
20const HUGE_PAGE_ALIGN: usize = 2 * 1024 * 1024;
21
22/// A 2 MiB-aligned f32 buffer backed by transparent huge pages on Linux.
23/// Reduces TLB misses during H2D DMA transfers to the GPU.
24struct HugePageF32 {
25    ptr: *mut f32,
26    len: usize,
27    /// Byte capacity (always `len * size_of::<f32>()` rounded up to alignment).
28    cap_bytes: usize,
29}
30
31// SAFETY: `HugePageF32` is a unique owner; rayon scoped threads borrow it
32// mutably via `as_mut_slice` within a guaranteed-over scope. The pointer is
33// never shared across threads without external synchronization (behind `Mutex`).
34unsafe impl Send for HugePageF32 {}
35unsafe impl Sync for HugePageF32 {}
36
37impl HugePageF32 {
38    fn new(len: usize) -> Self {
39        let byte_len = len * std::mem::size_of::<f32>();
40        let cap_bytes = byte_len;
41        let layout = Layout::from_size_align(cap_bytes, HUGE_PAGE_ALIGN).expect("huge-page layout");
42        let ptr = unsafe { std::alloc::alloc(layout) as *mut f32 };
43        if ptr.is_null() {
44            std::alloc::handle_alloc_error(layout);
45        }
46        // On Linux, advise the kernel to back this region with huge pages.
47        #[cfg(target_os = "linux")]
48        unsafe {
49            let ret = libc::madvise(ptr as *mut libc::c_void, cap_bytes, libc::MADV_HUGEPAGE);
50            if ret != 0 {
51                log::debug!("huge_page|madvise_failed|errno={ret}");
52            }
53        }
54        Self {
55            ptr,
56            len,
57            cap_bytes,
58        }
59    }
60
61    fn as_slice(&self) -> &[f32] {
62        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
63    }
64
65    fn as_mut_slice(&mut self) -> &mut [f32] {
66        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
67    }
68}
69
70impl Clone for HugePageF32 {
71    fn clone(&self) -> Self {
72        let buf = Self::new(self.len);
73        unsafe {
74            std::ptr::copy_nonoverlapping(self.ptr, buf.ptr, self.len);
75        }
76        buf
77    }
78}
79
80impl Drop for HugePageF32 {
81    fn drop(&mut self) {
82        let layout =
83            Layout::from_size_align(self.cap_bytes, HUGE_PAGE_ALIGN).expect("huge-page layout");
84        unsafe {
85            std::alloc::dealloc(self.ptr as *mut u8, layout);
86        }
87    }
88}
89
90#[derive(Clone)]
91struct WeightEntry {
92    n_in: usize,
93    n_out: usize,
94    /// Dense f32 row-major [n_out × n_in] (GGML convention: rows = out).
95    data: HugePageF32,
96    last_use: u64,
97}
98
99struct WeightCache {
100    entries: HashMap<u64, WeightEntry>,
101    clock: u64,
102}
103
104fn cache() -> &'static Mutex<WeightCache> {
105    static C: OnceLock<Mutex<WeightCache>> = OnceLock::new();
106    C.get_or_init(|| {
107        Mutex::new(WeightCache {
108            entries: HashMap::new(),
109            clock: 0,
110        })
111    })
112}
113
114/// FNV-1a over weight bytes + dims for cache key.
115pub fn weight_fingerprint(raw: &[u8], n_in: usize, n_out: usize) -> u64 {
116    let mut h: u64 = 0xcbf2_9ce4_8422_2325;
117    for &b in raw {
118        h ^= b as u64;
119        h = h.wrapping_mul(0x0100_0000_01b3);
120    }
121    h ^= n_in as u64;
122    h = h.wrapping_mul(0x0100_0000_01b3);
123    h ^= n_out as u64;
124    h.wrapping_mul(0x0100_0000_01b3)
125}
126
127/// True if a dense weight for `key` is already in the host TC cache.
128pub fn dense_weight_cached(key: u64) -> bool {
129    cache()
130        .lock()
131        .ok()
132        .map(|g| g.entries.contains_key(&key))
133        .unwrap_or(false)
134}
135
136/// Insert or refresh a dense weight matrix (row-major n_out × n_in).
137pub fn cache_dense_weight(key: u64, n_in: usize, n_out: usize, data: Vec<f32>) {
138    if data.len() != n_in.saturating_mul(n_out) {
139        return;
140    }
141    let len = data.len();
142    let buf = HugePageF32::new(len);
143    unsafe {
144        std::ptr::copy_nonoverlapping(data.as_ptr(), buf.ptr, len);
145    }
146    let Ok(mut g) = cache().lock() else {
147        return;
148    };
149    g.clock = g.clock.wrapping_add(1);
150    if g.entries.len() >= MAX_WEIGHT_ENTRIES && !g.entries.contains_key(&key) {
151        // Evict oldest.
152        if let Some(old_k) = g
153            .entries
154            .iter()
155            .min_by_key(|(_, e)| e.last_use)
156            .map(|(k, _)| *k)
157        {
158            g.entries.remove(&old_k);
159            log::info!("cuda_lane|weight_evict|key={old_k:#x}");
160        }
161    }
162    let clock = g.clock;
163    g.entries.insert(
164        key,
165        WeightEntry {
166            n_in,
167            n_out,
168            data: buf,
169            last_use: clock,
170        },
171    );
172    let n_ent = g.entries.len();
173    log::debug!("cuda_lane|weight_cache|key={key:#x}|n_in={n_in}|n_out={n_out}|entries={n_ent}");
174}
175
176/// Allocate a 2 MiB-aligned buffer and fill it via `fill` — skips the
177/// intermediate `Vec<f32>` that [`cache_dense_weight`] requires. The closure
178/// receives a row-major `n_out × n_in` slice to write into directly.
179pub fn cache_dense_weight_direct<F>(key: u64, n_in: usize, n_out: usize, fill: F) -> bool
180where
181    F: FnOnce(&mut [f32]) -> bool,
182{
183    let total = n_in.saturating_mul(n_out);
184    if total == 0 || total > MAX_DENSE_ELEMS {
185        return false;
186    }
187    let mut buf = HugePageF32::new(total);
188    let ok = fill(buf.as_mut_slice());
189    if !ok {
190        return false;
191    }
192    let Ok(mut g) = cache().lock() else {
193        return false;
194    };
195    g.clock = g.clock.wrapping_add(1);
196    if g.entries.len() >= MAX_WEIGHT_ENTRIES && !g.entries.contains_key(&key) {
197        if let Some(old_k) = g
198            .entries
199            .iter()
200            .min_by_key(|(_, e)| e.last_use)
201            .map(|(k, _)| *k)
202        {
203            g.entries.remove(&old_k);
204            log::info!("cuda_lane|weight_evict|key={old_k:#x}");
205        }
206    }
207    let clock = g.clock;
208    g.entries.insert(
209        key,
210        WeightEntry {
211            n_in,
212            n_out,
213            data: buf,
214            last_use: clock,
215        },
216    );
217    let n_ent = g.entries.len();
218    log::debug!(
219        "cuda_lane|weight_cache_direct|key={key:#x}|n_in={n_in}|n_out={n_out}|entries={n_ent}"
220    );
221    true
222}
223
224fn pad16(x: usize) -> usize {
225    ((x + 15) / 16) * 16
226}
227
228/// Batch dense GEMM: for each of `batch` rows of hidden (n_in), compute out (n_out).
229///
230/// `hidden` is packed `batch × n_in` contiguous. `weight` is row-major `n_out × n_in`
231/// (same layout as GGML / our GEMV rows). Writes `batch × n_out` into `out`.
232///
233/// Pads batch/n_in/n_out to multiples of 16 for WMMA; trims result.
234/// Returns false if mode/caps/dims ineligible or TC fails (caller keeps wgpu path).
235pub fn try_cuda_batch_gemv(
236    hidden: &[f32],
237    batch: usize,
238    n_in: usize,
239    n_out: usize,
240    weight: &[f32],
241    out: &mut [f32],
242) -> bool {
243    if !crate::inference_modes::prefer_tensor_core_gemm() {
244        return false;
245    }
246    if batch == 0 || n_in == 0 || n_out == 0 {
247        return false;
248    }
249    if hidden.len() < batch * n_in || weight.len() < n_out * n_in || out.len() < batch * n_out {
250        return false;
251    }
252    ensure_cuda_runtime_path();
253
254    // C[m×n] = A[m×k] · B[k×n]  with m=batch, k=n_in, n=n_out
255    // weight is n_out × n_in rows → need B as k×n = n_in × n_out (transpose of weight rows).
256    let m = pad16(batch);
257    let k = pad16(n_in);
258    let n = pad16(n_out);
259
260    let mut a = vec![0.0f32; m * k];
261    for b in 0..batch {
262        a[b * k..b * k + n_in].copy_from_slice(&hidden[b * n_in..b * n_in + n_in]);
263    }
264    // B[k×n]: B[i,j] = weight[j, i]  (weight row j, col i)
265    let mut bmat = vec![0.0f32; k * n];
266    for j in 0..n_out {
267        for i in 0..n_in {
268            bmat[i * n + j] = weight[j * n_in + i];
269        }
270    }
271
272    // Prefer f32-faithful TC/floor (`gemm_f32_tc`), not `gemm_f32_tc_reduced`.
273    // Reduced f16-WMMA on pad16 single-token decode was measured as incoherent garbage
274    // (2026-07-24). Lab may still force reduced via QUALIA_LLM_CUDA_TC_REDUCED=1.
275    let use_reduced = matches!(
276        std::env::var("QUALIA_LLM_CUDA_TC_REDUCED").ok().as_deref(),
277        Some("1") | Some("true")
278    );
279    let c = if use_reduced {
280        match gemm_f32_tc_reduced(m, k, n, &a, &bmat) {
281            Ok(v) => v,
282            Err(e) => {
283                log::warn!("cuda_lane|batch_gemv|tc_reduced_fail|{e:?}");
284                return false;
285            }
286        }
287    } else {
288        match crate::wgsl_forge::dispatch::gemm_f32_tc(m, k, n, &a, &bmat) {
289            Ok(v) => v,
290            Err(e) => {
291                log::warn!("cuda_lane|batch_gemv|tc_fail|{e:?}");
292                return false;
293            }
294        }
295    };
296    for b in 0..batch {
297        out[b * n_out..b * n_out + n_out].copy_from_slice(&c[b * n..b * n + n_out]);
298    }
299    true
300}
301
302/// Lookup-only: run GEMV if `key` is already densified in the cache.
303pub fn try_cuda_batch_gemv_cached_only(
304    key: u64,
305    hidden: &[f32],
306    batch: usize,
307    out: &mut [f32],
308) -> bool {
309    if !crate::inference_modes::prefer_tensor_core_gemm() {
310        return false;
311    }
312    let Ok(mut g) = cache().lock() else {
313        return false;
314    };
315    g.clock = g.clock.wrapping_add(1);
316    let clock = g.clock;
317    let Some(e) = g.entries.get_mut(&key) else {
318        return false;
319    };
320    e.last_use = clock;
321    let (wi, wo) = (e.n_in, e.n_out);
322    let w = e.data.clone();
323    drop(g);
324    try_cuda_batch_gemv(hidden, batch, wi, wo, w.as_slice(), out)
325}
326
327/// Like [`try_cuda_batch_gemv`] but uses a cached weight by fingerprint if present;
328/// otherwise caches `weight` under `key` then multiplies.
329pub fn try_cuda_batch_gemv_cached(
330    key: u64,
331    hidden: &[f32],
332    batch: usize,
333    n_in: usize,
334    n_out: usize,
335    weight: &[f32],
336    out: &mut [f32],
337) -> bool {
338    if try_cuda_batch_gemv_cached_only(key, hidden, batch, out) {
339        return true;
340    }
341    if weight.len() == n_in.saturating_mul(n_out) {
342        cache_dense_weight(key, n_in, n_out, weight.to_vec());
343        return try_cuda_batch_gemv(hidden, batch, n_in, n_out, weight, out);
344    }
345    false
346}
347
348/// How many weight matrices are currently cached.
349pub fn weight_cache_len() -> usize {
350    cache().lock().map(|g| g.entries.len()).unwrap_or(0)
351}
352
353/// Clear weight cache (tests / model swap).
354pub fn clear_weight_cache() {
355    if let Ok(mut g) = cache().lock() {
356        g.entries.clear();
357        g.clock = 0;
358    }
359}
qualia_core_db/inference/cuda_lane/weight_cache.rs

qualia_core_db/inference/cuda_lane/
weight_cache.rs
