cortiq_engine/

gpu.rs

//! Facade for GPU backends: a single call entry point for qtensor/pipeline/
//! linear_core. Job types and the threshold are canonical HERE; behind the
//! facade dispatch goes to a platform backend:
//!   - `gpu_metal` (Apple Silicon, unified memory + no-copy buffers);
//!   - `gpu_wgpu` (C1: Vulkan/DX12/Metal — NVIDIA/Radeon/Intel/Apple,
//!     weights resident in VRAM), available under `--features gpu`.
//!
//! Runtime selection via `CMF_GPU`: `1` — native Metal (macOS) or wgpu
//! (other OSes); `wgpu` — force wgpu (including for the local
//! Metal-via-wgpu parity test). Any backend refusal — `false` and the honest
//! CPU path, no partial results.

use cortiq_core::CmfModel;
use std::cell::Cell;
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, OnceLock};

/// Packed, f32-only Embryo model consumed by the resident Vulkan graph.
///
/// This is deliberately a separate representation from `CmfModel`: the
/// Embryo graph owns one contiguous device copy and never dereferences the
/// mmap or asks the ordinary per-op residency path for a tensor.  `meta`
/// contains byte-free element offsets into `weights`; `UINT_MAX` marks an
/// absent optional matrix.  The builder lives in `pipeline.rs`, where the
/// architecture-specific weight types are visible.
pub struct EmbryoGraphModel {
    pub id: u64,
    pub hidden: usize,
    pub intermediate: usize,
    pub vocab: usize,
    pub layers: usize,
    pub phase_heads: usize,
    pub nphase: usize,
    pub phase_dv: usize,
    pub anchor_q_heads: usize,
    pub anchor_kv_heads: usize,
    pub anchor_head_dim: usize,
    pub rotary_dim: usize,
    pub max_seq: usize,
    pub cluster_count: usize,
    pub cluster_size: usize,
    pub phase_state_len: usize,
    pub state_stride: usize,
    pub kv_stride: usize,
    pub norm_gemma: bool,
    pub phase_mass: f32,
    pub weights: Vec<f32>,
    pub meta: Vec<u32>,
    pub lm_head: Vec<f32>,
    pub clusters: Vec<f32>,
    pub final_norm: Vec<f32>,
    pub inv_freq: Vec<f32>,
    /// Natively bounded anchors (`swa_sink_v1`): every anchor layer owns a
    /// ring of `anchor_window` raw keys/values instead of a `max_seq` KV
    /// plane, so `kv_stride` is the ring size, `kv_layers` counts the
    /// anchors, and no position cap applies (the operator has none).
    pub bounded: bool,
    /// Layers that own a KV/ring slot of `kv_stride` f32 (all layers for
    /// the legacy full anchor — offsets are `layer·kv_stride` — or the
    /// bounded anchors only).
    pub kv_layers: usize,
    /// Recurrent (phase or GDN) layers — the state buffer holds one
    /// `state_stride` per such layer; anchors own none.
    pub state_layers: usize,
    pub anchor_window: usize,
    pub anchor_sink: usize,
    /// vmf_phase mixer layers (kind 0/1) in the stack.
    pub phase_layers: usize,
    /// GatedDeltaNet mixer layers (kind 4) in the stack; their geometry
    /// is `gdn_heads` (nv), `gdn_k_heads` (nk), `gdn_dk`, `gdn_dv`,
    /// `gdn_kk` (conv taps) — header words 24..29 of `meta`.
    pub gdn_layers: usize,
    pub gdn_heads: usize,
    pub gdn_k_heads: usize,
    pub gdn_dk: usize,
    pub gdn_dv: usize,
    pub gdn_kk: usize,
}

impl EmbryoGraphModel {
    /// Fused q/k/v projection width of the GDN mixer (`2·nk·dk + nv·dv`).
    pub fn gdn_c_dim(&self) -> usize {
        2 * self.gdn_k_heads * self.gdn_dk + self.gdn_heads * self.gdn_dv
    }
}

/// Words in the resident graph's `meta` header before the per-layer
/// records (64 words each). Shared by the host packer and the encoder.
pub const EMBRYO_META_HEADER: usize = 32;

/// Positions one chunked-prefill submit of the resident Embryo graph
/// covers (the shader's `CMAX`; the chunk window of the scratch buffer is
/// sized by it).
pub const EMBRYO_CHUNK_MAX: usize = 64;

/// `n` contiguous prompt positions (`rows` = n × hidden embeddings from
/// `position`) through the resident Embryo graph in one submit; logits of
/// the last position. False = refused before any device work (the caller
/// keeps the per-position path), or the sequence went cold.
pub fn forward_embryo_graph_chunk(
    model: &Arc<EmbryoGraphModel>,
    kv_id: u64,
    rows: &[f32],
    position: usize,
    n: usize,
    logits: &mut Vec<f32>,
) -> bool {
    #[cfg(feature = "gpu")]
    if backend() == Backend::Wgpu {
        return crate::gpu_wgpu::forward_embryo_graph_chunk(model, kv_id, rows, position, n, logits);
    }
    let _ = (model, kv_id, rows, position, n, logits);
    false
}

/// Bytes the resident Embryo graph holds on the device for a sequence:
/// `(recurrent state, anchor KV/ring)`. None when the sequence has no
/// device image (host path, or not started yet). This is the measured
/// long-context claim of the resident path — a bounded genome reports
/// the same two numbers at every context depth.
pub fn embryo_device_state_bytes(kv_id: u64) -> Option<(u64, u64)> {
    #[cfg(feature = "gpu")]
    if backend() == Backend::Wgpu {
        return crate::gpu_wgpu::embryo_device_state_bytes(kv_id);
    }
    let _ = kv_id;
    None
}

/// Does this directory entry count as WEIGHT in the device placement
/// heuristics (layer prefix, live-weight budget, `serve` replicas vs
/// split)? On a genome file only the trunk does: skill tensors replace
/// trunk tensors of the same shape inside their own lane, so counting
/// them would give the backbone of F1 another GPU prefix than F0's
/// (NF-7). Other files keep the historical whole-directory count.
pub fn counts_as_placement_weight(model: &CmfModel, name: &str) -> bool {
    model.header.genome.is_none() || cortiq_core::knowledge::is_trunk_tensor(name)
}

/// Weight bytes the placement heuristics budget for `model`: the trunk's
/// payload on a genome file ([`counts_as_placement_weight`]), the whole
/// mapped file otherwise (the historical measure).
pub fn placement_weight_bytes(model: &CmfModel) -> u64 {
    if model.header.genome.is_none() {
        return model.primary_bytes().len() as u64;
    }
    model
        .tensors
        .iter()
        .filter(|t| counts_as_placement_weight(model, &t.name))
        .map(|t| t.nbytes)
        .sum()
}

/// Next position of the resident Embryo sequence `kv_id` (`Some(n)` = the
/// device holds positions `0..n` of it), `None` when the device holds no
/// initialized image of that sequence — the host owns it (or none began).
pub fn embryo_device_next_position(kv_id: u64) -> Option<usize> {
    #[cfg(feature = "gpu")]
    if backend() == Backend::Wgpu {
        return crate::gpu_wgpu::embryo_device_next_position(kv_id);
    }
    let _ = kv_id;
    None
}

thread_local! {
    /// Index of the current forward layer (−1 = outside a numbered layer:
    /// lm_head/embed — always allowed). The pipeline sets it before
    /// each layer so that the GPU/CPU layer-split works.
    static CUR_LAYER: Cell<i64> = const { Cell::new(-1) };
    /// Inside `cpu_scope` every GPU gate reports disabled: the timed CPU
    /// arm of a probe (and a class that lost its probe) must run PURE
    /// CPU, or inner per-op hooks would re-enter the GPU and poison the
    /// comparison.
    static CPU_ONLY: Cell<bool> = const { Cell::new(false) };
    /// "This op paid a one-off cost" (weight upload / first pipeline
    /// build): backends set it, `probe_record` discards the sample so
    /// only steady-state timings compete.
    static PROBE_COLD: Cell<bool> = const { Cell::new(false) };
}

/// RAII form of `cpu_scope`, used when a device-prefix graph hands a whole
/// remainder of the forward pass to the host. Without a guard around that
/// tail, its ordinary per-op hooks re-entered the GPU and streamed the rest of
/// an over-size model through the residency arena, defeating the prefix's VRAM
/// bound at the driver-allocation level.
pub struct CpuScopeGuard(bool);

impl Drop for CpuScopeGuard {
    fn drop(&mut self) {
        CPU_ONLY.with(|c| c.set(self.0));
    }
}

pub fn enter_cpu_scope() -> CpuScopeGuard {
    let previous = CPU_ONLY.with(|c| c.replace(true));
    CpuScopeGuard(previous)
}

/// Run `f` with the GPU gates off on this thread (pure-CPU arm).
pub fn cpu_scope<R>(f: impl FnOnce() -> R) -> R {
    let _restore = enter_cpu_scope();
    f()
}

/// Capture CPU-only placement before dispatching whole operators to workers.
/// `cpu_scope` is thread-local, while a MoE panel worker calls QTensor again;
/// without inheritance, which expert happened to land on the caller changed
/// its precision/backend from run to run.
pub(crate) fn inherit_cpu_scope() -> impl Fn() -> Option<CpuScopeGuard> + Copy {
    let on = CPU_ONLY.get();
    move || on.then(enter_cpu_scope)
}

/// Backends: name the device once at init. The probe cache is keyed by
/// it, because a verdict is a property of THIS silicon and nothing else.
/// First writer wins: a process runs one backend, and on the rare host
/// where two initialize, the one that came up first is the one in use.
pub fn probe_set_device(label: &str) {
    let _ = DEVICE_LABEL.set(label.to_string());
}

fn device_label() -> &'static str {
    DEVICE_LABEL.get().map(String::as_str).unwrap_or("unknown")
}

static DEVICE_LABEL: std::sync::OnceLock<String> = std::sync::OnceLock::new();

/// Somewhere this process may write small caches.
///
/// `std::env::temp_dir()` is NOT that place on Android: with no `TMPDIR`
/// it answers `/tmp`, which does not exist in an app sandbox, and every
/// write fails silently — measured, after the pipeline cache appeared to
/// work in a shell (where `TMPDIR=/data/local/tmp`) and did nothing at
/// all in the app. The loader points this at the model's own directory,
/// which is somewhere the caller already writes.
static CACHE_DIR: std::sync::OnceLock<std::path::PathBuf> = std::sync::OnceLock::new();

/// Loader: name a directory this process can write to. First call wins.
pub fn set_cache_dir(dir: std::path::PathBuf) {
    let _ = CACHE_DIR.set(dir);
}

/// Same directory, for the backends.
pub fn cache_dir_pub() -> std::path::PathBuf {
    cache_dir()
}

fn cache_dir() -> std::path::PathBuf {
    if let Some(d) = CACHE_DIR.get() {
        return d.clone();
    }
    match std::env::var_os("TMPDIR") {
        Some(t) => std::path::PathBuf::from(t),
        None => std::env::temp_dir(),
    }
}

/// Where decided verdicts are remembered between runs. `CMF_PROBE_CACHE`
/// overrides the path; `0` disables the cache entirely.
fn probe_cache_path() -> Option<std::path::PathBuf> {
    match std::env::var("CMF_PROBE_CACHE") {
        Ok(v) if v == "0" => None,
        Ok(v) => Some(std::path::PathBuf::from(v)),
        Err(_) => Some(cache_dir().join("cortiq-gpu-probe.tsv")),
    }
}

/// One line per decided class: `version \t device \t class \t winner`.
/// A different engine build or a different device simply does not match,
/// so a stale file is inert rather than wrong.
fn probe_cache_key_named(class: &str) -> String {
    format!(
        "{}\t{}\t{}",
        env!("CARGO_PKG_VERSION"),
        device_label(),
        class
    )
}

const CLASS_NAMES: [&str; 7] = [
    "ffn",
    "matvec",
    "matmat",
    "qkv-batch",
    "matmat-wide",
    "lm-head",
    "gemm-nt",
];

/// Adopt every verdict this device already reached in an earlier run.
///
/// Probing is not cheap and it is not free of consequences: on a
/// Snapdragon 778G the three deciding classes took **three minutes of
/// wall clock** before the first token, every process, and in the phone
/// app that was the whole first answer — 209.6 s for 25 tokens against
/// 10.5 s on the CPU path. The verdict itself was the same every time.
/// Paying to rediscover it is the defect; the answer is to write it down.
fn probe_cache_load() {
    static ONCE: std::sync::Once = std::sync::Once::new();
    ONCE.call_once(|| {
        let Some(path) = probe_cache_path() else {
            return;
        };
        // Unit tests share this process and its default cache path; a
        // verdict left by an earlier run would decide a class before the
        // arbitration tests get to watch it alternate. Tests that mean to
        // exercise the cache point `CMF_PROBE_CACHE` at their own file.
        if cfg!(test) && std::env::var("CMF_PROBE_CACHE").is_err() {
            return;
        }
        let Ok(text) = std::fs::read_to_string(&path) else {
            return;
        };
        probe_cache_adopt(&text);
    });
}

/// Apply verdicts from a cache file's text. Split out from the file
/// reading so the adoption rule — including which lines must be IGNORED
/// — is testable without a filesystem.
fn probe_cache_adopt(text: &str) {
    for line in text.lines() {
        let Some((key, verdict)) = line.rsplit_once('\t') else {
            continue;
        };
        let winner = match verdict.trim() {
            "gpu" => 1u8,
            "cpu" => 2u8,
            _ => continue,
        };
        for (i, name) in CLASS_NAMES.iter().enumerate() {
            if probe_cache_key_named(name) == key {
                let _ = PROBES[i].state.compare_exchange(
                    0,
                    winner,
                    Ordering::Relaxed,
                    Ordering::Relaxed,
                );
                tracing::debug!("gpu probe [{name}]: remembered → {verdict}");
            }
        }
    }
}

/// Remember a verdict for the next run. Best-effort: a read-only cache
/// directory costs a re-probe, never a failure.
fn probe_cache_store(c: OpClass, winner: u8) {
    let Some(path) = probe_cache_path() else {
        return;
    };
    let line = format!(
        "{}\t{}\n",
        probe_cache_key_named(CLASS_NAMES[c as usize]),
        if winner == 1 { "gpu" } else { "cpu" }
    );
    use std::io::Write;
    if let Ok(mut f) = std::fs::OpenOptions::new()
        .create(true)
        .append(true)
        .open(&path)
    {
        let _ = f.write_all(line.as_bytes());
    }
}

/// Backends: note a one-off cost (weight upload, buffer-cache fill) so
/// the probe discards this sample.
/// Every buffer creation anywhere bumps this; the graph's bind-group
/// cache treats any cold event as total invalidation — a stale bind
/// group is silent corruption, a cleared cache is one re-encoded token.
pub fn cold_epoch() -> u64 {
    COLD_EPOCH.load(std::sync::atomic::Ordering::Relaxed)
}
static COLD_EPOCH: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

pub(crate) fn probe_note_cold() {
    COLD_EPOCH.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    PROBE_COLD.with(|c| c.set(true));
}

/// Peek the cold flag without consuming it (`probe_record` consumes).
/// Contention heuristics use this: a slow COLD op is a one-off build
/// cost, not evidence the device is busy.
pub(crate) fn probe_was_cold() -> bool {
    PROBE_COLD.with(|c| c.get())
}

/// Pipeline: mark the current layer (or −1 outside layers) for layer-split.
pub fn set_layer(l: i64) {
    CUR_LAYER.with(|c| c.set(l));
}

/// The layer `set_layer` last marked on this thread (−1 outside layers).
pub fn cur_layer() -> i64 {
    CUR_LAYER.with(|c| c.get())
}

/// Capacity-derived layer prefix for per-op walks. The explicit
/// `CMF_GPU_LAYERS` override is handled by the backend and takes precedence.
pub fn automatic_layer_prefix(
    model: &Arc<CmfModel>,
    num_layers: usize,
    physical_layers: usize,
) -> Option<usize> {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => {
            crate::gpu_wgpu::automatic_layer_prefix(model, num_layers, physical_layers)
        }
        _ => None,
    }
}

/// Parse `CMF_GPU_LAYERS` («0-19», «0,2,4», «0-9,30-39») once.
/// None = no restriction (all layers on GPU). Garbage → also no restriction.
fn layer_ranges() -> &'static Option<Vec<(i64, i64)>> {
    static R: OnceLock<Option<Vec<(i64, i64)>>> = OnceLock::new();
    R.get_or_init(|| {
        let s = std::env::var("CMF_GPU_LAYERS").ok()?;
        let mut v = Vec::new();
        for part in s.split(',') {
            let part = part.trim();
            match part.split_once('-') {
                Some((a, b)) => v.push((a.trim().parse().ok()?, b.trim().parse().ok()?)),
                None => {
                    let x: i64 = part.parse().ok()?;
                    v.push((x, x));
                }
            }
        }
        Some(v)
    })
}

fn layer_allowed() -> bool {
    match layer_ranges() {
        None => true,
        Some(ranges) => {
            let cur = CUR_LAYER.with(|c| c.get());
            cur < 0 || ranges.iter().any(|(a, b)| cur >= *a && cur <= *b)
        }
    }
}

/// GPU allowed FOR THE CURRENT LAYER: backend is initialized AND the layer
/// falls within `CMF_GPU_LAYERS` (GPU/CPU layer-split) AND we are not
/// inside a `cpu_scope`. Op gates call this.
pub fn enabled_here() -> bool {
    !CPU_ONLY.with(|c| c.get()) && enabled() && layer_allowed()
}

/// Descriptor-aware q2tp Vulkan kernels are kept behind an explicit opt-in
/// until the Prism full-graph/resident-weight path has a coherent generation
/// gate.  `CMF_GPU=1` alone must not silently turn synchronous per-op
/// readbacks into the default model path; callers and validation tests can
/// request the measured kernels with `CMF_Q2TP_GPU=1`.
pub fn q2tp_gpu_opt_in() -> bool {
    std::env::var("CMF_Q2TP_GPU").as_deref() == Ok("1")
}

// ── Runtime GPU-vs-CPU probe ────────────────────────────────────────────
// CMF_GPU=1 does not TRUST that the device wins — it MEASURES. For each
// op class the first calls alternate arms: GPU timed vs pure-CPU timed
// (under cpu_scope). Cold GPU calls (weight upload / cache fill) are
// discarded; after PROBE_SAMPLES clean samples per arm the faster arm is
// chosen for the rest of the process. Rationale: submit+poll latency
// differs by an order of magnitude across driver stacks (Metal/PCIe
// ~3-4 ms, Vulkan/4090 ~0.3 ms) — a static threshold cannot know whether
// per-op offload pays off HERE. CMF_GPU_PROBE=0 → always trust the GPU.

/// GPU-eligible op classes, each with an independent probe.
#[derive(Clone, Copy)]
pub enum OpClass {
    /// Whole FFN chain in one submission (dense / MoE block).
    Ffn = 0,
    /// Large hybrid CPU∥GPU matvec (lm_head class).
    Matvec = 1,
    /// Prefill GEMM (matmat).
    Matmat = 2,
    /// Batched matvecs of one input (QKV).
    Batch = 3,
    /// Prefill GEMM at image-diffusion widths (b ≥ 128). Probed apart
    /// from `Matmat`: one imagegen process runs BOTH populations
    /// (prompt encode b≈40 where the GPU wins big, DiT b≥256 where
    /// the CPU AMX arm is competitive) — a single shared verdict locks
    /// the wrong arm for whichever population samples second.
    MatmatWide = 4,
    /// The lm_head itself, apart from the merely-large matvecs. Same
    /// reasoning as `MatmatWide`, and DeepSeek-V4 is where it bit: its
    /// attention projections are 37M weights and its head is 529M, so
    /// the projections' verdict — CPU, honestly measured at 0.19 ms —
    /// decided for a matvec fourteen times their size that took 11 ms
    /// a token on the host.
    MatvecHead = 5,
    /// The blocked f32 GEMM (`fcd_ops::gemm_nt`): attention's QKᵀ and
    /// AV, and the VAE decoders' projections. It used to take every job
    /// over 4 M MACs on sight, with no CPU arm to lose to — which on
    /// the MiniMax-H3 video decoder was three times SLOWER than the
    /// host it displaced. Its population is per-head slices, nothing
    /// like the weight GEMMs above, so it probes on its own.
    GemmNt = 6,
}

/// Which probe a large matvec belongs to. The head is an order of
/// magnitude bigger than anything else that reaches this gate, and the
/// two populations do not have the same answer.
pub fn matvec_class(rows: usize, cols: usize) -> OpClass {
    if rows * cols >= 67_108_864 {
        OpClass::MatvecHead
    } else {
        OpClass::Matvec
    }
}

/// Probe verdict for one call.
pub enum ProbeArm {
    /// Run the GPU path (during probing: timed, recorded).
    Gpu,
    /// Probing: run the CPU path under `cpu_scope`, timed, recorded.
    CpuTimed,
    /// Decided: CPU won — run the CPU path (under `cpu_scope`).
    Cpu,
}

