cortiq_engine/

sam3.rs

//! SAM3 image inference on CPU (f32) over the sections written by
//! `cortiq-cli sam3-pack` (`trunk.*`, `fpn.*`, `text.*`, `det.*`).
//!
//! Semantics follow transformers 5.20.2 `Sam3Model` (fixtures: run
//! `tools/sam3_ref.py` on the pod): CLIP-L text tower + resizer, ViT trunk
//! with windowed axial RoPE (per-layer baked freqs_cis), SAM2-style FPN
//! neck, DETR encoder/decoder (presence token + 200 queries + log
//! box-RPB, POST-norm decoder layers) and the mask decoder. Geometry
//! prompts are not part of the image path. Batch size 1; `pixels` is a
//! preprocessed [3,1008,1008] tensor.

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;

use anyhow::{anyhow, Context, Result};
use cortiq_core::quant::dequant_tensor;
use cortiq_core::{CmfModel, TensorDtype};

use crate::inference::gelu_erf;
use crate::pool::Pool;
use crate::qtensor::QTensor;

/// Additive mask for "key is padded": saturates softmax to 0.
const NEG: f32 = -1e30;

/// The vendor image preprocess, exactly as `cfg.processor` spells it:
/// resize to 1008×1008 with `resample: 2` (bilinear), rescale 1/255,
/// normalize with mean/std 0.5 per channel, emit CHW f32. This produces
/// the parity fixtures' `pixels.f32` from the source image.
pub fn preprocess(frame: &crate::media::RgbFrame) -> Vec<f32> {
    use image::imageops::{self, FilterType};
    let img = image::RgbImage::from_raw(
        frame.width as u32,
        frame.height as u32,
        frame.data.clone(),
    )
    .expect("RgbFrame buffer matches its own dimensions");
    let rs = imageops::resize(&img, PIX as u32, PIX as u32, FilterType::Triangle);
    let raw = rs.into_raw();
    let mut out = vec![0f32; 3 * PIX * PIX];
    for i in 0..PIX * PIX {
        for c in 0..3 {
            out[c * PIX * PIX + i] = (raw[i * 3 + c] as f32 / 255.0 - 0.5) / 0.5;
        }
    }
    out
}

const PIX: usize = 1008;
const PATCH: usize = 14;
const GRID: usize = PIX / PATCH; // 72
const WIN: usize = 24;
const TD: usize = 1024;
const TH: usize = 16;
const TL: usize = 32;
const TI: usize = 4736;
const T_EPS: f32 = 1e-6;
const GLOBAL: [usize; 4] = [7, 15, 23, 31];

const SS: usize = 32;
const SD: usize = 1024;
const SL: usize = 24;
const SI: usize = 4096;
const S_EPS: f32 = 1e-5;
const PROJ: usize = 512;

const FD: usize = 256;
const FI: usize = 2048;
const FH: usize = 8;
const D_EPS: f32 = 1e-6;
const NL: usize = 6;
const NQ: usize = 200;
const MH: usize = 288;

/// One forward's outputs (batch 1), flat row-major f32.
pub struct Sam3Out {
    /// CLIP text projection of the pooled token, [512].
    pub text_embeds: Vec<f32>,
    /// Resizer output fed to the DETR stack, [32 * 256].
    pub text_feats: Vec<f32>,
    /// Trunk features, [5184 * 1024].
    pub trunk_feats: Vec<f32>,
    /// FPN levels 288/144/72/36, each [256 * H * W].
    pub fpn: Vec<Vec<f32>>,
    /// DETR encoder output, [5184 * 256].
    pub enc_hidden: Vec<f32>,
    /// Stacked decoder layer outputs, [6 * 200 * 256].
    pub dec_hidden: Vec<f32>,
    /// Final boxes (xyxy), [200 * 4].
    pub pred_boxes: Vec<f32>,
    /// Instance masks, [200 * 288 * 288].
    pub pred_masks: Vec<f32>,
    /// Per-query classification logits from the dot-product scorer
    /// (`det.scoring.*`), [200] — the final score of a query is
    /// `sigmoid(pred_logits) * sigmoid(presence_logit)` (HF's
    /// `post_process_object_detection`). Empty when `SAM3_ONLY`
    /// skips the decoder.
    pub pred_logits: Vec<f32>,
    /// Presence-token logit of the final decoder layer (clamped ±10).
    /// 0.0 when the decoder did not run.
    pub presence_logit: f32,
}

