1use cortiq_core::CmfModel;
14use std::cell::Cell;
15use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering};
16use std::sync::{Arc, OnceLock};
17
18thread_local! {
19 static CUR_LAYER: Cell<i64> = const { Cell::new(-1) };
23 static CPU_ONLY: Cell<bool> = const { Cell::new(false) };
28 static PROBE_COLD: Cell<bool> = const { Cell::new(false) };
32}
33
34pub fn cpu_scope<R>(f: impl FnOnce() -> R) -> R {
36 struct Restore(bool);
37 impl Drop for Restore {
38 fn drop(&mut self) {
39 CPU_ONLY.with(|c| c.set(self.0));
40 }
41 }
42 let previous = CPU_ONLY.with(|c| c.replace(true));
43 let _restore = Restore(previous);
44 f()
45}
46
47pub fn probe_set_device(label: &str) {
52 let _ = DEVICE_LABEL.set(label.to_string());
53}
54
55fn device_label() -> &'static str {
56 DEVICE_LABEL.get().map(String::as_str).unwrap_or("unknown")
57}
58
59static DEVICE_LABEL: std::sync::OnceLock<String> = std::sync::OnceLock::new();
60
61static CACHE_DIR: std::sync::OnceLock<std::path::PathBuf> = std::sync::OnceLock::new();
70
71pub fn set_cache_dir(dir: std::path::PathBuf) {
73 let _ = CACHE_DIR.set(dir);
74}
75
76pub fn cache_dir_pub() -> std::path::PathBuf {
78 cache_dir()
79}
80
81fn cache_dir() -> std::path::PathBuf {
82 if let Some(d) = CACHE_DIR.get() {
83 return d.clone();
84 }
85 match std::env::var_os("TMPDIR") {
86 Some(t) => std::path::PathBuf::from(t),
87 None => std::env::temp_dir(),
88 }
89}
90
91fn probe_cache_path() -> Option<std::path::PathBuf> {
94 match std::env::var("CMF_PROBE_CACHE") {
95 Ok(v) if v == "0" => None,
96 Ok(v) => Some(std::path::PathBuf::from(v)),
97 Err(_) => Some(cache_dir().join("cortiq-gpu-probe.tsv")),
98 }
99}
100
101fn probe_cache_key_named(class: &str) -> String {
105 format!("{}\t{}\t{}", env!("CARGO_PKG_VERSION"), device_label(), class)
106}
107
108const CLASS_NAMES: [&str; 7] = [
109 "ffn",
110 "matvec",
111 "matmat",
112 "qkv-batch",
113 "matmat-wide",
114 "lm-head",
115 "gemm-nt",
116];
117
118fn probe_cache_load() {
127 static ONCE: std::sync::Once = std::sync::Once::new();
128 ONCE.call_once(|| {
129 let Some(path) = probe_cache_path() else {
130 return;
131 };
132 if cfg!(test) && std::env::var("CMF_PROBE_CACHE").is_err() {
137 return;
138 }
139 let Ok(text) = std::fs::read_to_string(&path) else {
140 return;
141 };
142 probe_cache_adopt(&text);
143 });
144}
145
146fn probe_cache_adopt(text: &str) {
150 for line in text.lines() {
151 let Some((key, verdict)) = line.rsplit_once('\t') else {
152 continue;
153 };
154 let winner = match verdict.trim() {
155 "gpu" => 1u8,
156 "cpu" => 2u8,
157 _ => continue,
158 };
159 for (i, name) in CLASS_NAMES.iter().enumerate() {
160 if probe_cache_key_named(name) == key {
161 let _ =
162 PROBES[i]
163 .state
164 .compare_exchange(0, winner, Ordering::Relaxed, Ordering::Relaxed);
165 tracing::debug!("gpu probe [{name}]: remembered → {verdict}");
166 }
167 }
168 }
169}
170
171fn probe_cache_store(c: OpClass, winner: u8) {
174 let Some(path) = probe_cache_path() else {
175 return;
176 };
177 let line = format!(
178 "{}\t{}\n",
179 probe_cache_key_named(CLASS_NAMES[c as usize]),
180 if winner == 1 { "gpu" } else { "cpu" }
181 );
182 use std::io::Write;
183 if let Ok(mut f) = std::fs::OpenOptions::new()
184 .create(true)
185 .append(true)
186 .open(&path)
187 {
188 let _ = f.write_all(line.as_bytes());
189 }
190}
191
192pub fn cold_epoch() -> u64 {
198 COLD_EPOCH.load(std::sync::atomic::Ordering::Relaxed)
199}
200static COLD_EPOCH: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
201
202pub(crate) fn probe_note_cold() {
203 COLD_EPOCH.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
204 PROBE_COLD.with(|c| c.set(true));
205}
206
207pub(crate) fn probe_was_cold() -> bool {
211 PROBE_COLD.with(|c| c.get())
212}
213
214pub fn set_layer(l: i64) {
216 CUR_LAYER.with(|c| c.set(l));
217}
218
219pub fn cur_layer() -> i64 {
221 CUR_LAYER.with(|c| c.get())
222}
223
224fn layer_ranges() -> &'static Option<Vec<(i64, i64)>> {
227 static R: OnceLock<Option<Vec<(i64, i64)>>> = OnceLock::new();
228 R.get_or_init(|| {
229 let s = std::env::var("CMF_GPU_LAYERS").ok()?;
230 let mut v = Vec::new();
231 for part in s.split(',') {
232 let part = part.trim();
233 match part.split_once('-') {
234 Some((a, b)) => v.push((a.trim().parse().ok()?, b.trim().parse().ok()?)),
235 None => {
236 let x: i64 = part.parse().ok()?;
237 v.push((x, x));
238 }
239 }
240 }
241 Some(v)
242 })
243}
244
245fn layer_allowed() -> bool {
246 match layer_ranges() {
247 None => true,
248 Some(ranges) => {
249 let cur = CUR_LAYER.with(|c| c.get());
250 cur < 0 || ranges.iter().any(|(a, b)| cur >= *a && cur <= *b)
251 }
252 }
253}
254
255pub fn enabled_here() -> bool {
259 !CPU_ONLY.with(|c| c.get()) && enabled() && layer_allowed()
260}
261
262#[derive(Clone, Copy)]
274pub enum OpClass {
275 Ffn = 0,
277 Matvec = 1,
279 Matmat = 2,
281 Batch = 3,
283 MatmatWide = 4,
289 MatvecHead = 5,
296 GemmNt = 6,
303}
304
305pub fn matvec_class(rows: usize, cols: usize) -> OpClass {
309 if rows * cols >= 67_108_864 {
310 OpClass::MatvecHead
311 } else {
312 OpClass::Matvec
313 }
314}
315
316pub enum ProbeArm {
318 Gpu,
320 CpuTimed,
322 Cpu,
324}
325
326const PROBE_SAMPLES: u32 = 6;
328
329struct Probe {
330 state: AtomicU8,
332 flip: AtomicU32,
333 gpu_ns: AtomicU64,
334 gpu_n: AtomicU32,
335 cpu_ns: AtomicU64,
336 cpu_n: AtomicU32,
337 gpu_min: AtomicU64,
344 cpu_min: AtomicU64,
345}
346
347impl Probe {
348 const fn new() -> Self {
349 Self {
350 state: AtomicU8::new(0),
351 flip: AtomicU32::new(0),
352 gpu_ns: AtomicU64::new(0),
353 gpu_n: AtomicU32::new(0),
354 cpu_ns: AtomicU64::new(0),
355 cpu_n: AtomicU32::new(0),
356 gpu_min: AtomicU64::new(u64::MAX),
357 cpu_min: AtomicU64::new(u64::MAX),
358 }
359 }
360}
361
362static PROBES: [Probe; 7] = [
363 Probe::new(),
364 Probe::new(),
365 Probe::new(),
366 Probe::new(),
367 Probe::new(),
368 Probe::new(),
369 Probe::new(),
370];
371
372fn probe_on() -> bool {
373 static ON: OnceLock<bool> = OnceLock::new();
374 *ON.get_or_init(|| {
375 std::env::var("CMF_GPU_PROBE")
376 .map(|v| v != "0" && v != "off")
377 .unwrap_or(true)
378 })
379}
380
381pub fn q1_force() -> bool {
386 #[cfg(target_os = "macos")]
387 {
388 backend() == Backend::Metal
389 }
390 #[cfg(not(target_os = "macos"))]
391 {
392 false
393 }
394}
395
396pub fn fused_block_trusted() -> bool {
415 #[cfg(target_os = "macos")]
416 if backend() == Backend::Metal {
417 return true;
418 }
419 wgpu_graph_default()
420}
421
422pub fn weight_is_resident(model: &Arc<CmfModel>, idx: usize) -> bool {
434 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
435 {
436 return crate::gpu_wgpu::weight_is_resident(model, idx);
437 }
438 #[cfg(not(all(feature = "gpu", not(target_os = "macos"))))]
439 {
440 let _ = (model, idx);
441 true
442 }
443}
444
445pub fn probe_arm_cold_prefers_gpu(c: OpClass, weights_resident: bool) -> ProbeArm {
446 if !weights_resident && probe_deciding(c) {
447 return ProbeArm::Gpu;
448 }
449 probe_arm(c)
450}
451
452pub fn probe_arm(c: OpClass) -> ProbeArm {
453 PROBE_COLD.with(|f| f.set(false));
458 if !probe_on() {
459 return ProbeArm::Gpu;
460 }
461 probe_cache_load();
462 let p = &PROBES[c as usize];
463 match p.state.load(Ordering::Relaxed) {
464 1 => ProbeArm::Gpu,
465 2 => ProbeArm::Cpu,
466 _ => {
467 if p.flip.fetch_add(1, Ordering::Relaxed) % 2 == 0 {
468 ProbeArm::Gpu
469 } else {
470 ProbeArm::CpuTimed
471 }
472 }
473 }
474}
475
476pub fn probe_record(c: OpClass, gpu: bool, dur: std::time::Duration) {
479 let p = &PROBES[c as usize];
480 if p.state.load(Ordering::Relaxed) != 0 {
481 return;
482 }
483 if gpu && PROBE_COLD.with(|f| f.replace(false)) {
484 return; }
486 let ns = dur.as_nanos().min(u64::MAX as u128) as u64;
487 if gpu {
488 p.gpu_ns.fetch_add(ns, Ordering::Relaxed);
489 p.gpu_n.fetch_add(1, Ordering::Relaxed);
490 p.gpu_min.fetch_min(ns, Ordering::Relaxed);
491 } else {
492 p.cpu_ns.fetch_add(ns, Ordering::Relaxed);
493 p.cpu_n.fetch_add(1, Ordering::Relaxed);
494 p.cpu_min.fetch_min(ns, Ordering::Relaxed);
495 }
496 let (gn, cn) = (
497 p.gpu_n.load(Ordering::Relaxed),
498 p.cpu_n.load(Ordering::Relaxed),
499 );
500 if gn >= 2 && cn >= 2 {
501 let g = p.gpu_min.load(Ordering::Relaxed) as f64;
505 let cp = p.cpu_min.load(Ordering::Relaxed) as f64;
506 if (gn < PROBE_SAMPLES || cn < PROBE_SAMPLES) && g < cp * 2.0 && cp < g * 2.0 {
516 return;
517 }
518 let winner = if g <= cp { 1 } else { 2 };
519 if p.state
520 .compare_exchange(0, winner, Ordering::Relaxed, Ordering::Relaxed)
521 .is_ok()
522 {
523 tracing::info!(
524 "gpu probe [{}]: gpu {:.2} ms vs cpu {:.2} ms per op → {}",
525 CLASS_NAMES[c as usize],
526 g / 1e6,
527 cp / 1e6,
528 if winner == 1 { "gpu" } else { "cpu" },
529 );
530 probe_cache_store(c, winner);
531 }
532 }
533}
534
535pub fn probe_deciding(c: OpClass) -> bool {
538 probe_on() && PROBES[c as usize].state.load(Ordering::Relaxed) == 0
539}
540
541#[allow(unused_variables)]
551pub fn q8_resident_or_upload(model: &Arc<CmfModel>, idx: usize) -> bool {
552 static PROBE_UPLOADS: AtomicU32 = AtomicU32::new(0);
553 let may_upload = PROBE_UPLOADS.load(Ordering::Relaxed) < 4;
554 let resident = match backend() {
555 #[cfg(target_os = "macos")]
556 Backend::Metal => crate::gpu_metal::q8_resident_or_upload(model, idx, may_upload),
557 #[cfg(feature = "gpu")]
558 Backend::Wgpu => crate::gpu_wgpu::q8_resident_or_upload(model, idx, may_upload),
559 Backend::None => false,
560 };
561 if !resident && may_upload {
562 PROBE_UPLOADS.fetch_add(1, Ordering::Relaxed);
563 }
564 resident
565}
566
567#[cfg(test)]
569pub(crate) fn probe_reset() {
570 for p in &PROBES {
571 p.state.store(0, Ordering::Relaxed);
572 p.flip.store(0, Ordering::Relaxed);
573 p.gpu_ns.store(0, Ordering::Relaxed);
574 p.gpu_n.store(0, Ordering::Relaxed);
575 p.cpu_ns.store(0, Ordering::Relaxed);
576 p.cpu_n.store(0, Ordering::Relaxed);
577 }
578}
579
580#[cfg(test)]
581mod probe_tests {
582 use super::*;
583 use std::time::Duration;
584
585 #[test]
588 fn probe_alternates_discards_cold_and_decides() {
589 probe_reset();
590 assert!(matches!(probe_arm(OpClass::Ffn), ProbeArm::Gpu));
592 assert!(matches!(probe_arm(OpClass::Ffn), ProbeArm::CpuTimed));
593
594 probe_note_cold();
598 probe_record(OpClass::Ffn, true, Duration::from_secs(1000));
599 for _ in 0..PROBE_SAMPLES {
600 probe_record(OpClass::Ffn, true, Duration::from_millis(1));
601 probe_record(OpClass::Ffn, false, Duration::from_millis(4));
602 }
603 assert!(matches!(probe_arm(OpClass::Ffn), ProbeArm::Gpu));
604
605 for _ in 0..PROBE_SAMPLES {
607 probe_record(OpClass::Matmat, true, Duration::from_millis(4));
608 probe_record(OpClass::Matmat, false, Duration::from_millis(1));
609 }
610 assert!(matches!(probe_arm(OpClass::Matmat), ProbeArm::Cpu));
611
612 cpu_scope(|| CPU_ONLY.with(|c| assert!(c.get())));
614 CPU_ONLY.with(|c| assert!(!c.get()));
615 cpu_scope(|| {
616 cpu_scope(|| CPU_ONLY.with(|c| assert!(c.get())));
617 CPU_ONLY.with(|c| assert!(c.get()));
618 });
619 let _ = std::panic::catch_unwind(|| cpu_scope(|| panic!("scope test")));
620 CPU_ONLY.with(|c| assert!(!c.get()));
621 probe_reset();
622 }
623
624 #[test]
625 fn a_remembered_verdict_is_adopted_and_a_stranger_is_not() {
626 let mine = probe_cache_key_named("gemm-nt");
638 let state = || PROBES[OpClass::GemmNt as usize].state.load(Ordering::Relaxed);
639
640 probe_cache_adopt("SomeOtherGPU/Vulkan\tgemm-nt\tgpu\n");
642 assert_eq!(state(), 0);
643 let older = mine.replacen(env!("CARGO_PKG_VERSION"), "0.0.0-old", 1);
645 assert_ne!(older, mine);
646 probe_cache_adopt(&format!("{older}\tgpu\n"));
647 assert_eq!(state(), 0);
648 probe_cache_adopt(&format!("{mine}\tcpu\n"));
650 assert_eq!(state(), 2);
651
652 PROBES[OpClass::GemmNt as usize]
653 .state
654 .store(0, Ordering::Relaxed);
655 }
656}
657
658pub const GPU_MIN_ROWS: usize = 65_536;
661
662pub fn min_rows() -> usize {
669 if let Some(v) = std::env::var("CMF_GPU_MIN_ROWS")
670 .ok()
671 .and_then(|v| v.parse().ok())
672 {
673 return v;
674 }
675 if discrete() { 4096 } else { GPU_MIN_ROWS }
676}
677
678pub fn discrete() -> bool {
680 match backend() {
681 #[cfg(feature = "gpu")]
682 Backend::Wgpu => crate::gpu_wgpu::is_discrete(),
683 #[cfg(target_os = "macos")]
684 Backend::Metal => false, Backend::None => false,
686 }
687}
688
689pub struct MoeJob<'a> {
693 pub gate: (usize, usize, usize, &'a [f32]),
694 pub up: (usize, usize, usize, &'a [f32]),
695 pub down: (usize, usize, usize, &'a [f32]),
696 pub xs_gate: Vec<f32>,
697 pub xs_up: Vec<f32>,
698 pub down_col: &'a [f32],
699 pub w: f32,
700 pub q1: bool,
703 pub q4t: bool,
706 pub q4tp: bool,
710 pub gu_q2: bool,
714 pub swiglu_limit: f32,
719}
720
721pub struct BatchJob<'a> {
723 pub idx: usize,
724 pub rows: usize,
725 pub cols: usize,
726 pub row_scale: &'a [f32],
727 pub xs: Vec<f32>,
728 pub layout: BatchLayout,
732}
733
734#[derive(Clone, Copy, PartialEq, Eq, Debug)]
737pub enum BatchLayout {
738 Q8,
739 Q1,
740 Q4t,
741 Q4tp,
742}
743
744#[derive(Clone, Copy, PartialEq, Eq)]
745enum Backend {
746 None,
747 #[cfg(target_os = "macos")]
748 Metal,
749 #[cfg(feature = "gpu")]
750 Wgpu,
751}
752
753fn backend() -> Backend {
754 #[cfg(feature = "gpu")]
755 if crate::gpu_wgpu::selected() {
756 return if crate::gpu_wgpu::enabled() {
757 Backend::Wgpu
758 } else {
759 Backend::None
760 };
761 }
762 #[cfg(target_os = "macos")]
763 if crate::gpu_metal::enabled() {
764 return Backend::Metal;
765 }
766 Backend::None
767}
768
769pub fn backend_available() -> bool {
775 #[cfg(target_os = "macos")]
776 {
777 true
779 }
780 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
781 {
782 static AVAIL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
783 *AVAIL.get_or_init(crate::gpu_wgpu::adapter_probe)
784 }
785 #[cfg(all(not(feature = "gpu"), not(target_os = "macos")))]
786 {
787 false
788 }
789}
790
791pub fn enabled() -> bool {
792 backend() != Backend::None
793}
794
795pub fn wgpu_active() -> bool {
809 #[cfg(feature = "gpu")]
810 {
811 matches!(backend(), Backend::Wgpu)
812 }
813 #[cfg(not(feature = "gpu"))]
814 {
815 false
816 }
817}
818
819pub fn default_device() -> usize {
826 static D: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
827 *D.get_or_init(|| {
828 std::env::var("CMF_GPU_ADAPTER")
829 .ok()
830 .and_then(|v| v.trim().parse::<usize>().ok())
831 .unwrap_or(0)
832 })
833}
834
835thread_local! {
836 static CUR_DEV: std::cell::Cell<Option<usize>> = const { std::cell::Cell::new(None) };
837}
838
839pub fn current_device() -> usize {
841 CUR_DEV.with(|c| c.get()).unwrap_or_else(default_device)
842}
843
844pub fn set_current_device(i: usize) {
848 CUR_DEV.with(|c| c.set(Some(i)));
849}
850
851pub fn with_device<R>(dev: usize, f: impl FnOnce() -> R) -> R {
853 let prev = CUR_DEV.with(|c| c.replace(Some(dev)));
854 let r = f();
855 CUR_DEV.with(|c| c.set(prev));
856 r
857}
858
859pub fn device_count() -> usize {
862 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
863 {
864 return crate::gpu_wgpu::adapter_count();
865 }
866 #[cfg(not(all(feature = "gpu", not(target_os = "macos"))))]
867 {
868 usize::from(backend_available())
869 }
870}
871
872pub fn vram_budget() -> u64 {
876 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
877 {
878 return crate::gpu_wgpu::device_vram_budget();
879 }
880 #[cfg(not(all(feature = "gpu", not(target_os = "macos"))))]
881 {
882 if backend_available() { u64::MAX } else { 0 }
883 }
884}
885
886pub fn upload_bytes() -> u64 {
890 #[cfg(feature = "gpu")]
891 {
892 return crate::gpu_wgpu::UPLOAD_BYTES.load(std::sync::atomic::Ordering::Relaxed);
893 }
894 #[cfg(not(feature = "gpu"))]
895 0
896}
897
898#[derive(Clone, Copy, PartialEq, Eq, Debug)]
910pub enum GraphPhase {
911 Prefill,
912 Decode,
913}
914
915pub fn wgpu_graph_on(phase: GraphPhase) -> bool {
923 match std::env::var("CMF_GPU_WGPU_GRAPH").ok().as_deref() {
924 Some("0") => false,
925 Some("prefill") => phase == GraphPhase::Prefill,
926 Some(_) => true,
927 None => {
928 if wgpu_graph_default() {
929 return true;
930 }
931 let _ = phase;
936 false
937 }
938 }
939}
940
941pub fn wgpu_graph_default() -> bool {
942 #[cfg(feature = "gpu")]
943 {
944 matches!(backend(), Backend::Wgpu)
950 && (crate::gpu_wgpu::discrete_active()
951 || (cfg!(target_os = "macos") && crate::gpu_wgpu::adapter_up()))
952 }
953 #[cfg(not(feature = "gpu"))]
954 {
955 false
956 }
957}
958
959#[allow(clippy::too_many_arguments, unused_variables)]
961pub fn q8_matvec_range(
962 model: &Arc<CmfModel>,
963 idx: usize,
964 row0: usize,
965 row_scale: &[f32],
966 xs: &[f32],
967 rows: usize,
968 cols: usize,
969 out: &mut [f32],
970) -> bool {
971 match backend() {
972 #[cfg(target_os = "macos")]
973 Backend::Metal => {
974 crate::gpu_metal::q8_matvec_range(model, idx, row0, row_scale, xs, rows, cols, out)
975 }
976 #[cfg(feature = "gpu")]
977 Backend::Wgpu => {
978 crate::gpu_wgpu::q8_matvec_range(model, idx, row0, row_scale, xs, rows, cols, out)
979 }
980 Backend::None => false,
981 }
982}
983
984#[allow(clippy::too_many_arguments, unused_variables)]
987pub fn q8_matmat(
988 model: &Arc<CmfModel>,
989 idx: usize,
990 row_scale: &[f32],
991 pre: &[f32],
992 b: usize,
993 rows: usize,
994 cols: usize,
995 out: &mut [f32],
996) -> bool {
997 match backend() {
998 #[cfg(target_os = "macos")]
999 Backend::Metal => {
1000 crate::gpu_metal::q8_matmat(model, idx, row_scale, pre, b, rows, cols, out)
1001 }
1002 #[cfg(feature = "gpu")]
1003 Backend::Wgpu => crate::gpu_wgpu::q8_matmat(model, idx, row_scale, pre, b, rows, cols, out),
1004 Backend::None => false,
1005 }
1006}
1007
1008#[allow(unused_variables)]
1011pub fn q1_matvec(
1012 model: &Arc<CmfModel>,
1013 idx: usize,
1014 xs: &[f32],
1015 rows: usize,
1016 cols: usize,
1017 out: &mut [f32],
1018) -> bool {
1019 match backend() {
1020 #[cfg(target_os = "macos")]
1021 Backend::Metal => crate::gpu_metal::q1_matvec(model, idx, xs, rows, cols, out),
1022 #[cfg(feature = "gpu")]
1023 Backend::Wgpu => crate::gpu_wgpu::q1_matvec(model, idx, xs, rows, cols, out),
1024 Backend::None => false,
1025 }
1026}
1027
1028#[allow(clippy::too_many_arguments)]
1032pub fn attn_dropin(
1033 model: &Arc<CmfModel>,
1034 kv_id: u64,
1035 layer: usize,
1036 normed: &[f32],
1037 wq_idx: usize,
1038 wk_idx: usize,
1039 wv_idx: usize,
1040 wo_idx: usize,
1041 q_norm: Option<&[f32]>,
1042 k_norm: Option<&[f32]>,
1043 invf: &[f32],
1044 nh: usize,
1045 nkv: usize,
1046 hd: usize,
1047 rd: usize,
1048 hidden: usize,
1049 pos: usize,
1050 cap: usize,
1051 gemma: bool,
1052 eps: f32,
1053 cpu_k: &[Vec<f32>],
1054 cpu_v: &[Vec<f32>],
1055 out: &mut [f32],
1056) -> bool {
1057 match backend() {
1058 #[cfg(feature = "gpu")]
1059 Backend::Wgpu => crate::gpu_wgpu::attn_dropin_gpu(
1060 model, kv_id, layer, normed, wq_idx, wk_idx, wv_idx, wo_idx, q_norm, k_norm, invf, nh,
1061 nkv, hd, rd, hidden, pos, cap, gemma, eps, cpu_k, cpu_v, out,
1062 ),
1063 #[allow(unused_variables)]
1064 _ => false,
1065 }
1066}
1067
1068pub struct GraphW<'a> {
1072 pub idx: usize,
1073 pub kind: u8,
1074 pub row_scale: &'a [f32],
1075 pub data: &'a [f32],
1076}
1077
1078pub enum GraphAttn<'a> {
1081 Full {
1082 wq: GraphW<'a>,
1083 wk: GraphW<'a>,
1084 wv: GraphW<'a>,
1085 wo: GraphW<'a>,
1086 q_norm: Option<&'a [f32]>,
1087 k_norm: Option<&'a [f32]>,
1088 bias: Option<(&'a [f32], &'a [f32], &'a [f32])>,
1090 output_gate: bool,
1093 cpu_k: &'a [Vec<f32>],
1094 cpu_v: &'a [Vec<f32>],
1095 },
1096 Gdn {
1097 qkv: GraphW<'a>,
1098 z: GraphW<'a>,
1099 a: GraphW<'a>,
1100 b: GraphW<'a>,
1101 out: GraphW<'a>,
1102 conv1d: &'a [f32],
1103 a_log: &'a [f32],
1104 dt_bias: &'a [f32],
1105 norm: &'a [f32],
1106 nv: usize,
1107 nk: usize,
1108 dk: usize,
1109 dv: usize,
1110 kk: usize,
1111 cpu_state: &'a [f32],
1116 },
1117}
1118
1119pub struct GraphLayer<'a> {
1121 pub input_norm: &'a [f32],
1122 pub attn: GraphAttn<'a>,
1123 pub post_norm: &'a [f32],
1124 pub ffn: GraphFfn<'a>,
1125}
1126
1127pub enum GraphFfn<'a> {
1132 Dense {
1133 gate: GraphW<'a>,
1134 up: GraphW<'a>,
1135 down: GraphW<'a>,
1136 },
1137 Moe {
1138 router: GraphW<'a>,
1140 shared_gate: GraphW<'a>,
1142 experts: Vec<(usize, usize, usize)>,
1146 n_exp: usize,
1148 top_k: usize,
1149 inter: usize,
1150 norm_topk: bool,
1151 q4tp: bool,
1157 gu_q2: bool,
1161 },
1162}
1163
1164#[allow(clippy::too_many_arguments)]
1169pub fn forward_token_graph(
1170 model: &Arc<CmfModel>,
1171 kv_id: u64,
1172 layers: &[GraphLayer],
1173 o1: &[Option<Vec<crate::nystrom::O1DeviceView<'_>>>],
1176 o1_epoch: u64,
1177 invf: &[f32],
1178 h: &mut [f32],
1179 nh: usize,
1180 nkv: usize,
1181 hd: usize,
1182 rd: usize,
1183 hidden: usize,
1184 inter: usize,
1185 position: usize,
1186 cap: usize,
1187 gemma: bool,
1188 eps: f32,
1189 lm_head: Option<(&GraphW, usize)>,
1190 final_norm: &[f32],
1191 logits: &mut Vec<f32>,
1192 loop_norm_at: &[usize],
1193 steps: usize,
1194 embed: Option<(&GraphW, usize, f32)>,
1195 ids_out: Option<&mut Vec<u32>>,
1196 layers_run: Option<&mut usize>,
1199 layer_base: usize,
1203) -> bool {
1204 match backend() {
1205 #[cfg(feature = "gpu")]
1206 Backend::Wgpu => crate::gpu_wgpu::forward_token_graph(
1207 model,
1208 kv_id,
1209 layers,
1210 o1,
1211 o1_epoch,
1212 invf,
1213 h,
1214 nh,
1215 nkv,
1216 hd,
1217 rd,
1218 hidden,
1219 inter,
1220 position,
1221 cap,
1222 gemma,
1223 eps,
1224 lm_head,
1225 final_norm,
1226 logits,
1227 loop_norm_at,
1228 steps,
1229 embed,
1230 ids_out,
1231 layers_run,
1232 layer_base,
1233 ),
1234 #[allow(unused_variables)]
1235 _ => {
1236 let _ = (lm_head, final_norm, logits, loop_norm_at, layers_run, layer_base);
1237 false
1238 }
1239 }
1240}
1241
1242pub struct SpecTail<'a> {
1246 pub lm: GraphW<'a>,
1247 pub lm_rows: usize,
1248 pub final_norm: &'a [f32],
1249 pub logits_out: &'a mut Vec<f32>,
1250}
1251
1252#[allow(clippy::too_many_arguments)]
1256pub fn forward_batch_graph(
1257 model: &Arc<CmfModel>,
1258 kv_id: u64,
1259 layers: &[GraphLayer],
1260 invf: &[f32],
1261 h: &mut [f32],
1262 nh: usize,
1263 nkv: usize,
1264 hd: usize,
1265 rd: usize,
1266 hidden: usize,
1267 inter: usize,
1268 positions: &[usize],
1269 cap: usize,
1270 gemma: bool,
1271 eps: f32,
1272 k: usize,
1273 spec: Option<SpecTail<'_>>,
1274) -> bool {
1275 match backend() {
1276 #[cfg(feature = "gpu")]
1277 Backend::Wgpu => crate::gpu_wgpu::forward_batch_graph(
1278 model, kv_id, layers, invf, h, nh, nkv, hd, rd, hidden, inter, positions, cap, gemma,
1279 eps, k, spec,
1280 ),
1281 #[allow(unreachable_patterns)]
1282 _ => {
1283 let _ = spec;
1284 false
1285 }
1286 }
1287}
1288
1289pub fn gdn_spec_restore(kv_id: u64, slot: usize) -> bool {
1292 #[cfg(feature = "gpu")]
1293 if backend() == Backend::Wgpu {
1294 return crate::gpu_wgpu::gdn_spec_restore(kv_id, slot);
1295 }
1296 #[allow(unreachable_code)]
1297 {
1298 let _ = (kv_id, slot);
1299 false
1300 }
1301}
1302
1303pub fn graph_kv_reset(_kv_id: u64) {
1305 #[cfg(feature = "gpu")]
1306 if backend() == Backend::Wgpu {
1307 crate::gpu_wgpu::kv_mirror_reset(_kv_id);
1308 }
1309}
1310
1311pub fn q1t_matvec(
1315 model: &Arc<CmfModel>,
1316 idx: usize,
1317 xs: &[f32],
1318 rows: usize,
1319 cols: usize,
1320 out: &mut [f32],
1321) -> bool {
1322 match backend() {
1323 #[cfg(target_os = "macos")]
1324 Backend::Metal => {
1325 if metal_q1t_enabled() {
1326 crate::gpu_metal::q1t_matvec(model, idx, xs, rows, cols, out)
1327 } else {
1328 false
1329 }
1330 }
1331 #[cfg(feature = "gpu")]
1332 Backend::Wgpu => crate::gpu_wgpu::q1t_matvec(model, idx, xs, rows, cols, out),
1333 Backend::None => false,
1334 }
1335}
1336
1337#[allow(unused_variables)]
1340pub fn q4b_matvec(
1341 model: &Arc<CmfModel>,
1342 idx: usize,
1343 xs: &[f32],
1344 rows: usize,
1345 cols: usize,
1346 out: &mut [f32],
1347) -> bool {
1348 match backend() {
1349 #[cfg(target_os = "macos")]
1350 Backend::Metal => false,
1351 #[cfg(feature = "gpu")]
1352 Backend::Wgpu => crate::gpu_wgpu::q4b_matvec(model, idx, xs, rows, cols, out),
1353 Backend::None => false,
1354 }
1355}
1356
1357pub fn q1t_matmat(
1360 model: &Arc<CmfModel>,
1361 idx: usize,
1362 xs: &[f32],
1363 b: usize,
1364 rows: usize,
1365 cols: usize,
1366 out: &mut [f32],
1367) -> bool {
1368 match backend() {
1369 #[cfg(target_os = "macos")]
1370 Backend::Metal => crate::gpu_metal::q1t_matmat(model, idx, xs, b, rows, cols, out),
1374 #[cfg(feature = "gpu")]
1375 Backend::Wgpu => crate::gpu_wgpu::q1t_matmat(model, idx, xs, b, rows, cols, out),
1376 Backend::None => false,
1377 }
1378}
1379
1380#[cfg(target_os = "macos")]
1384pub(crate) fn metal_q1t_enabled() -> bool {
1385 std::env::var("CMF_METAL_Q1T")
1386 .map(|v| v != "0" && !v.eq_ignore_ascii_case("off"))
1387 .unwrap_or(true)
1388}
1389
1390pub fn q1_matmat(
1392 model: &Arc<CmfModel>,
1393 idx: usize,
1394 xs: &[f32],
1395 b: usize,
1396 rows: usize,
1397 cols: usize,
1398 out: &mut [f32],
1399) -> bool {
1400 match backend() {
1401 #[cfg(feature = "gpu")]
1402 Backend::Wgpu => crate::gpu_wgpu::q1_matmat(model, idx, xs, b, rows, cols, out),
1403 #[allow(unused_variables)]
1404 _ => false,
1405 }
1406}
1407
1408static MM_KILL: AtomicBool = AtomicBool::new(false);
1413pub(crate) fn mm_killed() -> bool {
1414 MM_KILL.load(Ordering::Relaxed)
1415}
1416pub(crate) fn mm_kill() {
1417 MM_KILL.store(true, Ordering::Relaxed);
1418}
1419
1420#[allow(unused_variables, clippy::too_many_arguments)]
1425pub fn chunk_attend(
1426 q: &[f32],
1427 k: &[&[f32]],
1428 v: &[&[f32]],
1429 b: usize,
1430 s0: usize,
1431 nh: usize,
1432 nkv: usize,
1433 hd: usize,
1434 scale: f32,
1435 out: &mut [f32],
1436) -> bool {
1437 match backend() {
1438 #[cfg(feature = "gpu")]
1439 Backend::Wgpu => crate::gpu_wgpu::chunk_attend(q, k, v, b, s0, nh, nkv, hd, scale, out),
1440 #[allow(unreachable_patterns)]
1441 _ => false,
1442 }
1443}
1444
1445#[allow(unused_variables, clippy::too_many_arguments)]
1449pub fn q4t_qkv(
1450 model: &Arc<CmfModel>,
1451 wq: usize,
1452 wk: usize,
1453 wv: usize,
1454 xs: &[f32],
1455 b: usize,
1456 cols: usize,
1457 rq: usize,
1458 rk: usize,
1459 rv: usize,
1460 out: &mut [f32],
1461) -> bool {
1462 match backend() {
1463 #[cfg(feature = "gpu")]
1464 Backend::Wgpu => crate::gpu_wgpu::q4t_qkv(model, wq, wk, wv, xs, b, cols, rq, rk, rv, out),
1465 #[allow(unreachable_patterns)]
1466 _ => false,
1467 }
1468}
1469
1470#[allow(unused_variables, clippy::too_many_arguments)]
1472#[allow(clippy::too_many_arguments, unused_variables)]
1476pub fn q4tp_ffn_packed(
1477 model: &Arc<CmfModel>,
1478 w1: usize,
1479 w2: usize,
1480 xs: &[f32],
1481 b: usize,
1482 hidden: usize,
1483 inter: usize,
1484 bias: Option<&[f32]>,
1485 out: &mut [f32],
1486) -> bool {
1487 match backend() {
1488 #[cfg(feature = "gpu")]
1489 Backend::Wgpu => {
1490 crate::gpu_wgpu::q4tp_ffn_packed(model, w1, w2, xs, b, hidden, inter, bias, out)
1491 }
1492 #[allow(unreachable_patterns)]
1493 _ => false,
1494 }
1495}
1496
1497pub fn q4tp_ffn(
1498 model: &Arc<CmfModel>,
1499 w1: usize,
1500 w3: usize,
1501 w2: usize,
1502 xs: &[f32],
1503 b: usize,
1504 hidden: usize,
1505 inter: usize,
1506 out: &mut [f32],
1507) -> bool {
1508 match backend() {
1509 #[cfg(target_os = "macos")]
1510 Backend::Metal => crate::gpu_metal::q4tp_ffn(model, w1, w3, w2, xs, b, hidden, inter, out),
1511 #[cfg(feature = "gpu")]
1512 Backend::Wgpu => crate::gpu_wgpu::q4tp_ffn(model, w1, w3, w2, xs, b, hidden, inter, out),
1513 #[allow(unreachable_patterns)]
1514 _ => false,
1515 }
1516}
1517
1518pub fn q4t_ffn(
1519 model: &Arc<CmfModel>,
1520 w1: usize,
1521 w3: usize,
1522 w2: usize,
1523 xs: &[f32],
1524 b: usize,
1525 hidden: usize,
1526 inter: usize,
1527 out: &mut [f32],
1528) -> bool {
1529 match backend() {
1530 #[cfg(target_os = "macos")]
1531 Backend::Metal => crate::gpu_metal::q4t_ffn(model, w1, w3, w2, xs, b, hidden, inter, out),
1532 #[cfg(feature = "gpu")]
1533 Backend::Wgpu => crate::gpu_wgpu::q4t_ffn(model, w1, w3, w2, xs, b, hidden, inter, out),
1534 #[allow(unreachable_patterns)]
1535 _ => false,
1536 }
1537}
1538
1539pub struct DitBlockArgs<'a> {
1544 pub n: usize,
1545 pub hidden: usize,
1546 pub inter: usize,
1547 pub nh: usize,
1548 pub nkv: usize,
1549 pub hd: usize,
1550 pub eps: f32,
1551 pub rope_cos: &'a [f32],
1552 pub rope_sin: &'a [f32],
1553 pub norm1: &'a [f32],
1554 pub norm2: &'a [f32],
1555 pub ffn_norm1: &'a [f32],
1556 pub ffn_norm2: &'a [f32],
1557 pub norm_q: &'a [f32],
1558 pub norm_k: &'a [f32],
1559 pub s_msa: &'a [f32],
1560 pub gate_msa: &'a [f32],
1561 pub s_mlp: &'a [f32],
1562 pub gate_mlp: &'a [f32],
1563 pub wq: usize,
1564 pub wk: usize,
1565 pub wv: usize,
1566 pub wo: usize,
1567 pub w1: usize,
1568 pub w3: usize,
1569 pub w2: usize,
1570 pub q4tp: bool,
1574 pub resident_in: bool,
1577 pub resident_out: bool,
1581}
1582
1583pub fn dit_chain_supported() -> bool {
1587 #[cfg(feature = "gpu")]
1588 {
1589 return matches!(backend(), Backend::Wgpu) && fused_dit_block_available();
1590 }
1591 #[allow(unreachable_code)]
1592 false
1593}
1594
1595pub fn dit_state_fetch(_x: &mut [f32]) -> bool {
1598 #[cfg(feature = "gpu")]
1599 {
1600 if matches!(backend(), Backend::Wgpu) {
1601 return crate::gpu_wgpu::dit_state_fetch(_x);
1602 }
1603 }
1604 false
1605}
1606
1607#[allow(unused_variables)]
1611#[allow(unused_variables, clippy::too_many_arguments)]
1615pub fn dit_qkv(
1616 model: &Arc<CmfModel>,
1617 wq: usize,
1618 wk: usize,
1619 wv: usize,
1620 xs: &[f32],
1621 b: usize,
1622 hidden: usize,
1623 qrows: usize,
1624 kvrows: usize,
1625 q_out: &mut [f32],
1626 k_out: &mut [f32],
1627 v_out: &mut [f32],
1628) -> bool {
1629 match backend() {
1630 #[cfg(feature = "gpu")]
1631 Backend::Wgpu => crate::gpu_wgpu::q4tp_qkv(
1632 model, wq, wk, wv, xs, b, hidden, qrows, kvrows, q_out, k_out, v_out,
1633 ),
1634 #[allow(unreachable_patterns)]
1635 _ => false,
1636 }
1637}
1638
1639pub fn fused_dit_block_available() -> bool {
1643 #[cfg(target_os = "macos")]
1644 {
1645 matches!(backend(), Backend::Metal) && fused_block_trusted()
1646 }
1647 #[cfg(not(target_os = "macos"))]
1648 {
1649 false
1650 }
1651}
1652
1653pub fn dit_block(model: &Arc<CmfModel>, a: &DitBlockArgs, x: &mut [f32]) -> bool {
1654 dit_block_seg(model, a, &[a.n], x)
1655}
1656
1657pub fn dit_block_seg(
1661 model: &Arc<CmfModel>,
1662 a: &DitBlockArgs,
1663 segs: &[usize],
1664 x: &mut [f32],
1665) -> bool {
1666 match backend() {
1667 #[cfg(target_os = "macos")]
1668 Backend::Metal if segs.len() <= 1 => crate::gpu_metal::dit_block(model, a, x),
1669 #[cfg(feature = "gpu")]
1676 Backend::Wgpu
1677 if match std::env::var("CMF_DIT_FUSED").ok().as_deref() {
1678 Some("0") => false,
1679 Some(_) => true,
1680 None => crate::gpu_wgpu::discrete_active(),
1681 } =>
1682 {
1683 crate::gpu_wgpu::dit_block_seg(model, a, segs, x)
1684 }
1685 #[allow(unreachable_patterns)]
1686 _ => false,
1687 }
1688}
1689
1690pub struct VaeResnetArgs<'a> {
1694 pub groups: usize,
1695 pub ic: usize,
1696 pub oc: usize,
1697 pub h: usize,
1698 pub w: usize,
1699 pub n1w: &'a [f32],
1700 pub n1b: &'a [f32],
1701 pub c1w: &'a [f32],
1702 pub c1b: &'a [f32],
1703 pub c1k: usize,
1704 pub n2w: &'a [f32],
1705 pub n2b: &'a [f32],
1706 pub c2w: &'a [f32],
1707 pub c2b: &'a [f32],
1708 pub c2k: usize,
1709 pub shortcut: Option<(&'a [f32], &'a [f32], usize)>,
1710}
1711
1712#[allow(unused_variables)]
1715pub fn vae_resnet(a: &VaeResnetArgs, x: &[f32], out: &mut [f32]) -> bool {
1716 match backend() {
1717 #[cfg(target_os = "macos")]
1718 Backend::Metal => crate::gpu_metal::vae_resnet(a, x, out),
1719 _ => false,
1720 }
1721}
1722
1723#[allow(unused_variables, clippy::too_many_arguments)]
1726pub fn vae_upsample_conv(
1727 w: &[f32],
1728 bias: &[f32],
1729 x: &[f32],
1730 ic: usize,
1731 oc: usize,
1732 h: usize,
1733 w_img: usize,
1734 k: usize,
1735 out: &mut [f32],
1736) -> bool {
1737 match backend() {
1738 #[cfg(target_os = "macos")]
1739 Backend::Metal => crate::gpu_metal::vae_upsample_conv(w, bias, x, ic, oc, h, w_img, k, out),
1740 #[cfg(feature = "gpu")]
1741 Backend::Wgpu => {
1742 crate::gpu_wgpu::vae_upsample_conv(w, bias, x, ic, oc, h, w_img, k, out)
1743 }
1744 #[allow(unreachable_patterns)]
1745 _ => false,
1746 }
1747}
1748
1749#[allow(unused_variables, clippy::too_many_arguments)]
1752pub fn vae_conv2d(
1753 w: &[f32],
1754 bias: &[f32],
1755 x: &[f32],
1756 ic: usize,
1757 oc: usize,
1758 h: usize,
1759 w_img: usize,
1760 k: usize,
1761 out: &mut [f32],
1762) -> bool {
1763 match backend() {
1764 #[cfg(target_os = "macos")]
1765 Backend::Metal => crate::gpu_metal::vae_conv2d(w, bias, x, ic, oc, h, w_img, k, out),
1766 #[cfg(feature = "gpu")]
1767 Backend::Wgpu => crate::gpu_wgpu::vae_conv2d(w, bias, x, ic, oc, h, w_img, k, out),
1768 #[allow(unreachable_patterns)]
1769 _ => false,
1770 }
1771}
1772
1773#[allow(unused_variables, clippy::too_many_arguments)]
1777#[allow(unused_variables)]
1781#[allow(clippy::too_many_arguments)]
1782#[allow(clippy::too_many_arguments, unused_variables)]
1785pub fn dit_qkv_attention(
1786 model: &Arc<CmfModel>,
1787 qkv_idx: usize,
1788 xn: &[f32],
1789 n: usize,
1790 hidden: usize,
1791 nh: usize,
1792 hd: usize,
1793 scale: f32,
1794 nr: (&[f32], &[f32], &[f32], f32),
1795 out: &mut [f32],
1796) -> bool {
1797 match backend() {
1798 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
1799 Backend::Wgpu => crate::gpu_wgpu::dit_qkv_attention(
1800 model, qkv_idx, xn, n, hidden, nh, hd, scale, nr, out,
1801 ),
1802 #[allow(unreachable_patterns)]
1803 _ => false,
1804 }
1805}
1806
1807#[allow(clippy::too_many_arguments)]
1810pub fn dit_qkv_attn_out(
1811 model: &Arc<CmfModel>,
1812 qkv_idx: usize,
1813 out_idx: usize,
1814 xn: &[f32],
1815 n: usize,
1816 hidden: usize,
1817 nh: usize,
1818 hd: usize,
1819 scale: f32,
1820 nr: (&[f32], &[f32], &[f32], f32),
1821 proj: &mut [f32],
1822) -> bool {
1823 match backend() {
1824 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
1825 Backend::Wgpu => crate::gpu_wgpu::dit_qkv_attn_out(
1826 model, qkv_idx, out_idx, xn, n, hidden, nh, hd, scale, nr, proj,
1827 ),
1828 #[allow(unreachable_patterns)]
1829 _ => false,
1830 }
1831}
1832
1833#[allow(clippy::too_many_arguments)]
1835pub fn vae_qkv_attn_out(
1836 model: &Arc<CmfModel>,
1837 qkv_idx: usize,
1838 out_idx: usize,
1839 xn: &[f32],
1840 n: usize,
1841 dim: usize,
1842 nh: usize,
1843 hd: usize,
1844 scale: f32,
1845 angles: &[f32],
1846 eps: f32,
1847 qkv_bias: &[f32],
1848 proj: &mut [f32],
1849) -> bool {
1850 match backend() {
1851 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
1852 Backend::Wgpu => crate::gpu_wgpu::vae_qkv_attn_out(
1853 model, qkv_idx, out_idx, xn, n, dim, nh, hd, scale, angles, eps, qkv_bias, proj,
1854 ),
1855 #[allow(unreachable_patterns)]
1856 _ => false,
1857 }
1858}
1859
1860#[allow(clippy::too_many_arguments)]
1861pub fn vae_attention_packed(
1862 qkv: &[f32],
1863 nh: usize,
1864 n: usize,
1865 hd: usize,
1866 scale: f32,
1867 angles: &[f32],
1868 eps: f32,
1869 out: &mut [f32],
1870) -> bool {
1871 vae_attention_packed_layout(qkv, nh, n, hd, scale, angles, eps, out, 1)
1872}
1873
1874#[allow(clippy::too_many_arguments)]
1875pub fn vae_attention_packed_layout(
1876 qkv: &[f32],
1877 nh: usize,
1878 n: usize,
1879 hd: usize,
1880 scale: f32,
1881 angles: &[f32],
1882 eps: f32,
1883 out: &mut [f32],
1884 layout: u32,
1885) -> bool {
1886 match backend() {
1887 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
1888 Backend::Wgpu => crate::gpu_wgpu::vae_attention_packed_layout(
1889 qkv, nh, n, hd, scale, angles, eps, out, layout,
1890 ),
1891 #[allow(unreachable_patterns)]
1892 _ => false,
1893 }
1894}
1895
1896#[allow(clippy::too_many_arguments)]
1897pub fn dit_split_only(
1898 qkv: &[f32],
1899 nh: usize,
1900 n: usize,
1901 hd: usize,
1902 layout: u32,
1903 norm: Option<(&[f32], f32)>,
1904 out_q: &mut [f32],
1905) -> bool {
1906 match backend() {
1907 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
1908 Backend::Wgpu => crate::gpu_wgpu::dit_split_only(qkv, nh, n, hd, layout, norm, out_q),
1909 #[allow(unreachable_patterns)]
1910 _ => false,
1911 }
1912}
1913
1914pub fn gemm_nt_f32(x: &[f32], w: &[f32], y: &mut [f32], n: usize, k: usize, m: usize) -> bool {
1919 match backend() {
1920 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
1921 Backend::Wgpu => crate::gpu_wgpu::gemm_nt_f32(x, w, y, n, k, m),
1922 #[allow(unreachable_patterns)]
1923 _ => false,
1924 }
1925}
1926
1927#[allow(clippy::too_many_arguments)]
1930pub fn music3_ffn(
1931 model: &std::sync::Arc<CmfModel>,
1932 idx_in: usize,
1933 idx_out: usize,
1934 h: &[f32],
1935 bias_in: &[f32],
1936 n: usize,
1937 hs: usize,
1938 inter: usize,
1939 out: &mut [f32],
1940) -> bool {
1941 match backend() {
1942 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
1943 Backend::Wgpu => {
1944 crate::gpu_wgpu::music3_ffn(model, idx_in, idx_out, h, bias_in, n, hs, inter, out)
1945 }
1946 #[allow(unreachable_patterns)]
1947 _ => false,
1948 }
1949}
1950
1951#[allow(clippy::too_many_arguments)]
1955pub fn conv1d_gemm(
1956 x: &[f32],
1957 w: &[f32],
1958 ic: usize,
1959 oc: usize,
1960 n: usize,
1961 k: usize,
1962 pad: usize,
1963 dil: usize,
1964 out_n: usize,
1965 yt: &mut [f32],
1966) -> bool {
1967 match backend() {
1968 #[cfg(target_os = "macos")]
1969 Backend::Metal => {
1970 crate::gpu_metal::conv1d_gemm(x, w, ic, oc, n, k, pad, dil, out_n, yt)
1971 }
1972 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
1973 Backend::Wgpu => {
1974 crate::gpu_wgpu::conv1d_gemm(x, w, ic, oc, n, k, pad, dil, out_n, yt)
1975 }
1976 #[allow(unreachable_patterns)]
1977 _ => false,
1978 }
1979}
1980
1981#[allow(clippy::too_many_arguments)]
1983pub fn vae_conv2d_coop(
1984 w: &[f32],
1985 bias: Option<&[f32]>,
1986 x: &[f32],
1987 ic: usize,
1988 oc: usize,
1989 h: usize,
1990 wi: usize,
1991 k: usize,
1992 out: &mut [f32],
1993) -> bool {
1994 match backend() {
1995 #[cfg(all(feature = "gpu", not(target_os = "macos")))]
1996 Backend::Wgpu => crate::gpu_wgpu::vae_conv2d_coop(w, bias, x, ic, oc, h, wi, k, out),
1997 #[allow(unreachable_patterns)]
1998 _ => false,
1999 }
2000}
2001
2002pub fn dit_attention_packed(
2003 qkv: &[f32],
2004 nh: usize,
2005 n: usize,
2006 hd: usize,
2007 scale: f32,
2008 nr: Option<(&[f32], &[f32], &[f32], f32)>,
2011 out: &mut [f32],
2012) -> bool {
2013 match backend() {
2014 #[cfg(feature = "gpu")]
2021 Backend::Wgpu => crate::gpu_wgpu::dit_attention_packed(qkv, nh, n, hd, scale, nr, out),
2022 #[allow(unreachable_patterns)]
2023 _ => false,
2024 }
2025}
2026
2027pub fn dit_attention_packed_available() -> bool {
2035 #[allow(unreachable_patterns)]
2036 match backend() {
2037 #[cfg(feature = "gpu")]
2038 Backend::Wgpu => crate::gpu_wgpu::dit_attention_packed_ready(),
2039 _ => false,
2040 }
2041}
2042
2043pub fn dit_attention(
2044 qh: &[f32],
2045 kh: &[f32],
2046 vh: &[f32],
2047 nh: usize,
2048 nkv: usize,
2049 n: usize,
2050 hd: usize,
2051 scale: f32,
2052 out: &mut [f32],
2053) -> bool {
2054 match backend() {
2055 #[cfg(target_os = "macos")]
2056 Backend::Metal => crate::gpu_metal::dit_attention(qh, kh, vh, nh, nkv, n, hd, scale, out),
2057 #[cfg(feature = "gpu")]
2058 Backend::Wgpu => crate::gpu_wgpu::dit_attention(qh, kh, vh, nh, nkv, n, hd, scale, out),
2059 #[allow(unreachable_patterns)]
2060 _ => false,
2061 }
2062}
2063
2064#[allow(unused_variables)]
2069pub fn q4tp_matmat(
2070 model: &Arc<CmfModel>,
2071 idx: usize,
2072 xs: &[f32],
2073 b: usize,
2074 rows: usize,
2075 cols: usize,
2076 out: &mut [f32],
2077) -> bool {
2078 match backend() {
2079 #[cfg(target_os = "macos")]
2080 Backend::Metal => crate::gpu_metal::q4tp_matmat(model, idx, xs, b, rows, cols, out),
2081 #[cfg(feature = "gpu")]
2082 Backend::Wgpu => crate::gpu_wgpu::q4tp_matmat(model, idx, xs, b, rows, cols, out),
2083 #[allow(unreachable_patterns)]
2084 _ => false,
2085 }
2086}
2087
2088pub fn q2tp_matmat(
2091 model: &Arc<CmfModel>,
2092 idx: usize,
2093 xs: &[f32],
2094 b: usize,
2095 rows: usize,
2096 cols: usize,
2097 out: &mut [f32],
2098) -> bool {
2099 match backend() {
2100 #[cfg(feature = "gpu")]
2101 Backend::Wgpu => crate::gpu_wgpu::q2tp_matmat(model, idx, xs, b, rows, cols, out),
2102 #[allow(unreachable_patterns)]
2103 _ => false,
2104 }
2105}
2106
2107pub fn q4tp_matvec(
2112 model: &Arc<CmfModel>,
2113 idx: usize,
2114 xs: &[f32],
2115 rows: usize,
2116 cols: usize,
2117 out: &mut [f32],
2118) -> bool {
2119 match backend() {
2120 #[cfg(target_os = "macos")]
2121 Backend::Metal => crate::gpu_metal::q4tp_matvec_for_test(model, idx, xs, rows, cols, out),
2122 #[cfg(feature = "gpu")]
2123 Backend::Wgpu => crate::gpu_wgpu::q4tp_matvec(model, idx, xs, rows, cols, out),
2124 #[allow(unreachable_patterns)]
2125 _ => false,
2126 }
2127}
2128
2129pub fn q4t_matvec(
2135 model: &Arc<CmfModel>,
2136 idx: usize,
2137 xs: &[f32],
2138 rows: usize,
2139 cols: usize,
2140 out: &mut [f32],
2141) -> bool {
2142 match backend() {
2143 #[cfg(target_os = "macos")]
2144 Backend::Metal => crate::gpu_metal::q4t_matvec_for_test(model, idx, xs, rows, cols, out),
2145 #[allow(unreachable_patterns)]
2146 _ => false,
2147 }
2148}
2149
2150pub fn q4t_matmat(
2151 model: &Arc<CmfModel>,
2152 idx: usize,
2153 xs: &[f32],
2154 b: usize,
2155 rows: usize,
2156 cols: usize,
2157 out: &mut [f32],
2158) -> bool {
2159 match backend() {
2160 #[cfg(target_os = "macos")]
2161 Backend::Metal => crate::gpu_metal::q4t_matmat(model, idx, xs, b, rows, cols, out),
2162 #[cfg(feature = "gpu")]
2163 Backend::Wgpu => crate::gpu_wgpu::q4t_matmat(model, idx, xs, b, rows, cols, out),
2164 #[allow(unreachable_patterns)]
2165 _ => false,
2166 }
2167}
2168
2169#[cfg(target_os = "macos")]
2171pub use crate::gpu_metal::{
2172 AttnDeviceParams, AttnGpuLayer, GdnGpuCfg, GdnGpuLayer, GpuMoe, GraphDims, MetalFfn,
2173 O1AttnParams, TokenGraph, kv_mirror_drop, kv_mirror_read_last, kv_mirror_take_imp,
2174};
2175
2176#[cfg(target_os = "macos")]
2178pub fn gdn_block(
2179 model: &Arc<CmfModel>,
2180 layers: &[GdnGpuLayer],
2181 states: &mut [&mut [f32]],
2182 cfg: &GdnGpuCfg,
2183 h: &mut [f32],
2184) -> bool {
2185 match backend() {
2186 Backend::Metal => crate::gpu_metal::gdn_block(model, layers, states, cfg, h),
2187 _ => false,
2188 }
2189}
2190
2191#[allow(unused_variables)]
2193pub fn moe_block(model: &Arc<CmfModel>, jobs: &[MoeJob], out: &mut [f32]) -> bool {
2194 match backend() {
2195 #[cfg(target_os = "macos")]
2196 Backend::Metal => crate::gpu_metal::moe_block(model, jobs, out),
2197 #[cfg(feature = "gpu")]
2198 Backend::Wgpu => crate::gpu_wgpu::moe_block(model, jobs, out),
2199 Backend::None => false,
2200 }
2201}
2202
2203#[allow(unused_variables)]
2205pub fn matvec_batch(model: &Arc<CmfModel>, jobs: &[BatchJob], out: &mut [&mut [f32]]) -> bool {
2206 match backend() {
2207 #[cfg(target_os = "macos")]
2208 Backend::Metal => crate::gpu_metal::matvec_batch(model, jobs, out),
2209 #[cfg(feature = "gpu")]
2210 Backend::Wgpu => crate::gpu_wgpu::matvec_batch(model, jobs, out),
2211 Backend::None => false,
2212 }
2213}
2214
2215static GRAPH_RACE_STATE: AtomicU8 = AtomicU8::new(0); static GRAPH_RACE_FLIP: AtomicU32 = AtomicU32::new(0);
2231static GRAPH_RACE_ARM_GRAPH: AtomicU8 = AtomicU8::new(0); static GRAPH_RACE_TOK: AtomicU32 = AtomicU32::new(0); static GRAPH_NS: [AtomicU64; 2] = [AtomicU64::new(0), AtomicU64::new(0)]; static GRAPH_N: [AtomicU32; 2] = [AtomicU32::new(0), AtomicU32::new(0)];
2235
2236const GRAPH_RACE_SAMPLES: u32 = 4;
2238
2239static GRAPH_UNSUPPORTED: AtomicBool = AtomicBool::new(false);
2249
2250pub fn graph_mark_unsupported() {
2255 if !GRAPH_UNSUPPORTED.swap(true, Ordering::Relaxed) {
2256 tracing::info!("wgpu token graph: unsupported for this model — not retrying");
2257 }
2258}
2259
2260pub fn graph_unsupported() -> bool {
2261 GRAPH_UNSUPPORTED.load(Ordering::Relaxed)
2262}
2263
2264pub fn graph_unsupported_reset() {
2266 GRAPH_UNSUPPORTED.store(false, Ordering::Relaxed);
2267}
2268
2269pub fn graph_race_begin_generation() {
2270 #[cfg(feature = "gpu")]
2275 {
2276 static FLUSHED: std::sync::Once = std::sync::Once::new();
2288 static FIRST: std::sync::atomic::AtomicBool =
2289 std::sync::atomic::AtomicBool::new(true);
2290 if FIRST.swap(false, Ordering::Relaxed) {
2291 } else {
2293 FLUSHED.call_once(crate::gpu_wgpu::pipeline_cache_flush);
2294 }
2295 }
2296 GRAPH_RACE_TOK.store(0, Ordering::Relaxed);
2297 if GRAPH_RACE_STATE.load(Ordering::Relaxed) != 0 {
2298 return;
2299 }
2300 let (gn, cn) = (
2301 GRAPH_N[1].load(Ordering::Relaxed),
2302 GRAPH_N[0].load(Ordering::Relaxed),
2303 );
2304 if gn >= GRAPH_RACE_SAMPLES && cn >= GRAPH_RACE_SAMPLES {
2305 let g_avg = GRAPH_NS[1].load(Ordering::Relaxed) / gn as u64;
2306 let c_avg = GRAPH_NS[0].load(Ordering::Relaxed) / cn as u64;
2307 let verdict = if g_avg < c_avg { 1 } else { 2 };
2308 GRAPH_RACE_STATE.store(verdict, Ordering::Relaxed);
2309 tracing::info!(
2310 "wgpu graph race: graph {:.2} ms/tok vs normal {:.2} ms/tok -> {}",
2311 g_avg as f64 / 1e6,
2312 c_avg as f64 / 1e6,
2313 if verdict == 1 { "graph" } else { "normal path" }
2314 );
2315 return;
2316 }
2317 let flip = GRAPH_RACE_FLIP.fetch_add(1, Ordering::Relaxed);
2318 GRAPH_RACE_ARM_GRAPH.store((flip % 2 == 1) as u8, Ordering::Relaxed);
2319}
2320
2321pub fn graph_race_use_graph(trusted: bool) -> bool {
2325 if trusted {
2326 return true;
2327 }
2328 match GRAPH_RACE_STATE.load(Ordering::Relaxed) {
2329 1 => true,
2330 2 => false,
2331 _ => GRAPH_RACE_ARM_GRAPH.load(Ordering::Relaxed) == 1,
2332 }
2333}
2334
2335pub fn graph_race_first_token_hopeless(dur: std::time::Duration) -> bool {
2340 if GRAPH_RACE_STATE.load(Ordering::Relaxed) != 0 {
2341 return false;
2342 }
2343 let first = GRAPH_RACE_TOK.load(Ordering::Relaxed) == 0;
2344 let cn = GRAPH_N[0].load(Ordering::Relaxed);
2345 if !first || cn == 0 {
2346 return false;
2347 }
2348 let c_avg = GRAPH_NS[0].load(Ordering::Relaxed) / cn as u64;
2349 let ns = dur.as_nanos() as u64;
2350 if ns > 1_000_000_000 && ns > 4 * c_avg {
2351 GRAPH_RACE_STATE.store(2, Ordering::Relaxed);
2352 tracing::info!(
2353 "wgpu graph race: first graph token {:.0} ms vs normal {:.2} ms/tok — hopeless, normal path wins",
2354 ns as f64 / 1e6,
2355 c_avg as f64 / 1e6
2356 );
2357 return true;
2358 }
2359 false
2360}
2361
2362pub fn graph_race_record(used_graph: bool, dur: std::time::Duration) {
2366 if GRAPH_RACE_STATE.load(Ordering::Relaxed) != 0 {
2367 return;
2368 }
2369 let tok = GRAPH_RACE_TOK.fetch_add(1, Ordering::Relaxed);
2370 if tok == 0 {
2371 return;
2372 }
2373 let i = used_graph as usize;
2374 GRAPH_NS[i].fetch_add(dur.as_nanos() as u64, Ordering::Relaxed);
2375 GRAPH_N[i].fetch_add(1, Ordering::Relaxed);
2376}
2377
2378pub(crate) fn fp_bytes(data: &[u8]) -> u64 {
2388 #[inline]
2389 fn fnv(mut h: u64, bytes: &[u8]) -> u64 {
2390 let (chunks, tail) = bytes.split_at(bytes.len() & !7);
2391 for c in chunks.chunks_exact(8) {
2392 h ^= u64::from_le_bytes(c.try_into().unwrap());
2393 h = h.wrapping_mul(0x100_0000_01b3);
2394 }
2395 for &b in tail {
2396 h ^= b as u64;
2397 h = h.wrapping_mul(0x100_0000_01b3);
2398 }
2399 h
2400 }
2401 let mut h = 0xcbf2_9ce4_8422_2325u64 ^ (data.len() as u64);
2402 if data.len() <= 4096 {
2403 return fnv(h, data);
2404 }
2405 let step = (data.len() - 64) / 63;
2406 for i in 0..64 {
2407 h = fnv(h, &data[i * step..i * step + 64]);
2408 }
2409 h
2410}
2411
2412pub(crate) fn fp_f32(data: &[f32]) -> u64 {
2415 let bytes =
2416 unsafe { std::slice::from_raw_parts(data.as_ptr() as *const u8, data.len() * 4) };
2417 fp_bytes(bytes)
2418}
2419
2420#[cfg(test)]
2421mod fp_tests {
2422 use super::fp_bytes;
2423
2424 #[test]
2429 fn fp_bytes_sees_a_change_anywhere_in_a_sampled_slice() {
2430 let n = 1 << 20; let base: Vec<u8> = (0..n).map(|i| (i * 31 + 7) as u8).collect();
2432 let h0 = fp_bytes(&base);
2433 assert_eq!(h0, fp_bytes(&base), "fingerprint must be deterministic");
2434 let mut dense = base.clone();
2437 for b in dense.iter_mut() {
2438 *b = b.wrapping_add(1);
2439 }
2440 assert_ne!(h0, fp_bytes(&dense), "a fully different tensor slipped through");
2441 assert_ne!(h0, fp_bytes(&base[..n - 64]));
2444 let mut small = vec![3u8; 4096];
2447 let hs = fp_bytes(&small);
2448 small[2048] ^= 1;
2449 assert_ne!(hs, fp_bytes(&small), "full hash missed a one-byte change");
2450 for n in [4097usize, 5000, 64 * 64, 1 << 16] {
2452 let v = vec![9u8; n];
2453 let _ = fp_bytes(&v); }
2455 }
2456}
2457
2458pub fn bake_release() {
2462 #[cfg(feature = "gpu")]
2463 crate::gpu_wgpu::bake_release();
2464}
2465
2466pub fn bake_precision_strict(on: bool) {
2470 #[cfg(feature = "gpu")]
2471 crate::gpu_wgpu::bake_precision_strict(on);
2472 #[cfg(not(feature = "gpu"))]
2473 let _ = on;
2474}
2475
2476
2477pub fn hostprof_encode_done(t0: std::time::Instant) {
2483 use std::sync::atomic::{AtomicU64, Ordering};
2484 static ENC: AtomicU64 = AtomicU64::new(0);
2485 static N: AtomicU64 = AtomicU64::new(0);
2486 if std::env::var("CMF_GRAPH_HOSTPROF").as_deref() != Ok("1") {
2487 return;
2488 }
2489 ENC.fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
2490 let n = N.fetch_add(1, Ordering::Relaxed) + 1;
2491 if n % 100 == 0 {
2492 eprintln!(
2493 "hostprof: encode {:.2} ms/token over {n} tokens",
2494 ENC.load(Ordering::Relaxed) as f64 / n as f64 / 1e6
2495 );
2496 }
2497}
2498
2499pub fn hostprof_total(t0: std::time::Instant) {
2500 use std::sync::atomic::{AtomicU64, Ordering};
2501 static TOT: AtomicU64 = AtomicU64::new(0);
2502 static N: AtomicU64 = AtomicU64::new(0);
2503 if std::env::var("CMF_GRAPH_HOSTPROF").as_deref() != Ok("1") {
2504 return;
2505 }
2506 TOT.fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
2507 let n = N.fetch_add(1, Ordering::Relaxed) + 1;
2508 if n % 100 == 0 {
2509 eprintln!(
2510 "hostprof: total {:.2} ms/token over {n} tokens",
2511 TOT.load(Ordering::Relaxed) as f64 / n as f64 / 1e6
2512 );
2513 }
2514}
2515
2516
2517pub fn stageprof(stage: u32, dt: std::time::Duration) {
2521 use std::sync::atomic::{AtomicU64, Ordering};
2522 static NS: [AtomicU64; 4] = [
2523 AtomicU64::new(0),
2524 AtomicU64::new(0),
2525 AtomicU64::new(0),
2526 AtomicU64::new(0),
2527 ];
2528 static N: AtomicU64 = AtomicU64::new(0);
2529 if std::env::var("CMF_GRAPH_HOSTPROF").as_deref() != Ok("1") {
2530 return;
2531 }
2532 NS[stage as usize % 4].fetch_add(dt.as_nanos() as u64, Ordering::Relaxed);
2533 if stage == 1 {
2534 let n = N.fetch_add(1, Ordering::Relaxed) + 1;
2535 if n % 200 == 0 {
2536 eprintln!(
2537 "stageprof: planning {:.2} ms/tok | gdn-item {:.2} ms/tok | attn-item {:.2} ms/tok ({n} tok)",
2538 NS[1].load(Ordering::Relaxed) as f64 / n as f64 / 1e6,
2539 NS[2].load(Ordering::Relaxed) as f64 / n as f64 / 1e6,
2540 NS[3].load(Ordering::Relaxed) as f64 / n as f64 / 1e6
2541 );
2542 }
2543 }
2544}
2545
2546
2547pub fn weight_bytes_dispatched() -> u64 {
2550 let mut total = 0u64;
2551 #[cfg(target_os = "macos")]
2552 {
2553 total += crate::gpu_metal::WEIGHT_BYTES.load(std::sync::atomic::Ordering::Relaxed);
2554 }
2555 #[cfg(feature = "gpu")]
2556 {
2557 total += crate::gpu_wgpu::WEIGHT_BYTES.load(std::sync::atomic::Ordering::Relaxed);
2558 }
2559 total
2560}
2561
2562
2563pub fn weight_bytes_by() -> [u64; 6] {
2566 #[cfg(target_os = "macos")]
2567 {
2568 let mut o = [0u64; 6];
2569 for (i, a) in crate::gpu_metal::WEIGHT_BYTES_BY.iter().enumerate() {
2570 o[i] = a.load(std::sync::atomic::Ordering::Relaxed);
2571 }
2572 return o;
2573 }
2574 #[allow(unreachable_code)]
2575 [0; 6]
2576}