Skip to main content

cortiq_engine/
gpu.rs

1//! Facade for GPU backends: a single call entry point for qtensor/pipeline/
2//! linear_core. Job types and the threshold are canonical HERE; behind the
3//! facade dispatch goes to a platform backend:
4//!   - `gpu_metal` (Apple Silicon, unified memory + no-copy buffers);
5//!   - `gpu_wgpu` (C1: Vulkan/DX12/Metal — NVIDIA/Radeon/Intel/Apple,
6//!     weights resident in VRAM), available under `--features gpu`.
7//!
8//! Runtime selection via `CMF_GPU`: `1` — native Metal (macOS) or wgpu
9//! (other OSes); `wgpu` — force wgpu (including for the local
10//! Metal-via-wgpu parity test). Any backend refusal — `false` and the honest
11//! CPU path, no partial results.
12
13use 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    /// Index of the current forward layer (−1 = outside a numbered layer:
20    /// lm_head/embed — always allowed). The pipeline sets it before
21    /// each layer so that the GPU/CPU layer-split works.
22    static CUR_LAYER: Cell<i64> = const { Cell::new(-1) };
23    /// Inside `cpu_scope` every GPU gate reports disabled: the timed CPU
24    /// arm of a probe (and a class that lost its probe) must run PURE
25    /// CPU, or inner per-op hooks would re-enter the GPU and poison the
26    /// comparison.
27    static CPU_ONLY: Cell<bool> = const { Cell::new(false) };
28    /// "This op paid a one-off cost" (weight upload / first pipeline
29    /// build): backends set it, `probe_record` discards the sample so
30    /// only steady-state timings compete.
31    static PROBE_COLD: Cell<bool> = const { Cell::new(false) };
32}
33
34/// Run `f` with the GPU gates off on this thread (pure-CPU arm).
35pub 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
47/// Backends: name the device once at init. The probe cache is keyed by
48/// it, because a verdict is a property of THIS silicon and nothing else.
49/// First writer wins: a process runs one backend, and on the rare host
50/// where two initialize, the one that came up first is the one in use.
51pub 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
61/// Somewhere this process may write small caches.
62///
63/// `std::env::temp_dir()` is NOT that place on Android: with no `TMPDIR`
64/// it answers `/tmp`, which does not exist in an app sandbox, and every
65/// write fails silently — measured, after the pipeline cache appeared to
66/// work in a shell (where `TMPDIR=/data/local/tmp`) and did nothing at
67/// all in the app. The loader points this at the model's own directory,
68/// which is somewhere the caller already writes.
69static CACHE_DIR: std::sync::OnceLock<std::path::PathBuf> = std::sync::OnceLock::new();
70
71/// Loader: name a directory this process can write to. First call wins.
72pub fn set_cache_dir(dir: std::path::PathBuf) {
73    let _ = CACHE_DIR.set(dir);
74}
75
76/// Same directory, for the backends.
77pub 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
91/// Where decided verdicts are remembered between runs. `CMF_PROBE_CACHE`
92/// overrides the path; `0` disables the cache entirely.
93fn 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
101/// One line per decided class: `version \t device \t class \t winner`.
102/// A different engine build or a different device simply does not match,
103/// so a stale file is inert rather than wrong.
104fn 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
118/// Adopt every verdict this device already reached in an earlier run.
119///
120/// Probing is not cheap and it is not free of consequences: on a
121/// Snapdragon 778G the three deciding classes took **three minutes of
122/// wall clock** before the first token, every process, and in the phone
123/// app that was the whole first answer — 209.6 s for 25 tokens against
124/// 10.5 s on the CPU path. The verdict itself was the same every time.
125/// Paying to rediscover it is the defect; the answer is to write it down.
126fn 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        // Unit tests share this process and its default cache path; a
133        // verdict left by an earlier run would decide a class before the
134        // arbitration tests get to watch it alternate. Tests that mean to
135        // exercise the cache point `CMF_PROBE_CACHE` at their own file.
136        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
146/// Apply verdicts from a cache file's text. Split out from the file
147/// reading so the adoption rule — including which lines must be IGNORED
148/// — is testable without a filesystem.
149fn 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
171/// Remember a verdict for the next run. Best-effort: a read-only cache
172/// directory costs a re-probe, never a failure.
173fn 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
192/// Backends: note a one-off cost (weight upload, buffer-cache fill) so
193/// the probe discards this sample.
194/// Every buffer creation anywhere bumps this; the graph's bind-group
195/// cache treats any cold event as total invalidation — a stale bind
196/// group is silent corruption, a cleared cache is one re-encoded token.
197pub 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
207/// Peek the cold flag without consuming it (`probe_record` consumes).
208/// Contention heuristics use this: a slow COLD op is a one-off build
209/// cost, not evidence the device is busy.
210pub(crate) fn probe_was_cold() -> bool {
211    PROBE_COLD.with(|c| c.get())
212}
213
214/// Pipeline: mark the current layer (or −1 outside layers) for layer-split.
215pub fn set_layer(l: i64) {
216    CUR_LAYER.with(|c| c.set(l));
217}
218
219/// The layer `set_layer` last marked on this thread (−1 outside layers).
220pub fn cur_layer() -> i64 {
221    CUR_LAYER.with(|c| c.get())
222}
223
224/// Parse `CMF_GPU_LAYERS` («0-19», «0,2,4», «0-9,30-39») once.
225/// None = no restriction (all layers on GPU). Garbage → also no restriction.
226fn 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
255/// GPU allowed FOR THE CURRENT LAYER: backend is initialized AND the layer
256/// falls within `CMF_GPU_LAYERS` (GPU/CPU layer-split) AND we are not
257/// inside a `cpu_scope`. Op gates call this.
258pub fn enabled_here() -> bool {
259    !CPU_ONLY.with(|c| c.get()) && enabled() && layer_allowed()
260}
261
262// ── Runtime GPU-vs-CPU probe ────────────────────────────────────────────
263// CMF_GPU=1 does not TRUST that the device wins — it MEASURES. For each
264// op class the first calls alternate arms: GPU timed vs pure-CPU timed
265// (under cpu_scope). Cold GPU calls (weight upload / cache fill) are
266// discarded; after PROBE_SAMPLES clean samples per arm the faster arm is
267// chosen for the rest of the process. Rationale: submit+poll latency
268// differs by an order of magnitude across driver stacks (Metal/PCIe
269// ~3-4 ms, Vulkan/4090 ~0.3 ms) — a static threshold cannot know whether
270// per-op offload pays off HERE. CMF_GPU_PROBE=0 → always trust the GPU.
271
272/// GPU-eligible op classes, each with an independent probe.
273#[derive(Clone, Copy)]
274pub enum OpClass {
275    /// Whole FFN chain in one submission (dense / MoE block).
276    Ffn = 0,
277    /// Large hybrid CPU∥GPU matvec (lm_head class).
278    Matvec = 1,
279    /// Prefill GEMM (matmat).
280    Matmat = 2,
281    /// Batched matvecs of one input (QKV).
282    Batch = 3,
283    /// Prefill GEMM at image-diffusion widths (b ≥ 128). Probed apart
284    /// from `Matmat`: one imagegen process runs BOTH populations
285    /// (prompt encode b≈40 where the GPU wins big, DiT b≥256 where
286    /// the CPU AMX arm is competitive) — a single shared verdict locks
287    /// the wrong arm for whichever population samples second.
288    MatmatWide = 4,
289    /// The lm_head itself, apart from the merely-large matvecs. Same
290    /// reasoning as `MatmatWide`, and DeepSeek-V4 is where it bit: its
291    /// attention projections are 37M weights and its head is 529M, so
292    /// the projections' verdict — CPU, honestly measured at 0.19 ms —
293    /// decided for a matvec fourteen times their size that took 11 ms
294    /// a token on the host.
295    MatvecHead = 5,
296    /// The blocked f32 GEMM (`fcd_ops::gemm_nt`): attention's QKᵀ and
297    /// AV, and the VAE decoders' projections. It used to take every job
298    /// over 4 M MACs on sight, with no CPU arm to lose to — which on
299    /// the MiniMax-H3 video decoder was three times SLOWER than the
300    /// host it displaced. Its population is per-head slices, nothing
301    /// like the weight GEMMs above, so it probes on its own.
302    GemmNt = 6,
303}
304
305/// Which probe a large matvec belongs to. The head is an order of
306/// magnitude bigger than anything else that reaches this gate, and the
307/// two populations do not have the same answer.
308pub 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
316/// Probe verdict for one call.
317pub enum ProbeArm {
318    /// Run the GPU path (during probing: timed, recorded).
319    Gpu,
320    /// Probing: run the CPU path under `cpu_scope`, timed, recorded.
321    CpuTimed,
322    /// Decided: CPU won — run the CPU path (under `cpu_scope`).
323    Cpu,
324}
325
326/// Clean samples per arm before a class decides.
327const PROBE_SAMPLES: u32 = 6;
328
329struct Probe {
330    /// 0 = probing, 1 = GPU won, 2 = CPU won.
331    state: AtomicU8,
332    flip: AtomicU32,
333    gpu_ns: AtomicU64,
334    gpu_n: AtomicU32,
335    cpu_ns: AtomicU64,
336    cpu_n: AtomicU32,
337    /// Best (minimum) sample per arm. The DECISION compares these:
338    /// means are poisoned by one-off cold costs the cold-flag cannot
339    /// see — e.g. the CPU arm's first mmap-cold expert matvec page
340    /// faults its weights in and reads 3× its steady state, which
341    /// locked the GPU arm on a 35B MoE at a 4× real-world loss. The
342    /// minimum is each arm's honest steady-state pace.
343    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
381/// q1 ops on the native Metal backend skip the probe entirely: the CPU
382/// q1 kernel is load-port-bound, the GPU one wins warm — and probe
383/// alternation itself cools the device between samples (measured: block
384/// times 5.8 ms warm vs 8.8 ms mixed). Other backends keep probing.
385pub 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
396/// Should a FUSED whole-block path trust the device instead of asking
397/// the per-op probe? True on native Metal and on discrete wgpu adapters.
398///
399/// The probe answers "is one wide matmat faster on the GPU", and for the
400/// DiT on Metal that is a coin flip — measured 2.62 ms GPU vs 2.56 ms
401/// CPU, a 2% spread that lands on either arm run to run. But the fused
402/// block's advantage is not per-op speed, it is that the hidden state,
403/// the packs and the attention panels never leave the device: end to end
404/// the whole-block path renders a 512² Lumina step in ~5.4 s against
405/// ~8.4 s when the probe happens to pick the CPU. Gating a fusion win on
406/// a per-op tie made every second render half-speed at random.
407///
408/// On a discrete card the verdict is never in doubt — an RTX 3090 against
409/// a 256-core EPYC measured 11.5 ms vs 31 ms per wide op, four runs out
410/// of four — so the probe's sampling phase is pure cost: it alone was 10%
411/// of a 512² render (74.3 s against 66.9 s with the probe off). Integrated
412/// and mobile adapters keep probing; there the submit latency is real and
413/// can genuinely lose.
414pub 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
422/// Which arm should this GPU-eligible call take? Consult AFTER the
423/// eligibility gates (`enabled_here` / `min_rows`) so only real
424/// candidates alternate.
425/// While a class is still probing, a call whose weights are NOT yet on
426/// the card should take the GPU arm anyway: the upload is work the next
427/// step needs regardless, and the sample it produces is discarded as
428/// cold — so handing that call to the CPU arm buys nothing and costs a
429/// host GEMM. Measured on a diffusion stack, where every layer is
430/// touched once per step and therefore EVERY first-step GPU sample is
431/// cold: one projection drew the CPU arm for the whole first step, 9.8 s
432/// against the 2.8 s it costs once the weights are warm.
433pub 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    // Every arbitrated call starts with a clean cold flag: both the
454    // sample discard in `probe_record` and the contention kill-switch
455    // read it AFTER the op, so a stale note from a previous call on
456    // this thread must not leak in.
457    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
476/// Record a timed arm sample; on the `PROBE_SAMPLES`-th clean sample of
477/// BOTH arms the class decides for the rest of the process.
478pub 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; // one-off cost in this call — not a steady-state sample
485    }
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        // Decide on each arm's BEST sample — the steady-state pace.
502        // Means carry one-off cold costs (mmap page-in on the CPU arm)
503        // that the cold-flag machinery cannot see.
504        let g = p.gpu_min.load(Ordering::Relaxed) as f64;
505        let cp = p.cpu_min.load(Ordering::Relaxed) as f64;
506        // Early verdict on a ≥2× gap — no reason to keep feeding the
507        // losing arm; close races take the full sample count. It was 3×,
508        // and the cost of that half-octave was measured: a DiT whose
509        // wide GEMMs run 11.4 ms on the device against 32.2 on the host
510        // (2.8×) kept ALTERNATING through the whole diffusion stack, and
511        // because the alternation counter is shared per class in call
512        // order, one projection drew the CPU arm every single time — 9.9
513        // seconds a step on a kernel that needs 0.4. Both arms are
514        // compared on their BEST sample, so a 2× gap is not noise.
515        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
535/// Is the class still collecting samples? (Call sites use this to route
536/// cold-weight calls away from the GPU arm during probing.)
537pub fn probe_deciding(c: OpClass) -> bool {
538    probe_on() && PROBES[c as usize].state.load(Ordering::Relaxed) == 0
539}
540
541/// Probing helper: true — tensor `idx`'s quant weights are ALREADY
542/// device-resident (a clean GPU sample is possible now); false — they
543/// were not (the upload starts within the VRAM budget, so a later call
544/// finds them warm) or the tensor cannot go to the GPU at all. Keeps the
545/// probe from billing a full cold dispatch+readback to a sample it will
546/// discard anyway. The verdict needs only a couple of warm tensors, so
547/// probe-driven uploads are capped — the losing-GPU machine should not
548/// pay for uploading the whole layer stack it will never use; if the GPU
549/// wins, the rest uploads lazily on demand, in the same first-touch order.
550#[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/// Test hook: reset all probes to the undecided state.
568#[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    // One test fn: PROBES is process-global and probe_reset touches all
586    // classes — parallel test threads would race.
587    #[test]
588    fn probe_alternates_discards_cold_and_decides() {
589        probe_reset();
590        // Probing: arms alternate.
591        assert!(matches!(probe_arm(OpClass::Ffn), ProbeArm::Gpu));
592        assert!(matches!(probe_arm(OpClass::Ffn), ProbeArm::CpuTimed));
593
594        // A cold GPU sample (upload noted) must be discarded: feed a
595        // catastrophic cold sample, then clean fast-GPU samples — GPU
596        // wins only if the cold one did not count.
597        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        // The reverse: a class where the CPU arm is faster decides CPU.
606        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: gates off inside, restored after.
613        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        // Probing is not free: on a Snapdragon 778G the deciding classes
627        // cost minutes of wall clock before the first token, every
628        // process, and reached the same verdict every time. The cache
629        // exists so that price is paid once.
630        //
631        // The key is built from THIS process's device, never a name this
632        // test sets: `probe_set_device` is first-writer-wins and on a Mac
633        // the Metal backend may already have named the silicon before the
634        // tests run — which is exactly how this test failed on CI while
635        // passing locally. GemmNt on purpose: the arbitration test never
636        // touches it, and both run in one process.
637        let mine = probe_cache_key_named("gemm-nt");
638        let state = || PROBES[OpClass::GemmNt as usize].state.load(Ordering::Relaxed);
639
640        // Another device's verdict is not mine, whatever it claims.
641        probe_cache_adopt("SomeOtherGPU/Vulkan\tgemm-nt\tgpu\n");
642        assert_eq!(state(), 0);
643        // Neither is one from another build of this engine.
644        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        // Mine is.
649        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
658/// Default row threshold: the GPU takes only larger matrices (lm_head
659/// class). Below it, the dispatch/readback cost does not pay off on unified memory.
660pub const GPU_MIN_ROWS: usize = 65_536;
661
662/// Effective threshold: `CMF_GPU_MIN_ROWS` overrides. Defaults differ
663/// by device class: on a DISCRETE card VRAM bandwidth pays off even for
664/// FFN/QKV-class matrices (4096), on unified memory only lm_head-class
665/// is worth the dispatch/readback (65536). Field case behind this: a
666/// 35B model on an RTX 4090 saw ~0 offload because every layer matrix
667/// sat below the old universal 65536.
668pub 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
678/// Is the active backend a discrete card (PCIe VRAM)?
679pub 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, // UMA by the init() guard
685        Backend::None => false,
686    }
687}
688
689/// A single MoE-FFN job (an expert with its own weight), executed in one
690/// submission: (rows, cols, idx, row_scale) for gate/up/down + prescaled
691/// inputs + the down θ-field + the blending weight.
692pub 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    /// q1 trio: scales live inside the 6-byte tiles (row_scale slices
701    /// empty, xs raw f32). Backends without a q1 kernel refuse the job.
702    pub q1: bool,
703    /// q4_tiled trio: scales inside the 18-byte tiles (row_scale
704    /// slices empty, xs raw f32) — the MoE-hybrid coder class.
705    pub q4t: bool,
706    /// q4tp trio: same raw-xs contract, 16-byte nibble stride and the scale
707    /// on a per-row ladder. Without this the experts of a q4tp MoE model fall
708    /// to the CPU while every other dtype rides the device.
709    pub q4tp: bool,
710    /// Mixed 2-bit profile: gate/up are q2tp (8-byte chunks, zero rung),
711    /// down stays q4tp. Set together with `q4tp`; a backend without the
712    /// 2-bit kernel must refuse the whole job.
713    pub gu_q2: bool,
714    /// The reference's `swiglu_limit`; 0 disables the clamp. A backend that
715    /// cannot apply it must REFUSE the job rather than drop it silently —
716    /// the difference only shows on saturating activations, which is the
717    /// hardest kind of divergence to notice.
718    pub swiglu_limit: f32,
719}
720
721/// A single independent batch matvec (GDN projections of one input).
722pub 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    /// Weight layout. Was a bare `q1: bool`, which could only ever spell two
729    /// of the four and silently sent everything else back to the CPU — the
730    /// GDN projections of a q4t/q4tp model never reached the device at all.
731    pub layout: BatchLayout,
732}
733
734/// Which kernel a batched matvec needs. q8 carries row scales in a side
735/// buffer; the rest embed them in the payload and differ in stride.
736#[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
769/// GPU enabled and initialized on the selected backend?
770/// Whether THIS build can bring a GPU up on THIS device: a compiled-in
771/// backend plus a live adapter. The mobile FFI exposes it so an app can
772/// tell "GPU off" from "GPU impossible" (a CPU-only .so ships no
773/// backend at all). Cached after the first call.
774pub fn backend_available() -> bool {
775    #[cfg(target_os = "macos")]
776    {
777        // The Metal path is always compiled on macOS.
778        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
795/// Default-on condition for the wgpu whole-token graph: the wgpu
796/// backend on a DISCRETE adapter. NOT plain `enabled()` (macOS/Metal
797/// must not pay a per-token layer scan for a graph its backend
798/// refuses), and NOT integrated adapters: the graph's ~300 barriered
799/// dispatches per token are cheap on desktop immediate-mode GPUs but
800/// tiled mobile GPUs (Adreno/Mali) drain the pipeline at every barrier
801/// — field report: 0.2 tok/s on-graph vs 15 tok/s on the CPU. On
802/// integrated adapters the per-op probe path arbitrates each op class
803/// against the CPU instead; CMF_GPU_WGPU_GRAPH=1 still forces the
804/// graph anywhere.
805/// Is the wgpu backend active at all (any adapter)? Eligibility gate
806/// for the whole-token graph — whether it actually RUNS is decided by
807/// `wgpu_graph_default` (trusted on discrete) or the generation race.
808pub 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
819/// Which GPU this thread's engine calls address. Multi-card hosts hold
820/// one wgpu context PER card (weights, KV mirrors and scratch live
821/// inside a context, so per-device contexts give per-device caches for
822/// free); this thread-local says which one is current. Default: the
823/// process pin (CMF_GPU_ADAPTER) or 0 — so single-card runs behave
824/// exactly as they always have.
825pub 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
839/// The device this thread is pinned to.
840pub fn current_device() -> usize {
841    CUR_DEV.with(|c| c.get()).unwrap_or_else(default_device)
842}
843
844/// Pin this thread to a device. Server slots call it once per request;
845/// the worker pool propagates it into its threads, so a dispatch begun
846/// on card 1 does not finish on card 0.
847pub fn set_current_device(i: usize) {
848    CUR_DEV.with(|c| c.set(Some(i)));
849}
850
851/// Run `f` with this thread pinned to `dev`, restoring the previous pin.
852pub 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
859/// How many GPUs this process can address (wgpu adapter count; 1 on
860/// Metal, 0 without a backend).
861pub 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
872/// Weight budget of the current GPU in bytes; 0 when there is none and
873/// u64::MAX on unified memory (where the OS pages shared RAM and the
874/// question "does the model fit the card" has no separate answer).
875pub 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
886/// Device weight bytes uploaded so far (wgpu; 0 on other backends).
887/// Steady-state windows must show a ZERO delta — growth mid-benchmark
888/// means eviction/re-upload and disqualifies the number.
889pub 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/// Which half of the run is asking.
899///
900/// The phase exists because the graph is plausibly two decisions, not
901/// one — but on the hardware measured so far it is only ever a decode
902/// decision. On an Adreno 642L with bonsai-1.7b, from identical clean
903/// starts and two repeats each: decode 11.6 tok/s without it and 0.72
904/// with, while prefill is 4.2 either way. A first reading of 3.4 -> 18.0
905/// for prefill did not survive a controlled re-run — it was a dirty
906/// probe cache between configurations, not the graph, and the prefill
907/// route through the graph is GDN-only in the first place, which this
908/// dense model never takes.
909#[derive(Clone, Copy, PartialEq, Eq, Debug)]
910pub enum GraphPhase {
911    Prefill,
912    Decode,
913}
914
915/// The one place that decides whether the whole-token graph runs.
916///
917/// `CMF_GPU_WGPU_GRAPH`: `0` off everywhere, `prefill` only for the
918/// prompt, anything else on everywhere. Unset: desktop-class GPUs take
919/// it for both phases; phone-class UMA takes it for PREFILL only, which
920/// is the measurement above rather than a guess — the per-op path keeps
921/// decode, where it is seventeen times better.
922pub 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            // Integrated/mobile keeps the per-op path for BOTH phases —
932            // unchanged, because the measurement that would have bought
933            // prefill a graph did not reproduce. `=prefill` is there for
934            // the device where it does; the default does not guess.
935            let _ = phase;
936            false
937        }
938    }
939}
940
941pub fn wgpu_graph_default() -> bool {
942    #[cfg(feature = "gpu")]
943    {
944        // Discrete cards always; Apple-silicon UMA on macOS too — desktop
945        // -class GPUs where the graph measured ~2x the CPU on the Qwen3.6
946        // family (M4: 13.3 tok/s against 7.3). Phone-class UMA (Android/
947        // iOS builds) keeps the per-op probe path: tiled mobile GPUs have
948        // turned the ~300-dispatch graph into seconds per token.
949        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/// q8_row/q8_2f matvec, rows [row0, row0+rows). `xs` — prescaled by the θ-field.
960#[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/// GEMM of a prefill batch: `pre` — prescaled inputs row-major [b, cols],
985/// out — row-major [b, rows].
986#[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/// q1 matvec: raw f32 activations, tile-embedded scales. Metal only
1009/// for now (wgpu q1 WGSL is queued); false = CPU fallback.
1010#[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/// Whole attention sub-block on the wgpu token graph (drop-in for
1029/// `qwen_attention`): normed hidden in, O-projection out, resident device
1030/// K/V mirror. false = refusal / not the wgpu backend → CPU path.
1031#[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
1068/// One weight in the whole-token graph: tensor idx + a codec tag (0=q8_row,
1069/// 1=q1, 2=q4_tiled, 3=q1t, 4=f32) + per-row scales (q8_row only) + the raw f32
1070/// data (kind 4 only — small unquantized projections like GDN in_proj_a/b).
1071pub struct GraphW<'a> {
1072    pub idx: usize,
1073    pub kind: u8,
1074    pub row_scale: &'a [f32],
1075    pub data: &'a [f32],
1076}
1077
1078/// A layer's token-mixing op: standard attention or a GDN (linear-attention)
1079/// block. The surrounding norms + SwiGLU FFN are common to both.
1080pub 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        /// (bq, bk, bv) attention biases (Qwen2). None ⇒ no bias.
1089        bias: Option<(&'a [f32], &'a [f32], &'a [f32])>,
1090        /// Qwen3.5 gated attention: wq emits 2·nh·hd (q||gate per head), the
1091        /// attention output is scaled by sigmoid(gate) before the O projection.
1092        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 recurrent state `[ring (kk-1)·cdim | S nv·dk·dv]` — seeds the
1112        /// device mirror when prefill ran on the host (o1 collection, CPU
1113        /// fallback): a zero-initialized device state at decode is exactly
1114        /// the "coherent but contextless" garble.
1115        cpu_state: &'a [f32],
1116    },
1117}
1118
1119/// Per-layer weights for the whole-token wgpu graph.
1120pub 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
1127/// The FFN of one graph layer: a dense SwiGLU trio, or a routed MoE —
1128/// router + top-k selection + all selected experts run ON DEVICE (the
1129/// routing decision depends on the resident hidden state, so a CPU
1130/// round-trip per layer would forfeit the one-submit design).
1131pub enum GraphFfn<'a> {
1132    Dense {
1133        gate: GraphW<'a>,
1134        up: GraphW<'a>,
1135        down: GraphW<'a>,
1136    },
1137    Moe {
1138        /// Router logits weight (f32, kind 4) `[n_exp, hidden]`.
1139        router: GraphW<'a>,
1140        /// Shared-expert sigmoid gate (f32) `[1, hidden]`.
1141        shared_gate: GraphW<'a>,
1142        /// Per-expert q4_tiled directory indices `(gate, up, down)`;
1143        /// the SHARED expert rides as the LAST entry — the select
1144        /// kernel pins it with the sigmoid weight.
1145        experts: Vec<(usize, usize, usize)>,
1146        /// Routed experts (shared excluded).
1147        n_exp: usize,
1148        top_k: usize,
1149        inter: usize,
1150        norm_topk: bool,
1151        /// Expert weight layout, uniform across the layer: `false` =
1152        /// q4_tiled (18 B tiles, inline f16 scale), `true` = q4tp
1153        /// (16 B nibbles + a per-row ladder plane). The two differ only
1154        /// in where the scale comes from, so they share every kernel
1155        /// but the weight-staging block.
1156        q4tp: bool,
1157        /// `true` = the gate/up experts are `q2tp` (2-bit plane) while
1158        /// `down` stays q4tp — the mixed profile a 2-bit-class checkpoint
1159        /// converts into. Only meaningful with `q4tp: true`.
1160        gu_q2: bool,
1161    },
1162}
1163
1164/// Whole-token decode graph on wgpu: the entire layer stack in ONE submit,
1165/// hidden resident, one readback. Updates `h` in place. false = refusal.
1166/// `loop_norm_at`: virtual layer indices after which `final_norm` is applied
1167/// (Looped Transformer mid-stack norm). Empty for standard models.
1168#[allow(clippy::too_many_arguments)]
1169pub fn forward_token_graph(
1170    model: &Arc<CmfModel>,
1171    kv_id: u64,
1172    layers: &[GraphLayer],
1173    // Per-layer sealed o1 (Nystrom) state; Some = replace this layer's
1174    // exact attention with the O(1) kernels. wgpu only.
1175    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    // How many leading layers the graph ran (see the wgpu twin) — smaller
1197    // than layers.len() when the expert budget ended the device prefix.
1198    layers_run: Option<&mut usize>,
1199    // Absolute index of layers[0] in the model — the KV/state mirrors key
1200    // on it, so a layer SPAN (network split segment) shares mirrors with
1201    // a full-stack run instead of colliding at slot 0.
1202    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
1242/// Speculative-verify tail for the batched graph: fold final-norm + lm_head
1243/// over every batch position and read all k logit rows back; the batch also
1244/// snapshots the GDN state per position for `gdn_spec_restore`.
1245pub 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/// Batched prefill: k contiguous positions through the whole graph in one submit
1253/// (projections/FFN as GEMMs, attention/GDN looped over scratch). `h` is
1254/// [k·hidden] in/out; `positions` len k. wgpu only.
1255#[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
1289/// After a partial speculative acceptance: restore every GDN layer's device
1290/// state to the snapshot after batch position `slot`. wgpu only.
1291pub 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
1303/// Drop the wgpu token graph's device K/V mirror for a pipeline.
1304pub 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
1311/// Ternary (q1t) BASE matvec on the GPU — fills `out` with the base dot; the
1312/// caller adds the sparse overlay on the CPU. Metal only for now (wgpu q1t not
1313/// yet written → CPU fallback).
1314pub 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/// q4_block matvec on the GPU — wgpu only (Metal drives q4_block through the
1338/// whole-token graph, not a standalone matvec).
1339#[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
1357/// q1t batched GEMM (prefill) — base + overlay on-device (Metal simdgroup or
1358/// wgpu register-blocked).
1359pub 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        // Batched prefill and single-token decode are both enabled. On the
1371        // real 14.8B Q1T model prefill PPL was within 0.3% of CPU (7.942 vs
1372        // 7.966), and the alignment-safe decode kernel reached 3.52e-6 max_rel.
1373        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/// Native Metal Q1T switch. Enabled by default after the byte-packed Q1T
1381/// fields were changed to alignment-safe loads; keep an explicit emergency
1382/// fallback for device/driver diagnostics.
1383#[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
1390/// Batched q1 GEMM (prefill). wgpu only — Metal has its own block path.
1391pub 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
1408/// Contention kill for the wide imagegen GEMM/FFN paths: one grossly
1409/// slow op under a work-proportional budget (fair-device ops are
1410/// ≤~100 ms even at 1024px) means another process owns the device —
1411/// verdicts are per-process, so CPU for the rest of this one.
1412static 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/// Fused DiT SwiGLU FFN on the device: g=X·W1ᵀ, u=X·W3ᵀ, silu(g)·u,
1421/// Causal chunk attention on the device: `b` queries against `s0 + b`
1422/// cached keys. wgpu only — Metal's chunk graph keeps attention inside
1423/// the resident block and never calls out.
1424#[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/// Fused QKV projection: one upload of the normed chunk, three GEMMs,
1446/// one readback of Q|K|V back to back. Metal has no twin yet — its
1447/// chunk graph keeps the whole layer resident and never surfaces QKV.
1448#[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/// y=·W2ᵀ — one command buffer, only X and Y cross the CPU boundary.
1471#[allow(unused_variables, clippy::too_many_arguments)]
1472/// SwiGLU FFN with a row-packed [gate|up] fc1 (MiniMax-H3's DiT), run
1473/// end to end on the device. wgpu only: Metal keeps the host loop until
1474/// its own packed kernel exists.
1475#[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
1539/// One whole modulated DiT block for `dit_block`: geometry, norm
1540/// weights, AdaLN scale/gate vectors (gates pre-tanh'd), a per-token
1541/// f32 RoPE cos/sin table, and the directory indices of the seven
1542/// q4t projections. `x` is in-out `[n, hidden]`.
1543pub 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    /// The projections' layout: q4tp (ladder scales) vs plain q4_tiled.
1571    /// The recommended Lumina file is q4tp, and a backend that only
1572    /// knows q4t must decline rather than decode with the wrong reader.
1573    pub q4tp: bool,
1574    /// The hidden state is already on the device from the previous block,
1575    /// so `x` need not be uploaded.
1576    pub resident_in: bool,
1577    /// Leave the result on the device instead of reading it back. The DiT
1578    /// loop does not touch `x` between blocks, so 27 of every 28 readbacks
1579    /// were moving 19 MB across PCIe and stalling on it for nothing.
1580    pub resident_out: bool,
1581}
1582
1583/// Can the selected backend keep the DiT's hidden state on the device
1584/// between blocks? Only the wgpu whole-block path; the Metal entry takes
1585/// and returns host memory every call.
1586pub 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
1595/// Pull the resident hidden state back to the host. For the caller that
1596/// chained blocks and then hit one the device declined.
1597pub 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/// One whole modulated DiT block on the device — norms, qkv, RoPE,
1608/// attention, residuals and the SwiGLU FFN in a single command
1609/// buffer; only `x` crosses the CPU boundary (in and out).
1610#[allow(unused_variables)]
1611/// The DiT's three projections in one submission (wgpu only; the
1612/// Metal path fuses the whole block instead). False = the caller keeps
1613/// its three separate calls.
1614#[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
1639/// Is a FUSED whole-block device path on offer? The batched-CFG shape
1640/// (two sequences in one tall batch) and the fused block (one sequence,
1641/// one command buffer) are alternatives, and the caller picks.
1642pub 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
1657/// The same block over a CONCATENATION of independent sequences:
1658/// attention per segment, everything position-wise batched. wgpu only —
1659/// the Metal path takes the single-sequence entry above.
1660pub 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        // The wgpu whole-block path. What it buys is host round trips —
1670        // six a block become one — so it defaults ON where those cost
1671        // real time (a discrete card across PCIe) and OFF on unified
1672        // memory, where the per-op path shares the same pages and the
1673        // fusion measured slightly slower on an M4. `CMF_DIT_FUSED=1`
1674        // forces it anywhere, `=0` forbids it.
1675        #[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
1690/// One VAE resnet block for `vae_resnet`: norm/conv weights and the
1691/// channel/shape geometry. `shortcut` is the 1×1 projection (w, b, k)
1692/// when in/out channels differ.
1693pub 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/// One whole VAE resnet block on the device (norm+silu → conv ×2 →
1713/// shortcut → add, one command buffer).
1714#[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/// Nearest-2× upsample fused with the following conv — the small
1724/// pre-upsample image is what crosses the CPU boundary.
1725#[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/// VAE conv2d on the device (implicit GEMM — the CPU path pays for a
1750/// multi-GB im2col matrix at high resolutions).
1751#[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/// DiT full bidirectional attention on the device (all heads:
1774/// scores GEMM → row softmax → P·V → panel unstack, one command
1775/// buffer). Head-major inputs; out is [n, nh·hd].
1776#[allow(unused_variables, clippy::too_many_arguments)]
1777/// Attention from an interleaved qkv panel, splitting into head-major
1778/// planes ON the device. wgpu only; `false` elsewhere so the caller
1779/// keeps its host repack.
1780#[allow(unused_variables)]
1781#[allow(clippy::too_many_arguments)]
1782/// qkv projection + attention with the panel never leaving the card.
1783/// wgpu only; `false` elsewhere and the caller keeps its host chain.
1784#[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/// The whole attention half of a DiT block on the card: qkv GEMM,
1808/// attention, output projection. Only `proj` comes home.
1809#[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/// The VAE decoder's attention half on the card. Only `proj` returns.
1834#[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
1914/// The backend's f32 NT GEMM: `y[n×m] = x[n×k] · wᵀ[m×k]`. Tensor
1915/// cores where the card has them. Refuses under `CMF_BAKE_GPU=0` or
1916/// strict f32, and for jobs below n·k·m = 4M, where the round trip
1917/// costs more than the arithmetic saves.
1918pub 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/// Music-3's FFN chain resident on the device — two GEMMs and the GLU
1928/// between them with no host round trip. `false` = refused, host runs.
1929#[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/// A 1D convolution as a GEMM whose column matrix is expanded on the
1952/// device instead of being built, transposed and uploaded by the host.
1953/// `yt` comes back `[out_n x oc]`. `false` = refused, caller runs host.
1954#[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/// The convolution as a GEMM on the matrix units. `false` = refused.
1982#[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    // (rope angles, q norm weights, k norm weights, eps) when the device
2009    // should apply qk-norm and RoPE itself; None when the host already did.
2010    nr: Option<(&[f32], &[f32], &[f32], f32)>,
2011    out: &mut [f32],
2012) -> bool {
2013    match backend() {
2014        // wgpu carries the only implementation, and it is not
2015        // platform-specific: `CMF_GPU=wgpu` on macOS runs it over Metal
2016        // like anywhere else. It used to be compiled out here on macOS,
2017        // which made the call a silent `false` — and the caller's
2018        // `assert!` turned that refusal into a panic on every
2019        // `cortiq animate` this platform ever ran.
2020        #[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
2027/// Whether `dit_attention_packed` has an implementation on the backend
2028/// that is actually selected.
2029///
2030/// The caller has to know BEFORE it skips the host qk-norm: deferring
2031/// the norm to a device that then refuses leaves q/k unnormalized with
2032/// no way back. Native Metal has no packed kernel, so on macOS this is
2033/// false unless `CMF_GPU=wgpu` picked the other backend.
2034pub 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/// Batched q4t GEMM on the device (imagegen DiT prefill shapes).
2065/// Metal: q4t_mul_mm decodes the mmap-resident tiles inside the
2066/// GEMM's K loop. wgpu (Vulkan/DX12 → NVIDIA/AMD/Intel/Adreno/Mali):
2067/// the register-blocked WGSL twin, weights cached in VRAM.
2068#[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
2088/// The same over a two-bit weight plane. Metal has no q2tp kernel, so
2089/// there it declines and the host takes it.
2090pub 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
2107/// Single-token q4tp matvec on the device — the lm_head class. Through the
2108/// DEDICATED matvec kernel: the batched GEMM at b=1 measured 11.73 ms
2109/// against the host's 9.51 on the release head, so the route that was
2110/// supposed to save eleven milliseconds a token lost its own probe instead.
2111pub 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
2129/// Single-token q4_tiled matvec on the device — the lm_head class (a
2130/// q4t checkpoint's head is its biggest host matvec, exactly like the
2131/// q4tp twin above). wgpu holds q4t_mv pipelines only inside the graph
2132/// encoder — the standalone arm stays an honest refusal until a
2133/// discrete-GPU q4t model reaches the bench.
2134pub 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/// Whole-block token-graph types re-exported from the Metal backend.
2170#[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/// A BLOCK of consecutive q1 GDN layers in one submission (Metal only).
2177#[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/// A layer's MoE-FFN in one submission (amortizing the dispatch cost).
2192#[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/// Independent matvecs of one input in a single submission (GDN projections).
2204#[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
2215// ── Whole-token wgpu graph race (generation granularity) ─────────────
2216// On integrated/mobile adapters the graph is neither trusted nor banned
2217// a priori — it RACES the normal path: generations alternate arms (the
2218// normal path first — known-good UX — then the graph), per-token wall
2219// times accumulate per arm, and once both arms have enough steady
2220// samples the faster one wins for the process. Arm switches happen ONLY
2221// at generation boundaries (`kv_cache.clear()` resets state), so the
2222// device KV mirror and the CPU cache never diverge mid-sequence. The
2223// single exception is the first-token bail: the very first decode token
2224// of a graph generation may be discarded and recomputed on the CPU
2225// path (the prompt KV is CPU-owned at that point, so this is safe) —
2226// a tiled mobile GPU that drains its pipeline at every barrier turns
2227// the ~300-dispatch graph into seconds per token (field report: 0.2
2228// tok/s vs 15 on the CPU), and one token is all it takes to see that.
2229static GRAPH_RACE_STATE: AtomicU8 = AtomicU8::new(0); // 0 racing, 1 graph won, 2 normal won
2230static GRAPH_RACE_FLIP: AtomicU32 = AtomicU32::new(0);
2231static GRAPH_RACE_ARM_GRAPH: AtomicU8 = AtomicU8::new(0); // this generation's arm
2232static GRAPH_RACE_TOK: AtomicU32 = AtomicU32::new(0); // token index within the generation
2233static GRAPH_NS: [AtomicU64; 2] = [AtomicU64::new(0), AtomicU64::new(0)]; // [normal, graph]
2234static GRAPH_N: [AtomicU32; 2] = [AtomicU32::new(0), AtomicU32::new(0)];
2235
2236/// Steady per-token samples per arm before the race decides.
2237const GRAPH_RACE_SAMPLES: u32 = 4;
2238
2239/// Called at every generation start (fresh KV). Applies a pending
2240/// verdict and picks this generation's arm while racing.
2241/// A graph that cannot be built for THIS model will never build: the
2242/// refusal is a property of the weights, not of the moment. Retrying it
2243/// per token is not free — the builder walks every layer and asks each
2244/// tensor for a graph view before giving up at layer 0 — and on an
2245/// Adreno 642L that retry cost 3x: forcing the graph on a model it
2246/// refuses measured 0.3 tok/s against 0.905 for the per-op path it falls
2247/// back to. Remembered once, the fallback runs at its own speed.
2248static GRAPH_UNSUPPORTED: AtomicBool = AtomicBool::new(false);
2249
2250/// The builder refused for a STRUCTURAL reason — an unsupported weight
2251/// or layer kind. Callers must NOT report the transient refusals (an
2252/// unsealed o1 state during prefill, a softcap): those clear on their
2253/// own and marking them would disable the graph for good.
2254pub 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
2264/// A different model in the same process starts with a clean slate.
2265pub fn graph_unsupported_reset() {
2266    GRAPH_UNSUPPORTED.store(false, Ordering::Relaxed);
2267}
2268
2269pub fn graph_race_begin_generation() {
2270    // One generation has now compiled whatever this model needs; keep it
2271    // for the next process. Once per run: the blob does not grow after
2272    // the pipelines exist, and the write is megabytes against the ~200 s
2273    // of compiling it saves on the device that needed this.
2274    #[cfg(feature = "gpu")]
2275    {
2276        // Save once, at the start of the SECOND generation: the first
2277        // has dispatched, so there is something to keep, and nothing is
2278        // saved before any work (the driver compiles at first use, not
2279        // at pipeline creation — the context comes up in 1.5 s while the
2280        // compiling costs minutes).
2281        //
2282        // Flushing again on 4, 8, 16 … was tried on the theory that a
2283        // chat turn compiles shapes the first one did not. It buys
2284        // nothing: a fresh app process still spent 49.0 s, then 58.7,
2285        // then 61.3 on its first answer with the backoff in place. One
2286        // flush it is.
2287        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            // Nothing dispatched yet.
2292        } 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
2321/// Should this decode token try the graph? `trusted` (discrete adapter,
2322/// explicit env, or a GDN hybrid whose state lives on the device) skips
2323/// the race entirely.
2324pub 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
2335/// First decode token of a racing graph generation: hopeless already?
2336/// (>4x the normal path's per-token average AND over a second.) Settles
2337/// the race immediately; the caller discards the graph result and
2338/// recomputes this token on the normal path.
2339pub 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
2362/// Record one decode-token wall time for the racing arm. The first
2363/// token of each generation is discarded (KV-mirror upload / cold
2364/// caches on the graph arm; cold mmap on the normal arm).
2365pub 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
2378/// Bounded-cost content fingerprint for the backends' pointer-keyed device
2379/// caches: FNV over the whole slice up to 4 KiB, over 64 spread 64-byte
2380/// windows (plus the length) above. An address-keyed hit must also prove
2381/// the bytes are still the ones it uploaded — the allocator reuses heap
2382/// and mmap addresses freely, so a reloaded model or a re-dequantized
2383/// layer lands where the old bytes were — and sampling keeps that proof at
2384/// ~a microsecond even for a 126 MB matrix. Real replacements (another
2385/// model's tensor, an Adam-updated master) differ densely, so a 4 KiB
2386/// spread cannot miss them.
2387pub(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
2412/// `fp_bytes` over an f32 slice without a bytemuck dependency (the Metal
2413/// backend builds with no GPU feature flags).
2414pub(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    /// The pointer-keyed caches survive on `fp_bytes` telling two different
2425    /// tensors apart at a reused address. Its sampling must therefore see a
2426    /// change ANYWHERE — head, tail, and the stretches between windows are
2427    /// the places a cheaper hash would go blind.
2428    #[test]
2429    fn fp_bytes_sees_a_change_anywhere_in_a_sampled_slice() {
2430        let n = 1 << 20; // 1 MiB — far above the 4 KiB full-hash threshold
2431        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        // A DENSE change (every requantized/redequantized tensor is one)
2435        // must flip the fingerprint no matter how the windows fall.
2436        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        // Length participates: the same prefix at a shorter length is a
2442        // different key AND a different fingerprint.
2443        assert_ne!(h0, fp_bytes(&base[..n - 64]));
2444        // Below the threshold the hash is exact: a single flipped byte in
2445        // a norm-sized vector must be seen.
2446        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        // And the sampled windows land within bounds on awkward sizes.
2451        for n in [4097usize, 5000, 64 * 64, 1 << 16] {
2452            let v = vec![9u8; n];
2453            let _ = fp_bytes(&v); // must not panic on window math
2454        }
2455    }
2456}
2457
2458/// Hand the card back after a bake: drop its resident weights, planes and
2459/// pools so the ordinary engine (the runtime gate, a serve that follows)
2460/// starts from a clean budget. No-op off the wgpu backend.
2461pub fn bake_release() {
2462    #[cfg(feature = "gpu")]
2463    crate::gpu_wgpu::bake_release();
2464}
2465
2466/// Strict-f32 for the bake's GEMMs (phase A mask training): the mask
2467/// selects neurons by a gradient signal, and f16 operand rounding on
2468/// that signal closes the wrong ones. No-op off the wgpu backend.
2469pub 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
2477/// CMF_GRAPH_HOSTPROF=1: how a graph token's wall splits between the
2478/// host encoding the command stream and the tail the GPU still owes
2479/// after encode. Fifteen GPU-side suspects measured null while the
2480/// bench counted 17.7k allocations a token — this is the instrument
2481/// that says whether the thief was on the host all along.
2482pub 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
2517/// Per-stage host-encode accumulator for the Metal token loop
2518/// (CMF_GRAPH_HOSTPROF=1). Stage 0 = GDN-run encode; everything else
2519/// falls out by subtraction from hostprof's encode total.
2520pub 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
2547/// Active weight bytes dispatched so far (Metal decode path); 0 where
2548/// the backend does not count. The honest floor's numerator.
2549pub 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
2563/// The per-stage split of `weight_bytes_dispatched`:
2564/// [misc, dense-ffn, moe, attn, gdn, head].
2565pub 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}