/// Additive attention masks (added to scores before softmax).
#[derive(Clone, Copy)]
pub(crate) enum Mask<'a> {
    None,
    /// One value per key, shared by all queries/heads (padding).
    Key(&'a [f32]),
    /// [n_q, n_k] matrix, shared across heads (CLIP causal).
    Add(&'a [f32]),
    /// [heads, n_q, n_k] per-head bias (box RPB).
    Head(&'a [f32]),
}

/// Attention sub-timers (ns), accumulated inside `run_attn` (qkv/proj
/// GEMMs vs the softmax core) and printed by the session timer's Drop.
static ATT_PROJ_NS: AtomicU64 = AtomicU64::new(0);
static ATT_CORE_NS: AtomicU64 = AtomicU64::new(0);

/// `SAM3_STAGES=1` — in-code wall-clock counters: `lap(name)` closes the
/// current segment under `name` (laps with the same name merge), `add`
/// accumulates inside loops. On drop the merged totals print as one
/// `[sam3-prof]` block on stderr. Off by default (one env read per run).
struct StageTimer {
    on: bool,
    main: bool,
    t0: std::time::Instant,
    acc: Vec<(&'static str, f64)>,
}

impl StageTimer {
    fn new() -> Self {
        Self {
            on: std::env::var("SAM3_STAGES").is_ok_and(|v| v != "0"),
            main: false,
            t0: std::time::Instant::now(),
            acc: Vec::new(),
        }
    }

    /// The forward-level timer: resets the per-run attention statics so
    /// multi-frame (video) runs report per-frame numbers.
    fn new_session() -> Self {
        ATT_PROJ_NS.store(0, Ordering::Relaxed);
        ATT_CORE_NS.store(0, Ordering::Relaxed);
        Self {
            on: std::env::var("SAM3_STAGES").is_ok_and(|v| v != "0"),
            main: true,
            t0: std::time::Instant::now(),
            acc: Vec::new(),
        }
    }

    fn lap(&mut self, name: &'static str) {
        if self.on {
            self.acc.push((name, self.t0.elapsed().as_secs_f64() * 1e3));
        }
        self.t0 = std::time::Instant::now();
    }
}

impl Drop for StageTimer {
    fn drop(&mut self) {
        if !self.on || self.acc.is_empty() {
            return;
        }
        let mut merged: Vec<(&'static str, f64)> = Vec::new();
        for (n, ms) in &self.acc {
            match merged.iter_mut().find(|(m, _)| m == n) {
                Some((_, t)) => *t += *ms,
                None => merged.push((*n, *ms)),
            }
        }
        let total: f64 = merged.iter().map(|(_, t)| t).sum();
        for (n, ms) in &merged {
            eprintln!("[sam3-prof] {n:>14} {ms:>10.1} ms  ({:.1}%)", ms / total * 100.0);
        }
        eprintln!("[sam3-prof] {:>14} {total:>10.1} ms", "TOTAL");
        if self.main {
            eprintln!(
                "[sam3-prof] {:>14} {:>10.1} ms",
                "attn_proj",
                ATT_PROJ_NS.load(Ordering::Relaxed) as f64 / 1e6
            );
            eprintln!(
                "[sam3-prof] {:>14} {:>10.1} ms",
                "attn_core",
                ATT_CORE_NS.load(Ordering::Relaxed) as f64 / 1e6
            );
            eprintln!(
                "[sam3-prof] {:>14} hit/gate={}/{}",
                "matmat_gpu",
                crate::qtensor::MATMAT_GPU_HIT.load(Ordering::Relaxed),
                crate::qtensor::MATMAT_GATE.load(Ordering::Relaxed)
            );
            #[cfg(feature = "gpu")]
            eprintln!(
                "[sam3-prof] {:>14} noctx/shape/entry/bounds/disp={}/{}/{}/{}/{}",
                "m2f_refuse",
                crate::gpu_wgpu::M2F_NO_CTX.load(Ordering::Relaxed),
                crate::gpu_wgpu::M2F_SHAPE.load(Ordering::Relaxed),
                crate::gpu_wgpu::M2F_ENTRY.load(Ordering::Relaxed),
                crate::gpu_wgpu::M2F_BOUNDS.load(Ordering::Relaxed),
                crate::gpu_wgpu::M2F_DISPATCH.load(Ordering::Relaxed),
            );
        }
    }
}

pub struct Sam3 {
    model: Arc<CmfModel>,
    pool: Option<Arc<Pool>>,
    /// `fpn.convs` (image detector neck) or `fpn.sam2_convs` (SAM3_FPN).
    fpn_prefix: &'static str,
    /// `trunk.pos_embed` is [1,577,1024]: Meta's leading cls row is
    /// dropped by the HF conversion (SAM3_POS=first keeps rows 0..576).
    pos_skip_cls: bool,
}

impl Sam3 {
    pub fn open(path: &str) -> Result<Self> {
        let model = CmfModel::open(path)
            .with_context(|| format!("open {path}"))?
            .into();
        let fpn_prefix = match std::env::var("SAM3_FPN").as_deref() {
            Ok("sam2_convs") => "fpn.sam2_convs",
            _ => "fpn.convs",
        };
        let pos_skip_cls = std::env::var("SAM3_POS").as_deref() != Ok("first");
        Ok(Self { model, pool: Pool::from_env(), fpn_prefix, pos_skip_cls })
    }

    /// The shared CMF handle (the video tracker reads its `trk.*` section
    /// through it).
    pub(crate) fn cmf(&self) -> &Arc<CmfModel> { &self.model }

    /// The GEMM thread pool the detector was opened with.
    pub(crate) fn pool_handle(&self) -> Option<&Arc<Pool>> { self.pool.as_ref() }

    /// Debug: `"name dtype shape"` lines for every tensor in the CMF
    /// directory (the parity tests filter them by section).
    pub fn directory(&self) -> Vec<String> {
        self.model
            .tensors
            .iter()
            .map(|t| format!("{} {:?} {:?}", t.name, t.dtype, t.shape))
            .collect()
    }

    /// Debug: raw payload bytes of one directory tensor.
    pub fn raw_tensor(&self, name: &str) -> Option<Vec<u8>> {
        let e = self.model.tensor(name)?;
        Some(self.model.entry_bytes(e).to_vec())
    }

    /// Debug: every `q8_row` row scale and every `f16` payload element
    /// whose bit pattern is non-finite, as `"name@idx bad=0x…"` lines.
    pub fn bad_scales(&self) -> Vec<String> {
        let f16_bad = |h: u16| h & 0x7C00 == 0x7C00;
        let mut out = Vec::new();
        for t in self.model.tensors.iter() {
            let bytes = self.model.entry_bytes(t);
            match t.dtype {
                TensorDtype::Q8Row if t.shape.len() == 2 => {
                    let (rows, cols) = (t.shape[0], t.shape[1]);
                    let off = rows * cols;
                    if bytes.len() < off + rows * 2 {
                        out.push(format!("{}@payload short", t.name));
                        continue;
                    }
                    for r in 0..rows {
                        let h = u16::from_le_bytes([bytes[off + r * 2], bytes[off + r * 2 + 1]]);
                        if f16_bad(h) {
                            out.push(format!("{}@row{r} scale_f16=0x{h:04x}", t.name));
                        }
                    }
                }
                TensorDtype::F16 => {
                    for (i, c) in bytes.chunks_exact(2).enumerate() {
                        let h = u16::from_le_bytes([c[0], c[1]]);
                        if f16_bad(h) {
                            out.push(format!("{}@{i} val_f16=0x{h:04x}", t.name));
                            if out.len() > 200 { return out; }
                        }
                    }
                }
                _ => {}
            }
        }
        out
    }

    /// `pixels` [3,1008,1008] preprocessed; `ids`/`attn` length 32
    /// (Sam3Processor, padding to max_length=32).
    ///
    /// Debug knobs (unset in normal runs): `SAM3_ONLY=trunk,...` runs only
    /// the listed stages and returns empty tensors for the rest (the
    /// parity test skips while it is set); `SAM3_DUMP_DIR` makes `trunk`
    /// write its stage outputs (`trunk_patch.f32`, `trunk_pre.f32`,
    /// `trunk_b{l}.f32`) for reference bisection; `SAM3_TEXT_DEBUG=1`
    /// prints per-stage finiteness stats of the text tower.
    pub fn forward(&mut self, pixels: &[f32], ids: &[u32], attn: &[u8]) -> Result<Sam3Out> {
        let _prof = crate::cpuprof::spans::dump_on_drop("image forward");
        if pixels.len() != 3 * PIX * PIX {
            return Err(anyhow!("pixels {} != 3*{}*{}", pixels.len(), PIX, PIX));
        }
        if ids.len() != SS || attn.len() != SS {
            return Err(anyhow!("ids/attn must be {SS} long"));
        }
        let only = std::env::var("SAM3_ONLY").ok();
        let want = |k: &str| {
            only
                .as_deref()
                .map_or(true, |o| o.split(',').any(|s| s.trim() == k))
        };
        let pool = self.pool.as_deref();
        let mut stg = StageTimer::new_session();
        let _p = crate::cpuprof::spans::span("text");
        let (text_embeds, text_feats) = self.text(ids, attn, pool)?;
        drop(_p);
        stg.lap("text");
        let blank = |te: Vec<f32>, tf: Vec<f32>| Sam3Out {
            text_embeds: te,
            text_feats: tf,
            trunk_feats: Vec::new(),
            fpn: Vec::new(),
            enc_hidden: Vec::new(),
            dec_hidden: Vec::new(),
            pred_boxes: Vec::new(),
            pred_masks: Vec::new(),
            pred_logits: Vec::new(),
            presence_logit: 0.0,
        };
        if !want("trunk") {
            return Ok(blank(text_embeds, text_feats));
        }
        let trunk_feats = {
            let _p = crate::cpuprof::spans::span("trunk");
            self.trunk(pixels, pool)?
        };
        stg.lap("trunk");
        if !want("fpn") {
            let mut o = blank(text_embeds, text_feats);
            o.trunk_feats = trunk_feats;
            return Ok(o);
        }
        let fpn = self.fpn(&trunk_feats, pool)?;
        stg.lap("fpn");
        let pos = sine_pos(GRID, GRID, FD / 2); // level-2 (72) encoding, rows
        let vision = chw_to_rows(&fpn[2], GRID, GRID);
        let enc_hidden = {
            let _p = crate::cpuprof::spans::span("detr_enc");
            self.detr_encoder(&vision, &pos, &text_feats, attn, pool)?
        };
        stg.lap("enc");
        if !want("dec") {
            let mut o = blank(text_embeds, text_feats);
            o.trunk_feats = trunk_feats;
            o.fpn = fpn;
            o.enc_hidden = enc_hidden;
            return Ok(o);
        }
        let (dec_hidden, final_refs, presence_logit) = {
            let _p = crate::cpuprof::spans::span("detr_dec");
            self.detr_decoder(&enc_hidden, &pos, &text_feats, attn, pool)?
        };
        stg.lap("dec");
        // Dot-product scoring (det.scoring.*): the FINAL layer's DETR
        // queries against the mean-pooled prompt features. Logits, not
        // probabilities — the caller multiplies both sigmoids, as HF's
        // post-processor does (`all_pred_logits[-1]`).
        let pred_logits = {
            let _p = crate::cpuprof::spans::span("score");
            self.score(
                &text_feats,
                attn,
                &dec_hidden[(NL - 1) * NQ * FD..],
                pool,
            )?
        };
        stg.lap("score");
        if !want("masks") {
            let mut o = blank(text_embeds, text_feats);
            o.trunk_feats = trunk_feats;
            o.fpn = fpn;
            o.enc_hidden = enc_hidden;
            o.dec_hidden = dec_hidden.clone();
            o.pred_boxes = xyxy(&final_refs);
            o.pred_logits = pred_logits;
            o.presence_logit = presence_logit;
            return Ok(o);
        }
        let pred_boxes = xyxy(&final_refs);
        let pred_masks = {
            let _p = crate::cpuprof::spans::span("mask_dec");
            self.mask_decoder(
                &dec_hidden[(NL - 1) * NQ * FD..],
                &enc_hidden,
                &fpn,
                &text_feats,
                attn,
                pool,
            )?
        };
        stg.lap("masks");
        Ok(Sam3Out {
            text_embeds,
            text_feats,
            trunk_feats,
            fpn,
            enc_hidden,
            dec_hidden,
            pred_boxes,
            pred_masks,
            pred_logits,
            presence_logit,
        })
    }

    // ── text tower (CLIP-L, pre-norm transformer) ───────────────────

    fn text(&self, ids: &[u32], attn: &[u8], pool: Option<&Pool>) -> Result<(Vec<f32>, Vec<f32>)> {
        let m = &self.model;
        let tok = lin(m, "text.encoder.token_embedding.weight")?;
        let pos = load_vec(m, "text.encoder.positional_embedding")?;
        let mut x = vec![0f32; SS * SD];
        let mut row = vec![0f32; SD];
        for (i, &id) in ids.iter().enumerate() {
            tok.row_f32(id as usize, &mut row);
            let r = &mut x[i * SD..(i + 1) * SD];
            r.copy_from_slice(&row);
            for (v, p) in r.iter_mut().zip(&pos[i * SD..(i + 1) * SD]) {
                *v += *p;
            }
        }
        dbg_stats("text.emb", &x);
        // causal + padding additive mask [SS, SS]
        let mut cm = vec![0f32; SS * SS];
        for q in 0..SS {
            for k in 0..=q {
                if attn[k] == 0 {
                    cm[q * SS + k] = NEG;
                }
            }
            for k in (q + 1)..SS {
                cm[q * SS + k] = NEG;
            }
        }
        for l in 0..SL {
            let b = format!("text.encoder.transformer.resblocks.{l}.");
            let (lw, lb) = (load_vec(m, &format!("{b}ln_1.weight"))?, load_vec(m, &format!("{b}ln_1.bias"))?);
            // Pre-norm: branch reads a normalized CLONE, x keeps the residual.
            let mut h = x.clone();
            layer_norm_rows(&mut h, &lw, &lb, S_EPS);
            let qkv = lin(m, &format!("{b}attn.in_proj_weight"))?;
            let ib = load_vec(m, &format!("{b}attn.in_proj_bias"))?;
            let op = lin(m, &format!("{b}attn.out_proj.weight"))?;
            let ob = load_vec(m, &format!("{b}attn.out_proj.bias"))?;
            let mut ao = vec![0f32; SS * SD];
            run_attn(
                &qkv, &ib, &op, &ob,
                &h, &h, &h, SS, SS, SD, 16, None, Mask::Add(&cm), pool, &mut ao,
            );
            for (a, c) in x.iter_mut().zip(&ao) { *a += *c; }
            let (lw, lb) = (load_vec(m, &format!("{b}ln_2.weight"))?, load_vec(m, &format!("{b}ln_2.bias"))?);
            let mut h = x.clone();
            layer_norm_rows(&mut h, &lw, &lb, S_EPS);
            let f1 = lin(m, &format!("{b}mlp.c_fc.weight"))?;
            let f1b = load_vec(m, &format!("{b}mlp.c_fc.bias"))?;
            let f2 = lin(m, &format!("{b}mlp.c_proj.weight"))?;
            let f2b = load_vec(m, &format!("{b}mlp.c_proj.bias"))?;
            let mut ff = vec![0f32; SS * SI];
            f1.matmat(&h, SS, &mut ff, pool);
            add_bias(&mut ff, SI, &f1b);
            ff.iter_mut().for_each(|v| *v = gelu_erf(*v));
            let mut fo = vec![0f32; SS * SD];
            f2.matmat(&ff, SS, &mut fo, pool);
            add_bias(&mut fo, SD, &f2b);
            for (a, c) in x.iter_mut().zip(&fo) { *a += *c; }
            dbg_stats(&format!("text.l{l}"), &x);
        }
        let (lw, lb) = (load_vec(m, "text.encoder.ln_final.weight")?, load_vec(m, "text.encoder.ln_final.bias")?);
        layer_norm_rows(&mut x, &lw, &lb, S_EPS);
        dbg_stats("text.ln_final", &x);
        if let Some(d) = std::env::var("SAM3_DUMP_DIR").ok() {
            dump_f32(&d, "text_lh", &x);
        }
        // CLIP pools at the first eos token: (ids == eos).argmax()
        let p = ids.iter().position(|&v| v == 49407).ok_or_else(|| anyhow!("no eos token"))?;
        let pooled = &x[p * SD..(p + 1) * SD];
        // text_embeds = pooled @ projection; projection is stored [1024, 512]
        let proj = load_vec(m, "text.encoder.text_projection")?;
        dbg_stats("text.proj", &proj);
        let mut embeds = vec![0f32; PROJ];
        for j in 0..PROJ {
            let mut acc = 0f32;
            for i in 0..SD { acc += pooled[i] * proj[i * PROJ + j]; }
            embeds[j] = acc;
        }
        dbg_stats("text.pooled", pooled);
        dbg_stats("text.embeds", &embeds);
        // DETR text features: Linear(1024 -> 256) over the full sequence
        let rz = lin(m, "text.resizer.weight")?;
        let rb = load_vec(m, "text.resizer.bias")?;
        let mut feats = vec![0f32; SS * FD];
        rz.matmat(&x, SS, &mut feats, pool);
        add_bias(&mut feats, FD, &rb);
        dbg_stats("text.feats", &feats);
        Ok((embeds, feats))
    }

    // ── trunk (ViT, pre-norm layers, windowed/global axial RoPE) ────

    fn trunk(&self, pixels: &[f32], pool: Option<&Pool>) -> Result<Vec<f32>> {
        let _patch = crate::cpuprof::spans::span("trunk.patch");
        let m = &self.model;
        let np = GRID * GRID;
        let kc = 3 * PATCH * PATCH;
        let dump = std::env::var("SAM3_DUMP_DIR").ok();
        let dlayers: Option<Vec<usize>> = std::env::var("SAM3_DUMP_LAYERS")
            .ok()
            .map(|s| s.split(',').filter_map(|t| t.trim().parse().ok()).collect());
        let w = load_vec(m, "trunk.patch_embed.proj.weight")?; // [1024,3,14,14]
        let mut pv = vec![0f32; np * kc];
        let pvb = pv.as_mut_ptr() as usize;
        let fill = |s: usize, e: usize| {
            let pv = unsafe { std::slice::from_raw_parts_mut(pvb as *mut f32, np * kc) };
            for pt in s..e {
                let (py, px) = (pt / GRID, pt % GRID);
                let d = &mut pv[pt * kc..(pt + 1) * kc];
                let mut o = 0;
                for c in 0..3 {
                    for ky in 0..PATCH {
                        let y = py * PATCH + ky;
                        for kx in 0..PATCH {
                            d[o] = pixels[c * PIX * PIX + y * PIX + px * PATCH + kx];
                            o += 1;
                        }
                    }
                }
            }
        };
        match pool { Some(p) => p.run_rows(np, &fill), None => fill(0, np) };
        let mut x = vec![0f32; np * TD];
        let xb = x.as_mut_ptr() as usize;
        let gemm = |s: usize, e: usize| {
            let x = unsafe { std::slice::from_raw_parts_mut(xb as *mut f32, np * TD) };
            for ch in s..e {
                let wr = &w[ch * kc..(ch + 1) * kc];
                for pt in 0..np {
                    let xv = &pv[pt * kc..(pt + 1) * kc];
                    let mut acc = 0f32;
                    for k in 0..kc { acc += wr[k] * xv[k]; }
                    x[pt * TD + ch] = acc;
                }
            }
        };
        match pool { Some(p) => p.run_rows(TD, &gemm), None => gemm(0, TD) };
        drop(_patch);
        if let Some(d) = &dump { dump_f32(d, "trunk_patch", &x); }
        // position embeddings [1,577,1024]: drop Meta's cls row, tile 24x24 -> 72x72
        let pos = load_vec(m, "trunk.pos_embed")?;
        let base = if self.pos_skip_cls { 1 } else { 0 };
        if pos.len() < (base + 576) * TD { return Err(anyhow!("pos_embed too small")); }
        for pt in 0..np {
            let (py, px) = (pt / GRID, pt % GRID);
            let src = (base + (py % 24) * 24 + px % 24) * TD;
            for d in 0..TD { x[pt * TD + d] += pos[src + d]; }
        }
        let (lw, lb) = (load_vec(m, "trunk.ln_pre.weight")?, load_vec(m, "trunk.ln_pre.bias")?);
        layer_norm_rows(&mut x, &lw, &lb, T_EPS);
        if let Some(d) = &dump { dump_f32(d, "trunk_pre", &x); }

        let mut tst = StageTimer::new();
        for l in 0..TL {
            let b = format!("trunk.blocks.{l}.");
            let _attn = crate::cpuprof::spans::span("trunk.attn");
            let (nw, nb) = (load_vec(m, &format!("{b}norm1.weight"))?, load_vec(m, &format!("{b}norm1.bias"))?);
            let mut h = x.clone();
            layer_norm_rows(&mut h, &nw, &nb, T_EPS);
            let dbg0 = dump.is_some() && l == 0;
            if let (Some(d), true) = (&dump, dbg0) { dump_f32(d, "b0_ln1", &h); }
            drop(_attn);
            let qkv = lin(m, &format!("{b}attn.qkv.weight"))?;
            let qb = load_vec(m, &format!("{b}attn.qkv.bias"))?;
            let proj = lin(m, &format!("{b}attn.proj.weight"))?;
            let pb = load_vec(m, &format!("{b}attn.proj.bias"))?;
            let fc = load_vec(m, &format!("{b}attn.freqs_cis"))?;
            tst.lap("trk_attn_ln");
            let attn_out = if GLOBAL.contains(&l) {
                let mut o = vec![0f32; np * TD];
                let _pa = crate::cpuprof::spans::span("trunk.attn.run");
                run_attn(&qkv, &qb, &proj, &pb, &h, &h, &h, np, np, TD, TH, Some(&fc), Mask::None, pool, &mut o);
                drop(_pa);
                o
            } else {
                let _pa = crate::cpuprof::spans::span("trunk.attn.run");
                let o = self.vit_windows(&h, &qkv, &qb, &proj, &pb, &fc, pool)?;
                drop(_pa);
                o
            };
            if let (Some(d), true) = (&dump, dbg0) { dump_f32(d, "b0_attn", &attn_out); }
            for (a, c) in x.iter_mut().zip(&attn_out) { *a += *c; }
            tst.lap("trk_attn_run");
            if let (Some(d), true) = (&dump, dbg0) { dump_f32(d, "b0_x2", &x); }
            // pre-norm MLP branch: norm a CLONE — x keeps the residual
            // (norming x in place made the block add mlp(ln2(x)) instead
            // of x + mlp(ln2(x))).
            let (nw, nb) = (load_vec(m, &format!("{b}norm2.weight"))?, load_vec(m, &format!("{b}norm2.bias"))?);
            let mut h = x.clone();
            layer_norm_rows(&mut h, &nw, &nb, T_EPS);
            if let (Some(d), true) = (&dump, dbg0) { dump_f32(d, "b0_ln2", &h); }
            let f1 = lin(m, &format!("{b}mlp.fc1.weight"))?;
            let f1b = load_vec(m, &format!("{b}mlp.fc1.bias"))?;
            let f2 = lin(m, &format!("{b}mlp.fc2.weight"))?;
            let f2b = load_vec(m, &format!("{b}mlp.fc2.bias"))?;
            let mut ff = vec![0f32; np * TI];
            f1.matmat(&h, np, &mut ff, pool);
            add_bias(&mut ff, TI, &f1b);
            gelu_slice(&mut ff);
            let mut fo = vec![0f32; np * TD];
            f2.matmat(&ff, np, &mut fo, pool);
            add_bias(&mut fo, TD, &f2b);
            if let (Some(d), true) = (&dump, dbg0) { dump_f32(d, "b0_mlp", &fo); }
            for (a, c) in x.iter_mut().zip(&fo) { *a += *c; }
            tst.lap("trk_mlp");
            if let (Some(d), true) = (&dump, dlayers.as_ref().map_or(true, |ls| ls.contains(&l))) {
                dump_f32(d, &format!("trunk_b{l}"), &x);
            }
        }
        Ok(x)
    }

    /// Windowed trunk attention: 3x3 windows of 24x24 tokens (freqs [576]).
    fn vit_windows(
        &self,
        h: &[f32],
        qkv: &QTensor,
        qb: &[f32],
        proj: &QTensor,
        pb: &[f32],
        fc: &[f32],
        pool: Option<&Pool>,
    ) -> Result<Vec<f32>> {
        let np = GRID * GRID;
        let nw = GRID / WIN;
        let wl = WIN * WIN;
        let mut wh = vec![0f32; wl * TD];
        let mut ao = vec![0f32; wl * TD];
        let mut out = vec![0f32; np * TD];
        for wy in 0..nw {
            for wx in 0..nw {
                for y in 0..WIN {
                    for x2 in 0..WIN {
                        let s = ((wy * WIN + y) * GRID + wx * WIN + x2) * TD;
                        let d = (y * WIN + x2) * TD;
                        wh[d..d + TD].copy_from_slice(&h[s..s + TD]);
                    }
                }
                run_attn(qkv, qb, proj, pb, &wh, &wh, &wh, wl, wl, TD, TH, Some(fc), Mask::None, pool, &mut ao);
                for y in 0..WIN {
                    for x2 in 0..WIN {
                        let d = ((wy * WIN + y) * GRID + wx * WIN + x2) * TD;
                        let s = (y * WIN + x2) * TD;
                        out[d..d + TD].copy_from_slice(&ao[s..s + TD]);
                    }
                }
            }
        }
        Ok(out)
    }

    // ── FPN neck ────────────────────────────────────────────────────

    /// Four pyramid levels from the trunk map: 288/144/72/36, each
    /// [256, H, W] channels-first.
    fn fpn(&self, trunk: &[f32], pool: Option<&Pool>) -> Result<Vec<Vec<f32>>> {
        fpn_levels(&self.model, self.fpn_prefix, trunk, pool, 4)
    }

    // ── DETR encoder (pre-norm layers) ──────────────────────────────

    #[allow(clippy::too_many_arguments)]
    fn detr_encoder(
        &self,
        vision: &[f32],
        pos: &[f32],
        text: &[f32],
        am: &[u8],
        pool: Option<&Pool>,
    ) -> Result<Vec<f32>> {
        let m = &self.model;
        let n = vision.len() / FD;
        let mut xs = vision.to_vec();
        let mut km = vec![0f32; SS];
        for (k, &v) in am.iter().enumerate() { km[k] = if v == 0 { NEG } else { 0.0 }; }
        let mut tse = StageTimer::new();
        for l in 0..NL {
            let b = format!("det.transformer.encoder.layers.{l}.");
            // pre-norm self attention: q=k=normed+pos, v=normed
            let (nw, nb) = (load_vec(m, &format!("{b}norm1.weight"))?, load_vec(m, &format!("{b}norm1.bias"))?);
            let mut h = xs.clone();
            layer_norm_rows(&mut h, &nw, &nb, D_EPS);
            let hp: Vec<f32> = h
                .chunks_exact(FD)
                .zip(pos.chunks_exact(FD))
                .flat_map(|(a, c)| a.iter().zip(c).map(|(x, y)| x + y))
                .collect();
            let qkv = lin(m, &format!("{b}self_attn.in_proj_weight"))?;
            let ib = load_vec(m, &format!("{b}self_attn.in_proj_bias"))?;
            let op = lin(m, &format!("{b}self_attn.out_proj.weight"))?;
            let ob = load_vec(m, &format!("{b}self_attn.out_proj.bias"))?;
            let mut ao = vec![0f32; n * FD];
            run_attn(&qkv, &ib, &op, &ob, &hp, &hp, &h, n, n, FD, FH, None, Mask::None, pool, &mut ao);
            for (a, c) in xs.iter_mut().zip(&ao) { *a += *c; }
            tse.lap("enc_self_attn");
            // cross attention to the text features (query has no pos add)
            let (nw, nb) = (load_vec(m, &format!("{b}norm2.weight"))?, load_vec(m, &format!("{b}norm2.bias"))?);
            let mut h = xs.clone();
            layer_norm_rows(&mut h, &nw, &nb, D_EPS);
            let qkv = lin(m, &format!("{b}cross_attn_image.in_proj_weight"))?;
            let ib = load_vec(m, &format!("{b}cross_attn_image.in_proj_bias"))?;
            let op = lin(m, &format!("{b}cross_attn_image.out_proj.weight"))?;
            let ob = load_vec(m, &format!("{b}cross_attn_image.out_proj.bias"))?;
            let mut ao = vec![0f32; n * FD];
            run_attn(&qkv, &ib, &op, &ob, &h, text, text, n, SS, FD, FH, None, Mask::Key(&km), pool, &mut ao);
            for (a, c) in xs.iter_mut().zip(&ao) { *a += *c; }
            tse.lap("enc_cross_text");
            // relu MLP
            let (nw, nb) = (load_vec(m, &format!("{b}norm3.weight"))?, load_vec(m, &format!("{b}norm3.bias"))?);
            let mut h = xs.clone();
            layer_norm_rows(&mut h, &nw, &nb, D_EPS);
            let l1 = lin(m, &format!("{b}linear1.weight"))?;
            let l1b = load_vec(m, &format!("{b}linear1.bias"))?;
            let l2 = lin(m, &format!("{b}linear2.weight"))?;
            let l2b = load_vec(m, &format!("{b}linear2.bias"))?;
            let mut mid = vec![0f32; n * FI];
            l1.matmat(&h, n, &mut mid, pool);
            add_bias(&mut mid, FI, &l1b);
            mid.iter_mut().for_each(|v| *v = v.max(0.0));
            let mut fo = vec![0f32; n * FD];
            l2.matmat(&mid, n, &mut fo, pool);
            add_bias(&mut fo, FD, &l2b);
            for (a, c) in xs.iter_mut().zip(&fo) { *a += *c; }
            tse.lap("enc_ffn");
        }
        Ok(xs)
    }

    // ── DETR decoder (POST-norm layers) ─────────────────────────────

    /// Returns (stacked layer outputs [6*200*256], final reference boxes
    /// in cxcywh — the box_head refinement of the last layer, which
    /// equals Sam3Model's `all_pred_boxes[-1]` pre-xyxy conversion — and
    /// the FINAL layer's presence logit: the presence row through
    /// `presence_token_out_norm` + the 3-layer presence head, clamped
    /// ±10 as HF clamps it).
    fn detr_decoder(
        &self,
        enc: &[f32],
        pos: &[f32],
        text: &[f32],
        am: &[u8],
        pool: Option<&Pool>,
    ) -> Result<(Vec<f32>, Vec<[f32; 4]>, f32)> {
        let m = &self.model;
        let n = enc.len() / FD;
        let db = "det.transformer.decoder.";
        let query = load_vec(m, &format!("{db}query_embed.weight"))?;
        let presence = load_vec(m, &format!("{db}presence_token.weight"))?;
        let ref0 = load_vec(m, &format!("{db}reference_points.weight"))?;
        let mut hidden = vec![0f32; (NQ + 1) * FD];
        hidden[..FD].copy_from_slice(&presence);
        hidden[FD..].copy_from_slice(&query);
        let mut refs: Vec<[f32; 4]> = ref0
            .chunks_exact(4)
            .map(|c| [sigm(c[0]), sigm(c[1]), sigm(c[2]), sigm(c[3])])
            .collect();
        let mut km = vec![0f32; SS];
        for (k, &v) in am.iter().enumerate() { km[k] = if v == 0 { NEG } else { 0.0 }; }
        let mut tsd = StageTimer::new();
        // heads (all loaded once, biases included)
        let rh0 = lin(m, &format!("{db}ref_point_head.layers.0.weight"))?;
        let rh0b = load_vec(m, &format!("{db}ref_point_head.layers.0.bias"))?;
        let rh1 = lin(m, &format!("{db}ref_point_head.layers.1.weight"))?;
        let rh1b = load_vec(m, &format!("{db}ref_point_head.layers.1.bias"))?;
        let rx = DecoderMlp::load(m, &format!("{db}boxRPB_embed_x."))?;
        let ry = DecoderMlp::load(m, &format!("{db}boxRPB_embed_y."))?;
        let bbox = DecoderMlp3::load3(m, &format!("{db}bbox_embed."))?;
        let olnw = load_vec(m, &format!("{db}norm.weight"))?;
        let olnb = load_vec(m, &format!("{db}norm.bias"))?;
        let pres_w = load_vec(m, &format!("{db}presence_token_out_norm.weight"))?;
        let pres_b = load_vec(m, &format!("{db}presence_token_out_norm.bias"))?;
        let pres_head = DecoderMlp3::load3(m, &format!("{db}presence_token_head."))?;
        let key = add_rows(enc, pos); // enc + pos, [n, FD]
        let mut inter: Vec<Vec<f32>> = Vec::with_capacity(NL);
        for l in 0..NL {
            let b = format!("{db}layers.{l}.");
            // query position from the sine encoding of the reference boxes
            let sine = sine_boxes(&refs);
            let mut qp = vec![0f32; NQ * FD];
            rh0.matmat(&sine, NQ, &mut qp, pool);
            add_bias(&mut qp, FD, &rh0b);
            qp.iter_mut().for_each(|v| *v = v.max(0.0)); // DecoderMLP: relu between
            let mut qp2 = vec![0f32; NQ * FD];
            rh1.matmat(&qp, NQ, &mut qp2, pool);
            add_bias(&mut qp2, FD, &rh1b);
            let qp = qp2;
            let mut qp_full = vec![0f32; (NQ + 1) * FD];
            qp_full[FD..].copy_from_slice(&qp);
            let rpb = rpb_bias(&refs, &rx, &ry, n, pool)?;
            // self attention: q=k=h+qp (pad row stays qp=0), v=h ; POST-norm
            let hqp = add_rows(&hidden, &qp_full);
            let qkv = lin(m, &format!("{b}self_attn.in_proj_weight"))?;
            let ib = load_vec(m, &format!("{b}self_attn.in_proj_bias"))?;
            let op = lin(m, &format!("{b}self_attn.out_proj.weight"))?;
            let ob = load_vec(m, &format!("{b}self_attn.out_proj.bias"))?;
            let mut ao = vec![0f32; (NQ + 1) * FD];
            run_attn(&qkv, &ib, &op, &ob, &hqp, &hqp, &hidden, NQ + 1, NQ + 1, FD, FH, None, Mask::None, pool, &mut ao);
            for (a, c) in hidden.iter_mut().zip(&ao) { *a += *c; }
            tsd.lap("dec_self_attn");
            let (nw, nb) = (load_vec(m, &format!("{b}norm1.weight"))?, load_vec(m, &format!("{b}norm1.bias"))?);
            layer_norm_rows(&mut hidden, &nw, &nb, D_EPS);
            // text cross attention
            let hqp = add_rows(&hidden, &qp_full);
            let qkv = lin(m, &format!("{b}ca_text.in_proj_weight"))?;
            let ib = load_vec(m, &format!("{b}ca_text.in_proj_bias"))?;
            let op = lin(m, &format!("{b}ca_text.out_proj.weight"))?;
            let ob = load_vec(m, &format!("{b}ca_text.out_proj.bias"))?;
            let mut ao = vec![0f32; (NQ + 1) * FD];
            run_attn(&qkv, &ib, &op, &ob, &hqp, text, text, NQ + 1, SS, FD, FH, None, Mask::Key(&km), pool, &mut ao);
            for (a, c) in hidden.iter_mut().zip(&ao) { *a += *c; }
            tsd.lap("dec_cross_text");
            let (nw, nb) = (load_vec(m, &format!("{b}catext_norm.weight"))?, load_vec(m, &format!("{b}catext_norm.bias"))?);
            layer_norm_rows(&mut hidden, &nw, &nb, D_EPS);
            // vision cross attention (keys = enc + pos, values = enc, RPB bias)
            let hqp = add_rows(&hidden, &qp_full);
            let qkv = lin(m, &format!("{b}cross_attn.in_proj_weight"))?;
            let ib = load_vec(m, &format!("{b}cross_attn.in_proj_bias"))?;
            let op = lin(m, &format!("{b}cross_attn.out_proj.weight"))?;
            let ob = load_vec(m, &format!("{b}cross_attn.out_proj.bias"))?;
            let mut ao = vec![0f32; (NQ + 1) * FD];
            run_attn(&qkv, &ib, &op, &ob, &hqp, &key, enc, NQ + 1, n, FD, FH, None, Mask::Head(&rpb), pool, &mut ao);
            for (a, c) in hidden.iter_mut().zip(&ao) { *a += *c; }
            tsd.lap("dec_cross_enc");
            let (nw, nb) = (load_vec(m, &format!("{b}norm2.weight"))?, load_vec(m, &format!("{b}norm2.bias"))?);
            layer_norm_rows(&mut hidden, &nw, &nb, D_EPS);
            // relu MLP (Sam3MLP)
            let l1 = lin(m, &format!("{b}linear1.weight"))?;
            let l1b = load_vec(m, &format!("{b}linear1.bias"))?;
            let l2 = lin(m, &format!("{b}linear2.weight"))?;
            let l2b = load_vec(m, &format!("{b}linear2.bias"))?;
            let mut mid = vec![0f32; (NQ + 1) * FI];
            l1.matmat(&hidden, NQ + 1, &mut mid, pool);
            add_bias(&mut mid, FI, &l1b);
            mid.iter_mut().for_each(|v| *v = v.max(0.0));
            let mut fo = vec![0f32; (NQ + 1) * FD];
            l2.matmat(&mid, NQ + 1, &mut fo, pool);
            add_bias(&mut fo, FD, &l2b);
            for (a, c) in hidden.iter_mut().zip(&fo) { *a += *c; }
            tsd.lap("dec_ffn");
            let (nw, nb) = (load_vec(m, &format!("{b}norm3.weight"))?, load_vec(m, &format!("{b}norm3.bias"))?);
            layer_norm_rows(&mut hidden, &nw, &nb, D_EPS);
            // box refine + intermediate (query rows only, output_layer_norm)
            let mut qh = hidden[FD..].to_vec();
            layer_norm_rows(&mut qh, &olnw, &olnb, D_EPS);
            inter.push(qh.clone());
            let d = bbox.forward(&qh, NQ, pool);
            for (r, dd) in refs.iter_mut().zip(d.chunks_exact(4)) {
                for k in 0..4 { r[k] = sigm(inv_sigmoid(r[k]) + dd[k]); }
            }
            tsd.lap("dec_box_ref");
        }
        let dec = inter.iter().flat_map(|v| v.iter().copied()).collect();
        // Presence logit of the FINAL layer: row 0 (the presence token)
        // is post-normed by the layer stack like every other row — feed
        // it through presence_layer_norm + presence_head (relu, relu,
        // linear — 3 layers), clamp ±10 (HF `clamp_presence_logit_max_val`).
        let mut ph = hidden[0..FD].to_vec();
        layer_norm_rows(&mut ph, &pres_w, &pres_b, D_EPS);
        let pl = pres_head.forward(&ph, 1, pool);
        let presence_logit = pl[0].clamp(-10.0, 10.0);
        Ok((dec, refs, presence_logit))
    }

    /// Dot-product concept scoring — HF `Sam3DotProductScoring` over
    /// the `det.scoring.*` section: the prompt MLP (2 layers, relu
    /// between) + residual + out-norm on the resizer features,
    /// mean-pooled over the unpadded positions, projected by
    /// `prompt_proj`; the output-normed queries of the FINAL decoder
    /// layer projected by `hs_proj`; the dot product scaled by
    /// 1/sqrt(256) and clamped to ±12 (HF `clamp_max_val`). `dec` is
    /// one layer's `[nq * FD]` queries; returns `nq` logits.
    fn score(
        &self,
        text: &[f32],
        am: &[u8],
        dec: &[f32],
        pool: Option<&Pool>,
    ) -> Result<Vec<f32>> {
        let m = &self.model;
        let mlp = DecoderMlp::load(m, "det.scoring.prompt_mlp.")?;
        let mut h = mlp.forward(text, SS, pool, FD);
        for (a, b) in h.iter_mut().zip(text) {
            *a += *b;
        }
        let (ow, ob) = (
            load_vec(m, "det.scoring.prompt_mlp.out_norm.weight")?,
            load_vec(m, "det.scoring.prompt_mlp.out_norm.bias")?,
        );
        layer_norm_rows(&mut h, &ow, &ob, D_EPS);
        // Mean pool over attention_mask == 1 (HF `_pool_text_features`).
        let n_valid = am.iter().filter(|&&v| v == 1).count().max(1) as f32;
        let mut pooled = vec![0f32; FD];
        for (i, &v) in am.iter().enumerate().take(SS) {
            if v == 1 {
                for d in 0..FD {
                    pooled[d] += h[i * FD + d];
                }
            }
        }
        for v in pooled.iter_mut() {
            *v /= n_valid;
        }
        let pt = lin(m, "det.scoring.prompt_proj.weight")?;
        let pb = load_vec(m, "det.scoring.prompt_proj.bias")?;
        let mut tp = vec![0f32; FD];
        pt.matmat(&pooled, 1, &mut tp, pool);
        add_bias(&mut tp, FD, &pb);
        let nq = dec.len() / FD;
        let hs = lin(m, "det.scoring.hs_proj.weight")?;
        let hb = load_vec(m, "det.scoring.hs_proj.bias")?;
        let mut q = vec![0f32; nq * FD];
        hs.matmat(dec, nq, &mut q, pool);
        add_bias(&mut q, FD, &hb);
        let scale = 1.0 / (FD as f32).sqrt();
        Ok(q
            .chunks(FD)
            .map(|r| {
                let mut acc = 0f32;
                for d in 0..FD {
                    acc += r[d] * tp[d];
                }
                (acc * scale).clamp(-12.0, 12.0)
            })
            .collect())
    }

    // ── mask decoder ────────────────────────────────────────────────

    /// Prompt cross-attention over the encoder states, pixel decoder
    /// (nearest upsample + sum + conv3x3 + GroupNorm + ReLU), then
    /// pred_masks[q] = einsum(mask_embed[q], instance_embed).
    #[allow(clippy::too_many_arguments)]
    fn mask_decoder(
        &self,
        queries: &[f32],
        enc: &[f32],
        fpn: &[Vec<f32>],
        text: &[f32],
        am: &[u8],
        pool: Option<&Pool>,
    ) -> Result<Vec<f32>> {
        let m = &self.model;
        let db = "det.seghead.";
        let n = enc.len() / FD;
        let mut km = vec![0f32; SS];
        for (k, &v) in am.iter().enumerate() { km[k] = if v == 0 { NEG } else { 0.0 }; }
        let nw = load_vec(m, &format!("{db}cross_attn_norm.weight"))?;
        let nb = load_vec(m, &format!("{db}cross_attn_norm.bias"))?;
        let mut h = enc.to_vec();
        layer_norm_rows(&mut h, &nw, &nb, D_EPS);
        let qkv = lin(m, &format!("{db}cross_attend_prompt.in_proj_weight"))?;
        let ib = load_vec(m, &format!("{db}cross_attend_prompt.in_proj_bias"))?;
        let op = lin(m, &format!("{db}cross_attend_prompt.out_proj.weight"))?;
        let ob = load_vec(m, &format!("{db}cross_attend_prompt.out_proj.bias"))?;
        let mut tsm = StageTimer::new();
        let mut e = enc.to_vec();
        let mut ao = vec![0f32; n * FD];
        run_attn(&qkv, &ib, &op, &ob, &h, text, text, n, SS, FD, FH, None, Mask::Key(&km), pool, &mut ao);
        for (a, c) in e.iter_mut().zip(&ao) { *a += *c; }
        tsm.lap("msk_cross_attn");
        // pixel decoder: encoder states replace level 2 (72), then upsample
        let l72 = rows_to_chw(&e, GRID, GRID);
        let sizes = [(144usize, 144usize), (288, 288)];
        let mut prev = l72;
        for (li, &(sh, sw)) in sizes.iter().enumerate() {
            let _pp = crate::cpuprof::spans::span("mask.pixdec");
            prev = upsample_nearest(&prev, FD, sh / 2, sh / 2, sh, sw);
            let fw = &fpn[1 - li]; // reversed(fpn[:-1]): 144 then 288
            for (a, b) in prev.iter_mut().zip(fw.iter()) { *a += *b; }
            let cw = load_vec(m, &format!("{db}pixel_decoder.conv_layers.{li}.weight"))?;
            let cb = load_vec(m, &format!("{db}pixel_decoder.conv_layers.{li}.bias"))?;
            prev = conv2d(&cw, &cb, &prev, FD, sh, sw, 3, 3, 1, 1, pool);
            let gw = load_vec(m, &format!("{db}pixel_decoder.norms.{li}.weight"))?;
            let gb = load_vec(m, &format!("{db}pixel_decoder.norms.{li}.bias"))?;
            group_norm(&mut prev, 8, &gw, &gb, 1e-5);
            prev.iter_mut().for_each(|v| *v = v.max(0.0));
            drop(_pp);
        }
        tsm.lap("msk_pixel_dec");
        let iw = load_vec(m, &format!("{db}instance_seg_head.weight"))?;
        let ib2 = load_vec(m, &format!("{db}instance_seg_head.bias"))?;
        let inst = conv2d(&iw, &ib2, &prev, FD, MH, MH, 1, 1, 1, 0, pool);
        // mask embeddings from the (already output-normed) queries
        let mut q = queries.to_vec();
        for i in 0..3 {
            let w = lin(m, &format!("{db}mask_predictor.mask_embed.layers.{i}.weight"))?;
            let b = load_vec(m, &format!("{db}mask_predictor.mask_embed.layers.{i}.bias"))?;
            let mut o = vec![0f32; NQ * FD];
            w.matmat(&q, NQ, &mut o, pool);
            add_bias(&mut o, FD, &b);
            if i < 2 { o.iter_mut().for_each(|v| *v = v.max(0.0)); }
            q = o;
        }
        tsm.lap("msk_inst_head");
        let hw = MH * MH;
        let mut out = vec![0f32; NQ * hw];
        let ob2 = out.as_mut_ptr() as usize;
        let _pe = crate::cpuprof::spans::span("mask.einsum");
        let run = |s: usize, e: usize| {
            let o = unsafe { std::slice::from_raw_parts_mut(ob2 as *mut f32, NQ * hw) };
            for qi in s..e {
                let qq = &q[qi * FD..(qi + 1) * FD];
                let dst = &mut o[qi * hw..(qi + 1) * hw];
                for c in 0..FD {
                    let wgt = qq[c];
                    if wgt == 0.0 { continue; }
                    let ic = &inst[c * hw..(c + 1) * hw];
                    for p in 0..hw { dst[p] += wgt * ic[p]; }
                }
            }
        };
        match pool { Some(p) => p.run_rows(NQ, &run), None => run(0, NQ) }
        drop(_pe);
        tsm.lap("msk_einsum");
        Ok(out)
    }
}

/// `Sam3VisionNeck` (`Sam3FPNLayer` stack) over the trunk map: the first
/// `n_levels` pyramid levels at 288/144/72/36, each [256, H, W]
/// channels-first. The DETR path uses all four; the video tracker's neck
/// (`fpn.sam2_convs`) only consumes the first three.
pub(crate) fn fpn_levels(
    m: &Arc<CmfModel>,
    p: &str,
    trunk: &[f32],
    pool: Option<&Pool>,
    n_levels: usize,
) -> Result<Vec<Vec<f32>>> {
    let hidden = rows_to_chw(trunk, GRID, GRID); // [1024,72,72]
    let scales = [4.0f32, 2.0, 1.0, 0.5];
    let mut levels = Vec::with_capacity(n_levels);
    let _p = crate::cpuprof::spans::span("fpn.levels");
    for (i, &s) in scales.iter().enumerate().take(n_levels) {
        let (mut ci, mut ch, mut cw) = (TD, GRID, GRID);
        let mut cur = hidden.clone();
        match s {
            4.0 => {
                let w0 = load_vec(m, &format!("{p}.{i}.dconv_2x2_0.weight"))?;
                let b0 = load_vec(m, &format!("{p}.{i}.dconv_2x2_0.bias"))?;
                cur = convt2x2(&w0, &b0, &cur, ci, ch, cw, pool);
                (ch, cw, ci) = (ch * 2, cw * 2, b0.len());
                gelu_slice(&mut cur);
                let w1 = load_vec(m, &format!("{p}.{i}.dconv_2x2_1.weight"))?;
                let b1 = load_vec(m, &format!("{p}.{i}.dconv_2x2_1.bias"))?;
                cur = convt2x2(&w1, &b1, &cur, ci, ch, cw, pool);
                (ch, cw, ci) = (ch * 2, cw * 2, b1.len());
            }
            2.0 => {
                let w0 = load_vec(m, &format!("{p}.{i}.dconv_2x2.weight"))?;
                let b0 = load_vec(m, &format!("{p}.{i}.dconv_2x2.bias"))?;
                cur = convt2x2(&w0, &b0, &cur, ci, ch, cw, pool);
                (ch, cw, ci) = (ch * 2, cw * 2, b0.len());
            }
            1.0 => {}
            _ => {
                cur = maxpool2(&cur, ci, ch, cw);
                (ch, cw) = (ch / 2, cw / 2);
            }
        }
        let w1 = load_vec(m, &format!("{p}.{i}.conv_1x1.weight"))?;
        let b1 = load_vec(m, &format!("{p}.{i}.conv_1x1.bias"))?;
        cur = conv2d(&w1, &b1, &cur, ci, ch, cw, 1, 1, 1, 0, pool);
        ci = b1.len();
        let w3 = load_vec(m, &format!("{p}.{i}.conv_3x3.weight"))?;
        let b3 = load_vec(m, &format!("{p}.{i}.conv_3x3.bias"))?;
        cur = conv2d(&w3, &b3, &cur, ci, ch, cw, 3, 3, 1, 1, pool);
        levels.push(cur);
    }
    Ok(levels)
}

// ───────────────────── model pieces ─────────────────────

/// Sam3DecoderMLP, 2 layers: relu(l1), l2 (with biases).
struct DecoderMlp {
    w0: QTensor,
    b0: Vec<f32>,
    w1: QTensor,
    b1: Vec<f32>,
}

impl DecoderMlp {
    fn load(m: &Arc<CmfModel>, prefix: &str) -> Result<Self> {
        Ok(Self {
            w0: lin(m, &format!("{prefix}layers.0.weight"))?,
            b0: load_vec(m, &format!("{prefix}layers.0.bias"))?,
            w1: lin(m, &format!("{prefix}layers.1.weight"))?,
            b1: load_vec(m, &format!("{prefix}layers.1.bias"))?,
        })
    }
    fn forward(&self, x: &[f32], rows: usize, pool: Option<&Pool>, out_w: usize) -> Vec<f32> {
        let mut h = vec![0f32; rows * self.b0.len()];
        self.w0.matmat(x, rows, &mut h, pool);
        add_bias(&mut h, self.b0.len(), &self.b0);
        h.iter_mut().for_each(|v| *v = v.max(0.0));
        let mut o = vec![0f32; rows * out_w];
        self.w1.matmat(&h, rows, &mut o, pool);
        add_bias(&mut o, out_w, &self.b1);
        o
    }
}

/// Sam3DecoderMLP, 3 layers: relu(l1), relu(l2), l3 (box head).
struct DecoderMlp3 {
    w0: QTensor,
    b0: Vec<f32>,
    w1: QTensor,
    b1: Vec<f32>,
    w2: QTensor,
    b2: Vec<f32>,
}

impl DecoderMlp3 {
    fn load3(m: &Arc<CmfModel>, prefix: &str) -> Result<Self> {
        Ok(Self {
            w0: lin(m, &format!("{prefix}layers.0.weight"))?,
            b0: load_vec(m, &format!("{prefix}layers.0.bias"))?,
            w1: lin(m, &format!("{prefix}layers.1.weight"))?,
            b1: load_vec(m, &format!("{prefix}layers.1.bias"))?,
            w2: lin(m, &format!("{prefix}layers.2.weight"))?,
            b2: load_vec(m, &format!("{prefix}layers.2.bias"))?,
        })
    }
    fn forward(&self, x: &[f32], rows: usize, pool: Option<&Pool>) -> Vec<f32> {
        let mut d = vec![0f32; rows * self.b0.len()];
        self.w0.matmat(x, rows, &mut d, pool);
        add_bias(&mut d, self.b0.len(), &self.b0);
        d.iter_mut().for_each(|v| *v = v.max(0.0));
        let mut d2 = vec![0f32; rows * self.b1.len()];
        self.w1.matmat(&d, rows, &mut d2, pool);
        add_bias(&mut d2, self.b1.len(), &self.b1);
        d2.iter_mut().for_each(|v| *v = v.max(0.0));
        let mut o = vec![0f32; rows * self.b2.len()];
        self.w2.matmat(&d2, rows, &mut o, pool);
        add_bias(&mut o, self.b2.len(), &self.b2);
        o
    }
}

pub(crate) fn sigm(x: f32) -> f32 {
    1.0 / (1.0 + (-x).exp())
}

fn inv_sigmoid(x: f32) -> f32 {
    let x = x.clamp(0.0, 1.0);
    ((x.max(1e-3)) / ((1.0 - x).max(1e-3))).ln()
}

fn xyxy(boxes: &[[f32; 4]]) -> Vec<f32> {
    let mut out = vec![0f32; boxes.len() * 4];
    for (i, b) in boxes.iter().enumerate() {
        out[i * 4] = b[0] - 0.5 * b[2];
        out[i * 4 + 1] = b[1] - 0.5 * b[3];
        out[i * 4 + 2] = b[0] + 0.5 * b[2];
        out[i * 4 + 3] = b[1] + 0.5 * b[3];
    }
    out
}

pub(crate) fn lin(m: &Arc<CmfModel>, name: &str) -> Result<QTensor> {
    let _p = crate::cpuprof::spans::span("qtensor.wrap");
    QTensor::from_model(m, name).map_err(|e| anyhow!(e))
}

/// Dequantize a tensor to flat f32 (norms, biases, tables, convs).
pub(crate) fn load_vec(m: &Arc<CmfModel>, name: &str) -> Result<Vec<f32>> {
    let _p = crate::cpuprof::spans::span("load.decode");
    let e = m.tensor(name).ok_or_else(|| anyhow!("tensor '{name}' not found"))?;
    if matches!(e.dtype, TensorDtype::U8) { return Err(anyhow!("'{name}' is u8")); }
    let bytes = m.entry_bytes(e);
    let mut dst = vec![0f32; e.n_elems()];
    dequant_tensor(e, bytes, &mut dst).map_err(|er| anyhow!("{name}: {er}"))?;
    Ok(dst)
}

/// `gelu_erf` over a whole buffer, one `SAM3_PROF` span for the loop.
pub(crate) fn gelu_slice(v: &mut [f32]) {
    let _p = crate::cpuprof::spans::span("op.gelu");
    v.iter_mut().for_each(|x| *x = gelu_erf(*x));
}

/// Row-major [rows, cols] += bias (broadcast over rows).
pub(crate) fn add_bias(out: &mut [f32], cols: usize, bias: &[f32]) {
    debug_assert_eq!(bias.len(), cols);
    for r in out.chunks_exact_mut(cols) {
        for (v, b) in r.iter_mut().zip(bias) { *v += *b; }
    }
}

/// LayerNorm over rows of `x` (row length = w.len()).
pub(crate) fn layer_norm_rows(x: &mut [f32], w: &[f32], b: &[f32], eps: f32) {
    let _p = crate::cpuprof::spans::span("op.layernorm");
    let d = w.len();
    for r in x.chunks_exact_mut(d) {
        let mean = r.iter().sum::<f32>() / d as f32;
        let var = r.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / d as f32;
        let inv = 1.0 / (var + eps).sqrt();
        for (v, (&ww, &bb)) in r.iter_mut().zip(w.iter().zip(b.iter())) {
            *v = (*v - mean) * inv * ww + bb;
        }
    }
}

/// Axial 2D RoPE in place on a fused QKV block: rotates the q part
/// (dims [0, dim)) and the k part ([dim, 2*dim)); v untouched. freqs_cis
/// = [n, 32, 2] = (cos, sin) per pair; pairs 0..15 from w-position,
/// 16..31 from h-position; per head block of 64 dims (16 heads share the
/// same position-dependent table).
fn apply_rope(buf: &mut [f32], n: usize, stride: usize, width: usize, fc: &[f32]) {
    let npairs = 32;
    debug_assert_eq!(fc.len(), n * npairs * 2);
    let blocks = width / 64;
    for i in 0..n {
        let f = &fc[i * npairs * 2..(i + 1) * npairs * 2];
        let row = &mut buf[i * stride..i * stride + width];
        for b in 0..blocks {
            let base = b * 64;
            for j in 0..npairs {
                let (c, s) = (f[j * 2], f[j * 2 + 1]);
                let (a, d) = (row[base + j * 2], row[base + j * 2 + 1]);
                row[base + j * 2] = a * c - d * s;
                row[base + j * 2 + 1] = d * c + a * s;
            }
        }
    }
}

/// RAII bypass of the f16-staged coop kernel for the calling thread —
/// the attention arm needs the pure-f32 scalar GEMM to stay inside the
/// parity gate (the coop kernel rounds operands to f16: measured ~1e-3
/// per-op, which compounds through a 32-layer trunk to rel 0.36).
struct CoopBypass;

impl CoopBypass {
    fn new() -> Self {
        #[cfg(feature = "gpu")]
        crate::gpu_wgpu::COOP_BYPASS.with(|b| b.set(true));
        Self
    }
}

impl Drop for CoopBypass {
    fn drop(&mut self) {
        #[cfg(feature = "gpu")]
        crate::gpu_wgpu::COOP_BYPASS.with(|b| b.set(false));
    }
}

/// GPU arm for `run_attn`'s core (QKᵀ softmax A·V), owner GPU port.
///
/// Per head: Q/K/V rows are packed contiguous, QKᵀ and A·V run as f32
/// `gemm_nt_f32` GEMMs (transient both operands are activations), the
/// scale/mask/softmax step stays on the CPU in exactly the scalar arm's
/// order — only the f32 dot-product summation order differs (allowed:
/// cos >= 0.999 vs the CPU arm). Returns false on ANY refusal — backend
/// down, `CMF_GPU=0`, inside `cpu_scope`, job under the GEMM floor, a
/// device decline mid-loop — and leaves `out` untouched, so the caller
/// falls back to the scalar arm. `CMF_SAM3_GPU=0` disables the arm.
#[allow(clippy::too_many_arguments)]
fn attn_core_gpu(
    qkv_q: &[f32],
    k: &[f32],
    v: &[f32],
    n_q: usize,
    n_k: usize,
    dim: usize,
    heads: usize,
    mask: Mask<'_>,
    scale: f32,
    out: &mut [f32],
) -> bool {
    #[cfg(feature = "gpu")]
    {
        if std::env::var("CMF_SAM3_GPU").as_deref() == Ok("0") {
            return false;
        }
        if !crate::gpu::enabled_here() {
            return false;
        }
        let hd = dim / heads;
        if hd == 0 || out.len() < n_q * dim {
            return false;
        }
        // Same floor gemm_nt enforces, checked before any packing: both
        // GEMMs carry the identical n_q * n_k * hd MAC product.
        if n_q.saturating_mul(n_k).saturating_mul(hd) < (1 << 22) {
            return false;
        }
        let h3 = dim * 3;
        // Q/K packed [.., hd] rows; V packed transposed (rows = hd) so
        // A·V is the same NT GEMM: out[n_q, hd] = A[n_q, n_k]·Vᵀ[hd, n_k].
        let _byp = CoopBypass::new();
        let mut qpk = vec![0f32; n_q * hd];
        let mut kpk = vec![0f32; n_k * hd];
        let mut vtp = vec![0f32; hd * n_k];
        let mut raw = vec![0f32; n_q * n_k];
        let mut outt = vec![0f32; n_q * hd];
        let mut acc = vec![0f32; n_q * dim];
        // The pure-f32 scalar arm refuses grids over 2^22 workgroups (a
        // hang guard) — QKᵀ over a big square (n_q * n_k) is chunked
        // along query rows; output rows are independent, so chunking
        // changes nothing but the dispatch.
        let qk_chunk = ((1usize << 22) / n_k.max(1)).max(1);
        for t in 0..heads {
            let off = t * hd;
            for i in 0..n_q {
                qpk[i * hd..(i + 1) * hd]
                    .copy_from_slice(&qkv_q[i * h3 + off..i * h3 + off + hd]);
            }
            for j in 0..n_k {
                kpk[j * hd..(j + 1) * hd]
                    .copy_from_slice(&k[j * h3 + dim + off..j * h3 + dim + off + hd]);
            }
            for d in 0..hd {
                let rowd = &mut vtp[d * n_k..(d + 1) * n_k];
                for (j, rv) in rowd.iter_mut().enumerate() {
                    *rv = v[j * h3 + 2 * dim + off + d];
                }
            }
            // QKᵀ: raw[i, j] = qpk[i]·kpk[j]  (chunked over query rows)
            let mut i0 = 0usize;
            while i0 < n_q {
                let i1 = (i0 + qk_chunk).min(n_q);
                if !crate::gpu_wgpu::gemm_nt_f32_transient(
                    &qpk[i0 * hd..i1 * hd],
                    &kpk,
                    &mut raw[i0 * n_k..i1 * n_k],
                    i1 - i0,
                    hd,
                    n_k,
                ) {
                    return false;
                }
                i0 = i1;
            }
            // scale + mask + stable softmax — the scalar arm's order.
            let mut rowmax = vec![0f32; n_q];
            for i in 0..n_q {
                let base = i * n_k;
                let mut mx = f32::NEG_INFINITY;
                for j in 0..n_k {
                    let mut val = raw[base + j] * scale;
                    val += match mask {
                        Mask::None => 0.0,
                        Mask::Key(km) => km[j],
                        Mask::Add(am) => am[i * n_k + j],
                        Mask::Head(hm) => hm[(t * n_q + i) * n_k + j],
                    };
                    raw[base + j] = val;
                    if val > mx {
                        mx = val;
                    }
                }
                rowmax[i] = mx;
                let mut sum = 0.0f32;
                for j in 0..n_k {
                    let e = (raw[base + j] - mx).exp();
                    raw[base + j] = e;
                    sum += e;
                }
                let inv = 1.0 / sum;
                for j in 0..n_k {
                    raw[base + j] *= inv;
                }
            }
            // A·V: outt[i, d] = raw[i]·vtp[d]
            if !crate::gpu_wgpu::gemm_nt_f32_transient(&raw, &vtp, &mut outt, n_q, n_k, hd) {
                return false;
            }
            for i in 0..n_q {
                acc[i * dim + off..i * dim + off + hd]
                    .copy_from_slice(&outt[i * hd..(i + 1) * hd]);
            }
        }
        out.copy_from_slice(&acc);
        true
    }
    #[cfg(not(feature = "gpu"))]
    {
        let _ = (qkv_q, k, v, n_q, n_k, dim, heads, mask, scale, out);
        false
    }
}

/// Multi-head attention for one sequence: fused QKV projection
/// (`w_qkv` rows = 3*dim laid out [q | k | v] per row, `qb` its bias),
/// optional axial RoPE on q/k, additive mask, softmax, output projection
/// (`w_o` + `ob`). `out` is [n_q, dim]; scale = head_dim^-0.5.
#[allow(clippy::too_many_arguments)]
pub(crate) fn run_attn(
    w_qkv: &QTensor,
    qb: &[f32],
    w_o: &QTensor,
    ob: &[f32],
    xs_q: &[f32],
    xs_k: &[f32],
    xs_v: &[f32],
    n_q: usize,
    n_k: usize,
    dim: usize,
    heads: usize,
    rope: Option<&[f32]>,
    mask: Mask,
    pool: Option<&Pool>,
    out: &mut [f32],
) {
    let h3 = dim * 3;
    let hd = dim / heads;
    let scale = 1.0 / (hd as f32).sqrt();
    let _tp = std::time::Instant::now();
    let mut qkv_q = vec![0f32; n_q * h3];
    let _pg = crate::cpuprof::spans::span("attn.gemm");
    w_qkv.matmat(xs_q, n_q, &mut qkv_q, pool);
    add_bias(&mut qkv_q, h3, qb);
    let same_k = xs_k.as_ptr() == xs_q.as_ptr();
    let mut kbuf = Vec::new();
    if !same_k {
        kbuf = vec![0f32; n_k * h3];
        w_qkv.matmat(xs_k, n_k, &mut kbuf, pool);
        add_bias(&mut kbuf, h3, qb);
    }
    let same_v = xs_v.as_ptr() == xs_q.as_ptr();
    let mut vbuf = Vec::new();
    if !same_v {
        vbuf = vec![0f32; n_k * h3];
        w_qkv.matmat(xs_v, n_k, &mut vbuf, pool);
        add_bias(&mut vbuf, h3, qb);
    }
    drop(_pg);
    if let Some(fc) = rope {
        // rotate q (and k when they share the buffer) over the first 2*dim dims
        apply_rope(&mut qkv_q, n_q, h3, dim * 2, fc);
        if !same_k {
            apply_rope(&mut kbuf, n_k, h3, dim * 2, fc);
        }
    }
    ATT_PROJ_NS.fetch_add(_tp.elapsed().as_nanos() as u64, Ordering::Relaxed);
    let _tc = std::time::Instant::now();
    let qkv_k: &[f32] = if same_k { &qkv_q } else { &kbuf };
    let qkv_v: &[f32] = if same_v { &qkv_q } else { &vbuf };
    let mut attn = vec![0f32; n_q * dim];
    if attn_core_gpu(&qkv_q, qkv_k, qkv_v, n_q, n_k, dim, heads, mask, scale, &mut attn) {
        ATT_CORE_NS.fetch_add(_tc.elapsed().as_nanos() as u64, Ordering::Relaxed);
        let _tw = std::time::Instant::now();
        w_o.matmat(&attn, n_q, out, pool);
        add_bias(out, dim, ob);
        ATT_PROJ_NS.fetch_add(_tw.elapsed().as_nanos() as u64, Ordering::Relaxed);
        return;
    }
    let ab = attn.as_mut_ptr() as usize;
    let _pc = crate::cpuprof::spans::span("attn.core");
    // Attention core as two packed GEMMs per head over each row chunk —
    // QKᵀ then softmax then AV — instead of the per-(row, head) scalar
    // loops with a dot()/axpy() call per key. matrixmultiply packs K/V
    // once per call (the old code rebuilt a hd-wide dot and re-sliced a
    // strided V row n_q×heads×n_k times: 1.0G+ tiny calls/frame at
    // ~115 µs per (row, head) group). Per-element mask/softmax values are
    // unchanged; only the reduction ORDER differs (gates are numeric).
    let run = |s: usize, e: usize| {
        if s >= e {
            return;
        }
        let a = unsafe { std::slice::from_raw_parts_mut(ab as *mut f32, n_q * dim) };
        let rows = e - s;
        let mut scores = vec![0f32; rows * n_k];
        let prof = crate::cpuprof::spans::on();
        for t in 0..heads {
            let tq = prof.then(std::time::Instant::now);
            // C[rows, n_k] = scale · Q[rows, hd] × Kᵀ (K head rows, stride h3)
            crate::qtensor::sgemm_public(
                rows,
                n_k,
                hd,
                scale,
                &qkv_q[s * h3 + t * hd..],
                h3,
                &qkv_k[dim + t * hd..],
                h3,
                true,
                &mut scores,
                n_k,
            );
            match mask {
                Mask::None => {}
                Mask::Key(km) => {
                    for r in 0..rows {
                        let sr = &mut scores[r * n_k..(r + 1) * n_k];
                        for (v, &m) in sr.iter_mut().zip(km) {
                            *v += m;
                        }
                    }
                }
                Mask::Add(am) => {
                    for (r, i) in (s..e).enumerate() {
                        let sr = &mut scores[r * n_k..(r + 1) * n_k];
                        for (v, &m) in sr.iter_mut().zip(&am[i * n_k..(i + 1) * n_k]) {
                            *v += m;
                        }
                    }
                }
                Mask::Head(hm) => {
                    for (r, i) in (s..e).enumerate() {
                        let base = (t * n_q + i) * n_k;
                        let sr = &mut scores[r * n_k..(r + 1) * n_k];
                        for (v, &m) in sr.iter_mut().zip(&hm[base..base + n_k]) {
                            *v += m;
                        }
                    }
                }
            }
            let tm = prof.then(std::time::Instant::now);
            // Stable softmax: the trunk's ViT scores reach ~155, so a
            // raw exp() overflows to inf for the row-max and the sum
            // turns inf → inv = 0 → the whole attention row silently
            // collapses to zeros. Subtract the row max (`mx`).
            for r in 0..rows {
                let sr = &mut scores[r * n_k..(r + 1) * n_k];
                let mut mx = f32::NEG_INFINITY;
                for &v in sr.iter() {
                    if v > mx {
                        mx = v;
                    }
                }
                let mut sum = 0.0f32;
                for v in sr.iter_mut() {
                    *v = (*v - mx).exp();
                    sum += *v;
                }
                let inv = 1.0 / sum;
                for v in sr.iter_mut() {
                    *v *= inv;
                }
            }
            let ta = prof.then(std::time::Instant::now);
            // C[rows, hd] = S[rows, n_k] × V_head (rows stride h3) → out
            crate::qtensor::sgemm_public(
                rows,
                hd,
                n_k,
                1.0,
                &scores,
                n_k,
                &qkv_v[2 * dim + t * hd..],
                h3,
                false,
                &mut a[s * dim + t * hd..],
                dim,
            );
            if let (Some(t0), Some(t1), Some(t2)) = (tq, tm, ta) {
                let t3 = std::time::Instant::now();
                crate::cpuprof::spans::add("attn.qk", t1 - t0);
                crate::cpuprof::spans::add("attn.sm", t2 - t1);
                crate::cpuprof::spans::add("attn.av", t3 - t2);
            }
        }
    };
    match pool { Some(p) => p.run_rows(n_q, &run), None => run(0, n_q) }
    drop(_pc);
    let _pg = crate::cpuprof::spans::span("attn.gemm");
    ATT_CORE_NS.fetch_add(_tc.elapsed().as_nanos() as u64, Ordering::Relaxed);
    let _tw = std::time::Instant::now();
    w_o.matmat(&attn, n_q, out, pool);
    add_bias(out, dim, ob);
    ATT_PROJ_NS.fetch_add(_tw.elapsed().as_nanos() as u64, Ordering::Relaxed);
}

// ─────────── convolutions (NCHW f32, explicit shapes) ───────────

pub(crate) fn conv2d(
    w: &[f32],
    b: &[f32],
    x: &[f32],
    ci: usize,
    h: usize,
    wid: usize,
    kh: usize,
    kw: usize,
    stride: usize,
    pad: usize,
    pool: Option<&Pool>,
) -> Vec<f32> {
    let co = b.len();
    debug_assert_eq!(w.len(), co * ci * kh * kw);
    debug_assert_eq!(x.len(), ci * h * wid);
    let oh = (h + 2 * pad - kh) / stride + 1;
    let ow = (wid + 2 * pad - kw) / stride + 1;
    let _p = crate::cpuprof::spans::span("op.conv2d");
    let mut out = vec![0f32; co * oh * ow];
    let ob = out.as_mut_ptr() as usize;
    // One output ROW per unit instead of one pixel: the (c, dy, dx)
    // accumulation order per element is unchanged (results stay bit-exact),
    // but the innermost loop becomes a contiguous `orow[lo..hi] += w *
    // xrow[..]` slice that auto-vectorizes — the pixel loop fed LLVM a
    // strided load, a div/mod for (oy, ox) and two bounds checks per tap,
    // and ran ~50% of the SAM3 image forward scalar.
    let run = |s: usize, e: usize| {
        let o = unsafe { std::slice::from_raw_parts_mut(ob as *mut f32, co * oh * ow) };
        for r in s..e {
            let c_o = r / oh;
            let oy = r % oh;
            let obase = c_o * oh * ow + oy * ow;
            let orow = &mut o[obase..obase + ow];
            orow.fill(b[c_o]);
            for c in 0..ci {
                for dy in 0..kh {
                    let iy = oy * stride + dy;
                    if iy < pad || iy - pad >= h { continue; }
                    let xrow = &x[(c * h + iy - pad) * wid..][..wid];
                    for dx in 0..kw {
                        let wv = w[((c_o * ci + c) * kh + dy) * kw + dx];
                        // ox with 0 <= ox*stride + dx - pad < wid, 0 <= ox < ow
                        let lo = if pad > dx {
                            (pad - dx).div_ceil(stride)
                        } else {
                            0
                        };
                        if lo >= ow { continue; }
                        let hi_excl = (pad + wid - dx).div_ceil(stride).min(ow);
                        if lo >= hi_excl { continue; }
                        if stride == 1 {
                            // vectorizable: both slices are contiguous
                            let xo = lo + dx - pad;
                            let hi = hi_excl;
                            for (a, b) in orow[lo..hi].iter_mut().zip(&xrow[xo..xo + (hi - lo)]) {
                                *a += wv * *b;
                            }
                        } else {
                            for ox in lo..hi_excl {
                                orow[ox] += wv * xrow[ox * stride + dx - pad];
                            }
                        }
                    }
                }
            }
        }
    };
    match pool { Some(p) => p.run_rows(co * oh, &run), None => run(0, co * oh) }
    out
}

/// ConvTranspose2d(k=2, s=2): w is [ci, co, 2, 2]; x is [ci, h, wid].
pub(crate) fn convt2x2(w: &[f32], b: &[f32], x: &[f32], ci: usize, h: usize, wid: usize, pool: Option<&Pool>) -> Vec<f32> {
    let co = b.len();
    debug_assert_eq!(w.len(), ci * co * 4);
    let (oh, ow) = (h * 2, wid * 2);
    let _p = crate::cpuprof::spans::span("op.convt");
    let mut out = vec![0f32; co * oh * ow];
    let ob = out.as_mut_ptr() as usize;
    // Per output row: (kh, kx) = (oy & 1, ox & 1), so even columns take
    // w[.., 0] and odd columns w[.., 1] — accumulate into two contiguous
    // scratch rows (vectorizable axpy over x, per-channel order kept →
    // bit-exact) and interleave once at the end.
    let run = |s: usize, e: usize| {
        let o = unsafe { std::slice::from_raw_parts_mut(ob as *mut f32, co * oh * ow) };
        let mut t0 = vec![0f32; wid];
        let mut t1 = vec![0f32; wid];
        for r in s..e {
            let c_o = r / oh;
            let oy = r % oh;
            let (slot, iy) = (oy & 1, oy >> 1);
            t0.fill(b[c_o]);
            t1.fill(b[c_o]);
            for c in 0..ci {
                let xrow = &x[(c * h + iy) * wid..][..wid];
                let wb = &w[(c * co + c_o) * 4 + slot * 2..][..2];
                let (w0, w1) = (wb[0], wb[1]);
                for (t, &xv) in t0.iter_mut().zip(xrow) {
                    *t += w0 * xv;
                }
                for (t, &xv) in t1.iter_mut().zip(xrow) {
                    *t += w1 * xv;
                }
            }
            let obase = c_o * oh * ow + oy * ow;
            for j in 0..wid {
                o[obase + 2 * j] = t0[j];
                o[obase + 2 * j + 1] = t1[j];
            }
        }
    };
    match pool { Some(p) => p.run_rows(co * oh, &run), None => run(0, co * oh) }
    out
}

pub(crate) fn maxpool2(x: &[f32], ci: usize, h: usize, wid: usize) -> Vec<f32> {
    let (oh, ow) = (h / 2, wid / 2);
    let mut out = vec![0f32; ci * oh * ow];
    for c in 0..ci {
        for oy in 0..oh {
            for ox in 0..ow {
                let mut mx = f32::NEG_INFINITY;
                for dy in 0..2 {
                    for dx in 0..2 {
                        let v = x[(c * h + oy * 2 + dy) * wid + ox * 2 + dx];
                        if v > mx { mx = v; }
                    }
                }
                out[(c * oh + oy) * ow + ox] = mx;
            }
        }
    }
    out
}

/// Nearest-neighbour upsample of [ci, h, wid] to [ci, sh, sw].
pub(crate) fn upsample_nearest(x: &[f32], ci: usize, h: usize, wid: usize, sh: usize, sw: usize) -> Vec<f32> {
    let mut out = vec![0f32; ci * sh * sw];
    for c in 0..ci {
        for oy in 0..sh {
            let iy = oy * h / sh;
            for ox in 0..sw {
                let ix = ox * wid / sw;
                out[(c * sh + oy) * sw + ox] = x[(c * h + iy) * wid + ix];
            }
        }
    }
    out
}

/// GroupNorm(eps=1e-5) over channels of [c, h, w].
pub(crate) fn group_norm(x: &mut [f32], groups: usize, w: &[f32], b: &[f32], eps: f32) {
    let co = w.len();
    let gsz = co / groups;
    let rest = x.len() / co;
    for g in 0..groups {
        let s = g * gsz * rest;
        let e = s + gsz * rest;
        let mean = x[s..e].iter().sum::<f32>() / (e - s) as f32;
        let var = x[s..e].iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / (e - s) as f32;
        let inv = 1.0 / (var + eps).sqrt();
        for (i, v) in x[s..e].iter_mut().enumerate() {
            let ch = g * gsz + (i / rest);
            *v = (*v - mean) * inv * w[ch] + b[ch];
        }
    }
}

// ─────────── layout helpers ───────────

pub(crate) fn rows_to_chw(rows: &[f32], h: usize, w: usize) -> Vec<f32> {
    let _p = crate::cpuprof::spans::span("op.layout");
    let n = h * w;
    let c = rows.len() / n;
    let mut out = vec![0f32; c * h * w];
    for ch in 0..c {
        for y in 0..h {
            for x in 0..w { out[(ch * h + y) * w + x] = rows[(y * w + x) * c + ch]; }
        }
    }
    out
}

pub(crate) fn chw_to_rows(chw: &[f32], h: usize, w: usize) -> Vec<f32> {
    let _p = crate::cpuprof::spans::span("op.layout");
    let c = chw.len() / (h * w);
    let mut out = vec![0f32; h * w * c];
    for ch in 0..c {
        for y in 0..h {
            for x in 0..w { out[(y * w + x) * c + ch] = chw[(ch * h + y) * w + x]; }
        }
    }
    out
}

/// Normalized sine position encoding for an FPN level: rows [h*w, 2n]
/// (pos_y then pos_x), scale 2π, temperature 10000.
pub(crate) fn sine_pos(h: usize, w: usize, n: usize) -> Vec<f32> {
    let _p = crate::cpuprof::spans::span("op.sine_pos");
    let scale = std::f32::consts::TAU;
    let mut dim_t = vec![0f32; n];
    for i in 0..n { dim_t[i] = 10000f32.powf((2 * (i / 2)) as f32 / n as f32); }
    let mut out = vec![0f32; h * w * n * 2];
    let (hn, wn) = (h as f32, w as f32);
    for y in 0..h {
        let yy = (y + 1) as f32 / (hn + 1e-6) * scale;
        for x in 0..w {
            let xx = (x + 1) as f32 / (wn + 1e-6) * scale;
            let row = (y * w + x) * n * 2;
            for i in 0..n {
                let (py, px) = (yy / dim_t[i], xx / dim_t[i]);
                if i % 2 == 0 {
                    out[row + i] = py.sin();
                    out[row + n + i] = px.sin();
                } else {
                    out[row + i] = py.cos();
                    out[row + n + i] = px.cos();
                }
            }
        }
    }
    out
}

/// Sine box embedding (decoder query_pos input): [N, 4*128] =
/// cat(pos_y, pos_x, pos_w, pos_h) for boxes in (cx, cy, w, h),
/// scale 2π (Sam3SinePositionEmbedding.encode_boxes).
fn sine_boxes(boxes: &[[f32; 4]]) -> Vec<f32> {
    let nf = 128;
    let scale = std::f32::consts::TAU;
    let mut dim_t = vec![0f32; nf];
    for i in 0..nf { dim_t[i] = 10000f32.powf((2 * (i / 2)) as f32 / nf as f32); }
    let mut out = vec![0f32; boxes.len() * 4 * nf];
    for (q, b) in boxes.iter().enumerate() {
        let base = q * 4 * nf;
        for (ci, v) in [b[1], b[0], b[2], b[3]].iter().enumerate() {
            for i in 0..nf {
                let val = v * scale / dim_t[i];
                out[base + ci * nf + i] = if i % 2 == 0 { val.sin() } else { val.cos() };
            }
        }
    }
    out
}

/// Per-head log box-RPB additive bias: [heads, NQ+1, n]. Row 0 (the
/// presence token) is zero — HF pads the query dim from NQ to NQ+1.
fn rpb_bias(refs: &[[f32; 4]], rx: &DecoderMlp, ry: &DecoderMlp, n: usize, pool: Option<&Pool>) -> Result<Vec<f32>> {
    let gh = GRID;
    debug_assert_eq!(n, gh * gh);
    let l2_8 = 8.0f32.ln();
    let mut dy = vec![0f32; refs.len() * gh * 2];
    let mut dx = vec![0f32; refs.len() * gh * 2];
    for (q, r) in refs.iter().enumerate() {
        let x1 = r[0] - 0.5 * r[2];
        let y1 = r[1] - 0.5 * r[3];
        let x2 = r[0] + 0.5 * r[2];
        let y2 = r[1] + 0.5 * r[3];
        for hi in 0..gh {
            let c = hi as f32 / gh as f32;
            dy[(q * gh + hi) * 2] = log_compress(c - y1, l2_8);
            dy[(q * gh + hi) * 2 + 1] = log_compress(c - y2, l2_8);
        }
        for wi in 0..gh {
            let c = wi as f32 / gh as f32;
            dx[(q * gh + wi) * 2] = log_compress(c - x1, l2_8);
            dx[(q * gh + wi) * 2 + 1] = log_compress(c - x2, l2_8);
        }
    }
    let py = ry.forward(&dy, refs.len() * gh, pool, FH);
    let px = rx.forward(&dx, refs.len() * gh, pool, FH);
    let mut out = vec![0f32; FH * (NQ + 1) * n];
    for t in 0..FH {
        for q in 0..refs.len() {
            for hi in 0..gh {
                for wi in 0..gh {
                    out[(t * (NQ + 1) + (q + 1)) * n + hi * gh + wi] =
                        py[(q * gh + hi) * FH + t] + px[(q * gh + wi) * FH + t];
                }
            }
        }
    }
    Ok(out)
}

/// sign(d) * log2(|8d| + 1) / log2(8)  (box_rpb_mode = "log").
fn log_compress(d: f32, l2_8: f32) -> f32 {
    d.signum() * ((8.0 * d).abs() + 1.0).log2() / l2_8
}

/// elementwise a+b over equal-length buffers
fn add_rows(a: &[f32], b: &[f32]) -> Vec<f32> {
    a.iter().zip(b).map(|(x, y)| x + y).collect()
}

/// Finiteness/range stats for one debug stage (env `SAM3_TEXT_DEBUG`).
#[allow(dead_code)]
fn dbg_stats(tag: &str, v: &[f32]) {
    if std::env::var("SAM3_TEXT_DEBUG").is_err() {
        return;
    }
    let (mut bad, mut mn, mut mx) = (0usize, f32::INFINITY, f32::NEG_INFINITY);
    for &x in v {
        if !x.is_finite() {
            bad += 1;
            continue;
        }
        if x < mn { mn = x; }
        if x > mx { mx = x; }
    }
    eprintln!(
        "[sam3] {tag} n={} nonfinite={} min={:.4e} max={:.4e}",
        v.len(), bad, mn, mx
    );
}

/// Raw f32 dump for reference bisection (env `SAM3_DUMP_DIR`).
#[allow(dead_code)]
fn dump_f32(dir: &str, name: &str, v: &[f32]) {
    use std::io::Write;
    let p = std::path::Path::new(dir).join(format!("{name}.f32"));
    let mut b = Vec::with_capacity(v.len() * 4);
    for x in v {
        b.extend_from_slice(&x.to_le_bytes());
    }
    match std::fs::File::create(&p).and_then(|mut f| f.write_all(&b)) {
        Ok(()) => eprintln!("[sam3] dumped {} ({} elems)", p.display(), v.len()),
        Err(e) => eprintln!("[sam3] dump {} failed: {e}", p.display()),
    }
}

// ───────────────────────── tests (gpu feature) ─────────────────────

/// Isolated GPU-vs-CPU check of `run_attn` itself on deterministic
/// pseudo-random inputs — no model, no fixtures: if the owner gate
/// (rel/cos) fails at the Sam3Out level, this says whether the
/// attention GPU arm or another block is the diverger.
#[cfg(all(test, feature = "gpu"))]
mod attn_gpu_tests {
    use super::*;

    fn lcg(state: &mut u64) -> f32 {
        *state = state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        ((*state >> 40) as f32 / (1u64 << 24) as f32) - 0.5
    }

    fn rel_cos(g: &[f32], w: &[f32]) -> (f32, f32) {
        let mut md = 0f32;
        let mut mx = 0f32;
        let (mut dot, mut ng, mut nw) = (0f64, 0f64, 0f64);
        for (&g, &w) in g.iter().zip(w) {
            let d = (g - w).abs();
            if d > md { md = d; }
            if g.abs() > mx { mx = g.abs(); }
            dot += g as f64 * w as f64;
            ng += g as f64 * g as f64;
            nw += w as f64 * w as f64;
        }
        let rel = if mx > 0.0 { md / mx } else { md };
        let cos = if ng > 0.0 && nw > 0.0 {
            (dot / (ng.sqrt() * nw.sqrt())) as f32
        } else {
            1.0
        };
        (rel, cos)
    }

    #[test]
    fn attn_gpu_matches_cpu() {
        if std::env::var("CMF_GPU").as_deref() != Ok("1") {
            eprintln!("skip attn_gpu_matches_cpu: CMF_GPU != 1");
            return;
        }
        if !crate::gpu::backend_available() {
            eprintln!("skip attn_gpu_matches_cpu: no wgpu backend");
            return;
        }
        // Shapes that mirror the real trunk: 1024x1024 (small square),
        // 576x576x1024/16 (windowed layer) and 5184x5184x1024/16 (global
        // layer — the only shape that exercises QK query-chunking).
        let cases: [(usize, usize, usize, usize, bool); 3] = [
            (1024, 1024, 256, 4, true),
            (576, 576, 1024, 16, false),
            (5184, 5184, 1024, 16, false),
        ];
        for &(n_q, n_k, dim, heads, with_head) in &cases {
            let h3 = dim * 3;
            let mut st = 0x9e37_79b9_7f4a_7c15u64;
            let xs: Vec<f32> = (0..n_k * h3).map(|_| lcg(&mut st) * 2.0).collect();
            let mut st2 = 0xdead_beefu64.wrapping_mul(0x5bd1_e995);
            let w_qkv = QTensor::F32 {
                data: (0..h3 * dim).map(|_| lcg(&mut st2)).collect(),
                rows: h3,
                cols: dim,
            };
            let w_o = QTensor::F32 {
                data: (0..dim * dim).map(|_| lcg(&mut st2)).collect(),
                rows: dim,
                cols: dim,
            };
            let qb = vec![0.01f32; h3];
            let ob = vec![0.02f32; dim];
            let head_mask: Vec<f32>;
            let mut head_owned: Vec<f32> = Vec::new();
            if with_head {
                head_owned = vec![0f32; heads * n_q * n_k];
                for (i, v) in head_owned.iter_mut().enumerate() {
                    *v = (i % 97) as f32 * 0.001 - 0.048;
            }
            }
            head_mask = head_owned;
            let mask_kinds: [(Option<usize>, &str); 2] =
                [(None, "none"), (with_head.then_some(0usize), "head")];
            for (kind, label) in mask_kinds {
                if kind.is_none() && !matches!(label, "none") {
                    continue;
                }
                let mask = match kind {
                    Some(_) if with_head => Mask::Head(&head_mask),
                    Some(_) => continue,
                    None => Mask::None,
                };
                if label == "head" && !with_head {
                    continue;
                }
                let mut out_cpu = vec![0f32; n_q * dim];
                let mut out_gpu = vec![0f32; n_q * dim];
                crate::gpu::cpu_scope(|| {
                    run_attn(
                        &w_qkv, &qb, &w_o, &ob, &xs, &xs, &xs,
                        n_q, n_k, dim, heads, None, mask, None, &mut out_cpu,
                    )
                });
                run_attn(
                    &w_qkv, &qb, &w_o, &ob, &xs, &xs, &xs,
                    n_q, n_k, dim, heads, None, mask, None, &mut out_gpu,
                );
                let (rel, cos) = rel_cos(&out_gpu, &out_cpu);
                eprintln!(
                    "attn_gpu_vs_cpu[{n_q}x{n_k}x{dim}/{heads} {label}] rel={rel:.3e} cos={cos:.6}"
                );
                assert!(
                    rel <= 1e-2 && cos >= 0.999,
                    "attn {n_q}x{n_k} mask={label}: rel={rel:.3e} cos={cos:.6} outside gate"
                );
            }
        }
    }
}