1use super::*;
4
5#[cfg(not(target_arch = "wasm32"))]
7const F16_PROMOTE_KEY_TAG: u64 = 0xF16E_F16E_0000_00F1;
8
9impl QTensorEngine {
10 #[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 #[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 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 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 #[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; let n1 = info.dims[1] as usize; let nbytes = n0.checked_mul(n1)?.checked_mul(2)?;
105 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 #[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 #[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 #[cfg(all(not(target_arch = "wasm32"), feature = "cuda"))]
179 {
180 if crate::prefer_tensor_core_gemm() {
181 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 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 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 #[cfg(not(target_arch = "wasm32"))]
260 let resident = if crate::llm_bench::resident_weights_enabled() {
261 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(¶ms));
281
282 #[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 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 #[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 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 #[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 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 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 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#[cfg(test)]
560mod substrate_parity_tests {
561 use crate::gguf_sharder::GgufTensorInfo;
562 use crate::solvers::linear_algebra::gemm::{matvec, Transpose};
563
564 fn run(n_in: usize, n_out: usize, seed: u64) -> (f32, f32) {
569 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]; 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 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 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 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 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 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 assert!(
667 quant_err < 0.5,
668 "Q8 quantization error unexpectedly large: quant_err={quant_err}"
669 );
670 }
671 }
672}