/// Clean samples per arm before a class decides.
const PROBE_SAMPLES: u32 = 6;

/// Declines before a class gives the work to the host for good. High
/// enough that a transient refusal — an unsealed state during prefill, a
/// shape the kernel skips this once — cannot settle the question.
const PROBE_DECLINE_LIMIT: u32 = 16;

/// Device samples discarded before any count — see `Probe::gpu_burn`.
const PROBE_WARMUP: u32 = 1;

struct Probe {
    /// 0 = probing, 1 = GPU won, 2 = CPU won.
    state: AtomicU8,
    flip: AtomicU32,
    gpu_ns: AtomicU64,
    gpu_n: AtomicU32,
    /// Times the device arm was chosen and the device DECLINED.
    ///
    /// A decline carries no timing, so nothing is recorded — and a class
    /// whose device path always refuses therefore never reaches a
    /// verdict, alternates arms forever, and pays a failed device
    /// attempt on half of every token's calls. Measured on an M4 with
    /// LFM2.5-2.6B: `ffn` was still undecided after 9000 calls, and a
    /// token cost 83.55 ms against 41.85 with the device off — twice the
    /// price for work the host did anyway.
    declines: AtomicU32,
    /// GPU samples still to discard as warm-up.
    ///
    /// The cold flag catches buffer and weight uploads, but a compute
    /// pipeline is compiled on first use and not every creation site
    /// raises it — the wgpu path has 21 pipeline creations against 12
    /// cold notes. One uncaught shader compile is enough to lose a
    /// class for the whole process: `gemm-nt` on an A100 was recorded at
    /// 117.01 ms against the host's 3.19 and sent to the CPU, which
    /// parked a 27B bake on 2.6 cores with the card idle. The decision
    /// already uses each arm's BEST sample, so discarding the first
    /// GPU sample costs one extra round trip and removes the whole
    /// class of first-call artefacts.
    gpu_burn: AtomicU32,
    cpu_ns: AtomicU64,
    cpu_n: AtomicU32,
    /// Best (minimum) sample per arm. The DECISION compares these:
    /// means are poisoned by one-off cold costs the cold-flag cannot
    /// see — e.g. the CPU arm's first mmap-cold expert matvec page
    /// faults its weights in and reads 3× its steady state, which
    /// locked the GPU arm on a 35B MoE at a 4× real-world loss. The
    /// minimum is each arm's honest steady-state pace.
    gpu_min: AtomicU64,
    cpu_min: AtomicU64,
}

impl Probe {
    const fn new() -> Self {
        Self {
            state: AtomicU8::new(0),
            flip: AtomicU32::new(0),
            gpu_ns: AtomicU64::new(0),
            gpu_n: AtomicU32::new(0),
            declines: AtomicU32::new(0),
            gpu_burn: AtomicU32::new(PROBE_WARMUP),
            cpu_ns: AtomicU64::new(0),
            cpu_n: AtomicU32::new(0),
            gpu_min: AtomicU64::new(u64::MAX),
            cpu_min: AtomicU64::new(u64::MAX),
        }
    }
}

static PROBES: [Probe; 7] = [
    Probe::new(),
    Probe::new(),
    Probe::new(),
    Probe::new(),
    Probe::new(),
    Probe::new(),
    Probe::new(),
];

/// A caller that knows its loop is long, uniform and warm can say so: the
/// probe times ops in isolation and alternates arms to do it, which reads a
/// sustained diffusion step as slower on the device than it is. Measured on
/// an M4 at 672 video tokens: the probe picked the CPU at 1.25 ms against
/// 0.88 ms per op, and the loop it picked for ran 23.9 s a step against the
/// device's 19.7 s.
static TRUST_GPU: AtomicBool = AtomicBool::new(false);

/// Take the probe out of the loop until the guard drops.
pub fn trust_gpu() -> GpuTrust {
    let was = TRUST_GPU.swap(true, Ordering::Relaxed);
    GpuTrust(was)
}

pub struct GpuTrust(bool);

impl Drop for GpuTrust {
    fn drop(&mut self) {
        TRUST_GPU.store(self.0, Ordering::Relaxed);
    }
}

fn probe_on_for(c: OpClass) -> bool {
    // The trust is only for the *wide* class. A sustained diffusion step is
    // where the probe reads a warm device as cold; the narrow batches inside
    // the same loop — an audio stream of fifty-one tokens against the same
    // weights — are small enough that submit latency can genuinely beat the
    // arithmetic, and there the probe is right and should keep deciding.
    if TRUST_GPU.load(Ordering::Relaxed) && matches!(c, OpClass::MatmatWide | OpClass::Ffn) {
        return false;
    }
    probe_on()
}

/// Is the per-op GPU/CPU probe enabled (`CMF_GPU_PROBE`, default on)?
/// The native Metal decode route never consults it — `q1_force` routes
/// the token graph to the device outright — so it is reported, not used.
pub fn probe_enabled() -> bool {
    probe_on()
}

fn probe_on() -> bool {
    static ON: OnceLock<bool> = OnceLock::new();
    *ON.get_or_init(|| {
        std::env::var("CMF_GPU_PROBE")
            .map(|v| v != "0" && v != "off")
            .unwrap_or(true)
    })
}

/// q1 ops on the native Metal backend skip the probe entirely: the CPU
/// q1 kernel is load-port-bound, the GPU one wins warm — and probe
/// alternation itself cools the device between samples (measured: block
/// times 5.8 ms warm vs 8.8 ms mixed). Other backends keep probing.
pub fn q1_force() -> bool {
    #[cfg(target_os = "macos")]
    {
        backend() == Backend::Metal
    }
    #[cfg(not(target_os = "macos"))]
    {
        false
    }
}

/// Should a FUSED whole-block path trust the device instead of asking
/// the per-op probe? True on native Metal and on discrete wgpu adapters.
///
/// The probe answers "is one wide matmat faster on the GPU", and for the
/// DiT on Metal that is a coin flip — measured 2.62 ms GPU vs 2.56 ms
/// CPU, a 2% spread that lands on either arm run to run. But the fused
/// block's advantage is not per-op speed, it is that the hidden state,
/// the packs and the attention panels never leave the device: end to end
/// the whole-block path renders a 512² Lumina step in ~5.4 s against
/// ~8.4 s when the probe happens to pick the CPU. Gating a fusion win on
/// a per-op tie made every second render half-speed at random.
///
/// On a discrete card the verdict is never in doubt — an RTX 3090 against
/// a 256-core EPYC measured 11.5 ms vs 31 ms per wide op, four runs out
/// of four — so the probe's sampling phase is pure cost: it alone was 10%
/// of a 512² render (74.3 s against 66.9 s with the probe off). Integrated
/// and mobile adapters keep probing; there the submit latency is real and
/// can genuinely lose.
pub fn fused_block_trusted() -> bool {
    #[cfg(target_os = "macos")]
    if backend() == Backend::Metal {
        return true;
    }
    wgpu_graph_default()
}

/// Which arm should this GPU-eligible call take? Consult AFTER the
/// eligibility gates (`enabled_here` / `min_rows`) so only real
/// candidates alternate.
/// While a class is still probing, a call whose weights are NOT yet on
/// the card should take the GPU arm anyway: the upload is work the next
/// step needs regardless, and the sample it produces is discarded as
/// cold — so handing that call to the CPU arm buys nothing and costs a
/// host GEMM. Measured on a diffusion stack, where every layer is
/// touched once per step and therefore EVERY first-step GPU sample is
/// cold: one projection drew the CPU arm for the whole first step, 9.8 s
/// against the 2.8 s it costs once the weights are warm.
pub fn weight_is_resident(model: &Arc<CmfModel>, idx: usize) -> bool {
    #[cfg(all(feature = "gpu", not(target_os = "macos")))]
    {
        return crate::gpu_wgpu::weight_is_resident(model, idx);
    }
    #[cfg(not(all(feature = "gpu", not(target_os = "macos"))))]
    {
        let _ = (model, idx);
        true
    }
}

pub fn probe_arm_cold_prefers_gpu(c: OpClass, weights_resident: bool) -> ProbeArm {
    if !weights_resident && probe_deciding(c) {
        return ProbeArm::Gpu;
    }
    probe_arm(c)
}

pub fn probe_arm(c: OpClass) -> ProbeArm {
    // Every arbitrated call starts with a clean cold flag: both the
    // sample discard in `probe_record` and the contention kill-switch
    // read it AFTER the op, so a stale note from a previous call on
    // this thread must not leak in.
    PROBE_COLD.with(|f| f.set(false));
    if !probe_on_for(c) {
        return ProbeArm::Gpu;
    }
    probe_cache_load();
    let p = &PROBES[c as usize];
    match p.state.load(Ordering::Relaxed) {
        1 => ProbeArm::Gpu,
        2 => ProbeArm::Cpu,
        _ => {
            if p.flip.fetch_add(1, Ordering::Relaxed) % 2 == 0 {
                ProbeArm::Gpu
            } else {
                ProbeArm::CpuTimed
            }
        }
    }
}

/// The device arm was chosen and the device refused the work, so there
/// is no time to record. Callers that fall through to the host MUST say
/// so here, or the class can never decide.
pub fn probe_note_decline(c: OpClass) {
    let p = &PROBES[c as usize];
    if p.state.load(Ordering::Relaxed) != 0 {
        return;
    }
    let n = p.declines.fetch_add(1, Ordering::Relaxed) + 1;
    if n >= PROBE_DECLINE_LIMIT
        && p.state
            .compare_exchange(0, 2, Ordering::Relaxed, Ordering::Relaxed)
            .is_ok()
    {
        tracing::info!(
            "gpu probe [{}]: device declined {n} times → cpu",
            CLASS_NAMES[c as usize]
        );
    }
}

/// Record a timed arm sample; on the `PROBE_SAMPLES`-th clean sample of
/// BOTH arms the class decides for the rest of the process.
pub fn probe_record(c: OpClass, gpu: bool, dur: std::time::Duration) {
    probe_record_into(
        &PROBES[c as usize],
        CLASS_NAMES[c as usize],
        Some(c),
        gpu,
        dur,
    )
}

/// The body of `probe_record` over ONE probe, so the decision can be
/// driven in a test without touching the process-wide array.
fn probe_record_into(
    p: &Probe,
    class_name: &str,
    cache: Option<OpClass>,
    gpu: bool,
    dur: std::time::Duration,
) {
    if p.state.load(Ordering::Relaxed) != 0 {
        return;
    }
    if gpu && PROBE_COLD.with(|f| f.replace(false)) {
        return; // one-off cost in this call — not a steady-state sample
    }
    if gpu {
        // Load-then-store rather than fetch_sub: a blind decrement at
        // zero wraps a u32 to its maximum and mutes the arm forever.
        // A benign race here burns one extra sample, which is free.
        let left = p.gpu_burn.load(Ordering::Relaxed);
        if left > 0 {
            p.gpu_burn.store(left - 1, Ordering::Relaxed);
            return; // warm-up: the first device sample builds its pipeline
        }
    }
    let ns = dur.as_nanos().min(u64::MAX as u128) as u64;
    if gpu {
        p.gpu_ns.fetch_add(ns, Ordering::Relaxed);
        p.gpu_n.fetch_add(1, Ordering::Relaxed);
        p.gpu_min.fetch_min(ns, Ordering::Relaxed);
    } else {
        p.cpu_ns.fetch_add(ns, Ordering::Relaxed);
        p.cpu_n.fetch_add(1, Ordering::Relaxed);
        p.cpu_min.fetch_min(ns, Ordering::Relaxed);
    }
    let (gn, cn) = (
        p.gpu_n.load(Ordering::Relaxed),
        p.cpu_n.load(Ordering::Relaxed),
    );
    if gn >= 2 && cn >= 2 {
        // Decide on each arm's BEST sample — the steady-state pace.
        // Means carry one-off cold costs (mmap page-in on the CPU arm)
        // that the cold-flag machinery cannot see.
        let g = p.gpu_min.load(Ordering::Relaxed) as f64;
        let cp = p.cpu_min.load(Ordering::Relaxed) as f64;
        // Early verdict on a ≥2× gap — no reason to keep feeding the
        // losing arm; close races take the full sample count. It was 3×,
        // and the cost of that half-octave was measured: a DiT whose
        // wide GEMMs run 11.4 ms on the device against 32.2 on the host
        // (2.8×) kept ALTERNATING through the whole diffusion stack, and
        // because the alternation counter is shared per class in call
        // order, one projection drew the CPU arm every single time — 9.9
        // seconds a step on a kernel that needs 0.4. Both arms are
        // compared on their BEST sample, so a 2× gap is not noise.
        if (gn < PROBE_SAMPLES || cn < PROBE_SAMPLES) && g < cp * 2.0 && cp < g * 2.0 {
            return;
        }
        let winner = if g <= cp { 1 } else { 2 };
        if p.state
            .compare_exchange(0, winner, Ordering::Relaxed, Ordering::Relaxed)
            .is_ok()
        {
            tracing::info!(
                "gpu probe [{}]: gpu {:.2} ms vs cpu {:.2} ms per op → {}",
                class_name,
                g / 1e6,
                cp / 1e6,
                if winner == 1 { "gpu" } else { "cpu" },
            );
            if let Some(c) = cache {
                probe_cache_store(c, winner);
            }
        }
    }
}

/// Is the class still collecting samples? (Call sites use this to route
/// cold-weight calls away from the GPU arm during probing.)
pub fn probe_deciding(c: OpClass) -> bool {
    probe_on_for(c) && PROBES[c as usize].state.load(Ordering::Relaxed) == 0
}

/// Probing helper: true — tensor `idx`'s quant weights are ALREADY
/// device-resident (a clean GPU sample is possible now); false — they
/// were not (the upload starts within the VRAM budget, so a later call
/// finds them warm) or the tensor cannot go to the GPU at all. Keeps the
/// probe from billing a full cold dispatch+readback to a sample it will
/// discard anyway. The verdict needs only a couple of warm tensors, so
/// probe-driven uploads are capped — the losing-GPU machine should not
/// pay for uploading the whole layer stack it will never use; if the GPU
/// wins, the rest uploads lazily on demand, in the same first-touch order.
#[allow(unused_variables)]
pub fn q8_resident_or_upload(model: &Arc<CmfModel>, idx: usize) -> bool {
    static PROBE_UPLOADS: AtomicU32 = AtomicU32::new(0);
    let may_upload = PROBE_UPLOADS.load(Ordering::Relaxed) < 4;
    let resident = match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::q8_resident_or_upload(model, idx, may_upload),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q8_resident_or_upload(model, idx, may_upload),
        Backend::None => false,
    };
    if !resident && may_upload {
        PROBE_UPLOADS.fetch_add(1, Ordering::Relaxed);
    }
    resident
}

/// Test hook: reset all probes to the undecided state.
#[cfg(test)]
pub(crate) fn probe_reset() {
    for p in &PROBES {
        p.state.store(0, Ordering::Relaxed);
        p.flip.store(0, Ordering::Relaxed);
        p.gpu_ns.store(0, Ordering::Relaxed);
        p.gpu_n.store(0, Ordering::Relaxed);
        p.cpu_ns.store(0, Ordering::Relaxed);
        p.cpu_n.store(0, Ordering::Relaxed);
    }
}

/// The probe table is process-global by design, while these unit tests reset
/// and seed selected entries to exercise arbitration. Keep only those tests
/// out of each other's way; production callers still probe concurrently.
#[cfg(test)]
static PROBE_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());

#[cfg(test)]
fn probe_test_guard() -> std::sync::MutexGuard<'static, ()> {
    PROBE_TEST_LOCK
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
}

#[cfg(test)]
mod probe_tests {
    use super::*;
    use std::time::Duration;

    #[test]
    fn cpu_only_whole_operator_dispatch_inherits_and_restores_scope() {
        let pool = crate::pool::Pool::with_spin(3, 0);
        cpu_scope(|| {
            let inherit = inherit_cpu_scope();
            pool.run_rows(64, &|_, _| {
                let _guard = inherit();
                assert!(CPU_ONLY.get());
            });
        });
        pool.run_rows(64, &|_, _| assert!(!CPU_ONLY.get()));
    }

    // One test fn: PROBES is process-global and probe_reset touches all
    // classes — parallel test threads would race.
    #[test]
    fn probe_alternates_discards_cold_and_decides() {
        let _probe_guard = probe_test_guard();
        probe_reset();
        // Probing: arms alternate.
        assert!(matches!(probe_arm(OpClass::Ffn), ProbeArm::Gpu));
        assert!(matches!(probe_arm(OpClass::Ffn), ProbeArm::CpuTimed));

        // A cold GPU sample (upload noted) must be discarded: feed a
        // catastrophic cold sample, then clean fast-GPU samples — GPU
        // wins only if the cold one did not count.
        probe_note_cold();
        probe_record(OpClass::Ffn, true, Duration::from_secs(1000));
        for _ in 0..PROBE_SAMPLES {
            probe_record(OpClass::Ffn, true, Duration::from_millis(1));
            probe_record(OpClass::Ffn, false, Duration::from_millis(4));
        }
        assert!(matches!(probe_arm(OpClass::Ffn), ProbeArm::Gpu));

        // The reverse: a class where the CPU arm is faster decides CPU.
        for _ in 0..PROBE_SAMPLES {
            probe_record(OpClass::Matmat, true, Duration::from_millis(4));
            probe_record(OpClass::Matmat, false, Duration::from_millis(1));
        }
        assert!(matches!(probe_arm(OpClass::Matmat), ProbeArm::Cpu));

        // cpu_scope: gates off inside, restored after.
        cpu_scope(|| CPU_ONLY.with(|c| assert!(c.get())));
        CPU_ONLY.with(|c| assert!(!c.get()));
        cpu_scope(|| {
            cpu_scope(|| CPU_ONLY.with(|c| assert!(c.get())));
            CPU_ONLY.with(|c| assert!(c.get()));
        });
        let _ = std::panic::catch_unwind(|| cpu_scope(|| panic!("scope test")));
        CPU_ONLY.with(|c| assert!(!c.get()));
        probe_reset();
    }

    #[test]
    fn a_remembered_verdict_is_adopted_and_a_stranger_is_not() {
        let _probe_guard = probe_test_guard();
        // Probing is not free: on a Snapdragon 778G the deciding classes
        // cost minutes of wall clock before the first token, every
        // process, and reached the same verdict every time. The cache
        // exists so that price is paid once.
        //
        // The key is built from THIS process's device, never a name this
        // test sets: `probe_set_device` is first-writer-wins and on a Mac
        // the Metal backend may already have named the silicon before the
        // tests run — which is exactly how this test failed on CI while
        // passing locally. GemmNt on purpose: the arbitration test never
        // touches it, and both run in one process.
        let mine = probe_cache_key_named("gemm-nt");
        let state = || {
            PROBES[OpClass::GemmNt as usize]
                .state
                .load(Ordering::Relaxed)
        };

        // Another device's verdict is not mine, whatever it claims.
        probe_cache_adopt("SomeOtherGPU/Vulkan\tgemm-nt\tgpu\n");
        assert_eq!(state(), 0);
        // Neither is one from another build of this engine.
        let older = mine.replacen(env!("CARGO_PKG_VERSION"), "0.0.0-old", 1);
        assert_ne!(older, mine);
        probe_cache_adopt(&format!("{older}\tgpu\n"));
        assert_eq!(state(), 0);
        // Mine is.
        probe_cache_adopt(&format!("{mine}\tcpu\n"));
        assert_eq!(state(), 2);

        PROBES[OpClass::GemmNt as usize]
            .state
            .store(0, Ordering::Relaxed);
    }
}

/// Default row threshold: the GPU takes only larger matrices (lm_head
/// class). Below it, the dispatch/readback cost does not pay off on unified memory.
pub const GPU_MIN_ROWS: usize = 65_536;

/// Effective threshold: `CMF_GPU_MIN_ROWS` overrides. Defaults differ
/// by device class: on a DISCRETE card VRAM bandwidth pays off even for
/// FFN/QKV-class matrices (4096), on unified memory only lm_head-class
/// is worth the dispatch/readback (65536). Field case behind this: a
/// 35B model on an RTX 4090 saw ~0 offload because every layer matrix
/// sat below the old universal 65536.
pub fn min_rows() -> usize {
    if let Some(v) = std::env::var("CMF_GPU_MIN_ROWS")
        .ok()
        .and_then(|v| v.parse().ok())
    {
        return v;
    }
    if discrete() { 4096 } else { GPU_MIN_ROWS }
}

/// Is the active backend a discrete card (PCIe VRAM)?
pub fn discrete() -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::is_discrete(),
        #[cfg(target_os = "macos")]
        Backend::Metal => false, // UMA by the init() guard
        Backend::None => false,
    }
}

