Skip to main content

qualia_core_db/gguf_bridge/
gemm.rs

1//! GEMM: weight upload + raw/quantized dispatch + ternary-FFN dispatch
2//! Split from gguf_bridge/mod.rs (structural refactor; no behaviour change).
3use super::*;
4
5/// XOR into the resident-weight key so f16-promoted FFN blobs never alias quant blobs.
6#[cfg(not(target_arch = "wasm32"))]
7const F16_PROMOTE_KEY_TAG: u64 = 0xF16E_F16E_0000_00F1;
8
9impl QTensorEngine {
10    /// Dequant full 2-D weight to dense f32 (row-major n_out × n_in) for CUDA TC path.
11    /// Parallel over rows (rayon) — Q4 densify is cold-path once, then cached.
12    #[cfg(all(not(target_arch = "wasm32"), feature = "cuda"))]
13    pub(crate) fn dequant_weight_dense_f32(
14        info: &GgufTensorInfo,
15        raw: &[u8],
16        n_in: usize,
17        n_out: usize,
18    ) -> Option<Vec<f32>> {
19        use crate::ggml_quants::dequant_matrix_row_into;
20        use rayon::prelude::*;
21        let total = n_out.checked_mul(n_in)?;
22        if total > crate::inference::cuda_lane::MAX_DENSE_ELEMS {
23            return None;
24        }
25        let mut dense = vec![0.0f32; total];
26        let ok = dense
27            .par_chunks_mut(n_in)
28            .enumerate()
29            .map(|(r, row)| dequant_matrix_row_into(raw, info, r, row).is_ok())
30            .all(|b| b);
31        if ok {
32            Some(dense)
33        } else {
34            None
35        }
36    }
37
38    /// Pre-densify and cache a weight for mode=cuda / TC GEMM (no thrash on first token).
39    ///
40    /// Used by the CUDA_DECODE plan path to fill the host dense-weight cache so
41    /// `try_cuda_batch_gemv` can hit without a cold dequant on first use. Skips
42    /// when already cached, dims too small for WMMA pad-16, or densify fails.
43    #[cfg(all(not(target_arch = "wasm32"), feature = "cuda"))]
44    pub(crate) fn prewarm_cuda_weight(
45        info: &GgufTensorInfo,
46        raw: &[u8],
47        n_in: usize,
48        n_out: usize,
49    ) -> bool {
50        if !crate::prefer_tensor_core_gemm() {
51            return false;
52        }
53        if n_in < 16 || n_out < 16 {
54            return false;
55        }
56        // Cap densify cost: one matrix ≤ ~64 MiB f32 (fits 3B FFN rows on A2000 headroom).
57        if n_in.saturating_mul(n_out).saturating_mul(4) > 64 * 1024 * 1024 {
58            return false;
59        }
60        let key = crate::weight_fingerprint(raw, n_in, n_out);
61        if crate::dense_weight_cached(key) {
62            return true;
63        }
64        // Dequantize directly into the 2 MiB-aligned huge-page buffer — skips
65        // the intermediate Vec<f32> allocation + copy that cache_dense_weight
66        // would require. Rayon parallelizes over rows.
67        use crate::ggml_quants::dequant_matrix_row_into;
68        use rayon::prelude::*;
69        crate::inference::cuda_lane::cache_dense_weight_direct(key, n_in, n_out, |buf| {
70            buf.par_chunks_mut(n_in)
71                .enumerate()
72                .map(|(r, row)| dequant_matrix_row_into(raw, info, r, row).is_ok())
73                .all(|b| b)
74        })
75    }
76
77    /// Promote a 2-D quant weight (Q4_K / SoA / Q6_K / Q8_0) to a resident **f16** buffer
78    /// for the fast coop GEMV path. Returns `(buffer, ggml_type=F16, byte_len, row_elems)`.
79    /// `None` when disabled, unsupported type, or OOM/dequant failure — caller keeps quant.
80    #[cfg(not(target_arch = "wasm32"))]
81    pub(crate) fn promote_matrix_to_f16_resident(
82        &self,
83        info: &GgufTensorInfo,
84        raw: &[u8],
85    ) -> Option<(wgpu::Buffer, u32, u32, u32)> {
86        if !crate::llm_bench::ffn_f16_enabled() {
87            return None;
88        }
89        use crate::ggml_quants::{
90            dequant_matrix_row_into, GGML_TYPE_F16, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K_SOA,
91            GGML_TYPE_Q6_K, GGML_TYPE_Q8_0,
92        };
93        if !matches!(
94            info.ggml_type,
95            GGML_TYPE_Q4_K | GGML_TYPE_Q4_K_SOA | GGML_TYPE_Q6_K | GGML_TYPE_Q8_0
96        ) {
97            return None;
98        }
99        if info.n_dims < 2 || info.dims[0] == 0 || info.dims[1] == 0 {
100            return None;
101        }
102        let n0 = info.dims[0] as usize; // in (row width)
103        let n1 = info.dims[1] as usize; // out (rows)
104        let nbytes = n0.checked_mul(n1)?.checked_mul(2)?;
105        // Skip absurd expansions (e.g. accidental full-model promote) — FFN matrices are
106        // typically ≤ ~200 MiB each on 3B-class models.
107        if nbytes > 512 * 1024 * 1024 {
108            return None;
109        }
110        let mut f16_bytes = vec![0u8; nbytes];
111        let mut row = vec![0f32; n0];
112        for r in 0..n1 {
113            if dequant_matrix_row_into(raw, info, r, &mut row).is_err() {
114                return None;
115            }
116            let base = r * n0 * 2;
117            for (c, &v) in row.iter().enumerate() {
118                let bits = half::f16::from_f32(v).to_le_bytes();
119                f16_bytes[base + c * 2] = bits[0];
120                f16_bytes[base + c * 2 + 1] = bits[1];
121            }
122        }
123        let key = (raw.as_ptr() as u64) ^ F16_PROMOTE_KEY_TAG;
124        let buf = self.resident_weight_buffer(key, &f16_bytes)?;
125        Some((buf, GGML_TYPE_F16, nbytes as u32, n0 as u32))
126    }
127
128    pub(crate) fn write_weight_words(&self, raw: &[u8], max_bytes: usize) {
129        let weight_buf = self.gemm_weight_buf.as_ref().expect("gemm weight buf");
130        let upload = if raw.len() <= max_bytes {
131            raw
132        } else {
133            &raw[..max_bytes]
134        };
135        self.gpu_queue().write_buffer(weight_buf, 0, upload);
136    }
137
138    /// Phase 2: get-or-create the resident VRAM buffer for a weight byte-region, keyed by `key`
139    /// (the region's absolute mmap address — unique per distinct weight, stable across tokens). The
140    /// bytes (from the immutable mmap) are uploaded **once** on first use, then this returns a clone
141    /// of the resident buffer handle (wgpu buffers are Arc-backed) to bind in place of the shared
142    /// per-token `gemm_weight_buf` — eliminating the per-token weight re-upload. Buffer size is
143    /// 256-aligned ≥ `raw.len()`; the shader only reads `weight_byte_len`.
144    #[cfg(not(target_arch = "wasm32"))]
145    pub(crate) fn resident_weight_buffer(&self, key: u64, raw: &[u8]) -> Option<wgpu::Buffer> {
146        let mut map = self.gemm_resident_weights.lock().ok()?;
147        if let Some(b) = map.get(&key) {
148            return Some(b.clone());
149        }
150        let size = (((raw.len() + 255) & !255).max(4)) as wgpu::BufferAddress;
151        let buf = self.gpu_device().create_buffer(&wgpu::BufferDescriptor {
152            label: Some("ResidentWeight"),
153            size,
154            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
155            mapped_at_creation: false,
156        });
157        self.gpu_queue().write_buffer(&buf, 0, raw);
158        map.insert(key, buf.clone());
159        Some(buf)
160    }
161
162    /// Quantized GEMM from a pre-sliced weight byte range (chunk-local row indices).
163    #[cfg(not(target_arch = "wasm32"))]
164    pub(crate) fn dispatch_gemm_raw_into(
165        &self,
166        info: &GgufTensorInfo,
167        raw: &[u8],
168        input: &[f32],
169        out: &mut [f32],
170        n_in: usize,
171        n_out: usize,
172    ) -> bool {
173        if n_in > input.len() || n_out > out.len() {
174            return false;
175        }
176
177        // Mode=cuda acceleration (fail closed → wgpu path below).
178        #[cfg(all(not(target_arch = "wasm32"), feature = "cuda"))]
179        {
180            if crate::prefer_tensor_core_gemm() {
181                // M2b: on-device Q4_K SoA dequant-GEMV (no host densify) — preferred for .soa.p64.
182                // This path has CPU differential tests; keep it on for mode=cuda.
183                if info.ggml_type == crate::ggml_quants::GGML_TYPE_Q4_K_SOA
184                    && crate::try_q4k_soa_gemv(n_in, n_out, &input[..n_in], raw, &mut out[..n_out])
185                {
186                    return true;
187                }
188                // Dense densify + CUDA GEMM for F16/Q8/etc. was measured **incoherent** on
189                // single-token decode (pad-to-16 + f16-WMMA reduced path → garbage tokens,
190                // 2026-07-24 SmolLM f16 package). Default OFF for decode GEMV (always batch=1
191                // here). Opt-in: QUALIA_LLM_CUDA_TC_DECODE=1 (lab only). Prefill should use
192                // multi-token dispatch or the resident mega-pass, not this broken default.
193                let densify_tc_decode = matches!(
194                    std::env::var("QUALIA_LLM_CUDA_TC_DECODE").ok().as_deref(),
195                    Some("1") | Some("true")
196                );
197                if densify_tc_decode
198                    && n_in >= 16
199                    && n_out >= 16
200                    && n_in.saturating_mul(n_out) <= crate::inference::cuda_lane::MAX_DENSE_ELEMS
201                {
202                    let key = crate::weight_fingerprint(raw, n_in, n_out);
203                    if crate::try_cuda_batch_gemv_cached_only(
204                        key,
205                        &input[..n_in],
206                        1,
207                        &mut out[..n_out],
208                    ) {
209                        return true;
210                    }
211                    if let Some(dense) = Self::dequant_weight_dense_f32(info, raw, n_in, n_out) {
212                        if crate::try_cuda_batch_gemv_cached(
213                            key,
214                            &input[..n_in],
215                            1,
216                            n_in,
217                            n_out,
218                            &dense,
219                            &mut out[..n_out],
220                        ) {
221                            return true;
222                        }
223                    }
224                }
225            }
226        }
227
228        let weight_bytes = raw.len();
229        // GEMM shader supports a wider quant set than the legacy `ggml_gpu_quant_supported` (Q4_K/Q6_K)
230        // — notably Q8_0, which was silently falling back to the CPU `stack_gemm_quant` below (the FFN
231        // bottleneck for Q8_0 models). The guards (size/buffer caps) still fail closed → CPU fallback.
232        if ggml_gpu_gemm_supported(info.ggml_type)
233            && n_in <= MAX_STACK_GEMM_IN
234            && n_out <= self.gemm_max_out_dim as usize
235            && weight_bytes <= self.max_tensor_bytes
236            && self.gemm_input_buf.is_some()
237        {
238            let params = GemmGpuParams {
239                n_in: n_in as u32,
240                n_out: n_out as u32,
241                weight_ggml_type: info.ggml_type,
242                weight_row_elems: info.dims[0] as u32,
243                weight_byte_len: raw.len() as u32,
244                n_batch: 1,
245                in_row_stride: 0,
246                out_row_stride: 0,
247            };
248            let input_buf = self.gemm_input_buf.as_ref().unwrap();
249            let weight_buf = self.gemm_weight_buf.as_ref().unwrap();
250            let output_buf = self.gemm_output_buf.as_ref().unwrap();
251            let params_buf = self.gemm_params_buf.as_ref().unwrap();
252            let staging = self.gemm_output_staging.as_ref().unwrap();
253
254            self.gpu_queue()
255                .write_buffer(input_buf, 0, bytemuck::cast_slice(&input[..n_in]));
256            // Phase 2 (native): bind this tensor's resident VRAM buffer (uploaded once, keyed by
257            // byte_offset) instead of re-uploading the weight into the shared gemm_weight_buf every
258            // token. On wasm the resident path is the MC8 arena, so this stays the per-token upload.
259            #[cfg(not(target_arch = "wasm32"))]
260            let resident = if crate::llm_bench::resident_weights_enabled() {
261                // Key on the chunk's absolute mmap address, NOT `info.byte_offset`: the output
262                // projection passes the SAME whole-tensor `info` (byte_offset == header size) for
263                // every vocab chunk, so byte_offset aliases all chunks to one buffer (wrong logits).
264                // `raw` is a slice of the immutable, lifetime-stable mmap → its start address is
265                // unique per distinct weight region and identical across tokens.
266                self.resident_weight_buffer(raw.as_ptr() as u64, raw)
267            } else {
268                None
269            };
270            #[cfg(target_arch = "wasm32")]
271            let resident: Option<wgpu::Buffer> = None;
272            let weight_binding: &wgpu::Buffer = match resident.as_ref() {
273                Some(r) => r,
274                None => {
275                    self.write_weight_words(raw, self.max_tensor_bytes);
276                    weight_buf
277                }
278            };
279            self.gpu_queue()
280                .write_buffer(params_buf, 0, bytemuck::bytes_of(&params));
281
282            // 0.0.21: select the cooperative GEMV kernel (one workgroup/row, coalesced + shared-mem
283            // reduction) when enabled, else the naive 1-thread/row kernel. Same group-0 bindings, so
284            // only the pipeline + dispatch geometry differ. The bind group must be built from the
285            // ACTIVE pipeline's auto-layout.
286            #[cfg(not(target_arch = "wasm32"))]
287            let use_coop = crate::llm_bench::coop_gemv_enabled();
288            #[cfg(target_arch = "wasm32")]
289            let use_coop = false;
290            #[cfg(not(target_arch = "wasm32"))]
291            let use_mr = use_coop
292                && info.ggml_type == crate::ggml_quants::GGML_TYPE_Q4_K_SOA
293                && n_out >= 512;
294            #[cfg(not(target_arch = "wasm32"))]
295            let active_pipeline: &wgpu::ComputePipeline = if use_mr {
296                &self.coop_gemv_mr_pipeline
297            } else if use_coop {
298                &self.coop_gemv_pipeline
299            } else {
300                &self.pipeline
301            };
302            #[cfg(target_arch = "wasm32")]
303            let use_mmv_q8_0 =
304                info.ggml_type == crate::ggml_quants::GGML_TYPE_Q8_0 && (n_in % 32 == 0);
305            #[cfg(target_arch = "wasm32")]
306            let active_pipeline: &wgpu::ComputePipeline = if use_mmv_q8_0 {
307                &self.mmv_q8_0_pipeline
308            } else {
309                &self.pipeline
310            };
311            #[cfg(target_arch = "wasm32")]
312            let use_mr = false;
313
314            #[cfg(not(target_arch = "wasm32"))]
315            let bind_layout = self.native_gemm_bind_layout(use_coop).clone();
316            #[cfg(target_arch = "wasm32")]
317            let bind_layout = active_pipeline.get_bind_group_layout(0);
318            // CoopGemvBGL is 5-slot (binding 4 = residual). Dummy residual = input.
319            let bind_group = if use_coop {
320                self.gpu_device()
321                    .create_bind_group(&wgpu::BindGroupDescriptor {
322                        label: Some("LayerGemmBindGroup"),
323                        layout: &bind_layout,
324                        entries: &[
325                            wgpu::BindGroupEntry {
326                                binding: 0,
327                                resource: input_buf.as_entire_binding(),
328                            },
329                            wgpu::BindGroupEntry {
330                                binding: 1,
331                                resource: weight_binding.as_entire_binding(),
332                            },
333                            wgpu::BindGroupEntry {
334                                binding: 2,
335                                resource: params_buf.as_entire_binding(),
336                            },
337                            wgpu::BindGroupEntry {
338                                binding: 3,
339                                resource: output_buf.as_entire_binding(),
340                            },
341                            wgpu::BindGroupEntry {
342                                binding: 4,
343                                resource: input_buf.as_entire_binding(),
344                            },
345                        ],
346                    })
347            } else {
348                self.gpu_device()
349                    .create_bind_group(&wgpu::BindGroupDescriptor {
350                        label: Some("LayerGemmBindGroup"),
351                        layout: &bind_layout,
352                        entries: &[
353                            wgpu::BindGroupEntry {
354                                binding: 0,
355                                resource: input_buf.as_entire_binding(),
356                            },
357                            wgpu::BindGroupEntry {
358                                binding: 1,
359                                resource: weight_binding.as_entire_binding(),
360                            },
361                            wgpu::BindGroupEntry {
362                                binding: 2,
363                                resource: params_buf.as_entire_binding(),
364                            },
365                            wgpu::BindGroupEntry {
366                                binding: 3,
367                                resource: output_buf.as_entire_binding(),
368                            },
369                        ],
370                    })
371            };
372
373            let mut encoder =
374                self.device()
375                    .create_command_encoder(&wgpu::CommandEncoderDescriptor {
376                        label: Some("LayerGemmEncoder"),
377                    });
378            {
379                let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
380                    label: None,
381                    timestamp_writes: crate::llm_gpu_profiler::pass_writes_both(),
382                });
383                cpass.set_pipeline(active_pipeline);
384                cpass.set_bind_group(0, &bind_group, &[]);
385                if use_mr {
386                    cpass.dispatch_workgroups(
387                        crate::llm_bench::coop_gemv_workgroups(n_out as u32),
388                        1,
389                        1,
390                    );
391                } else if use_coop {
392                    cpass.dispatch_workgroups(n_out as u32, 1, 1);
393                } else {
394                    #[cfg(target_arch = "wasm32")]
395                    if use_mmv_q8_0 {
396                        cpass.dispatch_workgroups((n_out as u32 + 3) / 4, 1, 1);
397                    } else {
398                        cpass.dispatch_workgroups((n_out as u32 + 63) / 64, 1, 1);
399                    }
400                    #[cfg(not(target_arch = "wasm32"))]
401                    cpass.dispatch_workgroups((n_out as u32 + 63) / 64, 1, 1);
402                }
403            }
404            let out_bytes = (n_out * 4) as wgpu::BufferAddress;
405            encoder.copy_buffer_to_buffer(output_buf, 0, staging, 0, out_bytes);
406            crate::llm_gpu_profiler::resolve(&mut encoder);
407            self.gpu_queue().submit(Some(encoder.finish()));
408            crate::llm_gpu_profiler::accumulate(crate::llm_gpu_profiler::Phase::Gemm);
409
410            let slice = staging.slice(..out_bytes);
411            let (tx, rx) = futures_channel::oneshot::channel();
412            slice.map_async(wgpu::MapMode::Read, move |r| {
413                let _ = tx.send(r);
414            });
415            self.poll_wait();
416            #[cfg(not(target_arch = "wasm32"))]
417            if let Ok(handle) = tokio::runtime::Handle::try_current() {
418                if handle.block_on(rx).ok().map(|m| m.is_ok()).unwrap_or(false) {
419                    let data = slice
420                        .get_mapped_range()
421                        .expect("wgpu buffer map_range failed");
422                    let floats: &[f32] = bytemuck::cast_slice(&data);
423                    out[..n_out].copy_from_slice(&floats[..n_out]);
424                    drop(data);
425                    staging.unmap();
426                    return true;
427                }
428            }
429            let _ = staging.unmap();
430        }
431
432        stack_gemm_quant(raw, info, input, out, n_in, n_out)
433    }
434
435    /// Browser compatibility path for synchronous callers. WebGPU mapping must
436    /// be awaited, so the public browser inference route uses the async GEMM
437    /// dispatcher and this fallback executes the deterministic CPU kernel.
438    #[cfg(target_arch = "wasm32")]
439    pub(crate) fn dispatch_gemm_raw_into(
440        &self,
441        info: &GgufTensorInfo,
442        raw: &[u8],
443        input: &[f32],
444        out: &mut [f32],
445        n_in: usize,
446        n_out: usize,
447    ) -> bool {
448        if n_in > input.len() || n_out > out.len() {
449            return false;
450        }
451        stack_gemm_quant(raw, info, input, out, n_in, n_out)
452    }
453
454    /// A1b: dispatch one **ternary** FFN GEMM (`GGML_TYPE_TERNARY_158`). The resident 2-bit GPU
455    /// kernel when the toggle is on AND the tensor is resident; otherwise the CPU base-3 oracle on
456    /// the blob fetched from the mmap. Fail-closed (returns false, never garbage) on any mismatch.
457    pub(crate) fn dispatch_ternary_ffn(
458        &self,
459        info: &GgufTensorInfo,
460        input: &[f32],
461        out: &mut [f32],
462        n_in: usize,
463        n_out: usize,
464    ) -> bool {
465        if n_in > input.len() || n_out > out.len() {
466            return false;
467        }
468        // GPU resident path (toggle on): keyed by the P64 blob offset (== info.byte_offset).
469        #[cfg(not(target_arch = "wasm32"))]
470        if crate::llm_bench::ternary_ffn_enabled() {
471            if let Some(res) = self.ternary_ffn.as_ref() {
472                if res.gemv(
473                    self.gpu_device(),
474                    self.gpu_queue(),
475                    info.byte_offset,
476                    input,
477                    out,
478                    n_in,
479                    n_out,
480                ) {
481                    return true;
482                }
483            }
484        }
485        // CPU oracle fallback (toggle off, or GPU unavailable): the SAME ternary weights via the
486        // base-3 CPU GEMM — correct, slower; this is the toggle's OFF baseline.
487        let mmap = match self.gguf_mmap.as_deref() {
488            Some(m) => m,
489            None => return false,
490        };
491        let raw = match crate::ggml_quants::fetch_tensor_bytes(mmap, self.tensor_data_offset, info)
492        {
493            Ok(s) => s,
494            Err(_) => return false,
495        };
496        if raw.len() < 4 {
497            return false;
498        }
499        let scale = f32::from_le_bytes([raw[0], raw[1], raw[2], raw[3]]);
500        crate::ternary::ternary_gemm_cpu(
501            &input[..n_in],
502            &raw[4..],
503            scale,
504            n_in,
505            n_out,
506            1,
507            0,
508            0,
509            &mut out[..n_out],
510        );
511        true
512    }
513
514    /// Quantized GEMM into caller `out` using reused GPU buffers (Q6_K) or CPU dequant fallback.
515    pub fn dispatch_gemm_into(
516        &self,
517        index: &crate::gguf_sharder::GgufTensorIndex,
518        info: &GgufTensorInfo,
519        input: &[f32],
520        out: &mut [f32],
521        n_in: usize,
522        n_out: usize,
523    ) -> bool {
524        if n_in > input.len() || n_out > out.len() {
525            wlog(&format!(
526                "[gemm_into] GUARD n_in={n_in} n_out={n_out} input={} out={}",
527                input.len(),
528                out.len()
529            ));
530            return false;
531        }
532        // A1b: ternary FFN tensors are not row-block quantized — route them to the dedicated ternary
533        // dispatch (resident 2-bit GPU kernel / CPU oracle) before the standard fetch+GEMM path.
534        if info.ggml_type == crate::ternary::GGML_TYPE_TERNARY_158 {
535            return self.dispatch_ternary_ffn(info, input, out, n_in, n_out);
536        }
537        let mmap = match self.gguf_mmap.as_deref() {
538            Some(m) => m,
539            None => return false,
540        };
541        let raw = match crate::ggml_quants::fetch_tensor_bytes(mmap, index.tensor_data_start, info)
542        {
543            Ok(s) => s,
544            Err(_) => return false,
545        };
546        self.dispatch_gemm_raw_into(info, raw, input, out, n_in, n_out)
547    }
548}
549
550/// Substrate-parity proof (CPU, no GPU): the LLM's quantized GEMV is *the same linear
551/// operation* as the engine's dense `solvers::linear_algebra::gemm`.
552///
553/// The LLM forward path is not a bespoke "AI inference" kernel — its weight×activation
554/// step is matrix–vector multiplication `out[i] = Σ_j W[i][j]·x[j]`, with `W` dequantized
555/// on the fly. This proves it: dequantize the quantized weights to a dense matrix, run the
556/// engine's `matvec` on them, and show the LLM kernel (`stack_gemm_quant`) agrees to f32
557/// rounding. Together with the existing GPU↔CPU probe (`gemm_parity_probe`), this closes
558/// the chain  substrate GEMM ≡ LLM CPU GEMV ≡ LLM GPU GEMV.
559#[cfg(test)]
560mod substrate_parity_tests {
561    use crate::gguf_sharder::GgufTensorInfo;
562    use crate::solvers::linear_algebra::gemm::{matvec, Transpose};
563
564    /// Returns `(exact_err, quant_err)`:
565    /// - `exact_err` = max|LLM_kernel(Q8(W)) − substrate(dequant(Q8(W)))| — should be ~f32 ε,
566    ///   proving the two compute the *same* operation;
567    /// - `quant_err` = max|LLM_kernel(Q8(W)) − substrate(W_original)| — the Q8 quantization cost.
568    fn run(n_in: usize, n_out: usize, seed: u64) -> (f32, f32) {
569        // Deterministic LCG → values in [-1, 1) (mirrors gemm_parity_probe_blocking).
570        let mut s = seed | 1;
571        let mut rng = move || -> f32 {
572            s = s
573                .wrapping_mul(6364136223846793005)
574                .wrapping_add(1442695040888963407);
575            ((s >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0
576        };
577
578        let row_bytes = crate::llm_kernel_parity::q8_0_bytes(n_in);
579        let mut raw = vec![0u8; row_bytes * n_out];
580        let mut w_orig = vec![0f32; n_in * n_out]; // dense, pre-quantization (row-major n_out×n_in)
581        let mut row_f32 = vec![0f32; n_in];
582        for r in 0..n_out {
583            for x in row_f32.iter_mut() {
584                *x = rng();
585            }
586            w_orig[r * n_in..(r + 1) * n_in].copy_from_slice(&row_f32);
587            assert!(crate::llm_kernel_parity::quantize_q8_0_from_f32(
588                &row_f32,
589                &mut raw[r * row_bytes..(r + 1) * row_bytes],
590            ));
591        }
592        let input: Vec<f32> = (0..n_in).map(|_| rng()).collect();
593
594        let info = GgufTensorInfo {
595            dims: [n_in as u64, n_out as u64, 1, 1],
596            n_dims: 2,
597            ggml_type: crate::ggml_quants::GGML_TYPE_Q8_0,
598            byte_offset: 0,
599        };
600
601        // (1) The actual LLM CPU kernel.
602        let mut out_llm = vec![0f32; n_out];
603        assert!(crate::gguf_bridge::stack_gemm_quant(
604            &raw,
605            &info,
606            &input,
607            &mut out_llm,
608            n_in,
609            n_out
610        ));
611
612        // (2) Dequantize the same quantized weights to a dense matrix, then run the
613        //     engine's GEMM (matvec) on it. Same operands ⇒ must match the LLM kernel.
614        let mut w_deq = vec![0f64; n_in * n_out];
615        let mut deq_row = vec![0f32; n_in];
616        for i in 0..n_out {
617            let got = crate::ggml_quants::dequant_matrix_row_into(&raw, &info, i, &mut deq_row)
618                .unwrap_or(0);
619            assert_eq!(got, n_in, "dequant row {i}");
620            for j in 0..n_in {
621                w_deq[i * n_in + j] = deq_row[j] as f64;
622            }
623        }
624        let x_f64: Vec<f64> = input.iter().map(|&v| v as f64).collect();
625        let mut out_sub_deq = vec![0f64; n_out];
626        matvec(Transpose::No, n_out, n_in, &w_deq, &x_f64, &mut out_sub_deq).unwrap();
627
628        // (3) The engine GEMM on the ORIGINAL (pre-quant) weights — the Q8 cost reference.
629        let w_orig_f64: Vec<f64> = w_orig.iter().map(|&v| v as f64).collect();
630        let mut out_sub_orig = vec![0f64; n_out];
631        matvec(
632            Transpose::No,
633            n_out,
634            n_in,
635            &w_orig_f64,
636            &x_f64,
637            &mut out_sub_orig,
638        )
639        .unwrap();
640
641        let exact_err = (0..n_out)
642            .map(|i| (out_llm[i] as f64 - out_sub_deq[i]).abs() as f32)
643            .fold(0.0f32, f32::max);
644        let quant_err = (0..n_out)
645            .map(|i| (out_llm[i] as f64 - out_sub_orig[i]).abs() as f32)
646            .fold(0.0f32, f32::max);
647        (exact_err, quant_err)
648    }
649
650    #[test]
651    fn llm_quant_gemv_is_the_substrate_gemm() {
652        // Several shapes/seeds; n_in a multiple of 32 (Q8_0 block size).
653        for &(n_in, n_out, seed) in &[
654            (64usize, 32usize, 0xC0FFEEu64),
655            (128, 96, 7),
656            (256, 64, 0xBEEF),
657        ] {
658            let (exact_err, quant_err) = run(n_in, n_out, seed);
659            // Same operation: LLM kernel == engine GEMM on identical (dequantized) weights,
660            // to f32 accumulation rounding only.
661            assert!(
662                exact_err < 1e-4,
663                "LLM GEMV diverges from substrate GEMM on identical weights: exact_err={exact_err} (n_in={n_in}, n_out={n_out})"
664            );
665            // Quantization is the *only* extra divergence from exact math, and it is bounded.
666            assert!(
667                quant_err < 0.5,
668                "Q8 quantization error unexpectedly large: quant_err={quant_err}"
669            );
670        }
671    }
672}