/// A single MoE-FFN job (an expert with its own weight), executed in one
/// submission: (rows, cols, idx, row_scale) for gate/up/down + prescaled
/// inputs + the down column scale + the blending weight.
pub struct MoeJob<'a> {
    pub gate: (usize, usize, usize, &'a [f32]),
    pub up: (usize, usize, usize, &'a [f32]),
    pub down: (usize, usize, usize, &'a [f32]),
    pub xs_gate: Vec<f32>,
    pub xs_up: Vec<f32>,
    pub down_col: &'a [f32],
    pub w: f32,
    /// q1 trio: scales live inside the 6-byte tiles (row_scale slices
    /// empty, xs raw f32). Backends without a q1 kernel refuse the job.
    pub q1: bool,
    /// q4_tiled trio: scales inside the 18-byte tiles (row_scale
    /// slices empty, xs raw f32) — the MoE-hybrid coder class.
    pub q4t: bool,
    /// q4tp trio: same raw-xs contract, 16-byte nibble stride and the scale
    /// on a per-row ladder. Without this the experts of a q4tp MoE model fall
    /// to the CPU while every other dtype rides the device.
    pub q4tp: bool,
    /// Mixed 2-bit profile: gate/up are q2tp (8-byte chunks, zero rung),
    /// down stays q4tp. Set together with `q4tp`; a backend without the
    /// 2-bit kernel must refuse the whole job.
    pub gu_q2: bool,
    /// The reference's `swiglu_limit`; 0 disables the clamp. A backend that
    /// cannot apply it must REFUSE the job rather than drop it silently —
    /// the difference only shows on saturating activations, which is the
    /// hardest kind of divergence to notice.
    pub swiglu_limit: f32,
}

/// A single independent batch matvec (GDN projections of one input).
pub struct BatchJob<'a> {
    pub idx: usize,
    pub rows: usize,
    pub cols: usize,
    pub row_scale: &'a [f32],
    pub xs: Vec<f32>,
    /// Weight layout. Was a bare `q1: bool`, which could only ever spell two
    /// of the four and silently sent everything else back to the CPU — the
    /// GDN projections of a q4t/q4tp model never reached the device at all.
    pub layout: BatchLayout,
}

/// Which kernel a batched matvec needs. q8 carries row scales in a side
/// buffer; the rest embed them in the payload and differ in stride.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum BatchLayout {
    Q8,
    Q1,
    Q4t,
    Q4tp,
}

#[derive(Clone, Copy, PartialEq, Eq)]
enum Backend {
    None,
    #[cfg(target_os = "macos")]
    Metal,
    #[cfg(feature = "gpu")]
    Wgpu,
}

fn backend() -> Backend {
    #[cfg(feature = "gpu")]
    if crate::gpu_wgpu::selected() {
        return if crate::gpu_wgpu::enabled() {
            Backend::Wgpu
        } else {
            Backend::None
        };
    }
    #[cfg(target_os = "macos")]
    if crate::gpu_metal::enabled() {
        return Backend::Metal;
    }
    Backend::None
}

/// GPU enabled and initialized on the selected backend?
/// Whether THIS build can bring a GPU up on THIS device: a compiled-in
/// backend plus a live adapter. The mobile FFI exposes it so an app can
/// tell "GPU off" from "GPU impossible" (a CPU-only .so ships no
/// backend at all). Cached after the first call.
pub fn backend_available() -> bool {
    #[cfg(target_os = "macos")]
    {
        // The Metal path is always compiled on macOS.
        true
    }
    #[cfg(all(feature = "gpu", not(target_os = "macos")))]
    {
        static AVAIL: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
        *AVAIL.get_or_init(crate::gpu_wgpu::adapter_probe)
    }
    #[cfg(all(not(feature = "gpu"), not(target_os = "macos")))]
    {
        false
    }
}

/// A process-wide, phase-scoped GPU gate. `cpu_scope` is thread-local and
/// the pool's workers do not inherit it, so a caller that wants a whole
/// *phase* off the device — a prompt encoder whose weights live in a part of
/// the file the hot loop never touches, on a machine that cannot keep both
/// wired — has to say so globally.
static GPU_PAUSED: AtomicBool = AtomicBool::new(false);

/// Park the device for every thread until the returned guard drops.
pub fn pause_gpu() -> GpuPause {
    GPU_PAUSED.store(true, Ordering::Relaxed);
    GpuPause(())
}

pub struct GpuPause(());

impl Drop for GpuPause {
    fn drop(&mut self) {
        GPU_PAUSED.store(false, Ordering::Relaxed);
    }
}

pub fn enabled() -> bool {
    !GPU_PAUSED.load(Ordering::Relaxed) && backend() != Backend::None
}

/// Default-on condition for the wgpu whole-token graph: the wgpu
/// backend on a DISCRETE adapter. NOT plain `enabled()` (macOS/Metal
/// must not pay a per-token layer scan for a graph its backend
/// refuses), and NOT integrated adapters: the graph's ~300 barriered
/// dispatches per token are cheap on desktop immediate-mode GPUs but
/// tiled mobile GPUs (Adreno/Mali) drain the pipeline at every barrier
/// — field report: 0.2 tok/s on-graph vs 15 tok/s on the CPU. On
/// integrated adapters the per-op probe path arbitrates each op class
/// against the CPU instead; CMF_GPU_WGPU_GRAPH=1 still forces the
/// graph anywhere.
/// Is the wgpu backend active at all (any adapter)? Eligibility gate
/// for the whole-token graph — whether it actually RUNS is decided by
/// `wgpu_graph_default` (trusted on discrete) or the generation race.
pub fn wgpu_active() -> bool {
    #[cfg(feature = "gpu")]
    {
        matches!(backend(), Backend::Wgpu)
    }
    #[cfg(not(feature = "gpu"))]
    {
        false
    }
}

/// Which GPU this thread's engine calls address. Multi-card hosts hold
/// one wgpu context PER card (weights, KV mirrors and scratch live
/// inside a context, so per-device contexts give per-device caches for
/// free); this thread-local says which one is current. Default: the
/// process pin (CMF_GPU_ADAPTER) or 0 — so single-card runs behave
/// exactly as they always have.
pub fn default_device() -> usize {
    static D: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
    *D.get_or_init(|| {
        std::env::var("CMF_GPU_ADAPTER")
            .ok()
            .and_then(|v| v.trim().parse::<usize>().ok())
            .unwrap_or(0)
    })
}

thread_local! {
    static CUR_DEV: std::cell::Cell<Option<usize>> = const { std::cell::Cell::new(None) };
}

/// The device this thread is pinned to.
pub fn current_device() -> usize {
    CUR_DEV.with(|c| c.get()).unwrap_or_else(default_device)
}

/// Pin this thread to a device. Server slots call it once per request;
/// the worker pool propagates it into its threads, so a dispatch begun
/// on card 1 does not finish on card 0.
pub fn set_current_device(i: usize) {
    CUR_DEV.with(|c| c.set(Some(i)));
}

/// Run `f` with this thread pinned to `dev`, restoring the previous pin.
pub fn with_device<R>(dev: usize, f: impl FnOnce() -> R) -> R {
    let prev = CUR_DEV.with(|c| c.replace(Some(dev)));
    let r = f();
    CUR_DEV.with(|c| c.set(prev));
    r
}

/// How many GPUs this process can address (wgpu adapter count; 1 on
/// Metal, 0 without a backend).
pub fn device_count() -> usize {
    #[cfg(all(feature = "gpu", not(target_os = "macos")))]
    {
        return crate::gpu_wgpu::adapter_count();
    }
    #[cfg(not(all(feature = "gpu", not(target_os = "macos"))))]
    {
        usize::from(backend_available())
    }
}

/// Weight budget of the current GPU in bytes; 0 when there is none and
/// u64::MAX on unified memory (where the OS pages shared RAM and the
/// question "does the model fit the card" has no separate answer).
pub fn vram_budget() -> u64 {
    #[cfg(all(feature = "gpu", not(target_os = "macos")))]
    {
        return crate::gpu_wgpu::device_vram_budget();
    }
    #[cfg(not(all(feature = "gpu", not(target_os = "macos"))))]
    {
        if backend_available() { u64::MAX } else { 0 }
    }
}

/// Bytes currently accounted as resident weight buffers on the active wgpu
/// adapter.  This is the logical device-local weight set; physical driver
/// allocations are reported separately by the platform tools.
pub fn resident_bytes() -> u64 {
    #[cfg(feature = "gpu")]
    {
        if backend() == Backend::Wgpu {
            return crate::gpu_wgpu::resident_bytes();
        }
    }
    0
}

/// Sealed O(1) device mirror count and logical bytes for one pipeline id.
/// Zero is returned when wgpu is unavailable or the sequence has not reached
/// an O(1) seal yet.
pub fn o1_device_stats(kv_id: u64) -> (usize, u64) {
    #[cfg(feature = "gpu")]
    {
        if backend() == Backend::Wgpu {
            return crate::gpu_wgpu::o1_device_stats(kv_id);
        }
    }
    let _ = kv_id;
    (0, 0)
}

/// Device weight bytes uploaded so far (wgpu; 0 on other backends).
/// Steady-state windows must show a ZERO delta — growth mid-benchmark
/// means eviction/re-upload and disqualifies the number.
pub fn upload_bytes() -> u64 {
    #[cfg(feature = "gpu")]
    {
        return crate::gpu_wgpu::UPLOAD_BYTES.load(std::sync::atomic::Ordering::Relaxed);
    }
    #[cfg(not(feature = "gpu"))]
    0
}

/// Measure a transient host-to-device upload when the wgpu backend is
/// compiled in. CPU-only builds keep the benchmark command available and
/// report no device measurement instead of referring to the gated module.
pub fn upload_bandwidth_probe(block: usize, rounds: usize) -> Option<f64> {
    #[cfg(feature = "gpu")]
    {
        return crate::gpu_wgpu::upload_bandwidth_probe(block, rounds);
    }
    let _ = (block, rounds);
    None
}

/// Which half of the run is asking.
///
/// The phase exists because the graph is plausibly two decisions, not
/// one — but on the hardware measured so far it is only ever a decode
/// decision. On an Adreno 642L with bonsai-1.7b, from identical clean
/// starts and two repeats each: decode 11.6 tok/s without it and 0.72
/// with, while prefill is 4.2 either way. A first reading of 3.4 -> 18.0
/// for prefill did not survive a controlled re-run — it was a dirty
/// probe cache between configurations, not the graph, and the prefill
/// route through the graph is GDN-only in the first place, which this
/// dense model never takes.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum GraphPhase {
    Prefill,
    Decode,
}

/// The one place that decides whether the whole-token graph runs.
///
/// `CMF_GPU_WGPU_GRAPH`: `0` off everywhere, `prefill` only for the
/// prompt, anything else on everywhere. Unset: desktop-class GPUs take
/// it for both phases; phone-class UMA takes it for PREFILL only, which
/// is the measurement above rather than a guess — the per-op path keeps
/// decode, where it is seventeen times better.
pub fn wgpu_graph_on(phase: GraphPhase) -> bool {
    match std::env::var("CMF_GPU_WGPU_GRAPH").ok().as_deref() {
        Some("0") => false,
        Some("prefill") => phase == GraphPhase::Prefill,
        Some(_) => true,
        None => {
            if wgpu_graph_default() {
                return true;
            }
            // Integrated/mobile keeps the per-op path for BOTH phases —
            // unchanged, because the measurement that would have bought
            // prefill a graph did not reproduce. `=prefill` is there for
            // the device where it does; the default does not guess.
            let _ = phase;
            false
        }
    }
}

pub fn wgpu_graph_default() -> bool {
    #[cfg(feature = "gpu")]
    {
        // Discrete cards always; Apple-silicon UMA on macOS too — desktop
        // -class GPUs where the graph measured ~2x the CPU on the Qwen3.6
        // family (M4: 13.3 tok/s against 7.3). Phone-class UMA (Android/
        // iOS builds) keeps the per-op probe path: tiled mobile GPUs have
        // turned the ~300-dispatch graph into seconds per token.
        matches!(backend(), Backend::Wgpu)
            && (crate::gpu_wgpu::discrete_active()
                || (cfg!(target_os = "macos") && crate::gpu_wgpu::adapter_up()))
    }
    #[cfg(not(feature = "gpu"))]
    {
        false
    }
}

/// q8_row/q8_2f matvec, rows [row0, row0+rows). `xs` — prescaled by the column scale.
#[allow(clippy::too_many_arguments, unused_variables)]
pub fn q8_matvec_range(
    model: &Arc<CmfModel>,
    idx: usize,
    row0: usize,
    row_scale: &[f32],
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => {
            crate::gpu_metal::q8_matvec_range(model, idx, row0, row_scale, xs, rows, cols, out)
        }
        #[cfg(feature = "gpu")]
        Backend::Wgpu => {
            crate::gpu_wgpu::q8_matvec_range(model, idx, row0, row_scale, xs, rows, cols, out)
        }
        Backend::None => false,
    }
}

/// Decode-exact short q8_2f panel for the banked MiMo head.
pub(crate) fn q82_short_rows(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    #[cfg(feature = "gpu")]
    if enabled_here() && backend() == Backend::Wgpu {
        return crate::gpu_wgpu::q82_short_rows(model, idx, xs, b, rows, cols, out);
    }
    let _ = (model, idx, xs, b, rows, cols, out);
    false
}

/// GEMM of a prefill batch: `pre` — prescaled inputs row-major [b, cols],
/// out — row-major [b, rows].
#[allow(clippy::too_many_arguments, unused_variables)]
/// The two-field int8 GEMM with the column field left for the device.
/// wgpu only — Metal's int8 kernel takes a pre-scaled activation, so the
/// caller keeps that path when this returns `false`.
#[allow(clippy::too_many_arguments)]
pub fn q8_matmat_2f(
    model: &Arc<CmfModel>,
    idx: usize,
    row_scale: &[f32],
    col_field: &[f32],
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    #[allow(unreachable_patterns)]
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => {
            crate::gpu_wgpu::q8_matmat_2f(model, idx, row_scale, col_field, xs, b, rows, cols, out)
        }
        _ => false,
    }
}

pub fn q8_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    row_scale: &[f32],
    pre: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => {
            crate::gpu_metal::q8_matmat(model, idx, row_scale, pre, b, rows, cols, out)
        }
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q8_matmat(model, idx, row_scale, pre, b, rows, cols, out),
        Backend::None => false,
    }
}

/// q1 matvec: raw f32 activations, tile-embedded scales. Metal only
/// for now (wgpu q1 WGSL is queued); false = CPU fallback.
#[allow(unused_variables)]
pub fn q1_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::q1_matvec(model, idx, xs, rows, cols, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q1_matvec(model, idx, xs, rows, cols, out),
        Backend::None => false,
    }
}

/// Whole attention sub-block on the wgpu token graph (drop-in for
/// `qwen_attention`): normed hidden in, O-projection out, resident device
/// K/V mirror. false = refusal / not the wgpu backend → CPU path.
#[allow(clippy::too_many_arguments)]
pub fn attn_dropin(
    model: &Arc<CmfModel>,
    kv_id: u64,
    layer: usize,
    normed: &[f32],
    wq_idx: usize,
    wk_idx: usize,
    wv_idx: usize,
    wo_idx: usize,
    q_norm: Option<&[f32]>,
    k_norm: Option<&[f32]>,
    late_qk_norm: bool,
    invf: &[f32],
    nh: usize,
    nkv: usize,
    hd: usize,
    rd: usize,
    hidden: usize,
    pos: usize,
    cap: usize,
    gemma: bool,
    eps: f32,
    cpu_k: &[Vec<f32>],
    cpu_v: &[Vec<f32>],
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::attn_dropin_gpu(
            model, kv_id, layer, normed, wq_idx, wk_idx, wv_idx, wo_idx, q_norm, k_norm,
            late_qk_norm, invf, nh, nkv, hd, rd, hidden, pos, cap, gemma, eps, cpu_k, cpu_v, out,
        ),
        #[allow(unused_variables)]
        _ => false,
    }
}

/// Descriptor operation attached to a graph weight.  `None` is the default
/// for ordinary CMF files; Prism weights are admitted only when the token
/// graph carries this explicit transform contract.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum GraphPrismOp {
    None,
    Forward,
    InverseEmbedding,
}

/// One weight in the whole-token graph: tensor idx + a codec tag (0=q8_row,
/// 1=q1, 2=q4_tiled, 3=q1t, 4=f32) + per-row scales (q8_row only) + the raw f32
/// data (kind 4 only — small unquantized projections like GDN in_proj_a/b).
pub struct GraphW<'a> {
    pub idx: usize,
    pub kind: u8,
    pub row_scale: &'a [f32],
    pub data: &'a [f32],
    pub prism: GraphPrismOp,
    pub affine: bool,
}

/// A layer's token-mixing op: standard attention or a GDN (linear-attention)
/// block. The surrounding norms + SwiGLU FFN are common to both.
pub enum GraphAttn<'a> {
    Full {
        wq: GraphW<'a>,
        wk: GraphW<'a>,
        wv: GraphW<'a>,
        wo: GraphW<'a>,
        q_norm: Option<&'a [f32]>,
        k_norm: Option<&'a [f32]>,
        /// HunYuan dense: q/k norm after RoPE (rope-kernel flag bit 32).
        late_qk_norm: bool,
        /// (bq, bk, bv) attention biases (Qwen2). None ⇒ no bias.
        bias: Option<(&'a [f32], &'a [f32], &'a [f32])>,
        /// Qwen3.5 gated attention: wq emits 2·nh·hd (q||gate per head), the
        /// attention output is scaled by sigmoid(gate) before the O projection.
        output_gate: bool,
        cpu_k: &'a [Vec<f32>],
        cpu_v: &'a [Vec<f32>],
        /// Absolute position of `cpu_k`/`cpu_v` row 0
        /// (`LayerKvCache::base`): a sliding layer that keeps only its
        /// tail stores rows from there on. 0 for every untrimmed layer.
        cpu_base: usize,
        /// This layer's own attention geometry, when the model's layers do
        /// not share one (MiMo-V2: 4/8 KV heads, 128-wide V under 192-wide
        /// heads, sliding windows with learned sinks, two RoPE tables).
        /// None = the call-wide (nkv, hd, rd, invf), V as wide as K, full
        /// context and a plain softmax — the historical contract, whose
        /// kernels and dispatch are untouched.
        geom: Option<GraphAttnGeom<'a>>,
        /// Spark-X2.5 head-wise output gate: `self_attn.g_proj`
        /// `[num_heads, hidden]` against the attention's normed input; head
        /// h's output is scaled by sigmoid(g[h]) before the O projection.
        /// None = no such gate (every other model).
        head_gate: Option<GraphW<'a>>,
    },
    Gdn {
        qkv: GraphW<'a>,
        z: GraphW<'a>,
        a: GraphW<'a>,
        b: GraphW<'a>,
        out: GraphW<'a>,
        conv1d: &'a [f32],
        a_log: &'a [f32],
        dt_bias: &'a [f32],
        norm: &'a [f32],
        nv: usize,
        nk: usize,
        dk: usize,
        dv: usize,
        kk: usize,
        /// CPU recurrent state `[ring (kk-1)·cdim | S nv·dk·dv]` — seeds the
        /// device mirror when prefill ran on the host (o1 collection, CPU
        /// fallback): a zero-initialized device state at decode is exactly
        /// the "coherent but contextless" garble.
        cpu_state: &'a [f32],
    },
    /// LFM2 gated short convolution: a fused (B, C, x) projection, a
    /// depthwise causal conv over a (kernel−1)-deep per-channel ring,
    /// C-gating, and an output projection. This mixer is what most of an
    /// LFM2 stack is (22 of the 2.6B's 30 layers), and before it had a
    /// graph arm the whole model fell to the per-op path — ~100 submits
    /// a token, 22 tok/s on an A100 for a 1.4 GB file.
    ShortConv {
        /// [3·hidden, hidden] fused input projection.
        inp: GraphW<'a>,
        /// [hidden, hidden] output projection.
        out: GraphW<'a>,
        /// [hidden · kernel] depthwise taps, `[channel][tap]`, tap
        /// kernel−1 multiplying the current position.
        taps: &'a [f32],
        kernel: usize,
        /// CPU conv ring `[channel][kernel−1]`, slot 0 newest — seeds
        /// the device mirror when prefill ran on the host, which for
        /// this mixer is always (the batch graph declines it).
        cpu_state: &'a [f32],
    },
}

/// Per-layer attention geometry for the wgpu graphs (see
/// `GraphAttn::Full::geom`). The layer's CPU cache keeps K rows `hd` wide
/// and V rows zero-padded to `hd`; the device mirror stores V `dv` wide,
/// and a windowed layer keeps a ring of the last positions only.
#[derive(Clone, Copy)]
pub struct GraphAttnGeom<'a> {
    /// KV heads of this layer (divides the Q heads).
    pub nkv: usize,
    /// V head width, `4 <= dv <= head_dim`, a multiple of 4.
    pub dv: usize,
    /// Rotary width of this layer (NeoX half-split over `[0, rd)`).
    pub rd: usize,
    /// This layer's RoPE inverse frequencies (`rd / 2` of them).
    pub invf: &'a [f32],
    /// Positions a query sees, its own included (MiMo-V2 SWA: 128);
    /// None = the whole context.
    pub window: Option<usize>,
    /// Learned per-Q-head sink logits (gpt-oss / MiMo-V2): they join the
    /// softmax max and denominator and carry no value row.
    pub sink: Option<&'a [f32]>,
}

/// Activation of a graph layer's dense FFN, `act(gate)·up`. The code is
/// the word the kernels read (`silu_mul_pre`'s `_c`, the fused gate+up
/// kernels' `_p1`): 0 keeps the historical SiLU expression untouched.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum GraphAct {
    Silu,
    /// Exact erf GELU (HF `gelu`), erf by Abramowitz–Stegun 7.1.26 — the
    /// formula `inference::erf_f32` uses on the host.
    GeluErf,
}

impl GraphAct {
    pub fn code(self) -> u32 {
        match self {
            Self::Silu => 0,
            Self::GeluErf => 1,
        }
    }
}

/// Per-layer weights for the whole-token wgpu graph.
pub struct GraphLayer<'a> {
    pub input_norm: &'a [f32],
    pub attn: GraphAttn<'a>,
    pub post_norm: &'a [f32],
    pub ffn: GraphFfn<'a>,
}

/// The FFN of one graph layer: a dense SwiGLU trio, or a routed MoE —
/// router + top-k selection + all selected experts run ON DEVICE (the
/// routing decision depends on the resident hidden state, so a CPU
/// round-trip per layer would forfeit the one-submit design).
pub enum GraphFfn<'a> {
    /// A singleton attention-only batch graph. Returns the post-attention
    /// residual, allowing a dynamic expert bank to own the FFN separately.
    AttentionOnly,
    Dense {
        gate: GraphW<'a>,
        up: GraphW<'a>,
        down: GraphW<'a>,
        /// act(gate)·up. The builders refuse an activation with no arm here
        /// (before this field the graph computed SiLU for any dense FFN).
        act: GraphAct,
    },
    Moe {
        /// Router logits weight (f32, kind 4) `[n_exp, hidden]`.
        router: GraphW<'a>,
        /// Shared-expert sigmoid gate (f32) `[1, hidden]`.
        shared_gate: GraphW<'a>,
        /// Per-expert q4_tiled directory indices `(gate, up, down)`;
        /// the SHARED expert rides as the LAST entry — the select
        /// kernel pins it with the sigmoid weight.
        experts: Vec<(usize, usize, usize)>,
        /// Routed experts (shared excluded).
        n_exp: usize,
        top_k: usize,
        inter: usize,
        norm_topk: bool,
        /// Expert weight layout, uniform across the layer: `false` =
        /// q4_tiled (18 B tiles, inline f16 scale), `true` = q4tp
        /// (16 B nibbles + a per-row ladder plane). The two differ only
        /// in where the scale comes from, so they share every kernel
        /// but the weight-staging block.
        q4tp: bool,
        /// `true` = the gate/up experts are `q2tp` (2-bit plane) while
        /// `down` stays q4tp — the mixed profile a 2-bit-class checkpoint
        /// converts into. Only meaningful with `q4tp: true`.
        gu_q2: bool,
        /// LFM2-MoE / DeepSeek-V3 `noaux_tc` routing: per-expert sigmoid
        /// scores instead of a softmax, and `norm_topk` renormalises with
        /// the 1e-6 floor. The softmax arm is bit-identical to before.
        sigmoid: bool,
        /// Per-expert SELECTION bias: added to the score for the top-k
        /// choice only — the mixing weights stay unbiased (noaux_tc).
        bias: Option<&'a [f32]>,
        /// Whether a shared expert rides as the last `experts` entry.
        /// LFM2-MoE has none; the select kernel then leaves slot `top_k`
        /// unwritten and the expert loop runs `top_k` slots, not +1.
        has_shared: bool,
        /// The shared expert carries a sigmoid gate (Qwen2/3-MoE). `false`
        /// with `has_shared`: the shared expert enters with weight 1
        /// (DeepSeek-V3 / HunYuan hy_v3) and `shared_gate` is a stand-in
        /// the select kernels ignore.
        shared_gated: bool,
        /// Multiplier on the routed mixing weights after the optional
        /// renormalization (`routed_scaling_factor`); 1.0 = none.
        route_scale: f32,
    },
}

/// Outcome of one whole-token graph attempt. A failed attempt after sealed
/// O(1) state was admitted must not fall through to the stale CPU state.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum TokenGraphOutcome {
    /// No command was committed; the caller may use its ordinary path.
    Declined,
    /// The graph completed and its hidden/logits output is valid.
    Completed,
    /// Sealed O(1) state was admitted and a later graph operation failed.
    Failed,
}

/// Whole-token decode graph on wgpu: the entire layer stack in ONE submit,
/// hidden resident, one readback. Updates `h` in place.
/// `loop_norm_at`: virtual layer indices after which `final_norm` is applied
/// (Looped Transformer mid-stack norm). Empty for standard models.
#[allow(clippy::too_many_arguments)]
pub fn forward_token_graph(
    model: &Arc<CmfModel>,
    kv_id: u64,
    layers: &[GraphLayer],
    // Per-layer sealed o1 (Nystrom) state; Some = replace this layer's
    // exact attention with the O(1) kernels. wgpu only.
    o1: &[Option<Vec<crate::nystrom::O1DeviceView<'_>>>],
    o1_epoch: u64,
    invf: &[f32],
    h: &mut [f32],
    nh: usize,
    nkv: usize,
    hd: usize,
    attn_scale: f32,
    rd: usize,
    hidden: usize,
    inter: usize,
    position: usize,
    cap: usize,
    gemma: bool,
    eps: f32,
    lm_head: Option<(&GraphW, usize)>,
    final_norm: &[f32],
    logits: &mut Vec<f32>,
    loop_norm_at: &[usize],
    steps: usize,
    embed: Option<(&GraphW, usize, f32)>,
    ids_out: Option<&mut Vec<u32>>,
    // How many leading layers the graph ran (see the wgpu twin) — smaller
    // than layers.len() when the expert budget ended the device prefix.
    layers_run: Option<&mut usize>,
    // Absolute index of layers[0] in the model — the KV/state mirrors key
    // on it, so a layer SPAN (network split segment) shares mirrors with
    // a full-stack run instead of colliding at slot 0.
    layer_base: usize,
    // Read the final hidden back alongside the fused head's logits.
    hidden_too: bool,
) -> TokenGraphOutcome {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::forward_token_graph(
            model,
            kv_id,
            layers,
            o1,
            o1_epoch,
            invf,
            h,
            nh,
            nkv,
            hd,
            attn_scale,
            rd,
            hidden,
            inter,
            position,
            cap,
            gemma,
            eps,
            lm_head,
            final_norm,
            logits,
            loop_norm_at,
            steps,
            embed,
            ids_out,
            layers_run,
            layer_base,
            hidden_too,
        ),
        #[allow(unused_variables)]
        _ => {
            let _ = (
                attn_scale,
                lm_head,
                final_norm,
                logits,
                loop_norm_at,
                layers_run,
                layer_base,
                hidden_too,
            );
            TokenGraphOutcome::Declined
        }
    }
}

/// Whole-token resident graph for the Embryo working model.  Unlike the
/// generic graph this path accepts a packed f32 model and keeps both phase
/// recurrent state and the anchor KV cache on the device.  `false` is an
/// honest capability refusal; callers must retain the CPU executor.
pub fn forward_embryo_graph(
    model: &Arc<EmbryoGraphModel>,
    kv_id: u64,
    hidden: &[f32],
    position: usize,
    logits: &mut Vec<f32>,
) -> bool {
    #[cfg(feature = "gpu")]
    if backend() == Backend::Wgpu {
        return crate::gpu_wgpu::forward_embryo_graph(model, kv_id, hidden, position, logits);
    }
    let _ = (model, kv_id, hidden, position, logits);
    false
}

/// Speculative-verify tail for the batched graph: fold final-norm + lm_head
/// over every batch position and read all k logit rows back; the batch also
/// snapshots the GDN state per position for `gdn_spec_restore`.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum BatchGraphOutcome {
    /// The graph declined before mutating persistent device state. Callers may
    /// safely use the existing per-position path.
    Declined,
    /// The complete batch committed and its readback succeeded.
    Completed,
    /// A batch that had admitted sealed O(1) state failed after admission.
    /// Falling back to CPU would mix two state machines, so the caller must
    /// abort and clear the sequence instead.
    Failed,
}

pub struct SpecTail<'a> {
    pub lm: GraphW<'a>,
    pub lm_rows: usize,
    pub final_norm: &'a [f32],
    pub logits_out: &'a mut Vec<f32>,
}

/// Batched prefill: k contiguous positions through the whole graph in one submit
/// (projections/FFN as GEMMs, attention/GDN looped over scratch). `h` is
/// [k·hidden] in/out; `positions` len k. wgpu only.
#[allow(clippy::too_many_arguments)]
pub fn forward_batch_graph(
    model: &Arc<CmfModel>,
    kv_id: u64,
    layers: &[GraphLayer],
    invf: &[f32],
    h: &mut [f32],
    nh: usize,
    nkv: usize,
    hd: usize,
    rd: usize,
    hidden: usize,
    inter: usize,
    positions: &[usize],
    cap: usize,
    gemma: bool,
    eps: f32,
    attn_scale: f32,
    k: usize,
    // Per-layer sealed O(1) device views. An empty slice means the ordinary
    // exact-KV path; otherwise it must have one entry per graph layer.
    o1: &[Option<Vec<crate::nystrom::O1DeviceView<'_>>>],
    o1_epoch: u64,
    spec: Option<SpecTail<'_>>,
    // Device-prefix mode (plain prefill only): Some = when the whole stack
    // does not fit the weight budget, run the leading layers that do — the
    // same prefix rule as the token graph — leave the boundary hidden in
    // `h` and report the count here; the caller runs the rest on the host.
    // None = all layers or a decline, as before.
    layers_run: Option<&mut usize>,
) -> BatchGraphOutcome {
    forward_batch_graph_at(model, kv_id, 0, layers, invf, h, nh, nkv, hd, rd, hidden, inter, positions, cap, gemma, eps, attn_scale, k, o1, o1_epoch, spec, layers_run)
}

thread_local! {
    static MIMO_ATTN_SCRATCH: std::cell::Cell<bool> = const { std::cell::Cell::new(true) };
}

pub(crate) fn mimo_attention_scratch_enabled() -> bool {
    MIMO_ATTN_SCRATCH.with(std::cell::Cell::get)
}

/// Scoped diagnostic A/B switch; unlike process environment mutations it
/// cannot race a model's background workers. Restores on panic as well.
#[doc(hidden)]
pub fn mimo_attention_scratch_scope<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
    struct Restore(bool);
    impl Drop for Restore {
        fn drop(&mut self) {
            MIMO_ATTN_SCRATCH.with(|v| v.set(self.0));
        }
    }
    let _restore = Restore(MIMO_ATTN_SCRATCH.with(|v| v.replace(enabled)));
    f()
}

thread_local! {
    static MIMO_Q8_SHORT: std::cell::Cell<bool> = const { std::cell::Cell::new(true) };
}

pub(crate) fn mimo_q8_short_enabled() -> bool {
    MIMO_Q8_SHORT.with(std::cell::Cell::get)
}

/// Diagnostic A/B switch for row-exact short q8 graph kernels.
#[doc(hidden)]
pub fn mimo_q8_short_scope<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
    struct Restore(bool);
    impl Drop for Restore {
        fn drop(&mut self) {
            MIMO_Q8_SHORT.with(|v| v.set(self.0));
        }
    }
    let _restore = Restore(MIMO_Q8_SHORT.with(|v| v.replace(enabled)));
    f()
}

/// Batched graph over a span whose first absolute layer is `layer_base`.
#[allow(clippy::too_many_arguments)]
pub fn forward_batch_graph_at(
    model: &Arc<CmfModel>,
    kv_id: u64,
    layer_base: usize,
    layers: &[GraphLayer],
    invf: &[f32],
    h: &mut [f32],
    nh: usize,
    nkv: usize,
    hd: usize,
    rd: usize,
    hidden: usize,
    inter: usize,
    positions: &[usize],
    cap: usize,
    gemma: bool,
    eps: f32,
    attn_scale: f32,
    k: usize,
    // Per-layer sealed O(1) device views. An empty slice means the ordinary
    // exact-KV path; otherwise it must have one entry per graph layer.
    o1: &[Option<Vec<crate::nystrom::O1DeviceView<'_>>>],
    o1_epoch: u64,
    spec: Option<SpecTail<'_>>,
    // Device-prefix mode (plain prefill only): Some = when the whole stack
    // does not fit the weight budget, run the leading layers that do — the
    // same prefix rule as the token graph — leave the boundary hidden in
    // `h` and report the count here; the caller runs the rest on the host.
    // None = all layers or a decline, as before.
    layers_run: Option<&mut usize>,
) -> BatchGraphOutcome {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::forward_batch_graph_at(
            model, kv_id, layer_base, layers, invf, h, nh, nkv, hd, rd, hidden, inter, positions,
            cap, gemma, eps, attn_scale, k, o1, o1_epoch, spec, layers_run,
        ),
        #[allow(unreachable_patterns)]
        _ => {
            let _ = (o1, o1_epoch, spec, layers_run);
            BatchGraphOutcome::Declined
        }
    }
}

/// After a partial speculative acceptance: restore every GDN layer's device
/// state to the snapshot after batch position `slot`. `base_pos` is the
/// absolute position of the first verify row and `expected_layers` makes the
/// restore all-or-nothing across the model's recurrent layers. wgpu only.
pub fn gdn_spec_restore(kv_id: u64, slot: usize, base_pos: usize, expected_layers: usize) -> bool {
    #[cfg(feature = "gpu")]
    if backend() == Backend::Wgpu {
        return crate::gpu_wgpu::gdn_spec_restore(kv_id, slot, base_pos, expected_layers);
    }
    #[allow(unreachable_code)]
    {
        let _ = (kv_id, slot, base_pos, expected_layers);
        false
    }
}

/// Re-point one exact-attention device mirror after a speculative round has
/// discarded unaccepted rows. The rows beyond `stored` remain allocated and
/// are overwritten by the next append; only the logical cursor moves. This
/// is the wgpu twin of Metal's existing mirror cursor helper and keeps the
/// MTP graph's speculative/device cache coherent with its real anchor.
pub fn graph_kv_set_stored(kv_id: u64, layer: usize, stored: usize) -> bool {
    #[cfg(feature = "gpu")]
    if backend() == Backend::Wgpu {
        return crate::gpu_wgpu::kv_mirror_set_stored(kv_id, layer, stored);
    }
    #[cfg(target_os = "macos")]
    if backend() == Backend::Metal {
        crate::gpu_metal::kv_mirror_set_stored(kv_id, layer, stored);
        return true;
    }
    false
}

/// Rows the wgpu token graph's exact-attention mirror holds for one layer
/// (None: no wgpu mirror). Metal keeps its owner cache current per token
/// and reports None here.
pub fn graph_kv_stored(_kv_id: u64, _layer: usize) -> Option<usize> {
    #[cfg(feature = "gpu")]
    if backend() == Backend::Wgpu {
        return crate::gpu_wgpu::kv_mirror_stored(_kv_id, _layer);
    }
    None
}

/// Does the wgpu token graph hold a device-resident recurrent state for
/// this layer (one the host `linear_state` has not seen)?
pub fn graph_state_resident(_kv_id: u64, _layer: usize) -> bool {
    #[cfg(feature = "gpu")]
    if backend() == Backend::Wgpu {
        return crate::gpu_wgpu::graph_state_resident(_kv_id, _layer);
    }
    false
}

/// Copy rows back from the wgpu token graph's K/V mirrors in one submit:
/// for each `(layer, from, to)` the K and V rows `[from..to)`, position-major
/// (`[(to − from) × nkv × hd]` each).
pub fn graph_kv_read_rows(
    _kv_id: u64,
    _reqs: &[(usize, usize, usize)],
    _nkv: usize,
    _hd: usize,
) -> Option<Vec<(Vec<f32>, Vec<f32>)>> {
    #[cfg(feature = "gpu")]
    if backend() == Backend::Wgpu {
        return crate::gpu_wgpu::kv_mirror_read_rows(_kv_id, _reqs, _nkv, _hd);
    }
    None
}

/// Rows `[from, to)` of one wgpu exact-attention mirror in the host
/// cache's layout (V zero-padded to `hd`), whatever the mirror's geometry
/// (narrow V, a sliding layer's ring). The third value is the first
/// position actually read: a ring returns zeros below it. None: no wgpu
/// mirror holding those rows at (nkv, hd).
pub fn graph_kv_pull_host(
    _kv_id: u64,
    _layer: usize,
    _from: usize,
    _to: usize,
    _nkv: usize,
    _hd: usize,
) -> Option<(Vec<f32>, Vec<f32>, usize)> {
    #[cfg(feature = "gpu")]
    if backend() == Backend::Wgpu {
        return crate::gpu_wgpu::kv_mirror_pull_host(_kv_id, _layer, _from, _to, _nkv, _hd);
    }
    None
}

/// Drop the wgpu token graph's device K/V mirror for a pipeline.
pub fn graph_kv_reset(_kv_id: u64) {
    #[cfg(feature = "gpu")]
    if backend() == Backend::Wgpu {
        crate::gpu_wgpu::kv_mirror_reset(_kv_id);
        crate::gpu_wgpu::embryo_graph_reset(_kv_id);
    }
}

/// Ternary (q1t) BASE matvec on the GPU — fills `out` with the base dot; the
/// caller adds the sparse overlay on the CPU. Metal only for now (wgpu q1t not
/// yet written → CPU fallback).
pub fn q1t_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => {
            if metal_q1t_enabled() {
                crate::gpu_metal::q1t_matvec(model, idx, xs, rows, cols, out)
            } else {
                false
            }
        }
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q1t_matvec(model, idx, xs, rows, cols, out),
        Backend::None => false,
    }
}

/// q4_block matvec on the GPU — wgpu only (Metal drives q4_block through the
/// whole-token graph, not a standalone matvec).
#[allow(unused_variables)]
pub fn q4b_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => false,
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q4b_matvec(model, idx, xs, rows, cols, out),
        Backend::None => false,
    }
}

/// q1t batched GEMM (prefill) — base + overlay on-device (Metal simdgroup or
/// wgpu register-blocked).
pub fn q1t_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        // Batched prefill and single-token decode are both enabled. On the
        // real 14.8B Q1T model prefill PPL was within 0.3% of CPU (7.942 vs
        // 7.966), and the alignment-safe decode kernel reached 3.52e-6 max_rel.
        Backend::Metal => crate::gpu_metal::q1t_matmat(model, idx, xs, b, rows, cols, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q1t_matmat(model, idx, xs, b, rows, cols, out),
        Backend::None => false,
    }
}

/// Native Metal Q1T switch. Enabled by default after the byte-packed Q1T
/// fields were changed to alignment-safe loads; keep an explicit emergency
/// fallback for device/driver diagnostics.
#[cfg(target_os = "macos")]
pub(crate) fn metal_q1t_enabled() -> bool {
    std::env::var("CMF_METAL_Q1T")
        .map(|v| v != "0" && !v.eq_ignore_ascii_case("off"))
        .unwrap_or(true)
}

/// Batched q1 GEMM (prefill). wgpu only — Metal has its own block path.
pub fn q1_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q1_matmat(model, idx, xs, b, rows, cols, out),
        #[allow(unused_variables)]
        _ => false,
    }
}

/// Contention kill for the wide imagegen GEMM/FFN paths: one grossly
/// slow op under a work-proportional budget (fair-device ops are
/// ≤~100 ms even at 1024px) means another process owns the device —
/// verdicts are per-process, so CPU for the rest of this one.
static MM_KILL: AtomicBool = AtomicBool::new(false);
pub(crate) fn mm_killed() -> bool {
    MM_KILL.load(Ordering::Relaxed)
}
pub(crate) fn mm_kill() {
    MM_KILL.store(true, Ordering::Relaxed);
}

/// Consecutive over-budget ops. ONE slow op is not contention: on a
/// 24 GB Mac running the 25.7 GB fl2va file the first ops after the
/// prompt encode page their weights in from the SSD and take seconds —
/// a field report (hololabs, HF discussion #2) had to neuter the kill
/// to keep the denoise on the GPU, and then measured 48 s/step where the
/// CPU fallback took >60. Contention is persistent; a page-in is not.
static MM_STRIKES: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(0);
const MM_STRIKES_TO_KILL: u32 = 3;
/// Whether the kill is armed at all. A one-shot phase whose slowness is
/// expected and not contention — the video prompt encoder streaming
/// 12 GB off the SSD on a 24 GB Mac (HF discussion #4: users had to
/// gut `mm_kill` to keep the denoise loop on the GPU) — disarms it and
/// re-arms it when the phase is over; strikes taken meanwhile are
/// forgotten.
static MM_ARMED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(true);

/// Disarm / re-arm the contention kill around a phase whose GEMMs are
/// slow for reasons that are not another process (see `MM_ARMED`).
pub fn mm_kill_arm(on: bool) {
    MM_ARMED.store(on, Ordering::Relaxed);
    if on {
        MM_STRIKES.store(0, Ordering::Relaxed);
    }
}

/// The contention verdict for one wide op: `el` against its
/// work-proportional `budget`. `exempt` marks ops whose time is not
/// evidence — the cold probe, or a weight that was not resident before
/// the call and rode in with it. Kills after `MM_STRIKES_TO_KILL`
/// consecutive strikes; a within-budget op clears the count.
/// `CMF_MM_KILL=0` disables the kill entirely (the device is trusted).
pub(crate) fn mm_budget_check(
    what: &str,
    el: std::time::Duration,
    budget: std::time::Duration,
    exempt: bool,
) {
    if el <= budget {
        MM_STRIKES.store(0, Ordering::Relaxed);
        return;
    }
    if exempt || !MM_ARMED.load(Ordering::Relaxed) {
        return;
    }
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    let on = *ON.get_or_init(|| std::env::var("CMF_MM_KILL").as_deref() != Ok("0"));
    let n = MM_STRIKES.fetch_add(1, Ordering::Relaxed) + 1;
    if !on {
        tracing::info!(
            "gpu {what} took {el:?} (budget {budget:?}) — over budget, CMF_MM_KILL=0 keeps the device"
        );
        return;
    }
    if n >= MM_STRIKES_TO_KILL {
        tracing::warn!(
            "gpu {what} took {el:?} (budget {budget:?}), {n} in a row — \
             device contended, CPU for the rest of the process (CMF_MM_KILL=0 to override)"
        );
        mm_kill();
    } else {
        tracing::info!(
            "gpu {what} took {el:?} (budget {budget:?}) — strike {n} of {MM_STRIKES_TO_KILL}"
        );
    }
}

/// Fused DiT SwiGLU FFN on the device: g=X·W1ᵀ, u=X·W3ᵀ, silu(g)·u,
/// Causal chunk attention on the device: `b` queries against `s0 + b`
/// cached keys. wgpu only — Metal's chunk graph keeps attention inside
/// the resident block and never calls out.
#[allow(unused_variables, clippy::too_many_arguments)]
pub fn chunk_attend(
    q: &[f32],
    k: &[&[f32]],
    v: &[&[f32]],
    b: usize,
    s0: usize,
    nh: usize,
    nkv: usize,
    hd: usize,
    scale: f32,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::chunk_attend(q, k, v, b, s0, nh, nkv, hd, scale, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// `chunk_attend` under a sliding window (`window` keys, the query's own
/// included). wgpu only; see `chunk_attend_windowed`.
#[allow(unused_variables, clippy::too_many_arguments)]
pub fn chunk_attend_win(
    q: &[f32],
    k: &[&[f32]],
    v: &[&[f32]],
    b: usize,
    s0: usize,
    nh: usize,
    nkv: usize,
    hd: usize,
    scale: f32,
    window: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => {
            crate::gpu_wgpu::chunk_attend_win(q, k, v, b, s0, nh, nkv, hd, scale, window, out)
        }
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Several projections of one host activation, read back in one trip:
/// `out` = [P0 | P1 | …], each `idxs` entry (tensor, rows). wgpu only.
#[allow(unused_variables)]
pub fn gemm_many_keep(
    model: &Arc<CmfModel>,
    idxs: &[(usize, usize)],
    xs: &[f32],
    b: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::gemm_many_keep(model, idxs, xs, b, cols, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

thread_local! {
    /// Set by the pipeline around a layer of Spark-X2.5's batched prefill:
    /// its int8 matrix-unit GEMMs may take the `zi_mm` tile, and the
    /// projections that read one activation share its upload (wgpu).
    static PREFILL_FAST_GEMM: Cell<bool> = const { Cell::new(false) };
}

/// Restores the previous flag on drop.
pub struct PrefillFastGemmGuard(bool);

impl Drop for PrefillFastGemmGuard {
    fn drop(&mut self) {
        PREFILL_FAST_GEMM.with(|c| c.set(self.0));
    }
}

/// Let this thread's int8 matrix-unit GEMMs take the `zi_mm` tile for the
/// guard's lifetime (bit-identical to the 64x64 kernel; see
/// `gpu_wgpu::zi_gemm_route`).
pub fn enter_prefill_fast_gemm() -> PrefillFastGemmGuard {
    PrefillFastGemmGuard(PREFILL_FAST_GEMM.with(|c| c.replace(true)))
}

/// Whether `enter_prefill_fast_gemm` is in force on this thread.
pub fn prefill_fast_gemm() -> bool {
    PREFILL_FAST_GEMM.with(|c| c.get())
}

/// The largest finite f16. An operand the matrix-unit attention casts to
/// f16 must stay within it, or the cast overflows to inf.
pub const F16_MAX: f32 = 65504.0;

/// max |x| over `xs`, or +inf when an entry is inf or NaN: the value the
/// f16 operand guards compare with `F16_MAX`. Sixteen independent lanes
/// and a flag, so it vectorizes (a single `fold` chain does not).
pub fn abs_max_or_inf(xs: &[f32]) -> f32 {
    let mut acc = [0f32; 16];
    let mut bad = [false; 16];
    let mut it = xs.chunks_exact(16);
    for c in &mut it {
        for ((a, f), &v) in acc.iter_mut().zip(bad.iter_mut()).zip(c) {
            let v = v.abs();
            // NaN compares false both ways: flagged, never taken.
            *f |= !(v < f32::INFINITY);
            *a = if v > *a { v } else { *a };
        }
    }
    let mut m = 0f32;
    for &v in it.remainder() {
        let v = v.abs();
        if !(v < f32::INFINITY) {
            return f32::INFINITY;
        }
        m = m.max(v);
    }
    if bad.iter().any(|&f| f) {
        return f32::INFINITY;
    }
    acc.iter().fold(m, |m, &a| m.max(a))
}

/// The decode graph's device K/V mirror a prefill chunk of one layer also
/// writes to (wgpu): `kv_id` is the pipeline's graph id, `layer` the
/// virtual layer index, `limit` the cache's `max_seq_len` — the same key
/// and capacity ceiling the token graph uses, so the mirror the prefill
/// fills is the one the first decode token reads.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct PrefillMirror {
    pub kv_id: u64,
    pub layer: usize,
    pub limit: usize,
}

thread_local! {
    /// Set by the pipeline around one layer's batched attention (see
    /// `enter_prefill_mirror`); every other caller of the batched
    /// attention (MTP heads, tests) sees None and keeps the host upload.
    static PREFILL_MIRROR: Cell<Option<PrefillMirror>> = const { Cell::new(None) };
}

/// Restores the previous target on drop.
pub struct PrefillMirrorGuard(Option<PrefillMirror>);

impl Drop for PrefillMirrorGuard {
    fn drop(&mut self) {
        PREFILL_MIRROR.with(|c| c.set(self.0));
    }
}

/// Name the device mirror the current layer's prefill chunk appends to
/// and attends against, for the guard's lifetime.
pub fn enter_prefill_mirror(t: PrefillMirror) -> PrefillMirrorGuard {
    PrefillMirrorGuard(PREFILL_MIRROR.with(|c| c.replace(Some(t))))
}

/// The mirror `enter_prefill_mirror` named on this thread, if any.
pub fn prefill_mirror() -> Option<PrefillMirror> {
    PREFILL_MIRROR.with(|c| c.get())
}

/// `chunk_attend_win` against the decode graph's device mirror instead of
/// an upload of the whole prefix: the host rows the mirror lacks (the
/// chunk's own `b`, or more after a host-only chunk) are appended first,
/// then the attention binds the mirror rows in place. `s0` is the chunk's
/// first position and `base` the position of host row 0
/// (`LayerKvCache::base`: past 0 on a trimmed sliding tail). `ring` is the
/// layer's window as the mirror geometry (Some on every sliding layer),
/// `window` the window the softmax masks with (0 while it masks nothing).
/// `operand_max` = max |x| over `q` and every K/V row the attend reads
/// (`abs_max_or_inf`; +inf when unknown): past `F16_MAX` the f32 kernels
/// run instead of the matrix units. wgpu only; false = nothing attended
/// (the caller uploads as before).
#[allow(unused_variables, clippy::too_many_arguments)]
pub fn chunk_attend_mirror(
    t: PrefillMirror,
    cpu_k: &[Vec<f32>],
    cpu_v: &[Vec<f32>],
    base: usize,
    q: &[f32],
    b: usize,
    s0: usize,
    nh: usize,
    nkv: usize,
    hd: usize,
    scale: f32,
    ring: Option<usize>,
    window: usize,
    operand_max: f32,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        // QKᵀ and P·V on the matrix units where the device has them (f16
        // operands, f32 accumulators — not bit-identical to the f32
        // kernels: Spark-X2.5 4B q8_2f wiki ppl 10.677 -> 10.672);
        // `CMF_PREFILL_ATTN_COOP=0` keeps the f32 ones (A/B).
        Backend::Wgpu => crate::gpu_wgpu::chunk_attend_mirror(
            t.kv_id,
            t.layer,
            t.limit,
            cpu_k,
            cpu_v,
            base,
            q,
            b,
            s0,
            nh,
            nkv,
            hd,
            scale,
            ring,
            window,
            std::env::var("CMF_PREFILL_ATTN_COOP").as_deref() != Ok("0"),
            operand_max,
            out,
        ),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// `chunk_attend_mirror`, then the head gate (`gains`, b × nh, computed on
/// the host) and the O projection `wo` on the card: the attention output
/// never comes home. `out` = [b][hidden]. The card scans the O operand
/// itself, which is the host path's product bit for bit only while max
/// |gated output| <= 1000: the caller vouches for that (Spark-X2.5's
/// prefill, max |v| <= 990). wgpu only; false = nothing written to `out`
/// (the caller attends and projects as before).
#[allow(unused_variables, clippy::too_many_arguments)]
pub fn chunk_attend_mirror_wo(
    t: PrefillMirror,
    cpu_k: &[Vec<f32>],
    cpu_v: &[Vec<f32>],
    base: usize,
    q: &[f32],
    b: usize,
    s0: usize,
    nh: usize,
    nkv: usize,
    hd: usize,
    scale: f32,
    ring: Option<usize>,
    window: usize,
    operand_max: f32,
    gains: Option<&[f32]>,
    wo: (&Arc<CmfModel>, usize),
    hidden: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::chunk_attend_mirror_wo(
            t.kv_id,
            t.layer,
            t.limit,
            cpu_k,
            cpu_v,
            base,
            q,
            b,
            s0,
            nh,
            nkv,
            hd,
            scale,
            ring,
            window,
            std::env::var("CMF_PREFILL_ATTN_COOP").as_deref() != Ok("0"),
            operand_max,
            gains,
            wo,
            hidden,
            out,
        ),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Whether the active backend's chunk attend can apply a sliding window.
pub fn chunk_attend_windowed() -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => true,
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Fused QKV projection: one upload of the normed chunk, three GEMMs,
/// one readback of Q|K|V back to back. Metal has no twin yet — its
/// chunk graph keeps the whole layer resident and never surfaces QKV.
#[allow(unused_variables, clippy::too_many_arguments)]
pub fn q4t_qkv(
    model: &Arc<CmfModel>,
    wq: usize,
    wk: usize,
    wv: usize,
    xs: &[f32],
    b: usize,
    cols: usize,
    rq: usize,
    rk: usize,
    rv: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q4t_qkv(model, wq, wk, wv, xs, b, cols, rq, rk, rv, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// y=·W2ᵀ — one command buffer, only X and Y cross the CPU boundary.
#[allow(unused_variables, clippy::too_many_arguments)]
/// SwiGLU FFN with a row-packed [gate|up] fc1 (MiniMax-H3's DiT), run
/// end to end on the device. wgpu only: Metal keeps the host loop until
/// its own packed kernel exists.
#[allow(clippy::too_many_arguments, unused_variables)]
pub fn q4tp_ffn_packed(
    model: &Arc<CmfModel>,
    w1: usize,
    w2: usize,
    xs: &[f32],
    b: usize,
    hidden: usize,
    inter: usize,
    bias: Option<&[f32]>,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => {
            crate::gpu_wgpu::ffn_packed(model, w1, w2, xs, b, hidden, inter, bias, out)
        }
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

pub fn q4tp_ffn(
    model: &Arc<CmfModel>,
    w1: usize,
    w3: usize,
    w2: usize,
    xs: &[f32],
    b: usize,
    hidden: usize,
    inter: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::q4tp_ffn(model, w1, w3, w2, xs, b, hidden, inter, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q4tp_ffn(model, w1, w3, w2, xs, b, hidden, inter, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// `q4tp_ffn` (`q4tp` = true) / `q4t_ffn` with the activation named:
/// SiLU takes the plain entry points (every backend), the exact GELU of
/// Spark-X2.5 has a wgpu arm only — other backends decline it.
#[allow(clippy::too_many_arguments, unused_variables)]
pub fn q4_ffn_act(
    model: &Arc<CmfModel>,
    w1: usize,
    w3: usize,
    w2: usize,
    xs: &[f32],
    b: usize,
    hidden: usize,
    inter: usize,
    q4tp: bool,
    act: GraphAct,
    out: &mut [f32],
) -> bool {
    match act {
        GraphAct::Silu if q4tp => q4tp_ffn(model, w1, w3, w2, xs, b, hidden, inter, out),
        GraphAct::Silu => q4t_ffn(model, w1, w3, w2, xs, b, hidden, inter, out),
        GraphAct::GeluErf => match backend() {
            #[cfg(feature = "gpu")]
            Backend::Wgpu => crate::gpu_wgpu::q4_ffn_act(
                model,
                w1,
                w3,
                w2,
                xs,
                b,
                hidden,
                inter,
                q4tp,
                act.code(),
                out,
            ),
            #[allow(unreachable_patterns)]
            _ => false,
        },
    }
}

/// The LLM prefill's dense FFN for separate gate/up/down weights in any
/// codec with a device GEMM (int8, q4tp, or a mix of them): the panels
/// stay on the card between the projections. wgpu only; other backends
/// decline and the caller keeps its per-GEMM path.
#[allow(clippy::too_many_arguments, unused_variables)]
pub fn ffn_act_keep(
    model: &Arc<CmfModel>,
    w1: usize,
    w3: usize,
    w2: usize,
    xs: &[f32],
    b: usize,
    hidden: usize,
    inter: usize,
    act: GraphAct,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::ffn_act_keep(
            model,
            w1,
            w3,
            w2,
            xs,
            b,
            hidden,
            inter,
            act.code(),
            out,
        ),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Qwen Image's exact two-projection tanh-GELU FFN.  The WGPU arm keeps the
/// intermediate on the device; other backends decline so the caller retains
/// its bounded CPU path.  `bias_in` is applied before GELU and `bias_out`
/// after the second projection, matching the official transformer.
#[allow(clippy::too_many_arguments, unused_variables)]
pub fn q4tp_gelu_ffn(
    model: &Arc<CmfModel>,
    w_in: usize,
    w_out: usize,
    xs: &[f32],
    b: usize,
    hidden: usize,
    inter: usize,
    bias_in: &[f32],
    bias_out: &[f32],
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q4tp_gelu_ffn(
            model, w_in, w_out, xs, b, hidden, inter, bias_in, bias_out, out,
        ),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

pub fn q4t_ffn(
    model: &Arc<CmfModel>,
    w1: usize,
    w3: usize,
    w2: usize,
    xs: &[f32],
    b: usize,
    hidden: usize,
    inter: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::q4t_ffn(model, w1, w3, w2, xs, b, hidden, inter, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q4t_ffn(model, w1, w3, w2, xs, b, hidden, inter, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// One whole modulated DiT block for `dit_block`: geometry, norm
/// weights, AdaLN scale/gate vectors (gates pre-tanh'd), a per-token
/// f32 RoPE cos/sin table, and the directory indices of the seven
/// q4t projections. `x` is in-out `[n, hidden]`.
pub struct DitBlockArgs<'a> {
    pub n: usize,
    pub hidden: usize,
    pub inter: usize,
    pub nh: usize,
    pub nkv: usize,
    pub hd: usize,
    pub eps: f32,
    pub rope_cos: &'a [f32],
    pub rope_sin: &'a [f32],
    pub norm1: &'a [f32],
    pub norm2: &'a [f32],
    pub ffn_norm1: &'a [f32],
    pub ffn_norm2: &'a [f32],
    pub norm_q: &'a [f32],
    pub norm_k: &'a [f32],
    pub s_msa: &'a [f32],
    pub gate_msa: &'a [f32],
    pub s_mlp: &'a [f32],
    pub gate_mlp: &'a [f32],
    pub wq: usize,
    pub wk: usize,
    pub wv: usize,
    pub wo: usize,
    pub w1: usize,
    pub w3: usize,
    pub w2: usize,
    /// The projections' layout: q4tp (ladder scales) vs plain q4_tiled.
    /// The recommended Lumina file is q4tp, and a backend that only
    /// knows q4t must decline rather than decode with the wrong reader.
    pub q4tp: bool,
    /// The hidden state is already on the device from the previous block,
    /// so `x` need not be uploaded.
    pub resident_in: bool,
    /// Leave the result on the device instead of reading it back. The DiT
    /// loop does not touch `x` between blocks, so 27 of every 28 readbacks
    /// were moving 19 MB across PCIe and stalling on it for nothing.
    pub resident_out: bool,
}

/// Can the selected backend keep the DiT's hidden state on the device
/// between blocks? Only the wgpu whole-block path; the Metal entry takes
/// and returns host memory every call.
pub fn dit_chain_supported() -> bool {
    #[cfg(feature = "gpu")]
    {
        return matches!(backend(), Backend::Wgpu) && fused_dit_block_available();
    }
    #[allow(unreachable_code)]
    false
}

/// Pull the resident hidden state back to the host. For the caller that
/// chained blocks and then hit one the device declined.
pub fn dit_state_fetch(_x: &mut [f32]) -> bool {
    #[cfg(feature = "gpu")]
    {
        if matches!(backend(), Backend::Wgpu) {
            return crate::gpu_wgpu::dit_state_fetch(_x);
        }
    }
    false
}

/// One whole modulated DiT block on the device — norms, qkv, RoPE,
/// attention, residuals and the SwiGLU FFN in a single command
/// buffer; only `x` crosses the CPU boundary (in and out).
#[allow(unused_variables)]
/// The DiT's three projections in one submission (wgpu only; the
/// Metal path fuses the whole block instead). False = the caller keeps
/// its three separate calls.
#[allow(unused_variables, clippy::too_many_arguments)]
pub fn dit_qkv(
    model: &Arc<CmfModel>,
    wq: usize,
    wk: usize,
    wv: usize,
    xs: &[f32],
    b: usize,
    hidden: usize,
    qrows: usize,
    kvrows: usize,
    q_out: &mut [f32],
    k_out: &mut [f32],
    v_out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q4tp_qkv(
            model, wq, wk, wv, xs, b, hidden, qrows, kvrows, q_out, k_out, v_out,
        ),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// The Qwen Image double-stream attention half.  The WGPU implementation
/// keeps the six Q/K/V projections, the stream join, qk-norm/RoPE, joint
/// attention, and both output projections on the device; the caller only
/// supplies the two normalized streams and receives the two projected
/// streams.  A backend or codec that cannot satisfy the full contract
/// returns `false` before changing either output, so the native host path
/// remains the portable fallback.
pub struct QwenImageAttentionArgs<'a> {
    pub image: &'a [f32],
    pub text: &'a [f32],
    pub image_tokens: usize,
    pub text_tokens: usize,
    pub heads: usize,
    pub head_dim: usize,
    pub image_q: usize,
    pub image_k: usize,
    pub image_v: usize,
    pub text_q: usize,
    pub text_k: usize,
    pub text_v: usize,
    pub image_out: usize,
    pub text_out: usize,
    pub image_q_norm: &'a [f32],
    pub image_k_norm: &'a [f32],
    pub text_q_norm: &'a [f32],
    pub text_k_norm: &'a [f32],
    pub image_cos: &'a [f32],
    pub image_sin: &'a [f32],
    pub text_cos: &'a [f32],
    pub text_sin: &'a [f32],
    pub image_q_bias: &'a [f32],
    pub image_k_bias: &'a [f32],
    pub image_v_bias: &'a [f32],
    pub text_q_bias: &'a [f32],
    pub text_k_bias: &'a [f32],
    pub text_v_bias: &'a [f32],
    pub image_out_bias: &'a [f32],
    pub text_out_bias: &'a [f32],
    pub image_proj: &'a mut [f32],
    pub text_proj: &'a mut [f32],
}

/// The per-layer controls and Q4TP directory indices used by the native
/// Qwen block.  Keeping this descriptor separate from the stream buffers
/// lets a whole transformer forward reuse one explicit device state without
/// a global scratch slot or a hidden context label.
#[allow(clippy::too_many_fields)]
pub struct QwenImageChainBlock<'a> {
    pub image_mod: &'a [f32],
    pub text_mod: &'a [f32],
    pub image_q: usize,
    pub image_k: usize,
    pub image_v: usize,
    pub text_q: usize,
    pub text_k: usize,
    pub text_v: usize,
    pub image_out: usize,
    pub text_out: usize,
    pub image_q_norm: &'a [f32],
    pub image_k_norm: &'a [f32],
    pub text_q_norm: &'a [f32],
    pub text_k_norm: &'a [f32],
    pub image_q_bias: &'a [f32],
    pub image_k_bias: &'a [f32],
    pub image_v_bias: &'a [f32],
    pub text_q_bias: &'a [f32],
    pub text_k_bias: &'a [f32],
    pub text_v_bias: &'a [f32],
    pub image_out_bias: &'a [f32],
    pub text_out_bias: &'a [f32],
    pub image_attn_gate: &'a [f32],
    pub text_attn_gate: &'a [f32],
    pub image_mlp_in: usize,
    pub image_mlp_out: usize,
    pub text_mlp_in: usize,
    pub text_mlp_out: usize,
    pub image_mlp_in_bias: &'a [f32],
    pub image_mlp_out_bias: &'a [f32],
    pub text_mlp_in_bias: &'a [f32],
    pub text_mlp_out_bias: &'a [f32],
}

/// Complete Qwen Image transformer block contract. The first norm/mod
/// panels are supplied by the native caller; the WGPU arm keeps both streams
/// resident through QKV, QK/RoPE, joint attention, output projections, both
/// gated residuals, and the exact tanh-GELU MLPs. A backend that cannot
/// satisfy the whole graph returns `false` without changing either output.
#[allow(clippy::too_many_fields)]
pub struct QwenImageBlockArgs<'a> {
    /// Raw stream state is read for the first gated residual and overwritten
    /// with the block's final state after the one readback.
    pub image: &'a mut [f32],
    pub text: &'a mut [f32],
    pub image_norm: &'a [f32],
    pub text_norm: &'a [f32],
    pub image_tokens: usize,
    pub text_tokens: usize,
    pub heads: usize,
    pub head_dim: usize,
    pub image_cos: &'a [f32],
    pub image_sin: &'a [f32],
    pub text_cos: &'a [f32],
    pub text_sin: &'a [f32],
    pub image_q: usize,
    pub image_k: usize,
    pub image_v: usize,
    pub text_q: usize,
    pub text_k: usize,
    pub text_v: usize,
    pub image_out: usize,
    pub text_out: usize,
    pub image_q_norm: &'a [f32],
    pub image_k_norm: &'a [f32],
    pub text_q_norm: &'a [f32],
    pub text_k_norm: &'a [f32],
    pub image_q_bias: &'a [f32],
    pub image_k_bias: &'a [f32],
    pub image_v_bias: &'a [f32],
    pub text_q_bias: &'a [f32],
    pub text_k_bias: &'a [f32],
    pub text_v_bias: &'a [f32],
    pub image_out_bias: &'a [f32],
    pub text_out_bias: &'a [f32],
    pub image_attn_gate: &'a [f32],
    pub text_attn_gate: &'a [f32],
    pub image_mlp_in: usize,
    pub image_mlp_out: usize,
    pub text_mlp_in: usize,
    pub text_mlp_out: usize,
    pub image_mlp_in_bias: &'a [f32],
    pub image_mlp_out_bias: &'a [f32],
    pub text_mlp_in_bias: &'a [f32],
    pub text_mlp_out_bias: &'a [f32],
    pub image_mlp_mod: &'a [f32],
    pub text_mlp_mod: &'a [f32],
    pub image_mlp_gate: &'a [f32],
    pub text_mlp_gate: &'a [f32],
}

/// Explicit whole-forward Qwen state contract.  The WGPU backend uploads the
/// two initial streams once, encodes a bounded number of complete blocks per
/// submission, and reads the final state once.  `blocks` is immutable for the
/// call, while the two stream slices receive only the final readback.
pub struct QwenImageChainArgs<'a> {
    pub image: &'a mut [f32],
    pub text: &'a mut [f32],
    pub image_tokens: usize,
    pub text_tokens: usize,
    pub heads: usize,
    pub head_dim: usize,
    pub image_cos: &'a [f32],
    pub image_sin: &'a [f32],
    pub text_cos: &'a [f32],
    pub text_sin: &'a [f32],
    pub blocks: &'a [QwenImageChainBlock<'a>],
}

#[allow(unused_variables)]
pub fn qwen_image_attention(
    model: &Arc<CmfModel>,
    args: &mut QwenImageAttentionArgs<'_>,
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::qwen_image_attention(model, args),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

#[allow(unused_variables)]
pub fn qwen_image_block(model: &Arc<CmfModel>, args: &mut QwenImageBlockArgs<'_>) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::qwen_image_block(model, args),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Keep all Qwen transformer blocks on the selected WGPU device, with only
/// bounded chunk submissions and one final readback.  Other backends decline
/// so the native caller can use its exact portable block loop.
#[allow(unused_variables)]
pub fn qwen_image_chain(model: &Arc<CmfModel>, args: &mut QwenImageChainArgs<'_>) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::qwen_image_chain(model, args),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// The Qwen Image second sub-block on WGPU: affine-free LayerNorm,
/// shift/scale modulation, Q4TP input projection, exact tanh-GELU, output
/// projection, bias and gated residual.  `data` is updated in place after a
/// single final readback.  Backends/codecs that cannot keep this chain on the
/// device return `false` before changing `data`, leaving the caller's
/// portable per-op path intact.
#[allow(unused_variables, clippy::too_many_arguments)]
pub fn qwen_image_mlp_inplace(
    model: &Arc<CmfModel>,
    w_in: usize,
    w_out: usize,
    data: &mut [f32],
    batch: usize,
    hidden: usize,
    inter: usize,
    bias_in: &[f32],
    bias_out: &[f32],
    modulation: &[f32],
    gate: &[f32],
) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::qwen_image_mlp_inplace(
            model,
            w_in,
            w_out,
            data,
            batch,
            hidden,
            inter,
            bias_in,
            bias_out,
            modulation,
            gate,
        ),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Is a FUSED whole-block device path on offer? The batched-CFG shape
/// (two sequences in one tall batch) and the fused block (one sequence,
/// one command buffer) are alternatives, and the caller picks.
pub fn fused_dit_block_available() -> bool {
    #[cfg(target_os = "macos")]
    {
        matches!(backend(), Backend::Metal) && fused_block_trusted()
    }
    #[cfg(not(target_os = "macos"))]
    {
        false
    }
}

pub fn dit_block(model: &Arc<CmfModel>, a: &DitBlockArgs, x: &mut [f32]) -> bool {
    dit_block_seg(model, a, &[a.n], x)
}

/// The same block over a CONCATENATION of independent sequences:
/// attention per segment, everything position-wise batched. wgpu only —
/// the Metal path takes the single-sequence entry above.
pub fn dit_block_seg(
    model: &Arc<CmfModel>,
    a: &DitBlockArgs,
    segs: &[usize],
    x: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal if segs.len() <= 1 => crate::gpu_metal::dit_block(model, a, x),
        // The wgpu whole-block path. What it buys is host round trips —
        // six a block become one — so it defaults ON where those cost
        // real time (a discrete card across PCIe) and OFF on unified
        // memory, where the per-op path shares the same pages and the
        // fusion measured slightly slower on an M4. `CMF_DIT_FUSED=1`
        // forces it anywhere, `=0` forbids it.
        #[cfg(feature = "gpu")]
        Backend::Wgpu
            if match std::env::var("CMF_DIT_FUSED").ok().as_deref() {
                Some("0") => false,
                Some(_) => true,
                None => crate::gpu_wgpu::discrete_active(),
            } =>
        {
            crate::gpu_wgpu::dit_block_seg(model, a, segs, x)
        }
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// One VAE resnet block for `vae_resnet`: norm/conv weights and the
/// channel/shape geometry. `shortcut` is the 1×1 projection (w, b, k)
/// when in/out channels differ.
pub struct VaeResnetArgs<'a> {
    pub groups: usize,
    pub ic: usize,
    pub oc: usize,
    pub h: usize,
    pub w: usize,
    pub n1w: &'a [f32],
    pub n1b: &'a [f32],
    pub c1w: &'a [f32],
    pub c1b: &'a [f32],
    pub c1k: usize,
    pub n2w: &'a [f32],
    pub n2b: &'a [f32],
    pub c2w: &'a [f32],
    pub c2b: &'a [f32],
    pub c2k: usize,
    pub shortcut: Option<(&'a [f32], &'a [f32], usize)>,
}

/// One whole VAE resnet block on the device (norm+silu → conv ×2 →
/// shortcut → add, one command buffer).
#[allow(unused_variables)]
pub fn vae_resnet(a: &VaeResnetArgs, x: &[f32], out: &mut [f32]) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::vae_resnet(a, x, out),
        _ => false,
    }
}

/// Nearest-2× upsample fused with the following conv — the small
/// pre-upsample image is what crosses the CPU boundary.
#[allow(unused_variables, clippy::too_many_arguments)]
pub fn vae_upsample_conv(
    w: &[f32],
    bias: &[f32],
    x: &[f32],
    ic: usize,
    oc: usize,
    h: usize,
    w_img: usize,
    k: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::vae_upsample_conv(w, bias, x, ic, oc, h, w_img, k, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::vae_upsample_conv(w, bias, x, ic, oc, h, w_img, k, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// VAE conv2d on the device (implicit GEMM — the CPU path pays for a
/// multi-GB im2col matrix at high resolutions).
#[allow(unused_variables, clippy::too_many_arguments)]
pub fn vae_conv2d(
    w: &[f32],
    bias: &[f32],
    x: &[f32],
    ic: usize,
    oc: usize,
    h: usize,
    w_img: usize,
    k: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::vae_conv2d(w, bias, x, ic, oc, h, w_img, k, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::vae_conv2d(w, bias, x, ic, oc, h, w_img, k, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// DiT full bidirectional attention on the device (all heads:
/// scores GEMM → row softmax → P·V → panel unstack, one command
/// buffer). Head-major inputs; out is [n, nh·hd].
#[allow(unused_variables, clippy::too_many_arguments)]
/// Attention from an interleaved qkv panel, splitting into head-major
/// planes ON the device. wgpu only; `false` elsewhere so the caller
/// keeps its host repack.
#[allow(unused_variables)]
#[allow(clippy::too_many_arguments)]
/// qkv projection + attention with the panel never leaving the card.
/// wgpu only; `false` elsewhere and the caller keeps its host chain.
#[allow(clippy::too_many_arguments, unused_variables)]
pub fn dit_qkv_attention(
    model: &Arc<CmfModel>,
    qkv_idx: usize,
    xn: &[f32],
    n: usize,
    hidden: usize,
    nh: usize,
    hd: usize,
    scale: f32,
    nr: (&[f32], &[f32], &[f32], f32),
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(all(feature = "gpu", not(target_os = "macos")))]
        Backend::Wgpu => crate::gpu_wgpu::dit_qkv_attention(
            model, qkv_idx, xn, n, hidden, nh, hd, scale, nr, out,
        ),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// The whole attention half of a DiT block on the card: qkv GEMM,
/// attention, output projection. Only `proj` comes home.
#[allow(clippy::too_many_arguments)]
pub fn dit_qkv_attn_out(
    model: &Arc<CmfModel>,
    qkv_idx: usize,
    out_idx: usize,
    xn: &[f32],
    n: usize,
    hidden: usize,
    nh: usize,
    hd: usize,
    scale: f32,
    nr: (&[f32], &[f32], &[f32], f32),
    proj: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(all(feature = "gpu", not(target_os = "macos")))]
        Backend::Wgpu => crate::gpu_wgpu::dit_qkv_attn_out(
            model, qkv_idx, out_idx, xn, n, hidden, nh, hd, scale, nr, proj,
        ),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// The VAE decoder's attention half on the card. Only `proj` returns.
#[allow(clippy::too_many_arguments)]
pub fn vae_qkv_attn_out(
    model: &Arc<CmfModel>,
    qkv_idx: usize,
    out_idx: usize,
    xn: &[f32],
    n: usize,
    dim: usize,
    nh: usize,
    hd: usize,
    scale: f32,
    angles: &[f32],
    eps: f32,
    qkv_bias: &[f32],
    proj: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(all(feature = "gpu", not(target_os = "macos")))]
        Backend::Wgpu => crate::gpu_wgpu::vae_qkv_attn_out(
            model, qkv_idx, out_idx, xn, n, dim, nh, hd, scale, angles, eps, qkv_bias, proj,
        ),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

#[allow(clippy::too_many_arguments)]
pub fn vae_attention_packed(
    qkv: &[f32],
    nh: usize,
    n: usize,
    hd: usize,
    scale: f32,
    angles: &[f32],
    eps: f32,
    out: &mut [f32],
) -> bool {
    vae_attention_packed_layout(qkv, nh, n, hd, scale, angles, eps, out, 1)
}

#[allow(clippy::too_many_arguments)]
pub fn vae_attention_packed_layout(
    qkv: &[f32],
    nh: usize,
    n: usize,
    hd: usize,
    scale: f32,
    angles: &[f32],
    eps: f32,
    out: &mut [f32],
    layout: u32,
) -> bool {
    match backend() {
        #[cfg(all(feature = "gpu", not(target_os = "macos")))]
        Backend::Wgpu => crate::gpu_wgpu::vae_attention_packed_layout(
            qkv, nh, n, hd, scale, angles, eps, out, layout,
        ),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

#[allow(clippy::too_many_arguments)]
pub fn dit_split_only(
    qkv: &[f32],
    nh: usize,
    n: usize,
    hd: usize,
    layout: u32,
    norm: Option<(&[f32], f32)>,
    out_q: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(all(feature = "gpu", not(target_os = "macos")))]
        Backend::Wgpu => crate::gpu_wgpu::dit_split_only(qkv, nh, n, hd, layout, norm, out_q),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// The backend's f32 NT GEMM: `y[n×m] = x[n×k] · wᵀ[m×k]`. Tensor
/// cores where the card has them. Refuses under `CMF_BAKE_GPU=0` or
/// strict f32, and for jobs below n·k·m = 4M, where the round trip
/// costs more than the arithmetic saves.
/// `gemm_nt_f32` whose `w` is known to change every call (an
/// accumulation over fresh activations, not a weight): it skips the
/// resident ledger and its per-call fingerprint of the whole operand.
pub fn gemm_nt_f32_transient(
    x: &[f32],
    w: &[f32],
    y: &mut [f32],
    n: usize,
    k: usize,
    m: usize,
) -> bool {
    match backend() {
        #[cfg(all(feature = "gpu", not(target_os = "macos")))]
        Backend::Wgpu => crate::gpu_wgpu::gemm_nt_f32_transient(x, w, y, n, k, m),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

pub fn gemm_nt_f32(x: &[f32], w: &[f32], y: &mut [f32], n: usize, k: usize, m: usize) -> bool {
    match backend() {
        #[cfg(all(feature = "gpu", not(target_os = "macos")))]
        Backend::Wgpu => crate::gpu_wgpu::gemm_nt_f32(x, w, y, n, k, m),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Music-3's FFN chain resident on the device — two GEMMs and the GLU
/// between them with no host round trip. `false` = refused, host runs.
#[allow(clippy::too_many_arguments)]
pub fn music3_ffn(
    model: &std::sync::Arc<CmfModel>,
    idx_in: usize,
    idx_out: usize,
    h: &[f32],
    bias_in: &[f32],
    n: usize,
    hs: usize,
    inter: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(all(feature = "gpu", not(target_os = "macos")))]
        Backend::Wgpu => {
            crate::gpu_wgpu::music3_ffn(model, idx_in, idx_out, h, bias_in, n, hs, inter, out)
        }
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// A 1D convolution as a GEMM whose column matrix is expanded on the
/// device instead of being built, transposed and uploaded by the host.
/// `yt` comes back `[out_n x oc]`. `false` = refused, caller runs host.
#[allow(clippy::too_many_arguments)]
pub fn conv1d_gemm(
    x: &[f32],
    w: &[f32],
    ic: usize,
    oc: usize,
    n: usize,
    k: usize,
    pad: usize,
    dil: usize,
    out_n: usize,
    yt: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::conv1d_gemm(x, w, ic, oc, n, k, pad, dil, out_n, yt),
        #[cfg(all(feature = "gpu", not(target_os = "macos")))]
        Backend::Wgpu => crate::gpu_wgpu::conv1d_gemm(x, w, ic, oc, n, k, pad, dil, out_n, yt),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// The convolution as a GEMM on the matrix units. `false` = refused.
#[allow(clippy::too_many_arguments)]
pub fn vae_conv2d_coop(
    w: &[f32],
    bias: Option<&[f32]>,
    x: &[f32],
    ic: usize,
    oc: usize,
    h: usize,
    wi: usize,
    k: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(all(feature = "gpu", not(target_os = "macos")))]
        Backend::Wgpu => crate::gpu_wgpu::vae_conv2d_coop(w, bias, x, ic, oc, h, wi, k, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

pub fn dit_attention_packed(
    qkv: &[f32],
    nh: usize,
    n: usize,
    hd: usize,
    scale: f32,
    // (rope angles, q norm weights, k norm weights, eps) when the device
    // should apply qk-norm and RoPE itself; None when the host already did.
    nr: Option<(&[f32], &[f32], &[f32], f32)>,
    out: &mut [f32],
) -> bool {
    match backend() {
        // wgpu carries the only implementation, and it is not
        // platform-specific: `CMF_GPU=wgpu` on macOS runs it over Metal
        // like anywhere else. It used to be compiled out here on macOS,
        // which made the call a silent `false` — and the caller's
        // `assert!` turned that refusal into a panic on every
        // `cortiq animate` this platform ever ran.
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::dit_attention_packed(qkv, nh, n, hd, scale, nr, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Whether `dit_attention_packed` has an implementation on the backend
/// that is actually selected.
///
/// The caller has to know BEFORE it skips the host qk-norm: deferring
/// the norm to a device that then refuses leaves q/k unnormalized with
/// no way back. Native Metal has no packed kernel, so on macOS this is
/// false unless `CMF_GPU=wgpu` picked the other backend.
pub fn dit_attention_packed_available() -> bool {
    #[allow(unreachable_patterns)]
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::dit_attention_packed_ready(),
        _ => false,
    }
}

pub fn dit_attention(
    qh: &[f32],
    kh: &[f32],
    vh: &[f32],
    nh: usize,
    nkv: usize,
    n: usize,
    hd: usize,
    scale: f32,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::dit_attention(qh, kh, vh, nh, nkv, n, hd, scale, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::dit_attention(qh, kh, vh, nh, nkv, n, hd, scale, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Batched q4t GEMM on the device (imagegen DiT prefill shapes).
/// Metal: q4t_mul_mm decodes the mmap-resident tiles inside the
/// GEMM's K loop. wgpu (Vulkan/DX12 → NVIDIA/AMD/Intel/Adreno/Mali):
/// the register-blocked WGSL twin, weights cached in VRAM.
#[allow(unused_variables)]
pub fn q4tp_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::q4tp_matmat(model, idx, xs, b, rows, cols, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q4tp_matmat(model, idx, xs, b, rows, cols, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// The same over a two-bit weight plane. Native Metal uses the dedicated
/// q2tp tile; unsupported shapes return false and preserve the host fallback.
pub fn q2tp_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::q2tp_matmat(model, idx, xs, b, rows, cols, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q2tp_matmat(model, idx, xs, b, rows, cols, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Descriptor-aware q2tp GEMM. The affine center is selected only for a
/// validated q2tp_affine target; the raw dtype16 payload remains unchanged.
pub fn q2tp_affine_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::q2tp_affine_matmat(model, idx, xs, b, rows, cols, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q2tp_affine_matmat(model, idx, xs, b, rows, cols, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Single-token q2tp matvec through the ordinary (center=1.5) WGSL kernel.
pub fn q2tp_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::q2tp_matvec(model, idx, xs, rows, cols, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q2tp_matvec(model, idx, xs, rows, cols, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Single-token q2tp matvec with the explicit affine center=1 descriptor
/// operator. This is kept separate from ordinary q2tp to make accidental
/// center changes impossible at a call site.
pub fn q2tp_affine_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::q2tp_affine_matvec(model, idx, xs, rows, cols, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q2tp_affine_matvec(model, idx, xs, rows, cols, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Single-token q4tp matvec on the device — the lm_head class. Through the
/// DEDICATED matvec kernel: the batched GEMM at b=1 measured 11.73 ms
/// against the host's 9.51 on the release head, so the route that was
/// supposed to save eleven milliseconds a token lost its own probe instead.
pub fn q4tp_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::q4tp_matvec_for_test(model, idx, xs, rows, cols, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q4tp_matvec(model, idx, xs, rows, cols, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Single-token q4_tiled matvec on the device — the lm_head class (a
/// q4t checkpoint's head is its biggest host matvec, exactly like the
/// q4tp twin above). wgpu holds q4t_mv pipelines only inside the graph
/// encoder — the standalone arm stays an honest refusal until a
/// discrete-GPU q4t model reaches the bench.
pub fn q4t_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::q4t_matvec_for_test(model, idx, xs, rows, cols, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

pub fn q4t_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::q4t_matmat(model, idx, xs, b, rows, cols, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::q4t_matmat(model, idx, xs, b, rows, cols, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Whole-block token-graph types re-exported from the Metal backend.
#[cfg(target_os = "macos")]
pub use crate::gpu_metal::{
    AttnDeviceParams, AttnGpuLayer, GdnGpuCfg, GdnGpuLayer, GpuMoe, GraphDims, MetalFfn,
    O1AttnParams, TokenGraph, kv_mirror_drop, kv_mirror_read_last, kv_mirror_take_imp,
};

/// A BLOCK of consecutive q1 GDN layers in one submission (Metal only).
#[cfg(target_os = "macos")]
pub fn gdn_block(
    model: &Arc<CmfModel>,
    layers: &[GdnGpuLayer],
    states: &mut [&mut [f32]],
    cfg: &GdnGpuCfg,
    h: &mut [f32],
) -> bool {
    match backend() {
        Backend::Metal => crate::gpu_metal::gdn_block(model, layers, states, cfg, h),
        _ => false,
    }
}

/// A layer's MoE-FFN in one submission (amortizing the dispatch cost).
#[allow(unused_variables)]
pub fn moe_block(model: &Arc<CmfModel>, jobs: &[MoeJob], out: &mut [f32]) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::moe_block(model, jobs, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::moe_block(model, jobs, out),
        Backend::None => false,
    }
}

/// Independent matvecs of one input in a single submission (GDN projections).
#[allow(unused_variables)]
pub fn matvec_batch(model: &Arc<CmfModel>, jobs: &[BatchJob], out: &mut [&mut [f32]]) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::matvec_batch(model, jobs, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::matvec_batch(model, jobs, out),
        Backend::None => false,
    }
}

// ── Whole-token wgpu graph race (generation granularity) ─────────────
// On integrated/mobile adapters the graph is neither trusted nor banned
// a priori — it RACES the normal path: generations alternate arms (the
// normal path first — known-good UX — then the graph), per-token wall
// times accumulate per arm, and once both arms have enough steady
// samples the faster one wins for the process. Arm switches happen ONLY
// at generation boundaries (`kv_cache.clear()` resets state), so the
// device KV mirror and the CPU cache never diverge mid-sequence. The
// single exception is the first-token bail: the very first decode token
// of a graph generation may be discarded and recomputed on the CPU
// path (the prompt KV is CPU-owned at that point, so this is safe) —
// a tiled mobile GPU that drains its pipeline at every barrier turns
// the ~300-dispatch graph into seconds per token (field report: 0.2
// tok/s vs 15 on the CPU), and one token is all it takes to see that.
static GRAPH_RACE_STATE: AtomicU8 = AtomicU8::new(0); // 0 racing, 1 graph won, 2 normal won
static GRAPH_RACE_FLIP: AtomicU32 = AtomicU32::new(0);
static GRAPH_RACE_ARM_GRAPH: AtomicU8 = AtomicU8::new(0); // this generation's arm
static GRAPH_RACE_TOK: AtomicU32 = AtomicU32::new(0); // token index within the generation
static GRAPH_NS: [AtomicU64; 2] = [AtomicU64::new(0), AtomicU64::new(0)]; // [normal, graph]
static GRAPH_N: [AtomicU32; 2] = [AtomicU32::new(0), AtomicU32::new(0)];

/// Steady per-token samples per arm before the race decides.
const GRAPH_RACE_SAMPLES: u32 = 4;

// A graph that cannot be built for THIS model will never build: the
// refusal is a property of the weights, not of the moment. Retrying it
// per token is not free — the builder walks every layer and asks each
// tensor for a graph view before giving up at layer 0 — and on an
// Adreno 642L that retry cost 3x: forcing the graph on a model it
// refuses measured 0.3 tok/s against 0.905 for the per-op path it falls
// back to. Remembered once, the fallback runs at its own speed.
//
// The verdict is kept PER PIPELINE (`Pipeline::graph_refused` /
// `mark_graph_refused`), not in a process-wide flag: several pipelines
// of one file share a process in `serve` (backbone slots + skill lanes
// loaded mid-traffic), and a refusal in one lane — or the reset a new
// lane used to issue — must not move another lane's running sequence
// between the device and the host (R4/NF-2). Callers must NOT report
// transient refusals (an unsealed o1 state during prefill, a softcap).

/// Called at every generation start (fresh KV). Applies a pending
/// verdict and picks this generation's arm while racing.
pub fn graph_race_begin_generation() {
    // One generation has now compiled whatever this model needs; keep it
    // for the next process. Once per run: the blob does not grow after
    // the pipelines exist, and the write is megabytes against the ~200 s
    // of compiling it saves on the device that needed this.
    #[cfg(feature = "gpu")]
    {
        // Save once, at the start of the SECOND generation: the first
        // has dispatched, so there is something to keep, and nothing is
        // saved before any work (the driver compiles at first use, not
        // at pipeline creation — the context comes up in 1.5 s while the
        // compiling costs minutes).
        //
        // Flushing again on 4, 8, 16 … was tried on the theory that a
        // chat turn compiles shapes the first one did not. It buys
        // nothing: a fresh app process still spent 49.0 s, then 58.7,
        // then 61.3 on its first answer with the backoff in place. One
        // flush it is.
        static FLUSHED: std::sync::Once = std::sync::Once::new();
        static FIRST: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(true);
        if FIRST.swap(false, Ordering::Relaxed) {
            // Nothing dispatched yet.
        } else {
            FLUSHED.call_once(crate::gpu_wgpu::pipeline_cache_flush);
        }
    }
    GRAPH_RACE_TOK.store(0, Ordering::Relaxed);
    if GRAPH_RACE_STATE.load(Ordering::Relaxed) != 0 {
        return;
    }
    let (gn, cn) = (
        GRAPH_N[1].load(Ordering::Relaxed),
        GRAPH_N[0].load(Ordering::Relaxed),
    );
    if gn >= GRAPH_RACE_SAMPLES && cn >= GRAPH_RACE_SAMPLES {
        let g_avg = GRAPH_NS[1].load(Ordering::Relaxed) / gn as u64;
        let c_avg = GRAPH_NS[0].load(Ordering::Relaxed) / cn as u64;
        let verdict = if g_avg < c_avg { 1 } else { 2 };
        GRAPH_RACE_STATE.store(verdict, Ordering::Relaxed);
        tracing::info!(
            "wgpu graph race: graph {:.2} ms/tok vs normal {:.2} ms/tok -> {}",
            g_avg as f64 / 1e6,
            c_avg as f64 / 1e6,
            if verdict == 1 { "graph" } else { "normal path" }
        );
        return;
    }
    let flip = GRAPH_RACE_FLIP.fetch_add(1, Ordering::Relaxed);
    GRAPH_RACE_ARM_GRAPH.store((flip % 2 == 1) as u8, Ordering::Relaxed);
}

/// Should this decode token try the graph? `trusted` (discrete adapter,
/// explicit env, or a GDN hybrid whose state lives on the device) skips
/// the race entirely.
pub fn graph_race_use_graph(trusted: bool) -> bool {
    if trusted {
        return true;
    }
    match GRAPH_RACE_STATE.load(Ordering::Relaxed) {
        1 => true,
        2 => false,
        _ => GRAPH_RACE_ARM_GRAPH.load(Ordering::Relaxed) == 1,
    }
}

/// First decode token of a racing graph generation: hopeless already?
/// (>4x the normal path's per-token average AND over a second.) Settles
/// the race immediately; the caller discards the graph result and
/// recomputes this token on the normal path.
pub fn graph_race_first_token_hopeless(dur: std::time::Duration) -> bool {
    if GRAPH_RACE_STATE.load(Ordering::Relaxed) != 0 {
        return false;
    }
    let first = GRAPH_RACE_TOK.load(Ordering::Relaxed) == 0;
    let cn = GRAPH_N[0].load(Ordering::Relaxed);
    if !first || cn == 0 {
        return false;
    }
    let c_avg = GRAPH_NS[0].load(Ordering::Relaxed) / cn as u64;
    let ns = dur.as_nanos() as u64;
    if ns > 1_000_000_000 && ns > 4 * c_avg {
        GRAPH_RACE_STATE.store(2, Ordering::Relaxed);
        tracing::info!(
            "wgpu graph race: first graph token {:.0} ms vs normal {:.2} ms/tok — hopeless, normal path wins",
            ns as f64 / 1e6,
            c_avg as f64 / 1e6
        );
        return true;
    }
    false
}

/// Record one decode-token wall time for the racing arm. The first
/// token of each generation is discarded (KV-mirror upload / cold
/// caches on the graph arm; cold mmap on the normal arm).
pub fn graph_race_record(used_graph: bool, dur: std::time::Duration) {
    if GRAPH_RACE_STATE.load(Ordering::Relaxed) != 0 {
        return;
    }
    let tok = GRAPH_RACE_TOK.fetch_add(1, Ordering::Relaxed);
    if tok == 0 {
        return;
    }
    let i = used_graph as usize;
    GRAPH_NS[i].fetch_add(dur.as_nanos() as u64, Ordering::Relaxed);
    GRAPH_N[i].fetch_add(1, Ordering::Relaxed);
}

/// Bounded-cost content fingerprint for the backends' pointer-keyed device
/// caches: FNV over the whole slice up to 4 KiB, over 64 spread 64-byte
/// windows (plus the length) above. An address-keyed hit must also prove
/// the bytes are still the ones it uploaded — the allocator reuses heap
/// and mmap addresses freely, so a reloaded model or a re-dequantized
/// layer lands where the old bytes were — and sampling keeps that proof at
/// ~a microsecond even for a 126 MB matrix. Real replacements (another
/// model's tensor, an Adam-updated master) differ densely, so a 4 KiB
/// spread cannot miss them.
pub(crate) fn fp_bytes(data: &[u8]) -> u64 {
    #[inline]
    fn fnv(mut h: u64, bytes: &[u8]) -> u64 {
        let (chunks, tail) = bytes.split_at(bytes.len() & !7);
        for c in chunks.chunks_exact(8) {
            h ^= u64::from_le_bytes(c.try_into().unwrap());
            h = h.wrapping_mul(0x100_0000_01b3);
        }
        for &b in tail {
            h ^= b as u64;
            h = h.wrapping_mul(0x100_0000_01b3);
        }
        h
    }
    let mut h = 0xcbf2_9ce4_8422_2325u64 ^ (data.len() as u64);
    if data.len() <= 4096 {
        return fnv(h, data);
    }
    let step = (data.len() - 64) / 63;
    for i in 0..64 {
        h = fnv(h, &data[i * step..i * step + 64]);
    }
    h
}

/// `fp_bytes` over an f32 slice without a bytemuck dependency (the Metal
/// backend builds with no GPU feature flags).
pub(crate) fn fp_f32(data: &[f32]) -> u64 {
    let bytes = unsafe { std::slice::from_raw_parts(data.as_ptr() as *const u8, data.len() * 4) };
    fp_bytes(bytes)
}

#[cfg(test)]
mod fp_tests {
    use super::fp_bytes;

    /// The pointer-keyed caches survive on `fp_bytes` telling two different
    /// tensors apart at a reused address. Its sampling must therefore see a
    /// change ANYWHERE — head, tail, and the stretches between windows are
    /// the places a cheaper hash would go blind.
    #[test]
    fn fp_bytes_sees_a_change_anywhere_in_a_sampled_slice() {
        let n = 1 << 20; // 1 MiB — far above the 4 KiB full-hash threshold
        let base: Vec<u8> = (0..n).map(|i| (i * 31 + 7) as u8).collect();
        let h0 = fp_bytes(&base);
        assert_eq!(h0, fp_bytes(&base), "fingerprint must be deterministic");
        // A DENSE change (every requantized/redequantized tensor is one)
        // must flip the fingerprint no matter how the windows fall.
        let mut dense = base.clone();
        for b in dense.iter_mut() {
            *b = b.wrapping_add(1);
        }
        assert_ne!(
            h0,
            fp_bytes(&dense),
            "a fully different tensor slipped through"
        );
        // Length participates: the same prefix at a shorter length is a
        // different key AND a different fingerprint.
        assert_ne!(h0, fp_bytes(&base[..n - 64]));
        // Below the threshold the hash is exact: a single flipped byte in
        // a norm-sized vector must be seen.
        let mut small = vec![3u8; 4096];
        let hs = fp_bytes(&small);
        small[2048] ^= 1;
        assert_ne!(hs, fp_bytes(&small), "full hash missed a one-byte change");
        // And the sampled windows land within bounds on awkward sizes.
        for n in [4097usize, 5000, 64 * 64, 1 << 16] {
            let v = vec![9u8; n];
            let _ = fp_bytes(&v); // must not panic on window math
        }
    }
}

/// Hand the card back after a bake: drop its resident weights, planes and
/// pools so the ordinary engine (the runtime gate, a serve that follows)
/// starts from a clean budget. No-op off the wgpu backend.
pub fn bake_release() {
    #[cfg(feature = "gpu")]
    crate::gpu_wgpu::bake_release();
}

/// Strict-f32 for the bake's GEMMs (phase A mask training): the mask
/// selects neurons by a gradient signal, and f16 operand rounding on
/// that signal closes the wrong ones. No-op off the wgpu backend.
pub fn bake_precision_strict(on: bool) {
    #[cfg(feature = "gpu")]
    crate::gpu_wgpu::bake_precision_strict(on);
    #[cfg(not(feature = "gpu"))]
    let _ = on;
}

/// CMF_GRAPH_HOSTPROF=1: how a graph token's wall splits between the
/// host encoding the command stream and the tail the GPU still owes
/// after encode. Fifteen GPU-side suspects measured null while the
/// bench counted 17.7k allocations a token — this is the instrument
/// that says whether the thief was on the host all along.
pub fn hostprof_encode_done(t0: std::time::Instant) {
    use std::sync::atomic::{AtomicU64, Ordering};
    static ENC: AtomicU64 = AtomicU64::new(0);
    static N: AtomicU64 = AtomicU64::new(0);
    if std::env::var("CMF_GRAPH_HOSTPROF").as_deref() != Ok("1") {
        return;
    }
    ENC.fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
    let n = N.fetch_add(1, Ordering::Relaxed) + 1;
    if n % 100 == 0 {
        eprintln!(
            "hostprof: encode {:.2} ms/token over {n} tokens",
            ENC.load(Ordering::Relaxed) as f64 / n as f64 / 1e6
        );
    }
}

pub fn hostprof_total(t0: std::time::Instant) {
    use std::sync::atomic::{AtomicU64, Ordering};
    static TOT: AtomicU64 = AtomicU64::new(0);
    static N: AtomicU64 = AtomicU64::new(0);
    if std::env::var("CMF_GRAPH_HOSTPROF").as_deref() != Ok("1") {
        return;
    }
    TOT.fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
    let n = N.fetch_add(1, Ordering::Relaxed) + 1;
    if n % 100 == 0 {
        eprintln!(
            "hostprof: total {:.2} ms/token over {n} tokens",
            TOT.load(Ordering::Relaxed) as f64 / n as f64 / 1e6
        );
    }
}

/// Per-stage host-encode accumulator for the Metal token loop
/// (CMF_GRAPH_HOSTPROF=1). Stage 0 = GDN-run encode; everything else
/// falls out by subtraction from hostprof's encode total.
pub fn stageprof(stage: u32, dt: std::time::Duration) {
    use std::sync::atomic::{AtomicU64, Ordering};
    static NS: [AtomicU64; 4] = [
        AtomicU64::new(0),
        AtomicU64::new(0),
        AtomicU64::new(0),
        AtomicU64::new(0),
    ];
    static N: AtomicU64 = AtomicU64::new(0);
    if std::env::var("CMF_GRAPH_HOSTPROF").as_deref() != Ok("1") {
        return;
    }
    NS[stage as usize % 4].fetch_add(dt.as_nanos() as u64, Ordering::Relaxed);
    if stage == 1 {
        let n = N.fetch_add(1, Ordering::Relaxed) + 1;
        if n % 200 == 0 {
            eprintln!(
                "stageprof: planning {:.2} ms/tok | gdn-item {:.2} ms/tok | attn-item {:.2} ms/tok ({n} tok)",
                NS[1].load(Ordering::Relaxed) as f64 / n as f64 / 1e6,
                NS[2].load(Ordering::Relaxed) as f64 / n as f64 / 1e6,
                NS[3].load(Ordering::Relaxed) as f64 / n as f64 / 1e6
            );
        }
    }
}

/// Active weight bytes dispatched so far (Metal decode path); 0 where
/// the backend does not count. The honest floor's numerator.
pub fn weight_bytes_dispatched() -> u64 {
    let mut total = 0u64;
    #[cfg(target_os = "macos")]
    {
        total += crate::gpu_metal::WEIGHT_BYTES.load(std::sync::atomic::Ordering::Relaxed);
    }
    #[cfg(feature = "gpu")]
    {
        total += crate::gpu_wgpu::WEIGHT_BYTES.load(std::sync::atomic::Ordering::Relaxed);
    }
    total
}

/// The per-stage split of `weight_bytes_dispatched`:
/// [misc, dense-ffn, moe, attn, gdn, head].
pub fn weight_bytes_by() -> [u64; 6] {
    #[cfg(target_os = "macos")]
    {
        let mut o = [0u64; 6];
        for (i, a) in crate::gpu_metal::WEIGHT_BYTES_BY.iter().enumerate() {
            o[i] = a.load(std::sync::atomic::Ordering::Relaxed);
        }
        return o;
    }
    #[allow(unreachable_code)]
    [0; 6]
}

#[cfg(test)]
mod f16_guard_tests {
    use super::*;

    /// The f16 guard's maximum: the largest magnitude through the lanes and
    /// the remainder, +inf for an inf or a NaN anywhere.
    #[test]
    fn abs_max_or_inf_is_the_magnitude_maximum() {
        let mut xs: Vec<f32> = (0..37).map(|i| (i as f32 - 18.0) * 0.5).collect();
        assert_eq!(abs_max_or_inf(&xs), 9.0);
        xs[36] = -70000.0; // the remainder
        assert_eq!(abs_max_or_inf(&xs), 70000.0);
        xs[3] = f32::NAN; // a lane
        assert_eq!(abs_max_or_inf(&xs), f32::INFINITY);
        xs[3] = 0.0;
        xs[35] = f32::NEG_INFINITY;
        assert_eq!(abs_max_or_inf(&xs), f32::INFINITY);
        assert_eq!(abs_max_or_inf(&[]), 0.0);
        assert_eq!(abs_max_or_inf(&[-0.0, -1e-30]), 1e-30);
    }
}

#[cfg(test)]
mod probe_warmup_tests {
    use super::*;
    use std::time::Duration;

    fn ms(v: f64) -> Duration {
        Duration::from_nanos((v * 1e6) as u64)
    }

    /// The bug this pins, measured on an A100: the first device call for
    /// a class compiles its pipeline, was timed at 117.01 ms against the
    /// host's 3.19, and sent `gemm-nt` to the CPU for the whole process —
    /// which ran a 27B bake on 2.6 cores with the card idle.
    #[test]
    fn one_cold_first_sample_does_not_lose_the_class() {
        let p = Probe::new();
        // First device sample is the pipeline build. Then the truth.
        probe_record_into(&p, "gemm-nt", None, true, ms(117.01));
        probe_record_into(&p, "gemm-nt", None, true, ms(1.1));
        probe_record_into(&p, "gemm-nt", None, true, ms(1.0));
        probe_record_into(&p, "gemm-nt", None, false, ms(3.19));
        probe_record_into(&p, "gemm-nt", None, false, ms(3.20));
        assert_eq!(
            p.state.load(Ordering::Relaxed),
            1,
            "the device is 3x faster once warm and must win"
        );
    }

    /// The warm-up must not become a way to never decide, and must not
    /// underflow: a blind decrement at zero wraps a u32 to its maximum
    /// and mutes the arm for the life of the process.
    #[test]
    fn the_warmup_is_spent_once_and_never_underflows() {
        let p = Probe::new();
        for _ in 0..8 {
            probe_record_into(&p, "matmat", None, true, ms(10.0));
        }
        assert_eq!(p.gpu_burn.load(Ordering::Relaxed), 0, "spent, not wrapped");
        assert_eq!(
            p.gpu_n.load(Ordering::Relaxed),
            7,
            "one sample burned, the rest counted"
        );
    }

    /// A device path that always refuses records no timing, so without
    /// counting the refusals the class can never reach a verdict. On an
    /// M4 with LFM2.5-2.6B `ffn` was still undecided after 9000 calls,
    /// alternating arms and paying a failed device attempt on half of
    /// them.
    #[test]
    fn a_class_whose_device_always_declines_settles_on_the_host() {
        let _probe_guard = probe_test_guard();
        // A class no other test in this file touches: `probe_note_decline`
        // works on the process-wide probes by design, and the tests in
        // this binary share them.
        let c = OpClass::MatmatWide;
        let p = &PROBES[c as usize];
        p.state.store(0, Ordering::Relaxed);
        p.declines.store(0, Ordering::Relaxed);
        for _ in 0..(PROBE_DECLINE_LIMIT - 1) {
            probe_note_decline(c);
        }
        assert_eq!(
            p.state.load(Ordering::Relaxed),
            0,
            "one short of the limit is still a question, not an answer"
        );
        probe_note_decline(c);
        assert_eq!(p.state.load(Ordering::Relaxed), 2, "settled on the host");
        assert!(matches!(probe_arm(c), ProbeArm::Cpu));
        p.state.store(0, Ordering::Relaxed);
        p.declines.store(0, Ordering::Relaxed);
    }

    /// A genuinely slower device still loses — the warm-up removes an
    /// artefact, it does not put a thumb on the scale.
    #[test]
    fn a_slow_device_still_loses_after_the_warmup() {
        let p = Probe::new();
        for _ in 0..4 {
            probe_record_into(&p, "matvec", None, true, ms(40.0));
        }
        for _ in 0..4 {
            probe_record_into(&p, "matvec", None, false, ms(2.0));
        }
        assert_eq!(p.state.load(Ordering::Relaxed), 2, "host wins on merit");
    }
}

/// Scratch/weight lifetime for a synchronous image-pipeline stage. Declare
/// this before the stage model so the model drops before cache collection.
pub(crate) struct ImageStageGuard {
    #[cfg(target_os = "macos")]
    metal: Option<crate::gpu_metal::ImageStageGuard>,
    #[cfg(feature = "gpu")]
    wgpu: crate::gpu_wgpu::ImageStageGuard,
}

pub(crate) fn image_stage_scope() -> ImageStageGuard {
    ImageStageGuard {
        #[cfg(target_os = "macos")]
        metal: if matches!(backend(), Backend::Metal) {
            Some(crate::gpu_metal::image_stage_scope())
        } else {
            None
        },
        #[cfg(feature = "gpu")]
        wgpu: crate::gpu_wgpu::image_stage_scope(),
    }
}

impl ImageStageGuard {
    pub(crate) fn track_model(&mut self, uid: u64) {
        #[cfg(target_os = "macos")]
        if let Some(metal) = &mut self.metal {
            metal.track_model(uid);
        }
        #[cfg(feature = "gpu")]
        self.wgpu.track_model(uid);
        #[cfg(not(target_os = "macos"))]
        let _ = uid;
    }
}

// ════════════════════════════════════════════════════════════════════
// Z-Image-Turbo device contract (WP0 scaffold, plan §2.1). APPEND-ONLY.
//
// Owner of the contract: the WP1 lead. The backends implement it in their
// own child modules — `gpu_wgpu/zimage.rs` (WP2) and `gpu_metal/zimage.rs`
// (WP3) — and never edit the parent files. New fields are added only as
// `Option<…>` with agreed semantics; existing fields never change meaning.
//
// Convention (the same as every `gpu::*` entry): `false` = "not handled",
// nothing observable was changed, and the caller runs the CPU path
// (`zimage::ZImageDit::step_cpu` etc.), which is the bit-level reference.
//
// Sequence order everywhere is diffusers' [img rows…, cap rows…], with
// padded lengths n_img_p = ceil32(n_img) and n_cap_p = ceil32(L).
// ════════════════════════════════════════════════════════════════════

// ───────────────────── Qwen-Image-2.1 device contract ─────────────────────

/// Geometry of the Qwen-Image-2.1 denoiser.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Qi21Geom {
    pub hidden: usize,
    pub nh: usize,
    pub hd: usize,
    pub inter: usize,
    pub in_ch: usize,
    pub eps: f32,
}

/// One block's weights: tensor indices of `to_q, to_k, to_v, to_out.0,
/// gate_layer, proj (up), out (down)` and the per-head q/k norm weights.
pub struct Qi21BlockRef<'a> {
    pub w: [usize; 7],
    pub norm_q: &'a [f32],
    pub norm_k: &'a [f32],
}

/// Everything the device needs to build a program and run its prefix.
pub struct Qi21PrefillArgs<'a> {
    pub model: &'a Arc<CmfModel>,
    pub geom: Qi21Geom,
    pub blocks: &'a [Qi21BlockRef<'a>],
    /// `img_in` [hidden, 64] and `proj_out` [64, hidden], f32.
    pub img_in: &'a [f32],
    pub proj_out: &'a [f32],
    pub key: u64,
    /// Prefix rows after `txt_in` / `img_in`, `[lp, hidden]`.
    pub x: &'a [f32],
    pub lp: usize,
    /// RoPE cos/sin of the prefix rows and of the target rows, `[rows, 64]`.
    pub rope_p: (&'a [f32], &'a [f32]),
    pub rope_t: (&'a [f32], &'a [f32]),
    /// Keys each prefix row sees (a leading run; non-decreasing).
    pub vis: &'a [u32],
    /// The t = 0 modulation `[s1|g1|s2|g2]`.
    pub mods0: &'a [f32],
    /// Target rows.
    pub n: usize,
}

/// Build the program `a.key` and run its prefix on the device.
#[allow(unused_variables)]
pub fn qi21_prefill(a: &Qi21PrefillArgs) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::qi21::prefill(a),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::qi21::prefill(a),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// One denoiser call of program `key`: `xtok [n, 64]` → `out [n, 64]`.
#[allow(unused_variables)]
pub fn qi21_step(key: u64, xtok: &[f32], mods: &[f32], fs: &[f32], out: &mut [f32]) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::qi21::step(key, xtok, mods, fs, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::qi21::step(key, xtok, mods, fs, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Drop program `key`. The wgpu module is released directly (module-local
/// state only), so a release never brings a device up.
#[allow(unused_variables)]
pub fn qi21_release_key(key: u64) {
    #[cfg(target_os = "macos")]
    if matches!(backend(), Backend::Metal) {
        crate::gpu_metal::qi21::release_key(key);
    }
    #[cfg(feature = "gpu")]
    crate::gpu_wgpu::qi21::release_key(key);
}

/// Drop the denoiser's device state (see `qi21_release_key`), and the
/// resident VAE decoder's weights.
pub fn qi21_release() {
    #[cfg(target_os = "macos")]
    if matches!(backend(), Backend::Metal) {
        crate::gpu_metal::qi21::release();
    }
    #[cfg(feature = "gpu")]
    {
        crate::gpu_wgpu::qi21::release();
        crate::gpu_wgpu::qi21_vae::release();
    }
}

/// One conv of the Qwen-Image-2.1 VAE decoder: host f32 weights
/// `[co][ci][k][k]` and bias `[co]` (k = 1 or 3, stride 1, same padding).
#[derive(Clone, Copy)]
pub struct Qi21VaeConvRef<'a> {
    pub w: &'a [f32],
    pub b: &'a [f32],
    pub ci: usize,
    pub co: usize,
    pub k: usize,
}

/// A residual block: `conv2(φ(conv1(φ(x)))) + x` (through the 1×1
/// shortcut when ci ≠ co), φ = RMS_norm over channels (x/‖x‖·√C·γ), SiLU.
pub struct Qi21VaeResRef<'a> {
    pub g1: &'a [f32],
    pub c1: Qi21VaeConvRef<'a>,
    pub g2: &'a [f32],
    pub c2: Qi21VaeConvRef<'a>,
    pub shortcut: Option<Qi21VaeConvRef<'a>>,
}

/// An up block: its resnets, then (all but the last block) nearest-2× +
/// the 3×3 conv, plus the DupUp shortcut of the block's input (`in_dim` →
/// `out_dim` channels, temporal factor `ft`).
pub struct Qi21VaeUpRef<'a> {
    pub resnets: Vec<Qi21VaeResRef<'a>>,
    pub up: Option<(Qi21VaeConvRef<'a>, usize)>,
    pub in_dim: usize,
    pub out_dim: usize,
}

/// The Qwen-Image-2.1 VAE decoder (`qwen_image21_vae.rs`) for a resident
/// device decode.
pub struct Qi21VaeDecodeArgs<'a> {
    /// Identity of the weights: a backend caches its planes by it.
    pub key: u64,
    pub post_quant: Qi21VaeConvRef<'a>,
    pub conv_in: Qi21VaeConvRef<'a>,
    pub mid_res: [Qi21VaeResRef<'a>; 2],
    /// The mid attention: RMS_norm γ, `to_qkv` (1×1, 3c outputs: q, k, v),
    /// `proj` (1×1); single head, residual.
    pub attn_gamma: &'a [f32],
    pub attn_qkv: Qi21VaeConvRef<'a>,
    pub attn_proj: Qi21VaeConvRef<'a>,
    pub ups: Vec<Qi21VaeUpRef<'a>>,
    pub norm_out: &'a [f32],
    pub conv_out: Qi21VaeConvRef<'a>,
}

/// Resident VAE decode: `z` = raw latent planes `[z_dim, h·w]` → `out`
/// `[out_channels, 16h·16w]` (before the clamp). `false` = not handled,
/// `out` untouched: the per-conv path runs.
#[allow(unused_variables)]
pub fn qi21_vae_decode(a: &Qi21VaeDecodeArgs, z: &[f32], h: usize, w: usize, out: &mut [f32]) -> bool {
    match backend() {
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::qi21_vae::decode(a, z, h, w, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// One Z-Image transformer block's device inputs (noise refiner, context
/// refiner or main layer — all share this shape). Weights are tensor
/// indices into `model.tensors` (diffusers names under `dit.`); the codec
/// is whatever the container holds (F16/Bf16/Q8Row/Q8_2f/Q4TiledP…), and a
/// backend that cannot expand a codec declines (returns `false`).
/// Norm vectors are f32 host slices that live as long as the caller's
/// `ZImageDit`; a backend may cache them by pointer (they do not change).
#[derive(Clone, Copy)]
pub struct ZBlockRef<'a> {
    /// `attention.to_q/to_k/to_v/to_out.0.weight`, each [hidden, hidden].
    pub wq: usize,
    pub wk: usize,
    pub wv: usize,
    pub wo: usize,
    /// `feed_forward.w1` (gate) / `w3` (up) [inter, hidden], `w2` (down)
    /// [hidden, inter]. FFN = w2(silu(w1·x) ⊙ w3·x).
    pub w1: usize,
    pub w3: usize,
    pub w2: usize,
    /// `attention_norm1` / `attention_norm2`, [hidden] (plain-w RMSNorm).
    pub norm1: &'a [f32],
    pub norm2: &'a [f32],
    /// `ffn_norm1` / `ffn_norm2`, [hidden].
    pub ffn_norm1: &'a [f32],
    pub ffn_norm2: &'a [f32],
    /// `attention.norm_q` / `norm_k`, [hd] (per-head RMSNorm before RoPE).
    pub norm_q: &'a [f32],
    pub norm_k: &'a [f32],
}

/// Z-Image geometry. Turbo: hidden 3840, nh 30 (MHA, no GQA), hd 128,
/// inter 10240, eps 1e-5 (all RMSNorms incl. qk-norm), final_eps 1e-6
/// (the affine-free final LayerNorm), patch_dim 64 (2×2×16).
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct ZGeom {
    pub hidden: usize,
    pub nh: usize,
    pub hd: usize,
    pub inter: usize,
    pub eps: f32,
    pub final_eps: f32,
    pub patch_dim: usize,
}

/// Once per (prompt, resolution). The backend uploads/caches what it needs
/// keyed by `key`; weight planes are keyed by the MODEL (not by `key`) and
/// survive across prompts until `zimage_release`.
pub struct ZPrepareArgs<'a> {
    pub model: &'a Arc<CmfModel>,
    pub geom: ZGeom,
    /// Caller-chosen identity of this (prompt, resolution) state; every
    /// `ZStepArgs` of the same image carries the same key.
    pub key: u64,
    /// Image tokens (H/16 · W/16), padded count ceil32(n_img), caption
    /// padded count ceil32(L). S = n_img_p + n_cap_p.
    pub n_img: usize,
    pub n_img_p: usize,
    pub n_cap_p: usize,
    /// The patch grid (H/16, W/16); n_img = grid.0 · grid.1. Row-major
    /// token order `hp·grid.1 + wp`.
    pub grid: (usize, usize),
    /// [n_cap_p, hidden], ALREADY context-refined (host or device).
    pub cap: &'a [f32],
    /// Noise-refiner RoPE: [n_img_p · hd/2] cos, sin (complex-interleaved
    /// pairs, hd/2 angles per token).
    pub rope_img: (&'a [f32], &'a [f32]),
    /// Main-layer RoPE: [(n_img_p + n_cap_p) · hd/2], rows ordered [img, cap].
    pub rope_joint: (&'a [f32], &'a [f32]),
    /// `all_x_embedder.2-1.weight` [hidden, 64], `.bias` [hidden],
    /// `x_pad_token` [hidden] (replaces rows ≥ n_img after the embed).
    pub x_emb_w: &'a [f32],
    pub x_emb_b: &'a [f32],
    pub x_pad: &'a [f32],
    /// `all_final_layer.2-1.linear.weight` [64, hidden], `.bias` [64].
    pub final_w: &'a [f32],
    pub final_b: &'a [f32],
    /// 2 noise-refiner blocks (image rows only) and 30 main layers.
    pub noise_refiner: &'a [ZBlockRef<'a>],
    pub layers: &'a [ZBlockRef<'a>],
    /// OPTIONAL (backends may ignore): the modulation of EVERY step of this
    /// image, [steps][(2+30)·4·hidden] in the `ZStepArgs::mods` layout, and
    /// [steps][hidden] final scales, so a backend can upload them once per
    /// image and index them by `ZStepArgs::step`. `ZStepArgs::mods` is still
    /// always supplied and is authoritative.
    pub mods_all: Option<&'a [f32]>,
    pub final_scale_all: Option<&'a [f32]>,
    /// OPTIONAL (B2): the CFG negative item. When `Some`, the backend
    /// prepares ONE batch-2 program under `key` — item 0 is this prompt,
    /// item 1 the negative — and every `ZStepArgs` of that key must carry
    /// `out_neg`. A backend without batch 2 returns `false` (the caller
    /// then prepares the two items separately or runs the CPU path).
    pub neg: Option<ZNegArgs<'a>>,
}

/// The negative (unconditional) item of a CFG pair: its own refined
/// caption, padded caption length and joint RoPE table (the image ids sit
/// at axis-0 position L_p+1, so both tables depend on the item's L_p).
pub struct ZNegArgs<'a> {
    /// [n_cap_p, hidden], context-refined.
    pub cap: &'a [f32],
    pub n_cap_p: usize,
    /// [n_img_p · hd/2] cos, sin (noise refiner) of this item.
    pub rope_img: (&'a [f32], &'a [f32]),
    /// [(n_img_p + n_cap_p) · hd/2] cos, sin, rows [img, cap].
    pub rope_joint: (&'a [f32], &'a [f32]),
}

/// Once per denoising step.
pub struct ZStepArgs<'a> {
    /// The `ZPrepareArgs::key` this step belongs to. A key the backend has
    /// not prepared → `false`.
    pub key: u64,
    /// Step index into the schedule (0..steps); selects the row of
    /// `ZPrepareArgs::mods_all` when a backend uses it.
    pub step: usize,
    /// [n_img_p, 64] patchified latent, inner order (dy·2+dx)·16+c. Rows
    /// ≥ n_img are copies of the last row; the backend replaces them with
    /// `x_pad` after the embed.
    pub x_tok: &'a [f32],
    /// Per block (noise_refiner then layers) the RAW chunks
    /// [scale_msa, gate_msa, scale_mlp, gate_mlp] of Linear(temb) (no SiLU
    /// before it), [(2+30)·4·hidden]. The backend applies (1+s) and tanh(g).
    pub mods: &'a [f32],
    /// [hidden] = 1 + Linear(SiLU(temb)) — already includes the +1.
    pub final_scale: &'a [f32],
    /// [n_img, 64]: the model output v (before the pipeline's negation),
    /// image rows only, patchified order.
    pub out: &'a mut [f32],
    /// [n_img, 64]: the negative item's v — required (and only valid) for
    /// a key prepared with `ZPrepareArgs::neg`. Both items see `x_tok`.
    pub out_neg: Option<&'a mut [f32]>,
}

/// Prepare the per-(prompt, resolution) device state. Backends: wgpu →
/// `gpu_wgpu::zimage::prepare` (WP2), Metal → `gpu_metal::zimage::prepare`
/// (WP3).
#[allow(unused_variables)]
pub fn zimage_prepare(a: &ZPrepareArgs) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::zimage::prepare(a),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::zimage::prepare(a),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// One full DiT forward on the device: x_embed → pad rows → noise refiner
/// ×2 → concat [img, cap] → 30 layers → final LayerNorm·scale → Linear →
/// image rows into `a.out`.
#[allow(unused_variables)]
pub fn zimage_step(a: &mut ZStepArgs) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::zimage::step(a),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::zimage::step(a),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// OPTIONAL (B2): build the backend's weight planes for the per-step blocks
/// and the context refiner ahead of `zimage_prepare`, so the caller can
/// overlap the upload with the (CPU) text encoder. `false` = not done;
/// `zimage_prepare` builds whatever is missing either way.
#[allow(unused_variables)]
pub fn zimage_preload(
    model: &Arc<CmfModel>,
    geom: &ZGeom,
    noise_refiner: &[ZBlockRef],
    layers: &[ZBlockRef],
    context_refiner: &[ZBlockRef],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => {
            crate::gpu_metal::zimage::preload(model, geom, noise_refiner, layers, context_refiner)
        }
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::zimage::preload(model, geom, noise_refiner, layers, context_refiner),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Persist the driver's compiled pipelines after a Z-Image generation (the
/// chain's kernels are built at first use, after the context came up), so
/// the next process skips the compile. Best-effort, no-op off wgpu.
pub fn zimage_flush_pipelines() {
    #[cfg(feature = "gpu")]
    if matches!(backend(), Backend::Wgpu) {
        crate::gpu_wgpu::pipeline_cache_flush();
    }
}

/// OPTIONAL (B2): bring the device up and compile the Z-Image kernels, on
/// a helper thread at the start of a generation (the context and the
/// compiles cost ~1 s cold, beside the host-side loading). `false` = no
/// device path here.
pub fn zimage_warmup() -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::zimage::warmup(),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::zimage::warmup(),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// OPTIONAL (B2): upload the resident VAE's weights and compile its
/// kernels ahead of `vae_decode_chain` (the caller runs it on a helper
/// thread while the DiT steps keep the device busy). `false` = not done.
#[allow(unused_variables)]
pub fn vae_prewarm(a: &crate::vae::VaeChainArgs) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::zimage::vae_prewarm(a),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::zimage::vae_prewarm(a),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Drop the Z-Image DiT device state (planes, prepared programs) but keep
/// the VAE chain (B2: the generator frees the DiT before decoding).
pub fn zimage_release_dit() {
    #[cfg(target_os = "macos")]
    crate::gpu_metal::zimage::release_dit();
    #[cfg(feature = "gpu")]
    crate::gpu_wgpu::zimage::release_dit();
}

/// Drop every Z-Image device resource (planes, prepared states, VAE chain
/// buffers): stage change or process end. Calls each compiled backend's
/// release directly, without `backend()`, so it never brings a device up;
/// the child modules' `release` must touch module-local state only.
pub fn zimage_release() {
    #[cfg(target_os = "macos")]
    crate::gpu_metal::zimage::release();
    #[cfg(feature = "gpu")]
    crate::gpu_wgpu::zimage::release();
}

/// Optional device context refiner: the same block math with scale = 0 and
/// gate = 1 (unmodulated): x += norm2(attn(norm1(x))); x += ffn_norm2(ffn(
/// ffn_norm1(x))). `cap` is [n_cap_p, hidden] in/out (the cap_embedder
/// output with pad rows already = cap_pad_token); `rope_cap` is
/// [n_cap_p · hd/2] cos, sin. `false` = untouched, run the CPU refiner.
#[allow(unused_variables)]
pub fn zimage_refine_caption(
    model: &Arc<CmfModel>,
    geom: &ZGeom,
    blocks: &[ZBlockRef],
    rope_cap: (&[f32], &[f32]),
    cap: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => {
            crate::gpu_metal::zimage::refine_caption(model, geom, blocks, rope_cap, cap)
        }
        #[cfg(feature = "gpu")]
        Backend::Wgpu => {
            crate::gpu_wgpu::zimage::refine_caption(model, geom, blocks, rope_cap, cap)
        }
        #[allow(unreachable_patterns)]
        _ => false,
    }
}

/// Resident Flux-VAE decode (the whole decoder on the device, one latent
/// upload, one RGB readback). `a` comes from `VaeDecoder::chain_args()`.
/// `z` is [latent_channels, h, w] ALREADY de-normalised
/// (z/scaling_factor + shift_factor — the conv_in input); `out` is
/// [3, 8h, 8w], the raw decoder output (≈[-1, 1], before x/2+0.5).
#[allow(unused_variables)]
pub fn vae_decode_chain(
    a: &crate::vae::VaeChainArgs,
    z: &[f32],
    h: usize,
    w: usize,
    out: &mut [f32],
) -> bool {
    match backend() {
        #[cfg(target_os = "macos")]
        Backend::Metal => crate::gpu_metal::zimage::vae_decode_chain(a, z, h, w, out),
        #[cfg(feature = "gpu")]
        Backend::Wgpu => crate::gpu_wgpu::zimage::vae_decode_chain(a, z, h, w, out),
        #[allow(unreachable_patterns)]
        _ => false,
    }
}