cortiq-engine 0.7.0

//! Cross-platform GPU backend (C1): wgpu → Vulkan / DX12 / Metal
//! (NVIDIA, AMD Radeon, Intel Arc, Apple). Implements the same contract as
//! `gpu_metal.rs`, behind the `gpu.rs` facade — runtime call-sites do not change.
//!
//! Difference from the Metal path: a discrete card has no unified memory, so
//! the quantized weights are LOADED into VRAM ONCE (residency cache keyed by
//! tensor index) — that is where the win lives (VRAM bandwidth ×5–10 vs CPU). The math
//! is identical to CPU/Metal: y[o] = row_scale[o]·Σ q[o,i]·xs[i], where xs is already
//! prescaled by the column scale (the two-scale q8_2f folds into the input prescale).
//!
//! Enabling: `CMF_GPU=wgpu` (or `=1` on non-macOS, where wgpu is the only backend).
//! Any init/limit failure — `false` and an honest CPU path.

use crate::gpu::{BatchJob, MoeJob};
use cortiq_core::CmfModel;
use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};
use wgpu::util::DeviceExt;

/// Workgroup limit per dimension (WebGPU minimum; lm_head has more
/// rows — we use grid-stride in the shader).
const MAX_WG: u32 = 65_535;

/// Subgroup-accelerated MoE select: the top-k rounds ride subgroupMax /
/// subgroupMin (barrier-free within a subgroup) — two barriers per slot
/// against the tree version's eight. Lives in its OWN module: `enable
/// subgroups` fails validation on devices without the feature, and one
/// invalid function kills every entry point of a module (0.5.40's Metal
/// lesson, re-learned on Vulkan this afternoon).
// The subgroup-reduction decode matvec, in its OWN module: `enable
// subgroups` must never reach a device without the feature (the same
// rule the MoE select follows). Bindings mirror the q4tp module
// verbatim so the existing bind groups fit unchanged.
const MV_SG_SRC: &str = r#"
enable subgroups;

struct Q1Params { np: u32, rows: u32, _p0: u32, _p1: u32 };
@group(0) @binding(0) var<storage, read>       q1w : array<u32>;
@group(0) @binding(1) var<storage, read>       q1x : array<f32>;
@group(0) @binding(2) var<storage, read_write> q1y : array<f32>;
@group(0) @binding(3) var<uniform>             q1p : Q1Params;
@group(0) @binding(4) var<storage, read> q4v_w : array<vec4<u32>>;
@group(0) @binding(5) var<storage, read> q4v_x : array<vec4<f32>>;

var<workgroup> lad_q4v: array<f32, 256>;
var<workgroup> sg_part: array<f32, 16>;

fn q4tp_byte(off: u32) -> u32 {
    return (q1w[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
}
// A nibble as f32 WITHOUT an integer-to-float conversion: OR it into the
// mantissa of 2^23 and subtract 2^23 + 8. Exact — n − 8 for every n in
// 0..16 — so each kernel's arithmetic is unchanged to the bit; what
// changes is the instruction mix. I2F issues at 1/8 of the FMA rate on
// NVIDIA and there were eight of them per weight word, more than the
// eight FMAs the word is for. `CMF_MAGIC_UNPACK=0` restores the
// conversions for an A/B (the Rust side swaps the bodies).
fn q4v_nib(w: u32, sh: u32) -> f32 {
    return bitcast<f32>(((w >> sh) & 0xFu) | 0x4B000000u) - 8388616.0;
}
fn q4v_dot8(w: u32, a: vec4<f32>, b: vec4<f32>) -> f32 {
    return q4v_nib(w, 0u) * a.x
         + q4v_nib(w, 4u) * a.y
         + q4v_nib(w, 8u) * a.z
         + q4v_nib(w, 12u) * a.w
         + q4v_nib(w, 16u) * b.x
         + q4v_nib(w, 20u) * b.y
         + q4v_nib(w, 24u) * b.z
         + q4v_nib(w, 28u) * b.w;
}

@compute @workgroup_size(256)
fn q4tp_matvec4_sg(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(num_workgroups) nwg: vec3<u32>,
                @builtin(local_invocation_index) lid: u32,
                @builtin(subgroup_invocation_id) sg_inv: u32,
                @builtin(subgroup_size) ssz: u32) {
    // Plain `enable subgroups` has no subgroup_id builtin — that is the
    // cooperative-matrix extension's gift, and asking for it here is
    // the validation error that took the whole device down (the init
    // path had no error scope; it does now). Linear tiling makes the
    // index arithmetic: lanes [i*ssz, (i+1)*ssz) are subgroup i.
    let sg_id = lid / ssz;
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    // 8 rows per workgroup, register-blocked in pairs: sub-block `sub` owns
    // rows base+sub and base+sub+4, and every x vec4 fetched for a group
    // feeds BOTH rows' dot chains — the x side of the LSU load nearly
    // halves. Each row's group order and add order stay those of the
    // one-row kernel.
    // `_p0`: how many activation vectors share this weight. One is a matvec;
    // more is a batch. The BATCH is the fast axis of the dispatch, so the
    // workgroups that read the same weight rows are neighbours and meet in
    // L2; walking the whole output space instead put them `rows/16` apart,
    // which streams the weight once per batch element and defeats the point.
    // Reuse is still L2's to give — this is not a register-blocked B kernel —
    // so the win is a measurement, not a claim.
    let nb = max(q1p._p0, 1u);
    let blocks = (rows + 7u) / 8u;
    var wb = wid.x;
    loop {
        if (wb >= blocks * nb) { break; }
        let bi = wb % nb;
        let base = (wb / nb) * 8u;
        let bofs = bi * rows;
        {
            let r = base + (lid >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4v[lid] = exp2(pr.x + f32(lid & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let wrow_a = base + sub;
        let wrow_b = base + sub + 4u;
        let row_a = bofs + wrow_a;
        let row_b = bofs + wrow_b;
        let live_a = wrow_a < rows;
        let live_b = wrow_b < rows;
        // In vec4 units: (row / lora) * gpr * 32 floats.
        var xblk = 0u;
        if (nb > 1u) { xblk = bi * gpr * 8u; }
        else if (q1p._p1 > 0u) { xblk = (wrow_a / q1p._p1) * gpr * 8u; }
        var acc_a = 0.0;
        var acc_b = 0.0;
        if (live_a) {
            let crow_a = codes_b + wrow_a * cstride;
            let crow_b = codes_b + wrow_b * cstride;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv_a = q4tp_byte(crow_a + cbo);
                if (sh > 3u) { cv_a = cv_a | (q4tp_byte(crow_a + cbo + 1u) << 8u); }
                let v_a = q4v_w[wrow_a * gpr + g];
                // `_p1` is the low-rank group width. Set, it slides the
                // activation window with the row — which is the ONLY thing
                // the grouped output projection does differently, and the
                // reason it had a kernel of its own reading 3.82 ms against
                // this one's 1.24 on comparable weights. Rows base..base+7
                // share a window whenever the width is a multiple of 8, and
                // the caller only takes this path then.
                let xq = xblk + g * 8u;
                let x0 = q4v_x[xq];      let x1 = q4v_x[xq + 1u];
                let x2 = q4v_x[xq + 2u]; let x3 = q4v_x[xq + 3u];
                let x4 = q4v_x[xq + 4u]; let x5 = q4v_x[xq + 5u];
                let x6 = q4v_x[xq + 6u]; let x7 = q4v_x[xq + 7u];
                let sa = lad_q4v[(sub << 5u) + ((cv_a >> sh) & 31u)];
                acc_a = acc_a + sa
                    * (q4v_dot8(v_a.x, x0, x1) + q4v_dot8(v_a.y, x2, x3)
                     + q4v_dot8(v_a.z, x4, x5) + q4v_dot8(v_a.w, x6, x7));
                if (live_b) {
                    var cv_b = q4tp_byte(crow_b + cbo);
                    if (sh > 3u) { cv_b = cv_b | (q4tp_byte(crow_b + cbo + 1u) << 8u); }
                    let v_b = q4v_w[wrow_b * gpr + g];
                    let sb = lad_q4v[128u + (sub << 5u) + ((cv_b >> sh) & 31u)];
                    acc_b = acc_b + sb
                        * (q4v_dot8(v_b.x, x0, x1) + q4v_dot8(v_b.y, x2, x3)
                         + q4v_dot8(v_b.z, x4, x5) + q4v_dot8(v_b.w, x6, x7));
                }
                g = g + 64u;
            }
        }
        // Subgroup reduction: 64 lanes of a sub-block are exactly two
        // 32-wide subgroups. One barrier replaces the tree's eight —
        // which the bandwidth test convicted as the missing third of
        // the bus (stream 1570 GB/s, this kernel's pattern 1623, the
        // kernel itself ~1000).
        // Keep binding 1 in the auto layout: an entry point only owns
        // the bindings it READS, and a five-entry layout against the
        // caller's six-entry bind group is a validation error per
        // dispatch — which fell this kernel back to the CPU at 5.3
        // tok/s while producing bit-identical (host) output. The qknorm
        // module documents the same trap.
        if (q1p._p0 == 0xFFFFFFFFu) { q1y[0] = q1x[0]; }
        let sga = subgroupAdd(acc_a);
        let sgb = subgroupAdd(acc_b);
        if (sg_inv == 0u) {
            sg_part[sg_id * 2u] = sga;
            sg_part[sg_id * 2u + 1u] = sgb;
        }
        workgroupBarrier();
        if (l == 0u) {
            // Sub-block = 64 lanes = 64/ssz subgroups; sum whatever the
            // width made, so a 64-wide device stays correct.
            let per = 64u / ssz;
            let base_sg = lid / ssz;
            var suma = 0.0;
            var sumb = 0.0;
            var k = 0u;
            loop {
                if (k >= per) { break; }
                suma = suma + sg_part[(base_sg + k) * 2u];
                sumb = sumb + sg_part[(base_sg + k) * 2u + 1u];
                k = k + 1u;
            }
            if (wrow_a < rows) { q1y[row_a] = suma; }
            if (wrow_b < rows) { q1y[row_b] = sumb; }
        }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}
"#;

/// Optional subgroup reduction for the resident q2tp 16-row matvec.  This is
/// intentionally a separate module: adapters without SUBGROUP must retain
/// the validated tree-reduction pipeline, and a validation error in this
/// shader must never poison the ordinary q2tp module.  The workgroup still
/// owns four 64-lane row groups; subgroup leaders write partial vec4 sums and
/// one lane per row group combines them.  The CPU fallback is selected unless
/// CMF_Q2TP_SG=1 explicitly opts into the A/B after the probe records
/// CMF_Q2TP_SG_WIDTH=(32|64) and CMF_Q2TP_SG_LINEAR=1.
const Q2TP_SG_SRC: &str = r#"
struct Q1Params { np: u32, rows: u32, _p0: u32, _p1: u32 };
@group(0) @binding(0) var<storage, read>       q1w : array<u32>;
@group(0) @binding(2) var<storage, read_write> q1y : array<f32>;
@group(0) @binding(3) var<uniform>             q1p : Q1Params;
@group(0) @binding(5) var<storage, read>       q4v_x : array<vec4<f32>>;

var<workgroup> lad_q4w: array<f32, 512>;
// Four f32 components for at most eight subgroups (256 lanes / 32).
var<workgroup> sg_part: array<f32, 32>;

fn q4tp_byte(off: u32) -> u32 {
    return (q1w[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
}

fn q2v_c2(w: u32, sh: u32, affine: u32) -> f32 {
    return bitcast<f32>((((w >> sh) & 3u) << 1u) | 0x4B000000u)
        - select(8388611.0, 8388610.0, affine != 0u);
}

fn q2v_d16(w: u32, a: vec4<f32>, b: vec4<f32>, c: vec4<f32>, d: vec4<f32>, affine: u32) -> f32 {
    return (q2v_c2(w, 0u, affine) * a.x
         + q2v_c2(w, 2u, affine) * a.y
         + q2v_c2(w, 4u, affine) * a.z
         + q2v_c2(w, 6u, affine) * a.w
         + q2v_c2(w, 8u, affine) * b.x
         + q2v_c2(w, 10u, affine) * b.y
         + q2v_c2(w, 12u, affine) * b.z
         + q2v_c2(w, 14u, affine) * b.w
         + q2v_c2(w, 16u, affine) * c.x
         + q2v_c2(w, 18u, affine) * c.y
         + q2v_c2(w, 20u, affine) * c.z
         + q2v_c2(w, 22u, affine) * c.w
         + q2v_c2(w, 24u, affine) * d.x
         + q2v_c2(w, 26u, affine) * d.y
         + q2v_c2(w, 28u, affine) * d.z
         + q2v_c2(w, 30u, affine) * d.w) * 0.5;
}

@compute @workgroup_size(256)
fn q2tp_matvec16w_sg(@builtin(workgroup_id) wid: vec3<u32>,
                      @builtin(num_workgroups) nwg: vec3<u32>,
                      @builtin(local_invocation_index) lid: u32,
                      @builtin(subgroup_invocation_id) sg_inv: u32,
                      @builtin(subgroup_size) ssz: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 2u;
    let codes_b = rows * gpr * 8u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sg_id = lid / ssz;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let blocks = (rows + 15u) / 16u;
    var wb = wid.x;
    loop {
        if (wb >= blocks) { break; }
        let base = wb * 16u;
        for (var t = lid; t < 512u; t = t + 256u) {
            let r = base + (t >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                let rung = t & 31u;
                lad_q4w[t] = select(
                    exp2(pr.x + f32(max(rung, 1u) - 1u) * pr.y),
                    0.0,
                    rung == 0u,
                );
            }
        }
        workgroupBarrier();
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        var acc = vec4<f32>(0.0);
        if (r0 < rows) {
            let c0 = codes_b + r0 * cstride;
            let c1 = codes_b + r1 * cstride;
            let c2 = codes_b + r2 * cstride;
            let c3 = codes_b + r3 * cstride;
            let l1 = r1 < rows;
            let l2 = r2 < rows;
            let l3 = r3 < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                let x0 = g * 8u;
                let xa = q4v_x[x0];      let xb = q4v_x[x0 + 1u];
                let xc = q4v_x[x0 + 2u]; let xd = q4v_x[x0 + 3u];
                let xe = q4v_x[x0 + 4u]; let xf = q4v_x[x0 + 5u];
                let xg = q4v_x[x0 + 6u]; let xh = q4v_x[x0 + 7u];
                var cv = q4tp_byte(c0 + cbo);
                if (sh > 3u) { cv = cv | (q4tp_byte(c0 + cbo + 1u) << 8u); }
                var wi = (r0 * gpr + g) * 2u;
                acc.x = acc.x + lad_q4w[(sub << 5u) + ((cv >> sh) & 31u)]
                    * (q2v_d16(q1w[wi], xa, xb, xc, xd, q1p._p1)
                        + q2v_d16(q1w[wi + 1u], xe, xf, xg, xh, q1p._p1));
                if (l1) {
                    cv = q4tp_byte(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c1 + cbo + 1u) << 8u); }
                    wi = (r1 * gpr + g) * 2u;
                    acc.y = acc.y + lad_q4w[128u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q2v_d16(q1w[wi], xa, xb, xc, xd, q1p._p1)
                            + q2v_d16(q1w[wi + 1u], xe, xf, xg, xh, q1p._p1));
                }
                if (l2) {
                    cv = q4tp_byte(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c2 + cbo + 1u) << 8u); }
                    wi = (r2 * gpr + g) * 2u;
                    acc.z = acc.z + lad_q4w[256u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q2v_d16(q1w[wi], xa, xb, xc, xd, q1p._p1)
                            + q2v_d16(q1w[wi + 1u], xe, xf, xg, xh, q1p._p1));
                }
                if (l3) {
                    cv = q4tp_byte(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c3 + cbo + 1u) << 8u); }
                    wi = (r3 * gpr + g) * 2u;
                    acc.w = acc.w + lad_q4w[384u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q2v_d16(q1w[wi], xa, xb, xc, xd, q1p._p1)
                            + q2v_d16(q1w[wi + 1u], xe, xf, xg, xh, q1p._p1));
                }
                g = g + 64u;
            }
        }
        // The explicit path is admitted only after the host has observed a
        // 32/64-lane linear subgroup.  Keep all lanes participating so the
        // row tails and invalid rows reduce to zero deterministically.
        let sa = subgroupAdd(acc.x);
        let sb = subgroupAdd(acc.y);
        let sc = subgroupAdd(acc.z);
        let sd = subgroupAdd(acc.w);
        if (sg_inv == 0u) {
            let off = sg_id * 4u;
            sg_part[off] = sa;
            sg_part[off + 1u] = sb;
            sg_part[off + 2u] = sc;
            sg_part[off + 3u] = sd;
        }
        workgroupBarrier();
        if (l == 0u) {
            let per = 64u / ssz;
            let first = (lid >> 6u) * per;
            var a = 0.0; var b = 0.0; var c = 0.0; var d = 0.0;
            for (var k = 0u; k < per; k = k + 1u) {
                let off = (first + k) * 4u;
                a = a + sg_part[off]; b = b + sg_part[off + 1u];
                c = c + sg_part[off + 2u]; d = d + sg_part[off + 3u];
            }
            if (r0 < rows) { q1y[r0] = a; }
            if (r1 < rows) { q1y[r1] = b; }
            if (r2 < rows) { q1y[r2] = c; }
            if (r3 < rows) { q1y[r3] = d; }
        }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}
"#;

const SELECT_SG_SRC: &str = r#"
struct MoeSelP { n_exp: u32, top_k: u32, norm: u32, pk: u32, scale: f32, _s0: u32, _s1: u32, _s2: u32 };
@group(0) @binding(0) var<storage, read>       sg_logit : array<f32>;
@group(0) @binding(1) var<storage, read>       sg_slog  : array<f32>;
@group(0) @binding(2) var<storage, read_write> sg_sel   : array<u32>;
@group(0) @binding(3) var<storage, read_write> sg_w     : array<f32>;
@group(0) @binding(4) var<uniform>             sg_p     : MoeSelP;
@group(0) @binding(5) var<storage, read>       sg_sgw   : array<u32>;
@group(0) @binding(6) var<storage, read>       sg_x     : array<f32>;

var<workgroup> sgm_lg:  array<f32, 256>;
var<workgroup> sgm_red: array<f32, 256>;
var<workgroup> sgm_pv:  array<f32, 8>;
var<workgroup> sgm_pi:  array<u32, 8>;
var<workgroup> sgm_gate: f32;

@compute @workgroup_size(256)
fn moe_select_sg(@builtin(local_invocation_index) lid: u32,
                 @builtin(subgroup_invocation_id) sl: u32,
                 @builtin(subgroup_size) ssz: u32) {
    let sgid = lid / ssz;
    let n = sg_p.n_exp;
    let sg_kind = sg_p.pk & 0xFFu;
    let sg_hidden = sg_p.pk >> 8u;
    // shared-expert gate (same math as the tree kernel)
    if (sg_kind == 4u) {
        var d = 0.0;
        var i = lid;
        loop {
            if (i >= sg_hidden) { break; }
            d = d + bitcast<f32>(sg_sgw[i]) * sg_x[i];
            i = i + 256u;
        }
        sgm_red[lid] = d;
        workgroupBarrier();
        var st = 128u;
        loop {
            if (st == 0u) { break; }
            if (lid < st) { sgm_red[lid] = sgm_red[lid] + sgm_red[lid + st]; }
            workgroupBarrier();
            st = st >> 1u;
        }
        if (lid == 0u) { sgm_gate = sgm_red[0]; }
    } else {
        if (lid == 0u) { sgm_gate = sg_slog[0]; }
    }
    workgroupBarrier();
    var v = -3.0e38;
    if (lid < n) { v = sg_logit[lid]; }
    sgm_lg[lid] = v;
    // global max + softmax denom (subgroup sums, one barrier each)
    let m1 = subgroupMax(v);
    if (sl == 0u) { sgm_red[sgid] = m1; }
    workgroupBarrier();
    var mx = -3.0e38;
    if (lid < 8u) { mx = sgm_red[lid]; }
    mx = subgroupMax(mx);
    mx = subgroupBroadcast(mx, 0u);
    if (lid == 0u) { sgm_red[255] = mx; }
    workgroupBarrier();
    mx = sgm_red[255];
    let ev = select(0.0, exp(v - mx), lid < n);
    let s1 = subgroupAdd(ev);
    if (sl == 0u) { sgm_red[sgid] = s1; }
    workgroupBarrier();
    var denom = 0.0;
    if (lid < 8u) { denom = sgm_red[lid]; }
    denom = subgroupAdd(denom);
    denom = subgroupBroadcast(denom, 0u);
    if (lid == 0u) { sgm_red[254] = denom; }
    workgroupBarrier();
    denom = sgm_red[254];
    // top-k rounds: subgroup argmax (value then lowest index), then an
    // 8-wide final in subgroup 0.
    let k = sg_p.top_k;
    var wsum = 0.0;
    for (var slot = 0u; slot < k; slot = slot + 1u) {
        let lv = sgm_lg[lid];
        let sm = subgroupMax(lv);
        let cand = select(0xFFFFFFFFu, lid, lv == sm);
        let si = subgroupMin(cand);
        if (sl == 0u) {
            sgm_pv[sgid] = sm;
            sgm_pi[sgid] = si;
        }
        workgroupBarrier();
        if (sgid == 0u) {
            var pv = -3.0e38;
            var pi = 0xFFFFFFFFu;
            if (sl < 8u) {
                pv = sgm_pv[sl];
                pi = sgm_pi[sl];
            }
            let bm = subgroupMax(pv);
            let bc = select(0xFFFFFFFFu, pi, pv == bm);
            let bi = subgroupMin(bc);
            if (sl == 0u) {
                sgm_pv[0] = bm;
                sgm_pi[0] = bi;
            }
        }
        workgroupBarrier();
        let bi = sgm_pi[0];
        let w = exp(sgm_pv[0] - mx) / denom;
        if (lid == 0u) {
            sg_sel[slot] = bi;
            sg_w[slot] = w;
        }
        wsum = wsum + w;
        if (lid == bi) { sgm_lg[lid] = -3.0e38; }
        workgroupBarrier();
    }
    if (lid == 0u) {
        if (sg_p.norm != 0u) {
            for (var slot = 0u; slot < k; slot = slot + 1u) { sg_w[slot] = sg_w[slot] / wsum; }
        }
        sg_sel[k] = n;
        sg_w[k] = 1.0 / (1.0 + exp(-sgm_gate));
    }
}
"#;

/// Q4TP MoE over a genuine model-wide `(layer, expert)` slot pool.
///
/// Vulkan limits one storage-buffer binding to 4 GiB even on cards with far
/// more VRAM. The cache is therefore split into eight equal physical
/// segments, exposed to WGSL as binding arrays. `msel` remains one flat slot
/// id; the shader alone resolves `(segment, local slot)`, so experts selected
/// for one layer may live anywhere in the common pool and still execute in a
/// single gate/up pass and a single down pass. Parameters live in group 1:
/// wgpu forbids a uniform binding in a group that contains a binding array.
const DSV4_GLOBAL_MOE_SEGMENTS: usize = 8;
const DSV4_GLOBAL_MOE_SEGMENTS_S16: usize = 16;
const DSV4_GLOBAL_MOE_SRC: &str = r#"
enable wgpu_binding_array;

struct WordBank { words: array<u32> };
struct GGuP {
    gpr: u32, inter: u32, slots: u32, mat16: u32,
    lim: f32, segment_slots: u32, _p0: u32, _p1: u32,
};
@group(0) @binding(0) var<storage, read>       gg_gw  : binding_array<WordBank, 8>;
@group(0) @binding(1) var<storage, read>       gg_uw  : binding_array<WordBank, 8>;
@group(0) @binding(2) var<storage, read>       gg_x   : array<f32>;
@group(0) @binding(3) var<storage, read>       gg_sel : array<u32>;
@group(0) @binding(4) var<storage, read_write> gg_act : array<f32>;
// V4.1 applies the final route weight before its BF16 down-projection
// input cast.  Generic DSV4 leaves this binding unused and keeps the old
// f32 activation path.
@group(0) @binding(5) var<storage, read>       gg_wt  : array<f32>;
@group(1) @binding(0) var<uniform>             gg_p   : GGuP;
var<workgroup> gg_pg: array<f32, 64>;
var<workgroup> gg_pu: array<f32, 64>;

// WGSL has no portable BF16 storage scalar.  This is the same round-to-nearest
// even conversion used by the CPU V4.1 path, represented as an f32 value so
// the following quantized dot consumes the exact rounded value.
fn gg_bf16(v: f32) -> f32 {
    let bits = bitcast<u32>(v);
    if ((bits & 0x7F800000u) == 0x7F800000u) { return v; }
    let rounded = bits + 0x7FFFu + ((bits >> 16u) & 1u);
    return bitcast<f32>(rounded & 0xFFFF0000u);
}

fn gg_g32(seg: u32, o: u32) -> u32 { return gg_gw[seg].words[o]; }
fn gg_u32(seg: u32, o: u32) -> u32 { return gg_uw[seg].words[o]; }
fn gg_g16(seg: u32, o: u32) -> u32 {
    return (gg_g32(seg, o >> 1u) >> ((o & 1u) * 16u)) & 0xFFFFu;
}
fn gg_u16(seg: u32, o: u32) -> u32 {
    return (gg_u32(seg, o >> 1u) >> ((o & 1u) * 16u)) & 0xFFFFu;
}
fn gg_g8(seg: u32, o: u32) -> u32 {
    return (gg_g32(seg, o >> 2u) >> ((o & 3u) * 8u)) & 0xFFu;
}
fn gg_u8(seg: u32, o: u32) -> u32 {
    return (gg_u32(seg, o >> 2u) >> ((o & 3u) * 8u)) & 0xFFu;
}
fn gg_dot8(w: u32, xi: u32) -> f32 {
    return (f32(w & 0xFu) - 8.0) * gg_x[xi]
         + (f32((w >> 4u) & 0xFu) - 8.0) * gg_x[xi + 1u]
         + (f32((w >> 8u) & 0xFu) - 8.0) * gg_x[xi + 2u]
         + (f32((w >> 12u) & 0xFu) - 8.0) * gg_x[xi + 3u]
         + (f32((w >> 16u) & 0xFu) - 8.0) * gg_x[xi + 4u]
         + (f32((w >> 20u) & 0xFu) - 8.0) * gg_x[xi + 5u]
         + (f32((w >> 24u) & 0xFu) - 8.0) * gg_x[xi + 6u]
         + (f32((w >> 28u) & 0xFu) - 8.0) * gg_x[xi + 7u];
}

fn gg_dot16_q2(w: u32, xi: u32) -> f32 {
    return (f32(w & 3u) - 1.5) * gg_x[xi]
         + (f32((w >> 2u) & 3u) - 1.5) * gg_x[xi + 1u]
         + (f32((w >> 4u) & 3u) - 1.5) * gg_x[xi + 2u]
         + (f32((w >> 6u) & 3u) - 1.5) * gg_x[xi + 3u]
         + (f32((w >> 8u) & 3u) - 1.5) * gg_x[xi + 4u]
         + (f32((w >> 10u) & 3u) - 1.5) * gg_x[xi + 5u]
         + (f32((w >> 12u) & 3u) - 1.5) * gg_x[xi + 6u]
         + (f32((w >> 14u) & 3u) - 1.5) * gg_x[xi + 7u]
         + (f32((w >> 16u) & 3u) - 1.5) * gg_x[xi + 8u]
         + (f32((w >> 18u) & 3u) - 1.5) * gg_x[xi + 9u]
         + (f32((w >> 20u) & 3u) - 1.5) * gg_x[xi + 10u]
         + (f32((w >> 22u) & 3u) - 1.5) * gg_x[xi + 11u]
         + (f32((w >> 24u) & 3u) - 1.5) * gg_x[xi + 12u]
         + (f32((w >> 26u) & 3u) - 1.5) * gg_x[xi + 13u]
         + (f32((w >> 28u) & 3u) - 1.5) * gg_x[xi + 14u]
         + (f32((w >> 30u) & 3u) - 1.5) * gg_x[xi + 15u];
}

@compute @workgroup_size(64)
fn dsv4_global_gate_up_q4tp(@builtin(workgroup_id) wid: vec3<u32>,
                            @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let slot = wid.y;
    let batch = wid.z;
    let bslot = batch * gg_p.slots + slot;
    let flat = gg_sel[bslot];
    let seg = flat / gg_p.segment_slots;
    let local = flat - seg * gg_p.segment_slots;
    let gpr = gg_p.gpr;
    let rows = gg_p.inter;
    let base16 = local * gg_p.mat16;
    let nib16 = base16 + row * gpr * 8u;
    let par16 = base16 + rows * gpr * 8u + row * 2u;
    let cst = (gpr * 5u + 7u) / 8u;
    let cod8 = (base16 + rows * gpr * 8u + rows * 2u) * 2u + row * cst;
    let gl = unpack2x16float(gg_g16(seg, par16) | (gg_g16(seg, par16 + 1u) << 16u));
    let ul = unpack2x16float(gg_u16(seg, par16) | (gg_u16(seg, par16 + 1u) << 16u));
    var ag = 0.0;
    var au = 0.0;
    for (var g = lid; g < gpr; g = g + 64u) {
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cg = gg_g8(seg, cod8 + cb);
        var cu = gg_u8(seg, cod8 + cb);
        if (shf > 3u) {
            cg = cg | (gg_g8(seg, cod8 + cb + 1u) << 8u);
            cu = cu | (gg_u8(seg, cod8 + cb + 1u) << 8u);
        }
        let sg = exp2(gl.x + f32((cg >> shf) & 31u) * gl.y);
        let su = exp2(ul.x + f32((cu >> shf) & 31u) * ul.y);
        let t16 = nib16 + g * 8u;
        let xb = batch * gpr * 32u + g * 32u;
        var dg = 0.0;
        var du = 0.0;
        for (var k = 0u; k < 4u; k = k + 1u) {
            let wg = gg_g16(seg, t16 + 2u * k)
                   | (gg_g16(seg, t16 + 1u + 2u * k) << 16u);
            let wu = gg_u16(seg, t16 + 2u * k)
                   | (gg_u16(seg, t16 + 1u + 2u * k) << 16u);
            dg = dg + gg_dot8(wg, xb + 8u * k);
            du = du + gg_dot8(wu, xb + 8u * k);
        }
        ag = ag + sg * dg;
        au = au + su * du;
    }
    gg_pg[lid] = ag;
    gg_pu[lid] = au;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) {
            gg_pg[lid] = gg_pg[lid] + gg_pg[lid + stride];
            gg_pu[lid] = gg_pu[lid] + gg_pu[lid + stride];
        }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        var gate = gg_pg[0];
        var up = gg_pu[0];
        if (gg_p._p0 != 0u) {
            gate = gg_bf16(gate);
            up = gg_bf16(up);
        }
        if (gg_p.lim > 0.0) {
            up = clamp(up, -gg_p.lim, gg_p.lim);
            gate = min(gate, gg_p.lim);
        }
        var act = (gate / (1.0 + exp(-gate))) * up;
        if (gg_p._p0 != 0u) {
            // The reference multiplies by the selected route weight before
            // converting the down input back to the model BF16 dtype.  The
            // down shader therefore skips its normal weight multiplication.
            act = gg_bf16(gg_wt[bslot] * act);
        }
        gg_act[(bslot * gg_p.inter) + row] = act;
    }
}

// Mixed q2tp/q4tp profile: gate and up keep the exact q2tp layout while
// down remains q4tp. The descriptor-indexed global pool changes only the
// physical base address; its ladder, exact-zero rung and add order match the
// parity-proven local `moe_gate_up_q2tp` kernel above.
@compute @workgroup_size(64)
fn dsv4_global_gate_up_q2tp(@builtin(workgroup_id) wid: vec3<u32>,
                            @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let slot = wid.y;
    let batch = wid.z;
    let bslot = batch * gg_p.slots + slot;
    let flat = gg_sel[bslot];
    let seg = flat / gg_p.segment_slots;
    let local = flat - seg * gg_p.segment_slots;
    let gpr = gg_p.gpr;
    let rows = gg_p.inter;
    let base16 = local * gg_p.mat16;
    let plane16 = base16 + row * gpr * 4u;
    let par16 = base16 + rows * gpr * 4u + row * 2u;
    let cst = (gpr * 5u + 7u) / 8u;
    let cod8 = (base16 + rows * gpr * 4u + rows * 2u) * 2u + row * cst;
    let gl = unpack2x16float(gg_g16(seg, par16) | (gg_g16(seg, par16 + 1u) << 16u));
    let ul = unpack2x16float(gg_u16(seg, par16) | (gg_u16(seg, par16 + 1u) << 16u));
    var ag = 0.0;
    var au = 0.0;
    for (var g = lid; g < gpr; g = g + 64u) {
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cg = gg_g8(seg, cod8 + cb);
        var cu = gg_u8(seg, cod8 + cb);
        if (shf > 3u) {
            cg = cg | (gg_g8(seg, cod8 + cb + 1u) << 8u);
            cu = cu | (gg_u8(seg, cod8 + cb + 1u) << 8u);
        }
        let cgv = (cg >> shf) & 31u;
        let cuv = (cu >> shf) & 31u;
        let sg = select(exp2(gl.x + f32(max(cgv, 1u) - 1u) * gl.y), 0.0, cgv == 0u);
        let su = select(exp2(ul.x + f32(max(cuv, 1u) - 1u) * ul.y), 0.0, cuv == 0u);
        let w32 = (plane16 + g * 4u) >> 1u;
        let xb = batch * gpr * 32u + g * 32u;
        let dg = gg_dot16_q2(gg_g32(seg, w32), xb)
               + gg_dot16_q2(gg_g32(seg, w32 + 1u), xb + 16u);
        let du = gg_dot16_q2(gg_u32(seg, w32), xb)
               + gg_dot16_q2(gg_u32(seg, w32 + 1u), xb + 16u);
        ag = ag + sg * dg;
        au = au + su * du;
    }
    gg_pg[lid] = ag;
    gg_pu[lid] = au;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) {
            gg_pg[lid] = gg_pg[lid] + gg_pg[lid + stride];
            gg_pu[lid] = gg_pu[lid] + gg_pu[lid + stride];
        }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        var gate = gg_pg[0];
        var up = gg_pu[0];
        if (gg_p._p0 != 0u) {
            gate = gg_bf16(gate);
            up = gg_bf16(up);
        }
        if (gg_p.lim > 0.0) {
            up = clamp(up, -gg_p.lim, gg_p.lim);
            gate = min(gate, gg_p.lim);
        }
        var act = (gate / (1.0 + exp(-gate))) * up;
        if (gg_p._p0 != 0u) {
            act = gg_bf16(gg_wt[bslot] * act);
        }
        gg_act[(bslot * gg_p.inter) + row] = act;
    }
}

struct GDnP {
    gpr: u32, hidden: u32, slots: u32, mat16: u32,
    segment_slots: u32, _p0: u32, _p1: u32, _p2: u32,
};
@group(0) @binding(0) var<storage, read>       gd_w   : binding_array<WordBank, 8>;
@group(0) @binding(1) var<storage, read>       gd_act : array<f32>;
@group(0) @binding(2) var<storage, read>       gd_sel : array<u32>;
@group(0) @binding(3) var<storage, read>       gd_wt  : array<f32>;
@group(0) @binding(4) var<storage, read_write> gd_y   : array<f32>;
@group(1) @binding(0) var<uniform>             gd_p   : GDnP;
var<workgroup> gd_pt: array<f32, 64>;

fn gd_bf16(v: f32) -> f32 {
    let bits = bitcast<u32>(v);
    if ((bits & 0x7F800000u) == 0x7F800000u) { return v; }
    let rounded = bits + 0x7FFFu + ((bits >> 16u) & 1u);
    return bitcast<f32>(rounded & 0xFFFF0000u);
}

fn gd_32(seg: u32, o: u32) -> u32 { return gd_w[seg].words[o]; }
fn gd_16(seg: u32, o: u32) -> u32 {
    return (gd_32(seg, o >> 1u) >> ((o & 1u) * 16u)) & 0xFFFFu;
}
fn gd_8(seg: u32, o: u32) -> u32 {
    return (gd_32(seg, o >> 2u) >> ((o & 3u) * 8u)) & 0xFFu;
}
fn gd_dot8(w: u32, xi: u32) -> f32 {
    return (f32(w & 0xFu) - 8.0) * gd_act[xi]
         + (f32((w >> 4u) & 0xFu) - 8.0) * gd_act[xi + 1u]
         + (f32((w >> 8u) & 0xFu) - 8.0) * gd_act[xi + 2u]
         + (f32((w >> 12u) & 0xFu) - 8.0) * gd_act[xi + 3u]
         + (f32((w >> 16u) & 0xFu) - 8.0) * gd_act[xi + 4u]
         + (f32((w >> 20u) & 0xFu) - 8.0) * gd_act[xi + 5u]
         + (f32((w >> 24u) & 0xFu) - 8.0) * gd_act[xi + 6u]
         + (f32((w >> 28u) & 0xFu) - 8.0) * gd_act[xi + 7u];
}

@compute @workgroup_size(64)
fn dsv4_global_down_q4tp(@builtin(workgroup_id) wid: vec3<u32>,
                         @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let batch = wid.y;
    let gpr = gd_p.gpr;
    let rows = gd_p.hidden;
    let cst = (gpr * 5u + 7u) / 8u;
    let total = gd_p.slots * gpr;
    var acc = 0.0;
    for (var i = lid; i < total; i = i + 64u) {
        let slot = i / gpr;
        let g = i % gpr;
        let bslot = batch * gd_p.slots + slot;
        let flat = gd_sel[bslot];
        let seg = flat / gd_p.segment_slots;
        let local = flat - seg * gd_p.segment_slots;
        let base16 = local * gd_p.mat16;
        let par16 = base16 + rows * gpr * 8u + row * 2u;
        let cod8 = (base16 + rows * gpr * 8u + rows * 2u) * 2u + row * cst;
        let pl = unpack2x16float(gd_16(seg, par16) | (gd_16(seg, par16 + 1u) << 16u));
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cv = gd_8(seg, cod8 + cb);
        if (shf > 3u) { cv = cv | (gd_8(seg, cod8 + cb + 1u) << 8u); }
        let scale = exp2(pl.x + f32((cv >> shf) & 31u) * pl.y);
        let t16 = base16 + (row * gpr + g) * 8u;
        let xb = (bslot * gpr + g) * 32u;
        var d = 0.0;
        for (var k = 0u; k < 4u; k = k + 1u) {
            let w = gd_16(seg, t16 + 2u * k)
                  | (gd_16(seg, t16 + 1u + 2u * k) << 16u);
            d = d + gd_dot8(w, xb + 8u * k);
        }
        // V4.1's gate/up pass already folded the route weight into the BF16
        // activation.  Keep the generic path's historical f32 weighting.
        if (gd_p._p0 != 0u) {
            acc = acc + scale * d;
        } else {
            acc = acc + gd_wt[bslot] * scale * d;
        }
    }
    gd_pt[lid] = acc;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { gd_pt[lid] = gd_pt[lid] + gd_pt[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        gd_y[batch * gd_p.hidden + row] = select(gd_pt[0], gd_bf16(gd_pt[0]), gd_p._p0 != 0u);
    }
}
"#;

/// Return the global-MoE shader for one supported descriptor-array geometry.
/// The default source stays a static S8 literal; S16 is opt-in and generated
/// only after the adapter has passed the capability gate below.
fn dsv4_global_moe_shader_source(segments: usize) -> Option<String> {
    match segments {
        DSV4_GLOBAL_MOE_SEGMENTS => Some(DSV4_GLOBAL_MOE_SRC.to_owned()),
        DSV4_GLOBAL_MOE_SEGMENTS_S16 => {
            let src = DSV4_GLOBAL_MOE_SRC
                .replace("binding_array<WordBank, 8>", "binding_array<WordBank, 16>");
            (src.matches("binding_array<WordBank, 16>").count() == 3
                && !src.contains("binding_array<WordBank, 8>"))
            .then_some(src)
        }
        _ => None,
    }
}

#[cfg(test)]
mod global_moe_shader_tests {
    use super::{DSV4_GLOBAL_MOE_SRC, dsv4_global_moe_shader_source};

    fn validate(src: &str) {
        let module = wgpu::naga::front::wgsl::parse_str(src).expect("global WGSL must parse");
        wgpu::naga::valid::Validator::new(
            wgpu::naga::valid::ValidationFlags::all(),
            wgpu::naga::valid::Capabilities::all(),
        )
        .validate(&module)
        .expect("global WGSL must validate");
    }

    #[test]
    fn mixed_q2_q4_global_shader_validates() {
        validate(DSV4_GLOBAL_MOE_SRC);
    }

    #[test]
    fn opt_in_s16_global_shader_validates() {
        let src = dsv4_global_moe_shader_source(16).expect("S16 source");
        assert_eq!(src.matches("binding_array<WordBank, 16>").count(), 3);
        assert!(!src.contains("binding_array<WordBank, 8>"));
        validate(&src);
    }
}

const WGSL: &str = r#"
struct Params { cols4: u32, rows: u32, row0_words: u32, _pad: u32 };
@group(0) @binding(0) var<storage, read>       q  : array<u32>;   // 4×i8 packed into u32, row-major
@group(0) @binding(1) var<storage, read>       xs : array<f32>;   // cols, already prescaled by the column scale
@group(0) @binding(2) var<storage, read>       rs : array<f32>;   // row scales for the range
@group(0) @binding(3) var<storage, read_write> y  : array<f32>;   // output: rows
@group(0) @binding(4) var<uniform>             p  : Params;

var<workgroup> partial: array<f32, 64>;

// Exact unpack of 4 signed bytes from u32 (little-endian) — like char4→
// float4 on Metal, without snorm error.
fn i8x4(w: u32) -> vec4<f32> {
    let s = i32(w);
    let b0 = (s << 24u) >> 24u;
    let b1 = (s << 16u) >> 24u;
    let b2 = (s <<  8u) >> 24u;
    let b3 =  s          >> 24u;
    return vec4<f32>(f32(b0), f32(b1), f32(b2), f32(b3));
}

// Grid-stride over rows: the number of workgroups is capped at 65535/dimension,
// while rows (lm_head) number in the hundreds of thousands; one group processes rows
// wid.x, wid.x+nwg.x, … , reducing each with 64 threads.
@compute @workgroup_size(64)
fn q8_matvec(@builtin(workgroup_id) wid: vec3<u32>,
             @builtin(num_workgroups) nwg: vec3<u32>,
             @builtin(local_invocation_index) lid: u32) {
    var row = wid.x;
    loop {
        if (row >= p.rows) { break; }
        let base = p.row0_words + row * p.cols4;
        var acc = 0.0;
        var i = lid;
        loop {
            if (i >= p.cols4) { break; }
            let v = i8x4(q[base + i]);
            let xi = i * 4u;
            let xv = vec4<f32>(xs[xi], xs[xi + 1u], xs[xi + 2u], xs[xi + 3u]);
            acc = acc + dot(v, xv);
            i = i + 64u;
        }
        partial[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { partial[lid] = partial[lid] + partial[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) { y[row] = partial[0] * rs[row]; }
        workgroupBarrier(); // before partial is reused by the next row
        row = row + nwg.x;
    }
}

// q8_2f in the token graph: int8 body, then a per-ROW f16 plane, then a
// per-COLUMN one — w[o,i] = q·s_o·c_i. The per-op path folds the column
// field into the activations on the host and reuses the q8_row kernel;
// a graph has no host in the loop, so the kernel reads both planes where
// they lie. Multiply order follows that per-op arm — x·c first, the row
// scale last — so the two agree, which is the parity that matters.
//
// Without this the whole graph refused: one q8_2f tensor per layer (the
// FFN down_proj on a qwen3 that is otherwise 169 tensors of q1) made
// every layer unbuildable, and a refusal at layer 0 is a refusal for the
// model. Measured on the phone that model never once took the graph.
struct P82 { cols4: u32, rows: u32, cols: u32, _pad: u32 };
@group(0) @binding(0) var<storage, read>       q82  : array<u32>;
@group(0) @binding(1) var<storage, read>       xs82 : array<f32>;
@group(0) @binding(2) var<storage, read_write> y82  : array<f32>;
@group(0) @binding(3) var<uniform>             p82  : P82;

var<workgroup> partial82: array<f32, 64>;

// Four consecutive f16 values starting at an arbitrary half-word. q8_2f is
// packed without padding, so an odd row count starts the column-scale plane
// in the high half of the last row-scale word.
fn f16x4(half: u32) -> vec4<f32> {
    let w = half >> 1u;
    let a = unpack2x16float(q82[w]);
    let b = unpack2x16float(q82[w + 1u]);
    if ((half & 1u) == 0u) {
        return vec4<f32>(a.x, a.y, b.x, b.y);
    }
    let c = unpack2x16float(q82[w + 2u]);
    return vec4<f32>(a.y, b.x, b.y, c.x);
}

// A pruned model does not give this kernel word-aligned rows: bonsai's
// FFN is 6137 wide in one layer and 6130 in another, so a row of int8
// starts mid-word and the whole tensor shears if you index it by
// `row * cols/4`. That shear is what produced fluent nonsense — "the the
// the the" — on a graph that had just started building. Rows are
// addressed in BYTES here and assembled from two words when they
// straddle; only the two f16 planes need `rows*cols` to be a multiple of
// four, which `resolve` checks before offering the arm.
fn ldu(byte: u32) -> u32 {
    let w = byte >> 2u;
    let sh = (byte & 3u) * 8u;
    if (sh == 0u) { return q82[w]; }
    return (q82[w] >> sh) | (q82[w + 1u] << (32u - sh));
}

@compute @workgroup_size(64)
fn q8_2f_matvec(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(num_workgroups) nwg: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let qbytes = p82.rows * p82.cols;
    let rs0 = qbytes >> 2u;                 // per-row f16 plane, in words
    let cs0h = (qbytes >> 1u) + p82.rows;   // per-column plane, in f16 units
    let ngrp = (p82.cols + 3u) / 4u;        // 4 columns a step, last may be short
    var row = wid.x;
    loop {
        if (row >= p82.rows) { break; }
        let rowbase = row * p82.cols;
        var acc = 0.0;
        var i = lid;
        loop {
            if (i >= ngrp) { break; }
            let c0 = i * 4u;
            var v = i8x4(ldu(rowbase + c0));
            // Columns past the end contribute nothing.
            let rem = p82.cols - c0;
            if (rem < 4u) {
                if (rem < 2u) { v.y = 0.0; }
                if (rem < 3u) { v.z = 0.0; }
                v.w = 0.0;
            }
            let xv = vec4<f32>(
                xs82[c0],
                select(0.0, xs82[c0 + 1u], rem > 1u),
                select(0.0, xs82[c0 + 2u], rem > 2u),
                select(0.0, xs82[c0 + 3u], rem > 3u),
            );
            let cv = f16x4(cs0h + c0);
            acc = acc + dot(v, xv * cv);
            i = i + 64u;
        }
        partial82[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { partial82[lid] = partial82[lid] + partial82[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) {
            let rw = unpack2x16float(q82[rs0 + (row >> 1u)]);
            var sc = rw.x;
            if ((row & 1u) == 1u) { sc = rw.y; }
            y82[row] = partial82[0] * sc;
        }
        workgroupBarrier();
        row = row + nwg.x;
    }
}

// The same matvec for a mobile GPU, and the difference is where the
// activations come from. The kernel above re-reads xs out of GLOBAL
// memory for every row — on a desktop the L2 absorbs that, on an Adreno
// 642L it does not: a 2048-wide layer reads 8 KB back 2048 times, and
// measured end to end the whole decode moved ~0.5 GB/s on a bus the CPU
// itself drives at 8.5. Here the group stages xs into workgroup memory
// once and every row reads it from there. Arithmetic, order and output
// are unchanged — this is the same sum, fed from a closer shelf.
//
// 768 vec4 is 12 KB, inside the 16 KB workgroup-storage floor, so the
// arm is only offered for cols ≤ 3072; wider layers keep the old path.
var<workgroup> xsh: array<vec4<f32>, 768>;

@compute @workgroup_size(64)
fn q8_matvec_tiled(@builtin(workgroup_id) wid: vec3<u32>,
                   @builtin(num_workgroups) nwg: vec3<u32>,
                   @builtin(local_invocation_index) lid: u32) {
    var j = lid;
    loop {
        if (j >= p.cols4) { break; }
        let b = j * 4u;
        xsh[j] = vec4<f32>(xs[b], xs[b + 1u], xs[b + 2u], xs[b + 3u]);
        j = j + 64u;
    }
    workgroupBarrier();
    var row = wid.x;
    loop {
        if (row >= p.rows) { break; }
        let base = p.row0_words + row * p.cols4;
        var acc = 0.0;
        var i = lid;
        loop {
            if (i >= p.cols4) { break; }
            acc = acc + dot(i8x4(q[base + i]), xsh[i]);
            i = i + 64u;
        }
        partial[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { partial[lid] = partial[lid] + partial[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) { y[row] = partial[0] * rs[row]; }
        workgroupBarrier();
        row = row + nwg.x;
    }
}

// GEMM of the prefill batch: y[bi, o] = rs[o]·Σ q[o,i]·xs[bi,i]. One workgroup
// per (row, position); the quant row stays hot in cache across bi.
struct MMParams { cols4: u32, rows: u32, nb: u32, _pad: u32 };
@group(0) @binding(0) var<storage, read>       qm  : array<u32>;
@group(0) @binding(1) var<storage, read>       xsm : array<f32>;  // [nb, cols] row-major
@group(0) @binding(2) var<storage, read>       rsm : array<f32>;  // [rows]
@group(0) @binding(3) var<storage, read_write> ym  : array<f32>;  // [nb, rows] row-major
@group(0) @binding(4) var<uniform>             pm  : MMParams;

var<workgroup> partial_mm: array<f32, 64>;

@compute @workgroup_size(64)
fn q8_matmat(@builtin(workgroup_id) wid: vec3<u32>,
             @builtin(num_workgroups) nwg: vec3<u32>,
             @builtin(local_invocation_index) lid: u32) {
    let bi = wid.y;
    if (bi >= pm.nb) { return; }
    let xb = bi * pm.cols4 * 4u;
    var row = wid.x;
    loop {
        if (row >= pm.rows) { break; }
        let qb = row * pm.cols4;
        var acc = 0.0;
        var i = lid;
        loop {
            if (i >= pm.cols4) { break; }
            let v = i8x4(qm[qb + i]);
            let xi = xb + i * 4u;
            let xv = vec4<f32>(xsm[xi], xsm[xi + 1u], xsm[xi + 2u], xsm[xi + 3u]);
            acc = acc + dot(v, xv);
            i = i + 64u;
        }
        partial_mm[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { partial_mm[lid] = partial_mm[lid] + partial_mm[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) { ym[bi * pm.rows + row] = partial_mm[0] * rsm[row]; }
        workgroupBarrier();
        row = row + nwg.x;
    }
}

// q1: 6-byte tiles [f16 scale][4B sign bits] per 32-group; gpr is even,
// so a row is whole 12-byte tile PAIRS = 3 u32 each (same layout walk
// as the Metal kernel). Bit set → +x; np = gpr/2 tile-pairs/row (64 cols each).
//
// FAST kernel (the FFN q1 matvecs are ~59% of a 27B decode token): one
// workgroup owns Q1_ROWS output rows, 16 lanes/row. Activations are staged
// into shared memory in 1024-col tiles and reused across those rows. The
// desktop default is 8 rows/128 threads; Adreno selects 16 rows/256 threads.
// CMF_Q1_RPG=8|16 keeps the choice reproducible. Sign unpack is a branchless XOR sign-flip
// (bit clear ⇒ flip the f32 sign bit) instead of 32 vec4 selects.
struct Q1Params { np: u32, rows: u32, _p0: u32, _p1: u32 };
@group(0) @binding(0) var<storage, read>       q1w : array<u32>;
@group(0) @binding(1) var<storage, read>       q1x : array<f32>;   // raw f32 activations
@group(0) @binding(2) var<storage, read_write> q1y : array<f32>;
@group(0) @binding(3) var<uniform>             q1p : Q1Params;
// The grouped projection's weight tiles, 16 bytes at a time — the tile
// base (row·gpr+g)·4 words is vec4-aligned by construction. Only the
// batch o_lora kernel binds this view.
@group(0) @binding(4) var<storage, read>       q1wv : array<vec4<u32>>;
@group(0) @binding(5) var<storage, read>       q1xv : array<vec4<f32>>;

// q4b_dot8's eight terms, the activations arriving as two vec4 registers.
fn q1_dot8v(w: u32, a: vec4<f32>, b: vec4<f32>) -> f32 {
    return (f32(w & 0xFu) - 8.0) * a.x
         + (f32((w >> 4u) & 0xFu) - 8.0) * a.y
         + (f32((w >> 8u) & 0xFu) - 8.0) * a.z
         + (f32((w >> 12u) & 0xFu) - 8.0) * a.w
         + (f32((w >> 16u) & 0xFu) - 8.0) * b.x
         + (f32((w >> 20u) & 0xFu) - 8.0) * b.y
         + (f32((w >> 24u) & 0xFu) - 8.0) * b.z
         + (f32((w >> 28u) & 0xFu) - 8.0) * b.w;
}

var<workgroup> partial_q1: array<f32, 256>;   // up to 16 rows × 16 lanes
// 1024-col activation tile, PADDED to 33 slots per 32-col group. The read
// pattern is lane*64 + j*4 (all 16 lanes share bank (j*4) mod 32 with a flat
// 1024 tile => 16-way bank conflict, ~8x LSU penalty on the dominant inner
// loop). Padding to stride-33 spreads the lanes across 16 distinct banks
// (66 mod 32 = 2). Same math/accumulation order => token-identical.
var<workgroup> q1xs: array<f32, 1056>;        // 32 groups × 33

// Sum of ±x over one 32-weight group; x read from the shared tile at xbase.
// bit=1 → +x, bit=0 → -x, done by XORing the f32 sign bit (no select chain).
fn q1_tile_sum(bits: u32, xbase: u32) -> f32 {
    var s = vec4<f32>(0.0);
    let pb = (xbase >> 5u) * 33u;   // xbase is a multiple of 32 => padded group base
    for (var j = 0u; j < 8u; j = j + 1u) {
        let nib = bits >> (j * 4u);
        let o = pb + j * 4u;         // j*4+{0..3} stays in [0,32) < 33: no group crossing
        let x = vec4<f32>(q1xs[o], q1xs[o + 1u], q1xs[o + 2u], q1xs[o + 3u]);
        let m = vec4<u32>(
            ((nib & 1u) ^ 1u) << 31u,
            (((nib >> 1u) & 1u) ^ 1u) << 31u,
            (((nib >> 2u) & 1u) ^ 1u) << 31u,
            (((nib >> 3u) & 1u) ^ 1u) << 31u);
        s = s + bitcast<vec4<f32>>(bitcast<vec4<u32>>(x) ^ m);
    }
    return s.x + s.y + s.z + s.w;
}

override Q1_WG: u32 = 128u;
override Q1_ROWS: u32 = 8u;

@compute @workgroup_size(Q1_WG)
fn q1_matvec(@builtin(workgroup_id) wid: vec3<u32>,
             @builtin(num_workgroups) nwg: vec3<u32>,
             @builtin(local_invocation_index) lid: u32) {
    let cols = q1p.np * 64u;
    let r = lid / 16u;
    let lane = lid % 16u;   // which tile-pair lane within a column tile
    var row0 = wid.x * Q1_ROWS;
    loop {
        if (row0 >= q1p.rows) { break; }
        let row = row0 + r;
        var acc = 0.0;
        var ti = 0u;                       // column tile start, in tile-pairs
        loop {
            if (ti >= q1p.np) { break; }
            // Cooperatively stage 1024 activations (16 tile-pairs) into shared.
            let c0 = ti * 64u;
            var k = lid;
            loop {
                if (k >= 1024u) { break; }
                let c = c0 + k;
                q1xs[(k >> 5u) * 33u + (k & 31u)] = select(0.0, q1x[c], c < cols);
                k = k + Q1_WG;
            }
            workgroupBarrier();
            let pi = ti + lane;            // this lane's tile-pair
            if (row < q1p.rows && pi < q1p.np) {
                let base = row * q1p.np * 3u + pi * 3u;
                let a0 = q1w[base]; let a1 = q1w[base + 1u]; let a2 = q1w[base + 2u];
                let s0 = unpack2x16float(a0).x;
                let s1 = unpack2x16float(a1).y;
                let bits0 = (a0 >> 16u) | (a1 << 16u);
                let xb = lane * 64u;       // local offset of this pair in q1xs
                acc = acc + s0 * q1_tile_sum(bits0, xb) + s1 * q1_tile_sum(a2, xb + 32u);
            }
            workgroupBarrier();
            ti = ti + 16u;
        }
        partial_q1[lid] = acc;
        workgroupBarrier();
        // reduce the 16 lanes of each row (blocks of 16 in partial_q1)
        if (lane < 8u) { partial_q1[lid] = partial_q1[lid] + partial_q1[lid + 8u]; }
        workgroupBarrier();
        if (lane < 4u) { partial_q1[lid] = partial_q1[lid] + partial_q1[lid + 4u]; }
        workgroupBarrier();
        if (lane < 2u) { partial_q1[lid] = partial_q1[lid] + partial_q1[lid + 2u]; }
        workgroupBarrier();
        if (lane < 1u) { partial_q1[lid] = partial_q1[lid] + partial_q1[lid + 1u]; }
        workgroupBarrier();
        if (lane == 0u && row < q1p.rows) { q1y[row] = partial_q1[lid]; }
        workgroupBarrier();
        row0 = row0 + nwg.x * Q1_ROWS;
    }
}

// Tiled GEMM for wide prefill batches (the WGSL cousin of Metal's
// q8_mul_mm; WGSL has no subgroup matrices, so this is the classic
// register-blocked form): a 64(b)×64(rows) C-tile per 16×16 workgroup,
// each thread owning a 4×4 accumulator block; X and dequantized W stage
// through 8 KB of workgroup memory in K-steps of 16. The naive kernel
// above re-reads every W row per position — here W is read once per 64
// positions. Perf is hardware-dependent by design: the runtime probe
// decides per machine whether this beats the CPU, so a card where it
// loses simply keeps the CPU path.
var<workgroup> mm_at: array<f32, 64 * 16>;
var<workgroup> mm_wt: array<f32, 64 * 16>;

fn mm_store4(m: u32, n0: u32, v0: f32, v1: f32, v2: f32, v3: f32) {
    if (m >= pm.nb) { return; }
    let base = m * pm.rows + n0;
    if (n0 < pm.rows) { ymm[base] = v0; }
    if (n0 + 1u < pm.rows) { ymm[base + 1u] = v1; }
    if (n0 + 2u < pm.rows) { ymm[base + 2u] = v2; }
    if (n0 + 3u < pm.rows) { ymm[base + 3u] = v3; }
}

fn q8_store4(m: u32, n0: u32, v0: f32, v1: f32, v2: f32, v3: f32) {
    if (m >= pm.nb) { return; }
    let base = m * pm.rows + n0;
    if (n0 < pm.rows) { ym[base] = v0 * rsm[n0]; }
    if (n0 + 1u < pm.rows) { ym[base + 1u] = v1 * rsm[n0 + 1u]; }
    if (n0 + 2u < pm.rows) { ym[base + 2u] = v2 * rsm[n0 + 2u]; }
    if (n0 + 3u < pm.rows) { ym[base + 3u] = v3 * rsm[n0 + 3u]; }
}

fn q1m_store4(m: u32, n0: u32, v0: f32, v1: f32, v2: f32, v3: f32) {
    if (m >= pm.nb) { return; }
    let base = m * pm.rows + n0;
    if (n0 < pm.rows) { ym[base] = v0; }
    if (n0 + 1u < pm.rows) { ym[base + 1u] = v1; }
    if (n0 + 2u < pm.rows) { ym[base + 2u] = v2; }
    if (n0 + 3u < pm.rows) { ym[base + 3u] = v3; }
}

@compute @workgroup_size(16, 16)
fn q8_mul_mm(@builtin(workgroup_id) wid: vec3<u32>,
             @builtin(local_invocation_id) lid: vec3<u32>) {
    let cols = pm.cols4 * 4u;
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;
    let tid = lid.y * 16u + lid.x;
    // Sixteen named scalars, not array<array<f32,4>,4> — see q4t_mul_mm.
    var a00 = 0.0; var a01 = 0.0; var a02 = 0.0; var a03 = 0.0;
    var a10 = 0.0; var a11 = 0.0; var a12 = 0.0; var a13 = 0.0;
    var a20 = 0.0; var a21 = 0.0; var a22 = 0.0; var a23 = 0.0;
    var a30 = 0.0; var a31 = 0.0; var a32 = 0.0; var a33 = 0.0;
    var k0 = 0u;
    loop {
        if (k0 >= cols) { break; }
        // Stage X tile [64×16] (4 f32 per thread) and W tile [64×16]
        // (one u32 = 4 quants per thread per round).
        for (var t = tid; t < 64u * 4u; t = t + 256u) {
            let m = t / 4u;
            let k4 = t % 4u;
            var xv = vec4<f32>(0.0);
            if (m0 + m < pm.nb && (k0 / 4u) + k4 < pm.cols4) {
                let xi = (m0 + m) * cols + k0 + k4 * 4u;
                xv = vec4<f32>(xsm[xi], xsm[xi + 1u], xsm[xi + 2u], xsm[xi + 3u]);
            }
            let dst = m * 16u + k4 * 4u;
            mm_at[dst] = xv.x;
            mm_at[dst + 1u] = xv.y;
            mm_at[dst + 2u] = xv.z;
            mm_at[dst + 3u] = xv.w;
        }
        for (var t = tid; t < 64u * 4u; t = t + 256u) {
            let n = t / 4u;
            let k4 = t % 4u;
            var wv = vec4<f32>(0.0);
            if (n0 + n < pm.rows && (k0 / 4u) + k4 < pm.cols4) {
                wv = i8x4(qm[(n0 + n) * pm.cols4 + (k0 / 4u) + k4]);
            }
            let dst = n * 16u + k4 * 4u;
            mm_wt[dst] = wv.x;
            mm_wt[dst + 1u] = wv.y;
            mm_wt[dst + 2u] = wv.z;
            mm_wt[dst + 3u] = wv.w;
        }
        workgroupBarrier();
        // 4×4 outer-product accumulation over the 16 staged K values.
        let ab = lid.y * 64u;
        let wb = lid.x * 64u;
        for (var k = 0u; k < 16u; k = k + 1u) {
            let x0 = mm_at[ab + k];
            let x1 = mm_at[ab + 16u + k];
            let x2 = mm_at[ab + 32u + k];
            let x3 = mm_at[ab + 48u + k];
            let y0 = mm_wt[wb + k];
            let y1 = mm_wt[wb + 16u + k];
            let y2 = mm_wt[wb + 32u + k];
            let y3 = mm_wt[wb + 48u + k];
            a00 = a00 + x0 * y0; a01 = a01 + x0 * y1;
            a02 = a02 + x0 * y2; a03 = a03 + x0 * y3;
            a10 = a10 + x1 * y0; a11 = a11 + x1 * y1;
            a12 = a12 + x1 * y2; a13 = a13 + x1 * y3;
            a20 = a20 + x2 * y0; a21 = a21 + x2 * y1;
            a22 = a22 + x2 * y2; a23 = a23 + x2 * y3;
            a30 = a30 + x3 * y0; a31 = a31 + x3 * y1;
            a32 = a32 + x3 * y2; a33 = a33 + x3 * y3;
        }
        workgroupBarrier();
        k0 = k0 + 16u;
    }
    let mb = m0 + lid.y * 4u;
    let nb2 = n0 + lid.x * 4u;
    q8_store4(mb, nb2, a00, a01, a02, a03);
    q8_store4(mb + 1u, nb2, a10, a11, a12, a13);
    q8_store4(mb + 2u, nb2, a20, a21, a22, a23);
    q8_store4(mb + 3u, nb2, a30, a31, a32, a33);
}

// Tiled q1 GEMM for wide batches (prefill / speculative K-token decode): the
// q1 twin of q8_mul_mm. Reuses the mul_mm bindings (rsm is unused — q1's scale
// is per-32-group and folded into the staged weight). Decode a 4-wide run of
// weights for one output row: 4 cols in one 32-group share a bit-word + scale;
// bit set → +scale, clear → −scale (XOR the sign bit). cols4 = cols/4, so the
// row has np = cols4/16 six-byte tile-pairs (64 cols each, 2 groups of 32).
fn q1_w4(n: u32, k: u32, np: u32) -> vec4<f32> {
    let pi = k / 64u;
    let off = k % 64u;                 // 4-aligned ⇒ never straddles a 32-group
    let base = n * np * 3u + pi * 3u;
    let a0 = qm[base]; let a1 = qm[base + 1u]; let a2 = qm[base + 2u];
    var bits: u32;
    var scale: f32;
    if (off < 32u) { bits = (a0 >> 16u) | (a1 << 16u); scale = unpack2x16float(a0).x; }
    else           { bits = a2;                        scale = unpack2x16float(a1).y; }
    let bo = off & 31u;
    let m = vec4<u32>(
        (((bits >> bo)        & 1u) ^ 1u) << 31u,
        (((bits >> (bo + 1u)) & 1u) ^ 1u) << 31u,
        (((bits >> (bo + 2u)) & 1u) ^ 1u) << 31u,
        (((bits >> (bo + 3u)) & 1u) ^ 1u) << 31u);
    let sv = vec4<f32>(scale, scale, scale, scale);
    return bitcast<vec4<f32>>(bitcast<vec4<u32>>(sv) ^ m);
}

@compute @workgroup_size(16, 16)
fn q1_mul_mm(@builtin(workgroup_id) wid: vec3<u32>,
             @builtin(local_invocation_id) lid: vec3<u32>) {
    let cols = pm.cols4 * 4u;
    let np = pm.cols4 / 16u;
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;
    let tid = lid.y * 16u + lid.x;
    // Sixteen named scalars, not array<array<f32,4>,4> — see q4t_mul_mm.
    var a00 = 0.0; var a01 = 0.0; var a02 = 0.0; var a03 = 0.0;
    var a10 = 0.0; var a11 = 0.0; var a12 = 0.0; var a13 = 0.0;
    var a20 = 0.0; var a21 = 0.0; var a22 = 0.0; var a23 = 0.0;
    var a30 = 0.0; var a31 = 0.0; var a32 = 0.0; var a33 = 0.0;
    var k0 = 0u;
    loop {
        if (k0 >= cols) { break; }
        for (var t = tid; t < 64u * 4u; t = t + 256u) {
            let m = t / 4u;
            let k4 = t % 4u;
            var xv = vec4<f32>(0.0);
            if (m0 + m < pm.nb && (k0 / 4u) + k4 < pm.cols4) {
                let xi = (m0 + m) * cols + k0 + k4 * 4u;
                xv = vec4<f32>(xsm[xi], xsm[xi + 1u], xsm[xi + 2u], xsm[xi + 3u]);
            }
            let dst = m * 16u + k4 * 4u;
            mm_at[dst] = xv.x; mm_at[dst + 1u] = xv.y; mm_at[dst + 2u] = xv.z; mm_at[dst + 3u] = xv.w;
        }
        for (var t = tid; t < 64u * 4u; t = t + 256u) {
            let n = t / 4u;
            let k4 = t % 4u;
            var wv = vec4<f32>(0.0);
            if (n0 + n < pm.rows && (k0 / 4u) + k4 < pm.cols4) {
                wv = q1_w4(n0 + n, k0 + k4 * 4u, np);
            }
            let dst = n * 16u + k4 * 4u;
            mm_wt[dst] = wv.x; mm_wt[dst + 1u] = wv.y; mm_wt[dst + 2u] = wv.z; mm_wt[dst + 3u] = wv.w;
        }
        workgroupBarrier();
        let ab = lid.y * 64u;
        let wb = lid.x * 64u;
        for (var k = 0u; k < 16u; k = k + 1u) {
            let x0 = mm_at[ab + k];
            let x1 = mm_at[ab + 16u + k];
            let x2 = mm_at[ab + 32u + k];
            let x3 = mm_at[ab + 48u + k];
            let y0 = mm_wt[wb + k];
            let y1 = mm_wt[wb + 16u + k];
            let y2 = mm_wt[wb + 32u + k];
            let y3 = mm_wt[wb + 48u + k];
            a00 = a00 + x0 * y0; a01 = a01 + x0 * y1;
            a02 = a02 + x0 * y2; a03 = a03 + x0 * y3;
            a10 = a10 + x1 * y0; a11 = a11 + x1 * y1;
            a12 = a12 + x1 * y2; a13 = a13 + x1 * y3;
            a20 = a20 + x2 * y0; a21 = a21 + x2 * y1;
            a22 = a22 + x2 * y2; a23 = a23 + x2 * y3;
            a30 = a30 + x3 * y0; a31 = a31 + x3 * y1;
            a32 = a32 + x3 * y2; a33 = a33 + x3 * y3;
        }
        workgroupBarrier();
        k0 = k0 + 16u;
    }
    let mb = m0 + lid.y * 4u;
    let nb2 = n0 + lid.x * 4u;
    q1m_store4(mb, nb2, a00, a01, a02, a03);
    q1m_store4(mb + 1u, nb2, a10, a11, a12, a13);
    q1m_store4(mb + 2u, nb2, a20, a21, a22, a23);
    q1m_store4(mb + 3u, nb2, a30, a31, a32, a33);
}

// ── Element-wise kernels of the MoE block (silu·mul·col, axpy, zeroing) ──
struct N1 { n: u32, f: u32, lim: f32, _c: u32 };

@group(0) @binding(0) var<storage, read>       sg   : array<f32>;
@group(0) @binding(1) var<storage, read>       su   : array<f32>;
@group(0) @binding(2) var<storage, read>       scol : array<f32>;
@group(0) @binding(3) var<storage, read_write> sact : array<f32>;
@group(0) @binding(4) var<uniform>             snp  : N1;
@compute @workgroup_size(256)
fn silu_mul_pre(@builtin(global_invocation_id) gid: vec3<u32>,
                @builtin(num_workgroups) nwg: vec3<u32>) {
    let i = gid.x + gid.y * nwg.x * 256u;
    if (i >= snp.n) { return; }
    var gv = sg[i];
    var uv = su[i];
    // swiglu_limit, and its asymmetry is the reference's: `up` is clamped on
    // BOTH sides, `gate` only from above. A device that skips it diverges
    // from the CPU path exactly on the tokens that saturate.
    if (snp.lim > 0.0) {
        uv = clamp(uv, -snp.lim, snp.lim);
        gv = min(gv, snp.lim);
    }
    var v = (gv / (1.0 + exp(-gv))) * uv;
    if (snp.f == 1u) { v = v * scol[i]; }
    sact[i] = v;
}

// `set`: y = w·x[soff+i] rather than y += w·x[i]. Two callers wanted the
// assignment and were spending a whole zero-fill dispatch to get it, and one
// wanted a strided source and was spending a copy. `soff` is in floats.
// `asg`, not `set`: `set` is a WGSL reserved keyword, and a shader that
// fails to parse takes the WHOLE module with it — the context then does
// not come up and every op quietly walks the host.
struct AxpyP { w: f32, n: u32, asg: u32, soff: u32 };
@group(0) @binding(0) var<storage, read>       ad : array<f32>;
@group(0) @binding(1) var<storage, read_write> ay : array<f32>;
@group(0) @binding(2) var<uniform>             ap : AxpyP;
@compute @workgroup_size(256)
fn axpy(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= ap.n) { return; }
    let v = ap.w * ad[ap.soff + i];
    if (ap.asg != 0u) { ay[i] = v; } else { ay[i] = ay[i] + v; }
}

// The two-field codec's column scale, applied to a panel that is already on
// the card: `y[r][i] = x[r][i] · col[i]`. The host folds this field into the
// activation for free when it has the activation; when the panel is the
// previous kernel's output, this is what keeps the fusion intact instead of
// dragging it home to multiply.
struct ColP { n: u32, cols: u32, _a: u32, _b: u32 };
@group(0) @binding(0) var<storage, read>       csx : array<f32>;
@group(0) @binding(1) var<storage, read>       csc : array<f32>;
@group(0) @binding(2) var<storage, read_write> csy : array<f32>;
@group(0) @binding(3) var<uniform>             csp : ColP;
@compute @workgroup_size(256)
fn colscale(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= csp.n) { return; }
    csy[i] = csx[i] * csc[i % csp.cols];
}

// The same matvec with 256 threads and four rows to a workgroup.
//
// `f32_matvec` gives a row 64 threads and a workgroup, which for the
// hyper-connection mix — 24 rows over 16 384 columns — is 24 workgroups of
// 64 threads: fifteen hundred threads on a card that holds hundreds of
// thousands. The dispatch itself costs ~3 µs (measured), so the time was
// never the launch; it was the kernel using a rounding error of the machine.
struct F32WP { cols: u32, rows: u32, _a: u32, _b: u32 };
@group(0) @binding(0) var<storage, read>       fww : array<f32>;
@group(0) @binding(1) var<storage, read>       fwx : array<f32>;
@group(0) @binding(2) var<storage, read_write> fwy : array<f32>;
@group(0) @binding(3) var<uniform>             fwp : F32WP;
var<workgroup> fwpart: array<f32, 256>;
@compute @workgroup_size(256)
fn f32_matvec_w(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    if (row >= fwp.rows) { return; }
    let base = row * fwp.cols;
    var acc = 0.0;
    var i = lid;
    loop {
        if (i >= fwp.cols) { break; }
        acc = acc + fww[base + i] * fwx[i];
        i = i + 256u;
    }
    fwpart[lid] = acc;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { fwpart[lid] = fwpart[lid] + fwpart[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    if (lid == 0u) { fwy[row] = fwpart[0]; }
}

// The same again with a thousand threads, for the shapes with FEW rows. The
// hyper-connection mix is 24 rows over 16 384 columns: at 256 threads that
// is 24 workgroups of 256, six thousand threads. Rows are what give this
// kernel its workgroups, so when rows are scarce the only width left is
// inside one.

@group(0) @binding(0) var<storage, read>       fxw : array<f32>;
@group(0) @binding(1) var<storage, read>       fxx : array<f32>;
@group(0) @binding(2) var<storage, read_write> fxy : array<f32>;
@group(0) @binding(3) var<uniform>             fxp : F32WP;
var<workgroup> fxpart: array<f32, 1024>;
@compute @workgroup_size(1024)
fn f32_matvec_x(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    if (row >= fxp.rows) { return; }
    let base = row * fxp.cols;
    var acc = 0.0;
    var i = lid;
    loop {
        if (i >= fxp.cols) { break; }
        acc = acc + fxw[base + i] * fxx[i];
        i = i + 1024u;
    }
    fxpart[lid] = acc;
    workgroupBarrier();
    var stride = 512u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { fxpart[lid] = fxpart[lid] + fxpart[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    if (lid == 0u) { fxy[row] = fxpart[0]; }
}

// The f32 matvec split over COLUMNS, for the shapes with almost no rows.
//
// Rows are what give this kernel its workgroups, and the hyper-connection
// mix has 24 of them over 16 384 columns: even at 1024 threads that is
// twenty-four workgroups — six per cent of a card that holds hundreds of
// thousands of threads. Splitting the shared axis gives rows·splits
// workgroups and a cheap merge, the same trade the attention split makes.
struct F32SP { cols: u32, rows: u32, nsplit: u32, chunk: u32 };
@group(0) @binding(0) var<storage, read>       fmsw : array<f32>;
@group(0) @binding(1) var<storage, read>       fmsx : array<f32>;
@group(0) @binding(2) var<storage, read_write> fmsy : array<f32>;
@group(0) @binding(3) var<uniform>             fmsp : F32SP;
@group(0) @binding(4) var<storage, read_write> fmspart : array<f32>;  // rows*nsplit
var<workgroup> fmsred: array<f32, 256>;

@compute @workgroup_size(256)
fn f32_matvec_split(@builtin(workgroup_id) wid: vec3<u32>,
                    @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let sp = wid.y;
    if (row >= fmsp.rows || sp >= fmsp.nsplit) { return; }
    let lo = sp * fmsp.chunk;
    var hi = lo + fmsp.chunk;
    if (hi > fmsp.cols) { hi = fmsp.cols; }
    let base = row * fmsp.cols;
    var acc = 0.0;
    var i = lo + lid;
    loop {
        if (i >= hi) { break; }
        acc = acc + fmsw[base + i] * fmsx[i];
        i = i + 256u;
    }
    fmsred[lid] = acc;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { fmsred[lid] = fmsred[lid] + fmsred[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    if (lid == 0u) { fmspart[row * fmsp.nsplit + sp] = fmsred[0]; }
    // Same reason as in the merge: this entry point never writes the output,
    // so without a mention its derived layout is one binding short of the
    // merge's and no single bind group can serve both. Never runs.
    if (fmsp.rows == 0xFFFFFFFFu) { fmsy[0] = 0.0; }
}

@compute @workgroup_size(64)
fn f32_matvec_merge(@builtin(workgroup_id) wid: vec3<u32>,
                    @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    if (row >= fmsp.rows) { return; }
    // The splits are summed IN ORDER by one lane: the partials are few and
    // a fixed order keeps the answer reproducible.
    if (lid == 0u) {
        var s = 0.0;
        for (var t = 0u; t < fmsp.nsplit; t = t + 1u) {
            s = s + fmspart[row * fmsp.nsplit + t];
        }
        fmsy[row] = s;
    }
    // A layout is derived from the bindings an ENTRY POINT uses, not from
    // the module's globals — so without this the merge's layout has fewer
    // bindings than the split's and one bind group cannot serve both.
    // The branch never runs.
    if (fmsp.rows == 0xFFFFFFFFu) { fmsy[0] = fmsw[0] + fmsx[0]; }
}

// Qwen3.5 output gate: attn_out *= sigmoid(gate), element-wise over nh·hd.
@group(0) @binding(0) var<storage, read>       gm_g : array<f32>;
@group(0) @binding(1) var<storage, read_write> gm_o : array<f32>;
@group(0) @binding(2) var<uniform>             gm_p : N1;
@compute @workgroup_size(256)
fn gate_mul(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= gm_p.n) { return; }
    gm_o[i] = gm_o[i] * (1.0 / (1.0 + exp(-gm_g[i])));
}

@group(0) @binding(0) var<storage, read_write> zy  : array<f32>;
@group(0) @binding(1) var<uniform>             znp : N1;
@compute @workgroup_size(256)
fn fill_zero(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i < znp.n) { zy[i] = 0.0; }
}

// silu(g)·u in place on g — the glue pass of the fused imagegen FFN
// (w1/w3/silu/w2 in one submission, one readback).
@group(0) @binding(0) var<storage, read_write> fsg : array<f32>;
@group(0) @binding(1) var<storage, read>       fsu : array<f32>;
@group(0) @binding(2) var<uniform>             fsp : N1;
@compute @workgroup_size(256)
fn ffn_silu_mul(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.y * (65535u * 256u) + gid.x;
    if (i >= fsp.n) { return; }
    let g = fsg[i];
    fsg[i] = (g / (1.0 + exp(-g))) * fsu[i];
}

// Qwen Image's MLP middle is tanh-GELU, not the SwiGLU used by the
// neighbouring DiT paths.  Keep the exact add-then-GELU order from the
// reference, but leave the panel on the device between the two Q4TP GEMMs.
// `gelu` is 1 for the input projection bias and 0 for the output projection
// bias; using one entry point keeps the pipeline/cache surface small.
struct QwenGeluP { n: u32, width: u32, gelu: u32, _p: u32 };
@group(0) @binding(0) var<storage, read_write> qg_values: array<f32>;
@group(0) @binding(1) var<storage, read>       qg_bias: array<f32>;
@group(0) @binding(2) var<uniform>             qg_p: QwenGeluP;

@compute @workgroup_size(256)
fn qwen_gelu_bias(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.y * (65535u * 256u) + gid.x;
    if (i >= qg_p.n) { return; }
    let j = i - (i / qg_p.width) * qg_p.width;
    var x = qg_values[i] + qg_bias[j];
    if (qg_p.gelu != 0u) {
        let x3 = x * x * x;
        x = 0.5 * x * (1.0 + tanh(0.7978846 * (x + 0.044715 * x3)));
    }
    qg_values[i] = x;
}

// Qwen Image's affine-free LayerNorm followed by its per-stream
// shift/scale modulation.  The two reductions intentionally stay in one
// workgroup per token so the large hidden panel never crosses the host.
// `qnm_mod` is `[shift, scale]`, both width elements.
struct QwenNormModP { n: u32, width: u32, eps: f32, _p: u32 };
@group(0) @binding(0) var<storage, read>       qnm_x   : array<f32>;
@group(0) @binding(1) var<storage, read>       qnm_mod : array<f32>;
@group(0) @binding(2) var<storage, read_write> qnm_y   : array<f32>;
@group(0) @binding(3) var<uniform>             qnm_p   : QwenNormModP;
var<workgroup> qnm_part: array<f32, 256>;

@compute @workgroup_size(256)
fn qwen_layernorm_mod(
    @builtin(workgroup_id) wid: vec3<u32>,
    @builtin(local_invocation_index) lid: u32,
) {
    let row = wid.x;
    if (row >= qnm_p.n) { return; }
    let base = row * qnm_p.width;
    var sum = 0.0;
    var i = lid;
    loop {
        if (i >= qnm_p.width) { break; }
        sum = sum + qnm_x[base + i];
        i = i + 256u;
    }
    qnm_part[lid] = sum;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { qnm_part[lid] = qnm_part[lid] + qnm_part[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let mean = qnm_part[0] / f32(qnm_p.width);
    // Every invocation must finish reading the reduced mean before any lane
    // reuses qnm_part for the variance reduction below.
    workgroupBarrier();
    var var_sum = 0.0;
    i = lid;
    loop {
        if (i >= qnm_p.width) { break; }
        let d = qnm_x[base + i] - mean;
        var_sum = var_sum + d * d;
        i = i + 256u;
    }
    qnm_part[lid] = var_sum;
    workgroupBarrier();
    stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { qnm_part[lid] = qnm_part[lid] + qnm_part[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let inv = inverseSqrt(qnm_part[0] / f32(qnm_p.width) + qnm_p.eps);
    i = lid;
    loop {
        if (i >= qnm_p.width) { break; }
        let normed = (qnm_x[base + i] - mean) * inv;
        qnm_y[base + i] = normed * (1.0 + qnm_mod[qnm_p.width + i]) + qnm_mod[i];
        i = i + 256u;
    }
}

// Qwen's gated residual has no extra RMS normalization: the projection
// bias is added first, then the learned per-channel gate scales it.
struct QwenResidualP { n: u32, width: u32, _a: u32, _b: u32 };
@group(0) @binding(0) var<storage, read>       qgr_base : array<f32>;
@group(0) @binding(1) var<storage, read_write> qgr_delta: array<f32>;
@group(0) @binding(2) var<storage, read>       qgr_gate : array<f32>;
@group(0) @binding(3) var<uniform>             qgr_p    : QwenResidualP;

@compute @workgroup_size(256)
fn qwen_gated_residual(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.y * (65535u * 256u) + gid.x;
    if (i >= qgr_p.n * qgr_p.width) { return; }
    let j = i - (i / qgr_p.width) * qgr_p.width;
    qgr_delta[i] = qgr_base[i] + qgr_gate[j] * qgr_delta[i];
}

// The same, for a fc1 that emits gate and up PACKED IN ONE ROW
// ([gate|up] per token, as MiniMax-H3's DiT stores it) instead of two
// separate panels. `n` counts activations (b·inter), `fsp2.inter` is the
// half-width. Writes a compact [b][inter] panel the second GEMM reads.
struct FsP2 { n: u32, inter: u32, has_bias: u32, _b: u32 };
@group(0) @binding(0) var<storage, read>       fpgu : array<f32>;
@group(0) @binding(1) var<storage, read_write> fpact: array<f32>;
@group(0) @binding(2) var<uniform>             fsp2 : FsP2;
// Gate/up bias, [2·inter] — the 3D VAE's FFN carries one, the DiT's
// does not. A one-element dummy is bound when there is none.
@group(0) @binding(3) var<storage, read>       fpb  : array<f32>;

@compute @workgroup_size(256)
fn ffn_silu_mul_packed(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.y * (65535u * 256u) + gid.x;
    if (i >= fsp2.n) { return; }
    let p = i / fsp2.inter;
    let j = i - p * fsp2.inter;
    let base = p * 2u * fsp2.inter;
    var g = fpgu[base + j];
    var u = fpgu[base + fsp2.inter + j];
    if (fsp2.has_bias != 0u) {
        g = g + fpb[j];
        u = u + fpb[fsp2.inter + j];
    }
    fpact[i] = (g / (1.0 + exp(-g))) * u;
}

// Plain f32 matvec (for small unquantized projections like GDN in_proj_a/b):
// y[o] = Σ_i W[o,i]·x[i]. One workgroup per output row.
struct F32P { cols: u32, rows: u32, _a: u32, _b: u32 };
@group(0) @binding(0) var<storage, read>       f32w : array<f32>;
@group(0) @binding(1) var<storage, read>       f32x : array<f32>;
@group(0) @binding(2) var<storage, read_write> f32y : array<f32>;
@group(0) @binding(3) var<uniform>             f32p : F32P;
var<workgroup> f32part: array<f32, 64>;
@compute @workgroup_size(64)
fn f32_matvec(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    if (row >= f32p.rows) { return; }
    let base = row * f32p.cols;
    var acc = 0.0;
    var i = lid;
    loop {
        if (i >= f32p.cols) { break; }
        acc = acc + f32w[base + i] * f32x[i];
        i = i + 64u;
    }
    f32part[lid] = acc;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { f32part[lid] = f32part[lid] + f32part[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    if (lid == 0u) { f32y[row] = f32part[0]; }
}

// ── Two independent projections of the SAME input, in ONE dispatch.
//
// A GDN layer projects its input four ways (qkv, z, a, b); a MoE layer
// projects it twice (router, shared gate). Those are independent, but
// dispatches inside a compute pass are serialized by wgpu's
// memory-visibility guarantee — so each costs a full launch (~29 µs on
// this Vulkan stack) for work the card could overlap. Measured on a
// 5090, the 40-layer decode spends 15.3 of its 16.4 ms in
// submit+readback across ~520 dispatches, so collapsing pairs is worth
// more than any arithmetic in the kernels.
//
// The two jobs are laid end-to-end in one row space: workgroup r < rows0
// is job 0, the rest is job 1. Whole workgroups branch together, so the
// per-kind branch never diverges within a workgroup. Kinds: 4 = f32,
// 6 = q4tp; the caller keeps the unfused path for anything else.
struct MvP2 {
    rows0: u32, cols0: u32, kind0: u32, _pa: u32,
    rows1: u32, cols1: u32, kind1: u32, _pb: u32,
};
@group(0) @binding(0) var<storage, read>       m2w0 : array<u32>;
@group(0) @binding(1) var<storage, read>       m2w1 : array<u32>;
@group(0) @binding(2) var<storage, read>       m2x  : array<f32>;
@group(0) @binding(3) var<storage, read_write> m2y0 : array<f32>;
@group(0) @binding(4) var<storage, read_write> m2y1 : array<f32>;
@group(0) @binding(5) var<uniform>             m2p  : MvP2;
var<workgroup> m2part: array<f32, 64>;
var<workgroup> m2lad:  array<f32, 32>;

fn m2_dot8(w: u32, xi: u32) -> f32 {
    return (f32(w & 0xFu) - 8.0) * m2x[xi]
         + (f32((w >> 4u) & 0xFu) - 8.0) * m2x[xi + 1u]
         + (f32((w >> 8u) & 0xFu) - 8.0) * m2x[xi + 2u]
         + (f32((w >> 12u) & 0xFu) - 8.0) * m2x[xi + 3u]
         + (f32((w >> 16u) & 0xFu) - 8.0) * m2x[xi + 4u]
         + (f32((w >> 20u) & 0xFu) - 8.0) * m2x[xi + 5u]
         + (f32((w >> 24u) & 0xFu) - 8.0) * m2x[xi + 6u]
         + (f32((w >> 28u) & 0xFu) - 8.0) * m2x[xi + 7u];
}

// One partial dot over job 0's weights. Split per buffer because WGSL has
// no way to pick a binding at runtime.
fn m2_part0(kind: u32, row: u32, cols: u32, lid: u32) -> f32 {
    var acc = 0.0;
    if (kind == 4u) {
        let base = row * cols;
        var i = lid;
        loop {
            if (i >= cols) { break; }
            acc = acc + bitcast<f32>(m2w0[base + i]) * m2x[i];
            i = i + 64u;
        }
        return acc;
    }
    let gpr = cols / 32u;
    let rows = m2p.rows0;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    if (lid < 32u) {
        let pr = unpack2x16float(m2w0[params_w + row]);
        m2lad[lid] = exp2(pr.x + f32(lid) * pr.y);
    }
    workgroupBarrier();
    var g = lid;
    loop {
        if (g >= gpr) { break; }
        let bit = g * 5u;
        let cb = codes_b + row * cstride + (bit >> 3u);
        let sh = bit & 7u;
        var cv = (m2w0[cb >> 2u] >> ((cb & 3u) * 8u)) & 0xFFu;
        if (sh > 3u) {
            let cb1 = cb + 1u;
            cv = cv | (((m2w0[cb1 >> 2u] >> ((cb1 & 3u) * 8u)) & 0xFFu) << 8u);
        }
        let scale = m2lad[(cv >> sh) & 31u];
        let base = (row * gpr + g) * 4u;
        let xb = g * 32u;
        var gs = 0.0;
        for (var k = 0u; k < 4u; k = k + 1u) {
            gs = gs + m2_dot8(m2w0[base + k], xb + 8u * k);
        }
        acc = acc + scale * gs;
        g = g + 64u;
    }
    return acc;
}

fn m2_part1(kind: u32, row: u32, cols: u32, lid: u32) -> f32 {
    var acc = 0.0;
    if (kind == 4u) {
        let base = row * cols;
        var i = lid;
        loop {
            if (i >= cols) { break; }
            acc = acc + bitcast<f32>(m2w1[base + i]) * m2x[i];
            i = i + 64u;
        }
        return acc;
    }
    let gpr = cols / 32u;
    let rows = m2p.rows1;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    if (lid < 32u) {
        let pr = unpack2x16float(m2w1[params_w + row]);
        m2lad[lid] = exp2(pr.x + f32(lid) * pr.y);
    }
    workgroupBarrier();
    var g = lid;
    loop {
        if (g >= gpr) { break; }
        let bit = g * 5u;
        let cb = codes_b + row * cstride + (bit >> 3u);
        let sh = bit & 7u;
        var cv = (m2w1[cb >> 2u] >> ((cb & 3u) * 8u)) & 0xFFu;
        if (sh > 3u) {
            let cb1 = cb + 1u;
            cv = cv | (((m2w1[cb1 >> 2u] >> ((cb1 & 3u) * 8u)) & 0xFFu) << 8u);
        }
        let scale = m2lad[(cv >> sh) & 31u];
        let base = (row * gpr + g) * 4u;
        let xb = g * 32u;
        var gs = 0.0;
        for (var k = 0u; k < 4u; k = k + 1u) {
            gs = gs + m2_dot8(m2w1[base + k], xb + 8u * k);
        }
        acc = acc + scale * gs;
        g = g + 64u;
    }
    return acc;
}

@compute @workgroup_size(64)
fn matvec_pair(@builtin(workgroup_id) wid: vec3<u32>,
               @builtin(num_workgroups) nwg: vec3<u32>,
               @builtin(local_invocation_index) lid: u32) {
    let total = m2p.rows0 + m2p.rows1;
    var flat = wid.x;
    loop {
        if (flat >= total) { break; }
        var acc = 0.0;
        let second = flat >= m2p.rows0;
        if (second) {
            acc = m2_part1(m2p.kind1, flat - m2p.rows0, m2p.cols1, lid);
        } else {
            acc = m2_part0(m2p.kind0, flat, m2p.cols0, lid);
        }
        m2part[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { m2part[lid] = m2part[lid] + m2part[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) {
            if (second) { m2y1[flat - m2p.rows0] = m2part[0]; }
            else { m2y0[flat] = m2part[0]; }
        }
        // The next iteration rewrites m2lad/m2part; make sure every thread
        // is done reading them first.
        workgroupBarrier();
        flat = flat + nwg.x;
    }
}

// f32 matvec with a token axis for the batch graph: wid.y = token, and
// the PER-ROW math is f32_matvec verbatim — same lane stride, same tree
// reduction — so the logits it produces are bit-identical to k separate
// dispatches of the single-token kernel. That equivalence is what lets
// the batch prefill's router and GDN a/b projections collapse from one
// dispatch per token per layer (~3200 a chunk) to one per layer.
struct F32BP { cols: u32, rows: u32, _a: u32, _b: u32 };
@group(0) @binding(0) var<storage, read>       fb_w : array<f32>;
@group(0) @binding(1) var<storage, read>       fb_x : array<f32>;
@group(0) @binding(2) var<storage, read_write> fb_y : array<f32>;
@group(0) @binding(3) var<uniform>             fb_p : F32BP;
var<workgroup> fb_part: array<f32, 64>;
@compute @workgroup_size(64)
fn f32_matvec_b(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let t = wid.y;
    if (row >= fb_p.rows) { return; }
    let base = row * fb_p.cols;
    let xoff = t * fb_p.cols;
    var acc = 0.0;
    var i = lid;
    loop {
        if (i >= fb_p.cols) { break; }
        acc = acc + fb_w[base + i] * fb_x[xoff + i];
        i = i + 64u;
    }
    fb_part[lid] = acc;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { fb_part[lid] = fb_part[lid] + fb_part[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    if (lid == 0u) { fb_y[t * fb_p.rows + row] = fb_part[0]; }
}

// Backward GEMM for the skill-bake trainer: dx[n,k] = dy[n,m] · w[m,k].
// The forward kernel contracts over w's SECOND index; this one contracts
// over its first, which is the whole difference between y = W·x and the
// gradient that comes back through W. One workgroup per (output column,
// row), 64 lanes striding the contraction.
struct GdxP { m: u32, k: u32, _a: u32, _b: u32 };
@group(0) @binding(0) var<storage, read>       gx_w  : array<f32>;
@group(0) @binding(1) var<storage, read>       gx_dy : array<f32>;
@group(0) @binding(2) var<storage, read_write> gx_dx : array<f32>;
@group(0) @binding(3) var<uniform>             gx_p  : GdxP;
var<workgroup> gx_part: array<f32, 64>;
@compute @workgroup_size(64)
fn f32_gemm_dx(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_index) lid: u32) {
    let j = wid.x;
    let r = wid.y;
    if (j >= gx_p.k) { return; }
    var acc = 0.0;
    var o = lid;
    loop {
        if (o >= gx_p.m) { break; }
        acc = acc + gx_dy[r * gx_p.m + o] * gx_w[o * gx_p.k + j];
        o = o + 64u;
    }
    gx_part[lid] = acc;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { gx_part[lid] = gx_part[lid] + gx_part[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    if (lid == 0u) { gx_dx[r * gx_p.k + j] = gx_part[0]; }
}

// ── DiT fused-block kernels ─────────────────────────────────────────────
// Modulated RMSNorm: o = rms(x)·w·(1+s) when `has_s`, else rms(x)·w.
// One workgroup a row, 256-lane tree reduction — the shape every norm
// in this file uses.
struct DmP { n: u32, hs: u32, eps: f32, has_s: u32 };
@group(0) @binding(0) var<storage, read>       dm_x : array<f32>;
@group(0) @binding(1) var<storage, read>       dm_w : array<f32>;
@group(0) @binding(2) var<storage, read>       dm_s : array<f32>;
@group(0) @binding(3) var<storage, read_write> dm_o : array<f32>;
@group(0) @binding(4) var<uniform>             dm_p : DmP;
var<workgroup> dm_part: array<f32, 256>;
@compute @workgroup_size(256)
fn dit_rmsmod(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    if (row >= dm_p.n) { return; }
    let base = row * dm_p.hs;
    var acc = 0.0;
    var i = lid;
    loop {
        if (i >= dm_p.hs) { break; }
        let v = dm_x[base + i];
        acc = acc + v * v;
        i = i + 256u;
    }
    dm_part[lid] = acc;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { dm_part[lid] = dm_part[lid] + dm_part[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let inv = inverseSqrt(dm_part[0] / f32(dm_p.hs) + dm_p.eps);
    i = lid;
    loop {
        if (i >= dm_p.hs) { break; }
        var v = dm_x[base + i] * inv * dm_w[i];
        if (dm_p.has_s != 0u) { v = v * (1.0 + dm_s[i]); }
        dm_o[base + i] = v;
        i = i + 256u;
    }
}

// Gated residual: x += gate ⊙ rms(d)·w — the DiT's sandwich norm
// on the way out of a sub-block; the gate comes in already tanh'd.
struct GrP { n: u32, hs: u32, eps: f32, _p: u32 };
@group(0) @binding(0) var<storage, read_write> gr_x : array<f32>;
@group(0) @binding(1) var<storage, read>       gr_d : array<f32>;
@group(0) @binding(2) var<storage, read>       gr_w : array<f32>;
@group(0) @binding(3) var<storage, read>       gr_g : array<f32>;
@group(0) @binding(4) var<uniform>             gr_p : GrP;
var<workgroup> gr_part: array<f32, 256>;
@compute @workgroup_size(256)
fn dit_gated_residual(
    @builtin(workgroup_id) wid: vec3<u32>,
    @builtin(local_invocation_index) lid: u32,
) {
    let row = wid.x;
    if (row >= gr_p.n) { return; }
    let base = row * gr_p.hs;
    var acc = 0.0;
    var i = lid;
    loop {
        if (i >= gr_p.hs) { break; }
        let v = gr_d[base + i];
        acc = acc + v * v;
        i = i + 256u;
    }
    gr_part[lid] = acc;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { gr_part[lid] = gr_part[lid] + gr_part[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let inv = inverseSqrt(gr_part[0] / f32(gr_p.hs) + gr_p.eps);
    i = lid;
    loop {
        if (i >= gr_p.hs) { break; }
        // The gate arrives tanh'd — the host computes it once per block
        // for all rows, and doing it again here squashed it twice.
        gr_x[base + i] = gr_x[base + i] + gr_g[i] * (gr_d[base + i] * inv * gr_w[i]);
        i = i + 256u;
    }
}

// Per-head qk-norm + interleaved-pair RoPE, packing token-major input
// into head-major output. The DiT's rope table is (cos, sin) per token
// and pair; a head's dims share it. One workgroup per (token, head).
struct DrpP { n: u32, heads: u32, hd: u32, eps: f32, scale: f32, raw: u32, _b: u32, _c: u32 };
@group(0) @binding(0) var<storage, read>       drp_src : array<f32>;
@group(0) @binding(1) var<storage, read_write> drp_dst : array<f32>;
@group(0) @binding(2) var<storage, read>       drp_w   : array<f32>;
@group(0) @binding(3) var<storage, read>       drp_cos : array<f32>;
@group(0) @binding(4) var<storage, read>       drp_sin : array<f32>;
@group(0) @binding(5) var<uniform>             drp_p   : DrpP;
var<workgroup> drp_part: array<f32, 128>;
@compute @workgroup_size(128)
fn dit_ropepack(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_index) lid: u32) {
    let tokn = wid.x;
    let hh = wid.y;
    if (tokn >= drp_p.n || hh >= drp_p.heads) { return; }
    let hd = drp_p.hd;
    let src = (tokn * drp_p.heads + hh) * hd;
    var acc = 0.0;
    var i = lid;
    loop {
        if (i >= hd) { break; }
        let v = drp_src[src + i];
        acc = acc + v * v;
        i = i + 128u;
    }
    drp_part[lid] = acc;
    workgroupBarrier();
    var stride = 64u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { drp_part[lid] = drp_part[lid] + drp_part[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    // `raw`: transpose only — v needs the token-major to head-major move
    // without the qk-norm or the rotation, and a norm with a huge eps
    // would zero it rather than skip it.
    // `raw`: transpose only — v needs the token-major to head-major move
    // with neither the qk-norm nor the rotation, so the weight and the
    // rope table it is handed are ignored rather than faked.
    var inv = 1.0;
    if (drp_p.raw == 0u) { inv = inverseSqrt(drp_part[0] / f32(hd) + drp_p.eps); }
    let pairs = hd / 2u;
    let dst = (hh * drp_p.n + tokn) * hd;
    var j = lid;
    loop {
        if (j >= pairs) { break; }
        var a = drp_src[src + 2u * j];
        var b = drp_src[src + 2u * j + 1u];
        var cs = 1.0;
        var sn = 0.0;
        if (drp_p.raw == 0u) {
            a = a * inv * drp_w[2u * j];
            b = b * inv * drp_w[2u * j + 1u];
            cs = drp_cos[tokn * pairs + j];
            sn = drp_sin[tokn * pairs + j];
        }
        drp_dst[dst + 2u * j] = (a * cs - b * sn) * drp_p.scale;
        drp_dst[dst + 2u * j + 1u] = (a * sn + b * cs) * drp_p.scale;
        j = j + 128u;
    }
}

// ── VAE convolution (direct, same padding, stride 1) ────────────────────
// One thread per (output channel, pixel). The CPU path builds an im2col
// matrix — multi-GB at 512² — and the Metal path has had a device kernel
// since the image runtime shipped; without this twin every Vulkan/DX12/
// Android box decoded the latent on the CPU, which is 3.1 s of a 11.8 s
// 256² render and grows with the pixels.
// `up2` != 0 reads the source at (y/2, x/2): the nearest-2× upsample the
// decoder does before its conv, fused so only the SMALL image crosses.
struct VcP { ic: u32, oc: u32, h: u32, w: u32, k: u32, up2: u32, sh: u32, sw: u32 };
@group(0) @binding(0) var<storage, read>       vc_w : array<f32>;
@group(0) @binding(1) var<storage, read>       vc_b : array<f32>;
@group(0) @binding(2) var<storage, read>       vc_x : array<f32>;
@group(0) @binding(3) var<storage, read_write> vc_y : array<f32>;
@group(0) @binding(4) var<uniform>             vc_p : VcP;
@compute @workgroup_size(64)
fn vae_conv(@builtin(global_invocation_id) gid: vec3<u32>) {
    let hw = vc_p.h * vc_p.w;
    let pix = gid.x;
    let oc = gid.y;
    if (pix >= hw || oc >= vc_p.oc) { return; }
    let oy = pix / vc_p.w;
    let ox = pix % vc_p.w;
    let k = vc_p.k;
    let pad = k / 2u;
    let shw = vc_p.sh * vc_p.sw;
    var acc = vc_b[oc];
    for (var ci = 0u; ci < vc_p.ic; ci = ci + 1u) {
        let wbase = (oc * vc_p.ic + ci) * k * k;
        let xbase = ci * shw;
        for (var ky = 0u; ky < k; ky = ky + 1u) {
            // Signed arithmetic in u32: skip taps that fall outside.
            let sy = oy + ky;
            if (sy < pad) { continue; }
            var iy = sy - pad;
            if (vc_p.up2 != 0u) { iy = iy / 2u; }
            if (iy >= vc_p.sh) { continue; }
            for (var kx = 0u; kx < k; kx = kx + 1u) {
                let sx = ox + kx;
                if (sx < pad) { continue; }
                var ix = sx - pad;
                if (vc_p.up2 != 0u) { ix = ix / 2u; }
                if (ix >= vc_p.sw) { continue; }
                acc = acc + vc_w[wbase + ky * k + kx] * vc_x[xbase + iy * vc_p.sw + ix];
            }
        }
    }
    vc_y[oc * hw + pix] = acc;
}

// RMSNorm of one row (WGSL twin of Metal rmsnorm_k): o = x·rsqrt(mean(x²)+eps)·w',
// w' = w or (1+w) for gemma. One workgroup, 256-thread tree reduction — the
// building block that keeps the token graph's hidden resident across the norm.
struct RmsP { n: u32, gemma: u32, eps: f32, _p: u32 };
@group(0) @binding(0) var<storage, read>       rn_x : array<f32>;
@group(0) @binding(1) var<storage, read>       rn_w : array<f32>;
@group(0) @binding(2) var<storage, read_write> rn_o : array<f32>;
@group(0) @binding(3) var<uniform>             rn_p : RmsP;
// A THOUSAND threads, not 256. This kernel reduces, so it is one workgroup
// by construction — one SM of two hundred — and a dispatch costs 2.7 µs
// while this one takes ~23. There is nothing to hide the load latency
// behind except more threads on the same SM.
var<workgroup> rn_part: array<f32, 1024>;
@compute @workgroup_size(1024)
fn rmsnorm(@builtin(local_invocation_id) lid: vec3<u32>) {
    let tid = lid.x;
    let n = rn_p.n;
    var acc = 0.0;
    var i = tid;
    loop {
        if (i >= n) { break; }
        let v = rn_x[i];
        acc = acc + v * v;
        i = i + 1024u;
    }
    rn_part[tid] = acc;
    workgroupBarrier();
    var stride = 512u;
    loop {
        if (stride == 0u) { break; }
        if (tid < stride) { rn_part[tid] = rn_part[tid] + rn_part[tid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let inv = inverseSqrt(rn_part[0] / f32(n) + rn_p.eps);
    i = tid;
    loop {
        if (i >= n) { break; }
        var wv = rn_w[i];
        if (rn_p.gemma == 1u) { wv = 1.0 + wv; }
        rn_o[i] = rn_x[i] * inv * wv;
        i = i + 1024u;
    }
}

// GDN depthwise causal conv + SiLU over the ring buffer of the last kk-1
// positions plus the current qkv, then shift the ring (drop oldest, append
// current). One thread per conv channel. WGSL twin of the Metal gdn_conv.
struct GcP { cdim: u32, kk: u32, xoff: u32, _b: u32 };
@group(0) @binding(0) var<storage, read>       gc_qkv  : array<f32>;
@group(0) @binding(1) var<storage, read>       gc_taps : array<f32>;
@group(0) @binding(2) var<storage, read_write> gc_ring : array<f32>;
@group(0) @binding(3) var<storage, read_write> gc_cq   : array<f32>;
@group(0) @binding(4) var<uniform>             gc_p    : GcP;
@compute @workgroup_size(256)
fn gdn_conv(@builtin(global_invocation_id) gid: vec3<u32>) {
    let c = gid.x;
    let cdim = gc_p.cdim;
    if (c >= cdim) { return; }
    let kk = gc_p.kk;
    let tb = c * kk;
    var acc = gc_qkv[gc_p.xoff + c] * gc_taps[tb + kk - 1u];
    for (var j = 0u; j + 1u < kk; j = j + 1u) {
        acc = acc + gc_ring[j * cdim + c] * gc_taps[tb + j];
    }
    gc_cq[c] = acc / (1.0 + exp(-acc));
    // ring shift (columns are independent per thread c)
    for (var j = 0u; j + 2u < kk; j = j + 1u) {
        gc_ring[j * cdim + c] = gc_ring[(j + 1u) * cdim + c];
    }
    if (kk > 1u) {
        gc_ring[(kk - 2u) * cdim + c] = gc_qkv[gc_p.xoff + c];
    }
}

// LFM2 gated short-conv decode step (the mixer 22 of an LFM2.5-2.6B's 30
// layers are). Reuses gdn_conv's binding set so the module stays valid by
// construction: gc_qkv = bcx [B | C | x] (3h), gc_taps = [h·k] channel-major
// taps, gc_ring = the ring in the HOST layout ([channel][k-1], slot 0
// newest — the seed comes from kv_cache.linear_state without repacking),
// gc_cq = y out (h), gc_p: cdim = h, kk = k, xoff unused. One thread per
// channel; y = C · (tap[k-1]·B·x + Σ tap[k-2-s]·ring[s]), then the shift.
@compute @workgroup_size(256)
fn sconv_step(@builtin(global_invocation_id) gid: vec3<u32>) {
    let ch = gid.x;
    let h = gc_p.cdim;
    if (ch >= h) { return; }
    let k = gc_p.kk;
    let ring = k - 1u;
    let bx = gc_qkv[ch] * gc_qkv[2u * h + ch];
    let tb = ch * k;
    var acc = gc_taps[tb + k - 1u] * bx;
    for (var s = 0u; s < ring; s = s + 1u) {
        acc = acc + gc_taps[tb + k - 2u - s] * gc_ring[ch * ring + s];
    }
    gc_cq[ch] = gc_qkv[h + ch] * acc;
    var j = ring;
    while (j > 1u) {
        j = j - 1u;
        gc_ring[ch * ring + j] = gc_ring[ch * ring + j - 1u];
    }
    if (ring > 0u) { gc_ring[ch * ring] = bx; }
}

// ── GDN (gated DeltaNet / linear attention) decode step ──────────────────
// One workgroup per v-head. From the conv output cq it l2-norms q/k, forms the
// decay g and gate β, runs the delta-rule state recurrence S ← g·S + kf⊗β(v −
// kfᵀS) with o = qfᵀS, then the gated RMSNorm o·norm·silu(z). S ([nv,dk,dv])
// persists across tokens (device state buffer). WGSL twin of the Metal GDN
// state-update kernel; dk,dv ≤ 256.
struct GdnP { nv: u32, dk: u32, dv: u32, kd: u32, rep: u32, cdim: u32, eps: f32, tok: u32 };
@group(0) @binding(0) var<storage, read>       gd_cq   : array<f32>;
@group(0) @binding(1) var<storage, read>       gd_z    : array<f32>;
@group(0) @binding(2) var<storage, read>       gd_a    : array<f32>;
@group(0) @binding(3) var<storage, read>       gd_b    : array<f32>;
@group(0) @binding(4) var<storage, read>       gd_alog : array<f32>;
@group(0) @binding(5) var<storage, read>       gd_dtb  : array<f32>;
@group(0) @binding(6) var<storage, read>       gd_norm : array<f32>;
@group(0) @binding(7) var<storage, read_write> gd_S    : array<f32>;
@group(0) @binding(8) var<storage, read_write> gd_o    : array<f32>;
// S and o again as vec4 (same-slot rule): a state row is dv-contiguous,
// so a lane's four columns are ONE 16-byte access, not four scattered.
@group(0) @binding(7) var<storage, read_write> gd_S4   : array<vec4<f32>>;
@group(0) @binding(8) var<storage, read_write> gd_o4   : array<vec4<f32>>;
@group(0) @binding(9) var<uniform>             gd_p    : GdnP;
var<workgroup> gd_kf: array<f32, 256>;
var<workgroup> gd_qf: array<f32, 256>;
var<workgroup> gd_ov: array<f32, 256>;
var<workgroup> gd_red: array<f32, 256>;
var<workgroup> gd_red2: array<f32, 256>;
var<workgroup> gd_red3: array<f32, 256>;
var<workgroup> gd_red4: array<f32, 256>;
fn gd_softplus(x: f32) -> f32 {
    if (x > 20.0) { return x; }
    return log(1.0 + exp(x));
}
fn gd_reduce(t: u32) -> f32 {
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (t < stride) { gd_red[t] = gd_red[t] + gd_red[t + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    return gd_red[0];
}
@compute @workgroup_size(256)
fn gdn_step(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let t = lid.x;
    if (h >= gd_p.nv) { return; }
    let dk = gd_p.dk;
    let dv = gd_p.dv;
    let ko = h / gd_p.rep;
    let qs = ko * dk;
    let ks = gd_p.kd + ko * dk;
    // l2-norm of q then k over dk
    gd_red[t] = select(0.0, gd_cq[qs + t] * gd_cq[qs + t], t < dk);
    workgroupBarrier();
    let nq = gd_reduce(t);
    workgroupBarrier();
    gd_red[t] = select(0.0, gd_cq[ks + t] * gd_cq[ks + t], t < dk);
    workgroupBarrier();
    let nkn = gd_reduce(t);
    workgroupBarrier();
    let invq = 1.0 / (sqrt(nq + 1e-6) * sqrt(f32(dk)));
    let invk = 1.0 / sqrt(nkn + 1e-6);
    if (t < dk) {
        gd_qf[t] = gd_cq[qs + t] * invq;
        gd_kf[t] = gd_cq[ks + t] * invk;
    }
    workgroupBarrier();
    let abo = gd_p.tok * gd_p.nv;
    let g = exp(-exp(gd_alog[h]) * gd_softplus(gd_a[abo + h] + gd_dtb[h]));
    let beta = 1.0 / (1.0 + exp(-gd_b[abo + h]));
    let sbase = h * dk * dv;
    if (t < dv) {
        let dj = t;
        let vt = gd_cq[2u * gd_p.kd + h * dv + dj];
        var kv = 0.0;
        for (var di = 0u; di < dk; di = di + 1u) { kv = kv + gd_S[sbase + di * dv + dj] * gd_kf[di]; }
        let delta = (vt - g * kv) * beta;
        var o = 0.0;
        for (var di = 0u; di < dk; di = di + 1u) {
            let idx = sbase + di * dv + dj;
            let cell = g * gd_S[idx] + gd_kf[di] * delta;
            gd_S[idx] = cell;
            o = o + gd_qf[di] * cell;
        }
        gd_ov[dj] = o;
    }
    workgroupBarrier();
    // gated RMSNorm over dv
    gd_red[t] = select(0.0, gd_ov[t] * gd_ov[t], t < dv);
    workgroupBarrier();
    let ss = gd_reduce(t);
    workgroupBarrier();
    let inv = 1.0 / sqrt(ss / f32(dv) + gd_p.eps);
    if (t < dv) {
        let zo = gd_p.tok * gd_p.nv * dv;
        let zz = gd_z[zo + h * dv + t];
        gd_o[zo + h * dv + t] = gd_ov[t] * inv * gd_norm[t] * (zz / (1.0 + exp(-zz)));
    }
}

// ── GDN, k-looped twins for the batched verify: the position recurrence
// stays INSIDE the kernel, so a layer is two dispatches with one barrier
// between them instead of 2k dispatches with 2k barrier drains — which
// were 7-8 ms of a 3-position verify. Columns are independent in the
// conv (no barriers at all); heads are independent in the step, so the
// loop lives inside the workgroup. When snap_stride > 0 each position's
// (ring, S) lands in the snapshot buffer as it is produced — the rows a
// partial acceptance restores from, at zero extra passes.
struct GcKP { cdim: u32, kk: u32, kb: u32, snap: u32, stride: u32, p0: u32, p1: u32, p2: u32 };
@group(0) @binding(0) var<storage, read>       gck_qkv : array<f32>;
@group(0) @binding(1) var<storage, read>       gck_taps: array<f32>;
@group(0) @binding(2) var<storage, read_write> gck_ring: array<f32>;
@group(0) @binding(3) var<storage, read_write> gck_cq  : array<f32>;
@group(0) @binding(4) var<uniform>             gck_p   : GcKP;
@group(0) @binding(5) var<storage, read_write> gck_snap: array<f32>;
@compute @workgroup_size(256)
fn gdn_conv_k(@builtin(global_invocation_id) gid: vec3<u32>) {
    let c = gid.x;
    let cdim = gck_p.cdim;
    if (c >= cdim) { return; }
    let kk = gck_p.kk;
    let tb = c * kk;
    for (var i = 0u; i < gck_p.kb; i = i + 1u) {
        let x = gck_qkv[i * cdim + c];
        var acc = x * gck_taps[tb + kk - 1u];
        for (var j = 0u; j + 1u < kk; j = j + 1u) {
            acc = acc + gck_ring[j * cdim + c] * gck_taps[tb + j];
        }
        gck_cq[i * cdim + c] = acc / (1.0 + exp(-acc));
        for (var j = 0u; j + 2u < kk; j = j + 1u) {
            gck_ring[j * cdim + c] = gck_ring[(j + 1u) * cdim + c];
        }
        if (kk > 1u) {
            gck_ring[(kk - 2u) * cdim + c] = x;
        }
        if (gck_p.snap != 0u) {
            let off = i * gck_p.stride;
            for (var j = 0u; j + 1u < kk; j = j + 1u) {
                gck_snap[off + j * cdim + c] = gck_ring[j * cdim + c];
            }
        }
    }
}

struct GdKP {
    nv: u32, dk: u32, dv: u32, kd: u32,
    rep: u32, cdim: u32, eps: f32, kb: u32,
    stride: u32, ring_els: u32, p0: u32, p1: u32,
};
@group(0) @binding(0) var<storage, read>       gdk_cq   : array<f32>;
@group(0) @binding(1) var<storage, read>       gdk_z    : array<f32>;
@group(0) @binding(2) var<storage, read>       gdk_a    : array<f32>;
@group(0) @binding(3) var<storage, read>       gdk_b    : array<f32>;
@group(0) @binding(4) var<storage, read>       gdk_alog : array<f32>;
@group(0) @binding(5) var<storage, read>       gdk_dtb  : array<f32>;
@group(0) @binding(6) var<storage, read>       gdk_norm : array<f32>;
@group(0) @binding(7) var<storage, read_write> gdk_S    : array<f32>;
@group(0) @binding(8) var<storage, read_write> gdk_o    : array<f32>;
@group(0) @binding(9) var<uniform>             gdk_p    : GdKP;
@group(0) @binding(10) var<storage, read_write> gdk_snap: array<f32>;
var<workgroup> gdk_kf: array<f32, 256>;
var<workgroup> gdk_qf: array<f32, 256>;
var<workgroup> gdk_ov: array<f32, 256>;
var<workgroup> gdk_red: array<f32, 256>;
fn gdk_reduce(t: u32) -> f32 {
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (t < stride) { gdk_red[t] = gdk_red[t] + gdk_red[t + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    return gdk_red[0];
}
@compute @workgroup_size(256)
fn gdn_step_k(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let t = lid.x;
    if (h >= gdk_p.nv) { return; }
    let dk = gdk_p.dk;
    let dv = gdk_p.dv;
    let ko = h / gdk_p.rep;
    let sbase = h * dk * dv;
    for (var i = 0u; i < gdk_p.kb; i = i + 1u) {
        let cq0 = i * gdk_p.cdim;
        let qs = cq0 + ko * dk;
        let ks = cq0 + gdk_p.kd + ko * dk;
        gdk_red[t] = select(0.0, gdk_cq[qs + t] * gdk_cq[qs + t], t < dk);
        workgroupBarrier();
        let nq = gdk_reduce(t);
        workgroupBarrier();
        gdk_red[t] = select(0.0, gdk_cq[ks + t] * gdk_cq[ks + t], t < dk);
        workgroupBarrier();
        let nkn = gdk_reduce(t);
        workgroupBarrier();
        let invq = 1.0 / (sqrt(nq + 1e-6) * sqrt(f32(dk)));
        let invk = 1.0 / sqrt(nkn + 1e-6);
        if (t < dk) {
            gdk_qf[t] = gdk_cq[qs + t] * invq;
            gdk_kf[t] = gdk_cq[ks + t] * invk;
        }
        workgroupBarrier();
        let abo = i * gdk_p.nv;
        let g = exp(-exp(gdk_alog[h]) * gd_softplus(gdk_a[abo + h] + gdk_dtb[h]));
        let beta = 1.0 / (1.0 + exp(-gdk_b[abo + h]));
        if (t < dv) {
            let dj = t;
            let vt = gdk_cq[cq0 + 2u * gdk_p.kd + h * dv + dj];
            var kv = 0.0;
            for (var di = 0u; di < dk; di = di + 1u) {
                kv = kv + gdk_S[sbase + di * dv + dj] * gdk_kf[di];
            }
            let delta = (vt - g * kv) * beta;
            var o = 0.0;
            let snap0 = i * gdk_p.stride + gdk_p.ring_els + sbase;
            for (var di = 0u; di < dk; di = di + 1u) {
                let idx = sbase + di * dv + dj;
                let cell = g * gdk_S[idx] + gdk_kf[di] * delta;
                gdk_S[idx] = cell;
                if (gdk_p.stride != 0u) {
                    gdk_snap[snap0 + di * dv + dj] = cell;
                }
                o = o + gdk_qf[di] * cell;
            }
            gdk_ov[dj] = o;
        }
        workgroupBarrier();
        gdk_red[t] = select(0.0, gdk_ov[t] * gdk_ov[t], t < dv);
        workgroupBarrier();
        let ss = gdk_reduce(t);
        workgroupBarrier();
        let inv = 1.0 / sqrt(ss / f32(dv) + gdk_p.eps);
        if (t < dv) {
            let zo = i * gdk_p.nv * dv;
            let zz = gdk_z[zo + h * dv + t];
            gdk_o[zo + h * dv + t] = gdk_ov[t] * inv * gdk_norm[t] * (zz / (1.0 + exp(-zz)));
        }
        workgroupBarrier();
    }
}

// ── GDN step, parallel edition: one WORKGROUP PER (head, column). The
// one-workgroup-per-head kernel put 32 workgroups on a 188-SM card — 8%
// occupancy, 3.7 ms/token of a 12.5 ms frame on the 2-bit 35B. Column j
// is independent under the delta rule, and both dk-loops become 128-lane
// tree reductions. Reduction order differs from the serial kernel, so
// bits differ within the documented GPU tie class; CMF_GDN_PAR=0 keeps
// the old kernel for A/B. The raw o lands in gd_o and gdn_step_norm
// applies the gated RMSNorm in place.
@compute @workgroup_size(128)
fn gdn_step_par(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let dj4 = wid.y;   // FOUR columns per workgroup, vec4 access
    let t = lid.x;
    let dk = gd_p.dk;
    let dv = gd_p.dv;
    if (h >= gd_p.nv || dj4 * 4u >= dv) { return; }
    let ko = h / gd_p.rep;
    let qs = ko * dk;
    let ks = gd_p.kd + ko * dk;
    // q/k l2 norms over dk (identical formulas, tree order)
    gd_red[t] = select(0.0, gd_cq[qs + t] * gd_cq[qs + t], t < dk);
    workgroupBarrier();
    var stride = 64u;
    loop {
        if (stride == 0u) { break; }
        if (t < stride) { gd_red[t] = gd_red[t] + gd_red[t + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let nq = gd_red[0];
    workgroupBarrier();
    gd_red[t] = select(0.0, gd_cq[ks + t] * gd_cq[ks + t], t < dk);
    workgroupBarrier();
    stride = 64u;
    loop {
        if (stride == 0u) { break; }
        if (t < stride) { gd_red[t] = gd_red[t] + gd_red[t + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let nkn = gd_red[0];
    workgroupBarrier();
    let invq = 1.0 / (sqrt(nq + 1e-6) * sqrt(f32(dk)));
    let invk = 1.0 / sqrt(nkn + 1e-6);
    let abo = gd_p.tok * gd_p.nv;
    let g = exp(-exp(gd_alog[h]) * gd_softplus(gd_a[abo + h] + gd_dtb[h]));
    let beta = 1.0 / (1.0 + exp(-gd_b[abo + h]));
    let s4base = (h * dk * dv) >> 2u;
    let dv4 = dv >> 2u;
    let vto = 2u * gd_p.kd + h * dv + dj4 * 4u;
    let vt = vec4<f32>(gd_cq[vto], gd_cq[vto + 1u], gd_cq[vto + 2u], gd_cq[vto + 3u]);
    let kf_t = select(0.0, gd_cq[ks + t] * invk, t < dk);
    let qf_t = select(0.0, gd_cq[qs + t] * invq, t < dk);
    // kv = kfᵀ S[:, j..j+3] — the four column reductions ride together
    var kv4 = vec4<f32>(0.0);
    if (t < dk) {
        kv4 = gd_S4[s4base + t * dv4 + dj4] * kf_t;
    }
    gd_red[t] = kv4.x;
    gd_red2[t] = kv4.y;
    gd_red3[t] = kv4.z;
    gd_red4[t] = kv4.w;
    workgroupBarrier();
    stride = 64u;
    loop {
        if (stride == 0u) { break; }
        if (t < stride) {
            gd_red[t] = gd_red[t] + gd_red[t + stride];
            gd_red2[t] = gd_red2[t] + gd_red2[t + stride];
            gd_red3[t] = gd_red3[t] + gd_red3[t + stride];
            gd_red4[t] = gd_red4[t] + gd_red4[t + stride];
        }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let kv = vec4<f32>(gd_red[0], gd_red2[0], gd_red3[0], gd_red4[0]);
    workgroupBarrier();
    let delta = (vt - g * kv) * beta;
    var contrib = vec4<f32>(0.0);
    if (t < dk) {
        let idx = s4base + t * dv4 + dj4;
        let cell = g * gd_S4[idx] + kf_t * delta;
        gd_S4[idx] = cell;
        contrib = qf_t * cell;
    }
    gd_red[t] = contrib.x;
    gd_red2[t] = contrib.y;
    gd_red3[t] = contrib.z;
    gd_red4[t] = contrib.w;
    workgroupBarrier();
    stride = 64u;
    loop {
        if (stride == 0u) { break; }
        if (t < stride) {
            gd_red[t] = gd_red[t] + gd_red[t + stride];
            gd_red2[t] = gd_red2[t] + gd_red2[t + stride];
            gd_red3[t] = gd_red3[t] + gd_red3[t + stride];
            gd_red4[t] = gd_red4[t] + gd_red4[t + stride];
        }
        workgroupBarrier();
        stride = stride / 2u;
    }
    if (t == 0u) {
        let zo4 = (gd_p.tok * gd_p.nv * dv) >> 2u;
        gd_o4[zo4 + h * dv4 + dj4] =
            vec4<f32>(gd_red[0], gd_red2[0], gd_red3[0], gd_red4[0]);
    }
}

// v2 of the parallel pair: the conv is INLINE (the same taps math, the
// same order, computed per element from the PRE-shift ring), so the par
// kernel no longer waits on a conv dispatch — and the ring shift rides
// the norm kernel, which was going to run anyway. One dependent hop
// fewer per GDN layer, thirty layers a frame.
struct GciP { kk: u32, xoff: u32, _a: u32, _b: u32 };
@group(0) @binding(10) var<storage, read>       gi_qkv  : array<f32>;
@group(0) @binding(11) var<storage, read_write> gi_ring : array<f32>;
@group(0) @binding(12) var<storage, read>       gi_taps : array<f32>;
@group(0) @binding(13) var<uniform>             gi_p    : GciP;

fn gi_cq(c: u32) -> f32 {
    let kk = gi_p.kk;
    let tb = c * kk;
    var acc = gi_qkv[gi_p.xoff + c] * gi_taps[tb + kk - 1u];
    for (var j = 0u; j + 1u < kk; j = j + 1u) {
        acc = acc + gi_ring[j * gd_p.cdim + c] * gi_taps[tb + j];
    }
    return acc / (1.0 + exp(-acc));
}

@compute @workgroup_size(128)
fn gdn_step_par2(@builtin(workgroup_id) wid: vec3<u32>,
                 @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let dj = wid.y;
    let t = lid.x;
    let dk = gd_p.dk;
    let dv = gd_p.dv;
    if (h >= gd_p.nv || dj >= dv) { return; }
    let ko = h / gd_p.rep;
    let qs = ko * dk;
    let ks = gd_p.kd + ko * dk;
    let cq_q = select(0.0, gi_cq(qs + t), t < dk);
    let cq_k = select(0.0, gi_cq(ks + t), t < dk);
    gd_red[t] = cq_q * cq_q;
    workgroupBarrier();
    var stride = 64u;
    loop {
        if (stride == 0u) { break; }
        if (t < stride) { gd_red[t] = gd_red[t] + gd_red[t + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let nq = gd_red[0];
    workgroupBarrier();
    gd_red[t] = cq_k * cq_k;
    workgroupBarrier();
    stride = 64u;
    loop {
        if (stride == 0u) { break; }
        if (t < stride) { gd_red[t] = gd_red[t] + gd_red[t + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let nkn = gd_red[0];
    workgroupBarrier();
    let invq = 1.0 / (sqrt(nq + 1e-6) * sqrt(f32(dk)));
    let invk = 1.0 / sqrt(nkn + 1e-6);
    let abo = gd_p.tok * gd_p.nv;
    let g = exp(-exp(gd_alog[h]) * gd_softplus(gd_a[abo + h] + gd_dtb[h]));
    let beta = 1.0 / (1.0 + exp(-gd_b[abo + h]));
    let sbase = h * dk * dv;
    let vt = gi_cq(2u * gd_p.kd + h * dv + dj);
    let kf_t = cq_k * invk;
    let qf_t = cq_q * invq;
    gd_red[t] = select(0.0, gd_S[sbase + t * dv + dj] * kf_t, t < dk);
    workgroupBarrier();
    stride = 64u;
    loop {
        if (stride == 0u) { break; }
        if (t < stride) { gd_red[t] = gd_red[t] + gd_red[t + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let kv = gd_red[0];
    workgroupBarrier();
    let delta = (vt - g * kv) * beta;
    var contrib = 0.0;
    if (t < dk) {
        let idx = sbase + t * dv + dj;
        let cell = g * gd_S[idx] + kf_t * delta;
        gd_S[idx] = cell;
        contrib = qf_t * cell;
    }
    gd_red[t] = contrib;
    workgroupBarrier();
    stride = 64u;
    loop {
        if (stride == 0u) { break; }
        if (t < stride) { gd_red[t] = gd_red[t] + gd_red[t + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    if (t == 0u) {
        let zo = gd_p.tok * gd_p.nv * dv;
        gd_o[zo + h * dv + dj] = gd_red[0];
    }
}

// norm v2: the gated RMSNorm PLUS the ring shift the conv kernel used to
// do — its writers (par2's gi_cq readers) are all upstream in the pass.
@compute @workgroup_size(256)
fn gdn_step_norm2(@builtin(workgroup_id) wid: vec3<u32>,
                  @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let t = lid.x;
    let dv = gd_p.dv;
    if (h >= gd_p.nv) { return; }
    let zo = gd_p.tok * gd_p.nv * dv;
    gd_red[t] = select(0.0, gd_o[zo + h * dv + t] * gd_o[zo + h * dv + t], t < dv);
    workgroupBarrier();
    let ss = gd_reduce(t);
    workgroupBarrier();
    let inv = 1.0 / sqrt(ss / f32(dv) + gd_p.eps);
    if (t < dv) {
        let zz = gd_z[zo + h * dv + t];
        gd_o[zo + h * dv + t] =
            gd_o[zo + h * dv + t] * inv * gd_norm[t] * (zz / (1.0 + exp(-zz)));
    }
    // ring shift, strided over cdim across all norm workgroups
    let kk = gi_p.kk;
    let cdim = gd_p.cdim;
    var c = wid.x * 256u + t;
    loop {
        if (c >= cdim) { break; }
        for (var j = 0u; j + 2u < kk; j = j + 1u) {
            gi_ring[j * cdim + c] = gi_ring[(j + 1u) * cdim + c];
        }
        if (kk > 1u) {
            gi_ring[(kk - 2u) * cdim + c] = gi_qkv[gi_p.xoff + c];
        }
        c = c + gd_p.nv * 256u;
    }
}

// Gated RMSNorm tail of the parallel GDN step: in place over gd_o.
@compute @workgroup_size(256)
fn gdn_step_norm(@builtin(workgroup_id) wid: vec3<u32>,
                 @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let t = lid.x;
    let dv = gd_p.dv;
    if (h >= gd_p.nv) { return; }
    let zo = gd_p.tok * gd_p.nv * dv;
    gd_red[t] = select(0.0, gd_o[zo + h * dv + t] * gd_o[zo + h * dv + t], t < dv);
    workgroupBarrier();
    let ss = gd_reduce(t);
    workgroupBarrier();
    let inv = 1.0 / sqrt(ss / f32(dv) + gd_p.eps);
    if (t < dv) {
        let zz = gd_z[zo + h * dv + t];
        gd_o[zo + h * dv + t] =
            gd_o[zo + h * dv + t] * inv * gd_norm[t] * (zz / (1.0 + exp(-zz)));
    }
}

// Fused residual-add + RMSNorm (WGSL twin of Metal add_rmsnorm_rows): h += d
// in place, then o = rms(h)·w. Collapses an axpy + an rmsnorm dispatch into
// one — cuts two launches per layer off the token graph.
struct ArP { n: u32, gemma: u32, eps: f32, _p: u32 };
@group(0) @binding(0) var<storage, read_write> ar_h : array<f32>;
@group(0) @binding(1) var<storage, read>       ar_d : array<f32>;
@group(0) @binding(2) var<storage, read>       ar_w : array<f32>;
@group(0) @binding(3) var<storage, read_write> ar_o : array<f32>;
@group(0) @binding(4) var<uniform>             ar_p : ArP;
var<workgroup> ar_part: array<f32, 256>;
@compute @workgroup_size(256)
fn add_rmsnorm(@builtin(local_invocation_id) lid: vec3<u32>) {
    let tid = lid.x;
    let n = ar_p.n;
    var acc = 0.0;
    var i = tid;
    loop {
        if (i >= n) { break; }
        let v = ar_h[i] + ar_d[i];
        ar_h[i] = v;
        acc = acc + v * v;
        i = i + 256u;
    }
    ar_part[tid] = acc;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (tid < stride) { ar_part[tid] = ar_part[tid] + ar_part[tid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let inv = inverseSqrt(ar_part[0] / f32(n) + ar_p.eps);
    i = tid;
    loop {
        if (i >= n) { break; }
        var wv = ar_w[i];
        if (ar_p.gemma == 1u) { wv = 1.0 + wv; }
        ar_o[i] = ar_h[i] * inv * wv;
        i = i + 256u;
    }
}

// Batched RMSNorm for prefill: one workgroup per row (wid.x), row r reads/writes
// rn_x[r*n..] → rn_o[r*n..]; the weight rn_w[n] is shared. K prompt positions
// norm in one dispatch (twin of `rmsnorm`, strided by row).
@compute @workgroup_size(256)
fn rmsnorm_b(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let tid = lid.x;
    let n = rn_p.n;
    let base = wid.x * n;
    var acc = 0.0;
    var i = tid;
    loop { if (i >= n) { break; } let v = rn_x[base + i]; acc = acc + v * v; i = i + 256u; }
    rn_part[tid] = acc;
    workgroupBarrier();
    var stride = 128u;
    loop { if (stride == 0u) { break; } if (tid < stride) { rn_part[tid] = rn_part[tid] + rn_part[tid + stride]; } workgroupBarrier(); stride = stride / 2u; }
    let inv = inverseSqrt(rn_part[0] / f32(n) + rn_p.eps);
    i = tid;
    loop { if (i >= n) { break; } var wv = rn_w[i]; if (rn_p.gemma == 1u) { wv = 1.0 + wv; } rn_o[base + i] = rn_x[base + i] * inv * wv; i = i + 256u; }
}

// Batched fused residual-add + RMSNorm (one workgroup per row): ar_h[r] += ar_d[r]
// in place, then ar_o[r] = rms(ar_h[r])·w. Prefill twin of `add_rmsnorm`.
@compute @workgroup_size(256)
fn add_rmsnorm_b(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let tid = lid.x;
    let n = ar_p.n;
    let base = wid.x * n;
    var acc = 0.0;
    var i = tid;
    loop { if (i >= n) { break; } let v = ar_h[base + i] + ar_d[base + i]; ar_h[base + i] = v; acc = acc + v * v; i = i + 256u; }
    ar_part[tid] = acc;
    workgroupBarrier();
    var stride = 128u;
    loop { if (stride == 0u) { break; } if (tid < stride) { ar_part[tid] = ar_part[tid] + ar_part[tid + stride]; } workgroupBarrier(); stride = stride / 2u; }
    let inv = inverseSqrt(ar_part[0] / f32(n) + ar_p.eps);
    i = tid;
    loop { if (i >= n) { break; } var wv = ar_w[i]; if (ar_p.gemma == 1u) { wv = 1.0 + wv; } ar_o[base + i] = ar_h[base + i] * inv * wv; i = i + 256u; }
}

// RoPE + optional qk-norm + gate-split, one 32-thread workgroup per head
// (WGSL twin of Metal attn_rope_qkn; the qk-norm sum-of-squares reduces in
// workgroup memory — no subgroup ops, portable). Heads [0,nh)=Q (2·hd each
// when gated: q||gate), [nh,nh+nkv)=K. flags: 1=gate 2=qnorm 4=knorm 8=gemma
// 32=norm-after-rope (HunYuan dense).
struct RqP { nh: u32, nkv: u32, hd: u32, rd: u32, pos: u32, flags: u32, eps: f32, tok: u32 };
@group(0) @binding(0) var<storage, read>       rq_qraw : array<f32>;
@group(0) @binding(1) var<storage, read_write> rq_k    : array<f32>;
@group(0) @binding(2) var<storage, read_write> rq_qout : array<f32>;
@group(0) @binding(3) var<storage, read_write> rq_gout : array<f32>;
@group(0) @binding(4) var<storage, read>       rq_qnw  : array<f32>;
@group(0) @binding(5) var<storage, read>       rq_knw  : array<f32>;
@group(0) @binding(6) var<storage, read>       rq_invf : array<f32>;
@group(0) @binding(7) var<uniform>             rq_p    : RqP;
var<workgroup> rq_red: array<f32, 32>;
var<workgroup> rq_head: array<f32, 256>;
@compute @workgroup_size(32)
fn attn_rope_qkn(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let head = wid.x;
    let lane = lid.x;
    let nh = rq_p.nh;
    let hd = rq_p.hd;
    if (head >= nh + rq_p.nkv) { return; }
    let isq = head < nh;
    let gate = (rq_p.flags & 1u) != 0u;
    let src_base = select((head - nh) * hd, head * select(1u, 2u, gate) * hd, isq);
    // Batch-graph token offsets (0 in the token graph): q rows live in the
    // batched projection output, K is rotated IN PLACE in its batch slice.
    let qoff = rq_p.tok * nh * select(1u, 2u, gate) * hd;
    let koff = rq_p.tok * rq_p.nkv * hd;
    let nt = (hd + 31u) / 32u;  // ≤ 8 for head_dim ≤ 256 (Qwen3.5 uses 256)
    var xv: array<f32, 8>;
    var ss = 0.0;
    for (var t = 0u; t < nt; t = t + 1u) {
        let d = t * 32u + lane;
        var val = 0.0;
        if (d < hd) { val = select(rq_k[koff + src_base + d], rq_qraw[qoff + src_base + d], isq); }
        xv[t] = val;
        ss = ss + val * val;
    }
    rq_red[lane] = ss;
    workgroupBarrier();
    var stride = 16u;
    loop {
        if (stride == 0u) { break; }
        if (lane < stride) { rq_red[lane] = rq_red[lane] + rq_red[lane + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let normed = select((rq_p.flags & 4u) != 0u, (rq_p.flags & 2u) != 0u, isq);
    // HunYuan dense (flag 32): rotate FIRST, then norm. The rotation keeps
    // the head's sum of squares (rq_red[0] serves both orders); only the
    // elementwise norm weights must see the rotated vector. Both orders
    // run the same barrier sequence — the branches hold no barriers.
    let late = (rq_p.flags & 32u) != 0u;
    let hlf = rq_p.rd / 2u;
    // RoPE over the first rd dims, pairing dim i with dim i+hlf. Staged through
    // workgroup memory because the pair partner lands on a DIFFERENT lane when
    // hlf isn't a multiple of 32 (partial RoPE — Qwen3.5 rotates head_dim/4, so
    // hlf can be 16). The old register tiling (xv[t+toff], toff=hlf/32) silently
    // did nothing for hlf<32; here each lane ropes the pairs i=lane,lane+32,…
    for (var t = 0u; t < nt; t = t + 1u) {
        let d = t * 32u + lane;
        if (d < hd) { rq_head[d] = xv[t]; }
    }
    workgroupBarrier();
    if (late) {
        var ri0 = lane;
        loop {
            if (ri0 >= hlf) { break; }
            let angle0 = f32(rq_p.pos) * rq_invf[ri0];
            let cc0 = cos(angle0);
            let sf0 = sin(angle0);
            let y0 = rq_head[ri0];
            let y1 = rq_head[ri0 + hlf];
            rq_head[ri0] = y0 * cc0 - y1 * sf0;
            rq_head[ri0 + hlf] = y0 * sf0 + y1 * cc0;
            ri0 = ri0 + 32u;
        }
    }
    workgroupBarrier();
    for (var t = 0u; t < nt; t = t + 1u) {
        let d = t * 32u + lane;
        if (d < hd) { xv[t] = rq_head[d]; }
    }
    workgroupBarrier();
    if (normed) {
        let inv = 1.0 / sqrt(rq_red[0] / f32(hd) + rq_p.eps);
        let gemma = (rq_p.flags & 8u) != 0u;
        for (var t = 0u; t < nt; t = t + 1u) {
            let d = t * 32u + lane;
            if (d < hd) {
                var wd = select(rq_knw[d], rq_qnw[d], isq);
                if (gemma) { wd = 1.0 + wd; }
                xv[t] = xv[t] * inv * wd;
            }
        }
    }
    for (var t = 0u; t < nt; t = t + 1u) {
        let d = t * 32u + lane;
        if (d < hd) { rq_head[d] = xv[t]; }
    }
    workgroupBarrier();
    if (!late) {
        var ri = lane;
        loop {
            if (ri >= hlf) { break; }
            let angle = f32(rq_p.pos) * rq_invf[ri];
            let cc = cos(angle);
            let sfac = sin(angle);
            let x0 = rq_head[ri];
            let x1 = rq_head[ri + hlf];
            rq_head[ri] = x0 * cc - x1 * sfac;
            rq_head[ri + hlf] = x0 * sfac + x1 * cc;
            ri = ri + 32u;
        }
    }
    workgroupBarrier();
    let dst_base = select((head - nh) * hd, head * hd, isq);
    for (var t = 0u; t < nt; t = t + 1u) {
        let d = t * 32u + lane;
        if (d < hd) {
            if (isq) { rq_qout[dst_base + d] = rq_head[d]; } else { rq_k[koff + dst_base + d] = rq_head[d]; }
        }
    }
    if (isq && gate) {
        let gbase = head * 2u * hd + hd;
        for (var t = 0u; t < nt; t = t + 1u) {
            let d = t * 32u + lane;
            if (d < hd) { rq_gout[head * hd + d] = rq_qraw[qoff + gbase + d]; }
        }
    }
}

// Append this position's K/V rows into the device cache mirror ([nkv,cap,hd]
// each) at row `stored`. WGSL twin of Metal kv_append.
struct KvP { nkv: u32, hd: u32, cap: u32, stored: u32 };
@group(0) @binding(0) var<storage, read>       kv_k  : array<f32>;
@group(0) @binding(1) var<storage, read>       kv_v  : array<f32>;
@group(0) @binding(2) var<storage, read_write> kv_kb : array<f32>;
@group(0) @binding(3) var<storage, read_write> kv_vb : array<f32>;
@group(0) @binding(4) var<uniform>             kv_p  : KvP;
// `stored` carries the batch token index in its high bits (pos | tok<<20):
// the batch graph appends straight from its batched K/V buffers, the token
// graph passes tok=0 and reads from offset zero as before. Positions stay
// under 2^20, far above any cap the cache allows.
@compute @workgroup_size(256)
fn kv_append(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= kv_p.nkv * kv_p.hd) { return; }
    let stored = kv_p.stored & 0xFFFFFu;
    let toff = (kv_p.stored >> 20u) * kv_p.nkv * kv_p.hd;
    let h = i / kv_p.hd;
    let d = i % kv_p.hd;
    let dst = (h * kv_p.cap + stored) * kv_p.hd + d;
    kv_kb[dst] = kv_k[toff + i];
    kv_vb[dst] = kv_v[toff + i];
}

// Grouped decode attention, one 32-thread workgroup per Q-head. Dims sliced
// across lanes (dim d in lane d%32, slot d/32); online softmax over the n
// cached positions with the per-position q·k dot reduced in workgroup memory
// (portable — no subgroup ops). WGSL twin of Metal gqa_attend (output only;
// Attention-importance is handled on the CPU side when eviction is active).
struct AtP { nh: u32, hpk: u32, hd: u32, cap: u32, n: u32, scale: f32, _b: u32, _c: u32 };
@group(0) @binding(0) var<storage, read>       at_q : array<vec4<f32>>;
@group(0) @binding(1) var<storage, read>       at_k : array<vec4<f32>>;
@group(0) @binding(2) var<storage, read>       at_v : array<vec4<f32>>;
@group(0) @binding(3) var<storage, read_write> at_o : array<f32>;
@group(0) @binding(4) var<uniform>             at_p : AtP;
// Flash-decoding: split the n cached positions across the 32 lanes. Each lane
// runs an INDEPENDENT online softmax over positions lane, lane+32, … with NO
// barrier in the loop (the old kernel barriered twice PER position — O(ctx)
// serial chain), then a 5-step 32-way log-sum-exp merge. Serial steps: n → n/32.
// K/V/Q are vec4 bindings (hd % 4 == 0, gated by the Rust callers): each lane
// reads a DIFFERENT cache row, so f32 loads were 4B-used-per-32B-sector — the
// depth wall of the decode graph (4090, 1.7B q1 @ctx512: attend dominated the
// 15 ms/token submit). vec4 quarters the wasted sectors. The workgroup
// accumulator stays SCALAR at stride 257 — (lane·257 + d) mod 32 is unique per
// lane, bank-conflict-free; a vec4 accumulator array cannot be (stride must be
// ≡1 mod 32 AND a multiple of 4 — impossible).
var<workgroup> at_acc: array<f32, 8224>; // [lane*257 + d], stride 257 dodges 32-bank conflicts, hd ≤ 256 (Qwen3.5=256)
var<workgroup> at_m: array<f32, 32>;
var<workgroup> at_l: array<f32, 32>;
// Decode-regime attend: 256 threads per head instead of one warp. Lanes
// are POSITIONS for the score pass (dot over hd each) and DIMENSIONS for
// the value pass (coalesced v reads, one output dim per lane, hd <= 256).
// Online softmax over 256-position chunks; per-chunk stats via one tree.
// The 32-lane kernel above kept a 257-stride accumulator per lane and a
// five-level 256-wide merge — 137 us per layer at fifty positions of
// context. This shape does the same math in the natural order.
var<workgroup> ad_sc: array<f32, 256>;
var<workgroup> ad_red: array<f32, 256>;

@compute @workgroup_size(256)
fn gqa_attend_dec(@builtin(workgroup_id) wid: vec3<u32>,
                  @builtin(local_invocation_index) lid: u32) {
    let h = wid.x;
    if (h >= at_p.nh) { return; }
    let hd = at_p.hd;
    let hd4 = hd / 4u;
    let n = at_p.n;
    let kbase = (h / at_p.hpk) * at_p.cap * hd4;
    let qbase = h * hd4;
    let scale = at_p.scale;
    var m = -1.0e30;
    var l = 0.0;
    var acc = 0.0;                      // this lane's output dim (lid < hd)
    var c0 = 0u;
    loop {
        if (c0 >= n) { break; }
        let cn = min(256u, n - c0);
        // scores: lane p of the chunk
        var sc = -1.0e30;
        if (lid < cn) {
            let krow = kbase + (c0 + lid) * hd4;
            var dot4 = vec4<f32>(0.0);
            for (var d = 0u; d < hd4; d = d + 1u) {
                dot4 = dot4 + at_q[qbase + d] * at_k[krow + d];
            }
            sc = (dot4.x + dot4.y + dot4.z + dot4.w) * scale;
        }
        ad_sc[lid] = sc;
        ad_red[lid] = sc;
        workgroupBarrier();
        var st = 128u;
        loop {
            if (st == 0u) { break; }
            if (lid < st) { ad_red[lid] = max(ad_red[lid], ad_red[lid + st]); }
            workgroupBarrier();
            st = st >> 1u;
        }
        let cm = ad_red[0];
        workgroupBarrier();
        let mp = max(m, cm);
        let f = exp(m - mp);
        // weights into shared, denom via tree
        let w = select(0.0, exp(ad_sc[lid] - mp), lid < cn);
        ad_sc[lid] = w;
        ad_red[lid] = w;
        workgroupBarrier();
        st = 128u;
        loop {
            if (st == 0u) { break; }
            if (lid < st) { ad_red[lid] = ad_red[lid] + ad_red[lid + st]; }
            workgroupBarrier();
            st = st >> 1u;
        }
        l = l * f + ad_red[0];
        workgroupBarrier();
        // value pass: lane = output dim, coalesced across lanes
        if (lid < hd) {
            acc = acc * f;
            let dw = lid >> 2u;
            let dc = lid & 3u;
            for (var p = 0u; p < cn; p = p + 1u) {
                acc = acc + ad_sc[p] * at_v[kbase + (c0 + p) * hd4 + dw][dc];
            }
        }
        m = mp;
        c0 = c0 + 256u;
        workgroupBarrier();
    }
    if (lid < hd) {
        at_o[h * hd + lid] = acc / l;
    }
}

@compute @workgroup_size(32)
fn gqa_attend(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let lane = lid.x;
    if (h >= at_p.nh) { return; }
    let hd = at_p.hd;
    let hd4 = hd / 4u;
    let n = at_p.n;
    let kbase = (h / at_p.hpk) * at_p.cap * hd4;
    let qbase = h * hd4;
    let scale = at_p.scale;
    let base = lane * 257u;
    for (var d = 0u; d < hd; d = d + 1u) { at_acc[base + d] = 0.0; }
    var m = -1e30;
    var l = 0.0;
    var p = lane;
    loop {
        if (p >= n) { break; }
        let krow = kbase + p * hd4;
        var dot4 = vec4<f32>(0.0);
        for (var d = 0u; d < hd4; d = d + 1u) { dot4 = dot4 + at_q[qbase + d] * at_k[krow + d]; }
        let dot = (dot4.x + dot4.y + dot4.z + dot4.w) * scale;
        let mp = max(m, dot);
        let f = exp(m - mp);
        let w = exp(dot - mp);
        l = l * f + w;
        for (var d = 0u; d < hd4; d = d + 1u) {
            let vv = at_v[krow + d] * w;
            let a = base + d * 4u;
            at_acc[a]      = at_acc[a]      * f + vv.x;
            at_acc[a + 1u] = at_acc[a + 1u] * f + vv.y;
            at_acc[a + 2u] = at_acc[a + 2u] * f + vv.z;
            at_acc[a + 3u] = at_acc[a + 3u] * f + vv.w;
        }
        m = mp;
        p = p + 32u;
    }
    at_m[lane] = m;
    at_l[lane] = l;
    workgroupBarrier();
    var stride = 16u;
    loop {
        if (stride == 0u) { break; }
        if (lane < stride) {
            let o = lane + stride;
            let m1 = at_m[lane];
            let m2 = at_m[o];
            let mm = max(m1, m2);
            let f1 = exp(m1 - mm);
            let f2 = exp(m2 - mm);
            at_l[lane] = at_l[lane] * f1 + at_l[o] * f2;
            let bo = o * 257u;
            for (var d = 0u; d < hd; d = d + 1u) {
                at_acc[base + d] = at_acc[base + d] * f1 + at_acc[bo + d] * f2;
            }
            at_m[lane] = mm;
        }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let invl = select(0.0, 1.0 / at_l[0], at_l[0] > 0.0);
    for (var d = lane; d < hd; d = d + 32u) {
        at_o[h * hd + d] = at_acc[d] * invl;
    }
}

// head_dim ≤ 256 on a 32 KB device: same stride 257, HALF the lanes.
//
// The 32-lane kernel above needs 32·257·4 = 32 896 B of workgroup memory
// and cannot be created where the limit is 32 768 — wgpu-Metal and mobile.
// The stride cannot shrink (it must exceed hd, and 257 is what dodges the
// 32-bank conflicts), so the lane count is the only free dimension:
// 16·257·4 = 16 448 B fits with room to spare.
//
// Without this, `hd_cap` on Apple was 128 and the whole-token graph
// silently declined for the ENTIRE Qwen3.5/3.6 family (head_dim 256) —
// every layer walked the host on a machine whose GPU could have run it.
// Halving the lanes halves the position parallelism, which this kernel
// can afford: it is bound by the vec4 K/V reads, not by lane occupancy.
var<workgroup> at_acc16: array<f32, 4112>; // [lane*257 + d], 16 lanes
var<workgroup> at_m16: array<f32, 16>;
var<workgroup> at_l16: array<f32, 16>;
@compute @workgroup_size(16)
fn gqa_attend_w16(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let lane = lid.x;
    if (h >= at_p.nh) { return; }
    let hd = at_p.hd;
    let hd4 = hd / 4u;
    let n = at_p.n;
    let kbase = (h / at_p.hpk) * at_p.cap * hd4;
    let qbase = h * hd4;
    let scale = at_p.scale;
    let base = lane * 257u;
    for (var d = 0u; d < hd; d = d + 1u) { at_acc16[base + d] = 0.0; }
    var m = -1e30;
    var l = 0.0;
    var p = lane;
    loop {
        if (p >= n) { break; }
        let krow = kbase + p * hd4;
        var dot4 = vec4<f32>(0.0);
        for (var d = 0u; d < hd4; d = d + 1u) { dot4 = dot4 + at_q[qbase + d] * at_k[krow + d]; }
        let dot = (dot4.x + dot4.y + dot4.z + dot4.w) * scale;
        let mp = max(m, dot);
        let f = exp(m - mp);
        let w = exp(dot - mp);
        l = l * f + w;
        for (var d = 0u; d < hd4; d = d + 1u) {
            let vv = at_v[krow + d] * w;
            let a = base + d * 4u;
            at_acc16[a]      = at_acc16[a]      * f + vv.x;
            at_acc16[a + 1u] = at_acc16[a + 1u] * f + vv.y;
            at_acc16[a + 2u] = at_acc16[a + 2u] * f + vv.z;
            at_acc16[a + 3u] = at_acc16[a + 3u] * f + vv.w;
        }
        m = mp;
        p = p + 16u;
    }
    at_m16[lane] = m;
    at_l16[lane] = l;
    workgroupBarrier();
    var stride = 8u;
    loop {
        if (stride == 0u) { break; }
        if (lane < stride) {
            let o = lane + stride;
            let m1 = at_m16[lane];
            let m2 = at_m16[o];
            let mm = max(m1, m2);
            let f1 = exp(m1 - mm);
            let f2 = exp(m2 - mm);
            at_l16[lane] = at_l16[lane] * f1 + at_l16[o] * f2;
            let bo = o * 257u;
            for (var d = 0u; d < hd; d = d + 1u) {
                at_acc16[base + d] = at_acc16[base + d] * f1 + at_acc16[bo + d] * f2;
            }
            at_m16[lane] = mm;
        }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let invl = select(0.0, 1.0 / at_l16[0], at_l16[0] > 0.0);
    for (var d = lane; d < hd; d = d + 16u) {
        at_o[h * hd + d] = at_acc16[d] * invl;
    }
}

// hd <= 128 twin of gqa_attend at stride 129 — 16.5 KB of workgroup
// memory instead of 33 KB. Mobile GPUs (Adreno/Mali) and wgpu-Metal cap
// maxComputeWorkgroupStorageSize at 32768 B, where the 257-stride kernel
// cannot even be created: the invalid pipeline turned every dispatch
// into a no-op and the graph decoded garbage on phones. (lane*129 + d)
// mod 32 == (lane + d) mod 32 — still bank-conflict-free.
var<workgroup> at_acc_s: array<f32, 4128>;
@compute @workgroup_size(32)
fn gqa_attend_s(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let lane = lid.x;
    if (h >= at_p.nh) { return; }
    let hd = at_p.hd;
    let hd4 = hd / 4u;
    let n = at_p.n;
    let kbase = (h / at_p.hpk) * at_p.cap * hd4;
    let qbase = h * hd4;
    let scale = at_p.scale;
    let base = lane * 129u;
    for (var d = 0u; d < hd; d = d + 1u) { at_acc_s[base + d] = 0.0; }
    var m = -1e30;
    var l = 0.0;
    var p = lane;
    loop {
        if (p >= n) { break; }
        let krow = kbase + p * hd4;
        var dot4 = vec4<f32>(0.0);
        for (var d = 0u; d < hd4; d = d + 1u) { dot4 = dot4 + at_q[qbase + d] * at_k[krow + d]; }
        let dot = (dot4.x + dot4.y + dot4.z + dot4.w) * scale;
        let mp = max(m, dot);
        let f = exp(m - mp);
        let w = exp(dot - mp);
        l = l * f + w;
        for (var d = 0u; d < hd4; d = d + 1u) {
            let vv = at_v[krow + d] * w;
            let a = base + d * 4u;
            at_acc_s[a]      = at_acc_s[a]      * f + vv.x;
            at_acc_s[a + 1u] = at_acc_s[a + 1u] * f + vv.y;
            at_acc_s[a + 2u] = at_acc_s[a + 2u] * f + vv.z;
            at_acc_s[a + 3u] = at_acc_s[a + 3u] * f + vv.w;
        }
        m = mp;
        p = p + 32u;
    }
    at_m[lane] = m;
    at_l[lane] = l;
    workgroupBarrier();
    var stride = 16u;
    loop {
        if (stride == 0u) { break; }
        if (lane < stride) {
            let o = lane + stride;
            let m1 = at_m[lane];
            let m2 = at_m[o];
            let mm = max(m1, m2);
            let f1 = exp(m1 - mm);
            let f2 = exp(m2 - mm);
            at_l[lane] = at_l[lane] * f1 + at_l[o] * f2;
            let bo = o * 129u;
            for (var d = 0u; d < hd; d = d + 1u) {
                at_acc_s[base + d] = at_acc_s[base + d] * f1 + at_acc_s[bo + d] * f2;
            }
            at_m[lane] = mm;
        }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let invl = select(0.0, 1.0 / at_l[0], at_l[0] > 0.0);
    for (var d = lane; d < hd; d = d + 32u) {
        at_o[h * hd + d] = at_acc_s[d] * invl;
    }
}

// q1t (ternary base-3) + q4_block matvec — reuse the q1 bindings (q1w/q1x/q1y/
// q1p) and its 4-slot layout. Weights arrive as array<u32>, so bytes come out
// with shift+mask (q1t_byte). q1p fields are reinterpreted: np=gpr, _p0=cols.
var<workgroup> partial_q1t: array<f32, 64>;
fn q1t_byte(off: u32) -> u32 {
    return (q1w[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
}
const Q1T_LUT: array<u32, 243> = array<u32, 243>(
    0u, 1u, 2u, 4u, 5u, 6u, 8u, 9u, 10u, 16u, 17u, 18u, 20u, 21u, 22u, 24u,
    25u, 26u, 32u, 33u, 34u, 36u, 37u, 38u, 40u, 41u, 42u, 64u, 65u, 66u, 68u, 69u,
    70u, 72u, 73u, 74u, 80u, 81u, 82u, 84u, 85u, 86u, 88u, 89u, 90u, 96u, 97u, 98u,
    100u, 101u, 102u, 104u, 105u, 106u, 128u, 129u, 130u, 132u, 133u, 134u, 136u, 137u, 138u, 144u,
    145u, 146u, 148u, 149u, 150u, 152u, 153u, 154u, 160u, 161u, 162u, 164u, 165u, 166u, 168u, 169u,
    170u, 256u, 257u, 258u, 260u, 261u, 262u, 264u, 265u, 266u, 272u, 273u, 274u, 276u, 277u, 278u,
    280u, 281u, 282u, 288u, 289u, 290u, 292u, 293u, 294u, 296u, 297u, 298u, 320u, 321u, 322u, 324u,
    325u, 326u, 328u, 329u, 330u, 336u, 337u, 338u, 340u, 341u, 342u, 344u, 345u, 346u, 352u, 353u,
    354u, 356u, 357u, 358u, 360u, 361u, 362u, 384u, 385u, 386u, 388u, 389u, 390u, 392u, 393u, 394u,
    400u, 401u, 402u, 404u, 405u, 406u, 408u, 409u, 410u, 416u, 417u, 418u, 420u, 421u, 422u, 424u,
    425u, 426u, 512u, 513u, 514u, 516u, 517u, 518u, 520u, 521u, 522u, 528u, 529u, 530u, 532u, 533u,
    534u, 536u, 537u, 538u, 544u, 545u, 546u, 548u, 549u, 550u, 552u, 553u, 554u, 576u, 577u, 578u,
    580u, 581u, 582u, 584u, 585u, 586u, 592u, 593u, 594u, 596u, 597u, 598u, 600u, 601u, 602u, 608u,
    609u, 610u, 612u, 613u, 614u, 616u, 617u, 618u, 640u, 641u, 642u, 644u, 645u, 646u, 648u, 649u,
    650u, 656u, 657u, 658u, 660u, 661u, 662u, 664u, 665u, 666u, 672u, 673u, 674u, 676u, 677u, 678u,
    680u, 681u, 682u
);

@compute @workgroup_size(64)
fn q1t_matvec(@builtin(workgroup_id) wid: vec3<u32>,
              @builtin(num_workgroups) nwg: vec3<u32>,
              @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let base_len = rows * gpr * 9u;
    let ent_off = base_len + (rows + 1u) * 4u;
    var row = wid.x;
    loop {
        if (row >= rows) { break; }
        var acc = 0.0;
        var g = lid;
        loop {
            if (g >= gpr) { break; }
            let toff = (row * gpr + g) * 9u;
            let sc16 = q1t_byte(toff) | (q1t_byte(toff + 1u) << 8u);
            let scale = unpack2x16float(sc16).x;
            let codes = toff + 2u;
            let xb = g * 32u;
            var gsum = 0.0;
            // One byte carries FIVE base-3 codes: read (and LUT) it once
            // and spend it on all five, instead of re-reading per weight —
            // 7 byte loads a group against 32. Same k order, same adds.
            var k = 0u;
            for (var bi = 0u; bi < 7u; bi = bi + 1u) {
                let p = Q1T_LUT[q1t_byte(codes + bi)];
                let n = min(5u, 32u - k);
                for (var j = 0u; j < n; j = j + 1u) {
                    let code = (p >> (j * 2u)) & 3u;
                    let sgn = select(0.0, 1.0, code == 1u) - select(0.0, 1.0, code == 2u);
                    gsum = gsum + sgn * q1x[xb + k + j];
                }
                k = k + n;
            }
            acc = acc + scale * gsum;
            g = g + 64u;
        }
        partial_q1t[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { partial_q1t[lid] = partial_q1t[lid] + partial_q1t[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) {
            var corr = 0.0;
            let rp0 = base_len + row * 4u;
            let c0 = q1t_byte(rp0) | (q1t_byte(rp0 + 1u) << 8u) | (q1t_byte(rp0 + 2u) << 16u) | (q1t_byte(rp0 + 3u) << 24u);
            let rp1 = base_len + (row + 1u) * 4u;
            let c1 = q1t_byte(rp1) | (q1t_byte(rp1 + 1u) << 8u) | (q1t_byte(rp1 + 2u) << 16u) | (q1t_byte(rp1 + 3u) << 24u);
            for (var p = c0; p < c1; p = p + 1u) {
                let e = ent_off + p * 4u;
                let col = q1t_byte(e) | (q1t_byte(e + 1u) << 8u);
                let val16 = q1t_byte(e + 2u) | (q1t_byte(e + 3u) << 8u);
                corr = corr + unpack2x16float(val16).x * q1x[col];
            }
            q1y[row] = partial_q1t[0] + corr;
        }
        workgroupBarrier();
        row = row + nwg.x;
    }
}

// 8 nibbles from one u32 word dot 8 activations (fully unrolled FMA chain).
fn q4b_dot8(w: u32, xi: u32) -> f32 {
    return (f32(w & 0xFu) - 8.0) * q1x[xi]
         + (f32((w >> 4u) & 0xFu) - 8.0) * q1x[xi + 1u]
         + (f32((w >> 8u) & 0xFu) - 8.0) * q1x[xi + 2u]
         + (f32((w >> 12u) & 0xFu) - 8.0) * q1x[xi + 3u]
         + (f32((w >> 16u) & 0xFu) - 8.0) * q1x[xi + 4u]
         + (f32((w >> 20u) & 0xFu) - 8.0) * q1x[xi + 5u]
         + (f32((w >> 24u) & 0xFu) - 8.0) * q1x[xi + 6u]
         + (f32((w >> 28u) & 0xFu) - 8.0) * q1x[xi + 7u];
}

// q4b, tall edition: 8 rows per 256-thread workgroup in pairs with vec4
// activations — the same recipe as q4tp_matvec4, on the split layout
// (nibbles and f16 scales in two distant planes). Per-row group order
// and add order match the one-row kernel, so parity carries.
var<workgroup> p8a_q4b: array<f32, 256>;
var<workgroup> p8b_q4b: array<f32, 256>;

@compute @workgroup_size(256)
fn q4b_matvec8(@builtin(workgroup_id) wid: vec3<u32>,
               @builtin(num_workgroups) nwg: vec3<u32>,
               @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let scales_off = rows * gpr * 16u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    var base = wid.x * 8u;
    loop {
        if (base >= rows) { break; }
        let row_a = base + sub;
        let row_b = row_a + 4u;
        var acc_a = 0.0;
        var acc_b = 0.0;
        if (row_a < rows) {
            let live_b = row_b < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let xq = g * 8u;
                let x0 = q4v_x[xq];      let x1 = q4v_x[xq + 1u];
                let x2 = q4v_x[xq + 2u]; let x3 = q4v_x[xq + 3u];
                let x4 = q4v_x[xq + 4u]; let x5 = q4v_x[xq + 5u];
                let x6 = q4v_x[xq + 6u]; let x7 = q4v_x[xq + 7u];
                let ga = row_a * gpr + g;
                let sab = scales_off + ga * 2u;
                let sa = unpack2x16float((q1w[sab >> 2u] >> ((sab & 3u) * 8u)) & 0xFFFFu).x;
                let va = q4v_w[ga];
                acc_a = acc_a + sa
                    * (q4v_dot8(va.x, x0, x1) + q4v_dot8(va.y, x2, x3)
                     + q4v_dot8(va.z, x4, x5) + q4v_dot8(va.w, x6, x7));
                if (live_b) {
                    let gb = row_b * gpr + g;
                    let sbb = scales_off + gb * 2u;
                    let sb = unpack2x16float((q1w[sbb >> 2u] >> ((sbb & 3u) * 8u)) & 0xFFFFu).x;
                    let vb = q4v_w[gb];
                    acc_b = acc_b + sb
                        * (q4v_dot8(vb.x, x0, x1) + q4v_dot8(vb.y, x2, x3)
                         + q4v_dot8(vb.z, x4, x5) + q4v_dot8(vb.w, x6, x7));
                }
                g = g + 64u;
            }
        }
        p8a_q4b[lid] = acc_a;
        p8b_q4b[lid] = acc_b;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                p8a_q4b[lid] = p8a_q4b[lid] + p8a_q4b[lid + stride];
                p8b_q4b[lid] = p8b_q4b[lid] + p8b_q4b[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u) {
            if (row_a < rows) { q1y[row_a] = p8a_q4b[sub << 6u]; }
            if (row_b < rows) { q1y[row_b] = p8b_q4b[sub << 6u]; }
        }
        workgroupBarrier();
        base = base + nwg.x * 8u;
    }
}

@compute @workgroup_size(64)
fn q4b_matvec(@builtin(workgroup_id) wid: vec3<u32>,
              @builtin(num_workgroups) nwg: vec3<u32>,
              @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let scales_off = rows * gpr * 16u;
    var row = wid.x;
    loop {
        if (row >= rows) { break; }
        var acc = 0.0;
        var g = lid;
        loop {
            if (g >= gpr) { break; }
            let gi = row * gpr + g;
            // Scale: one u32 read instead of two byte reads.
            let sc_byte = scales_off + gi * 2u;
            let sc16 = (q1w[sc_byte >> 2u] >> ((sc_byte & 3u) * 8u)) & 0xFFFFu;
            let scale = unpack2x16float(sc16).x;
            // 4 u32 reads = 16 bytes = 32 weights (4× fewer array accesses
            // than the per-byte path, ~40% fewer ALU per group).
            let pk4 = gi * 4u;
            let xb = g * 32u;
            let gsum = q4b_dot8(q1w[pk4], xb)
                     + q4b_dot8(q1w[pk4 + 1u], xb + 8u)
                     + q4b_dot8(q1w[pk4 + 2u], xb + 16u)
                     + q4b_dot8(q1w[pk4 + 3u], xb + 24u);
            acc = acc + scale * gsum;
            g = g + 64u;
        }
        partial_q1t[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { partial_q1t[lid] = partial_q1t[lid] + partial_q1t[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) { q1y[row] = partial_q1t[0]; }
        workgroupBarrier();
        row = row + nwg.x;
    }
}

// q4_tiled matvec: 18-byte interleaved tiles [f16 scale][16B nibbles] — ONE
// stream per row (the split q4b layout above reads nibbles and scales from
// two distant regions; feeding TILED bytes to it produced garbage — caught by
// an end-to-end answer check on real Vulkan). Tiles are 2-aligned, so words
// assemble from u16 halves of the u32 weight array.
fn q4t_u16(off16: u32) -> u32 {
    return (q1w[off16 >> 1u] >> ((off16 & 1u) * 16u)) & 0xFFFFu;
}
@compute @workgroup_size(64)
fn q4t_matvec(@builtin(workgroup_id) wid: vec3<u32>,
              @builtin(num_workgroups) nwg: vec3<u32>,
              @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    var row = wid.x;
    loop {
        if (row >= rows) { break; }
        var acc = 0.0;
        var g = lid;
        loop {
            if (g >= gpr) { break; }
            let t16 = (row * gpr + g) * 9u;
            let scale = unpack2x16float(q4t_u16(t16)).x;
            let xb = g * 32u;
            var gsum = 0.0;
            for (var k = 0u; k < 4u; k = k + 1u) {
                let w = q4t_u16(t16 + 1u + 2u * k) | (q4t_u16(t16 + 2u + 2u * k) << 16u);
                gsum = gsum + q4b_dot8(w, xb + 8u * k);
            }
            acc = acc + scale * gsum;
            g = g + 64u;
        }
        partial_q1t[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { partial_q1t[lid] = partial_q1t[lid] + partial_q1t[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) { q1y[row] = partial_q1t[0]; }
        workgroupBarrier();
        row = row + nwg.x;
    }
}

// q4t, tall edition: 8 rows per 256-thread workgroup in pairs, vec4
// activations. The 18-byte tile stride is 2-aligned, not 4, so the
// weights stay u16-assembled (that is the layout's own cost) — but the
// activation side vectorizes exactly as in q4tp, and every x vec4 feeds
// two rows. Per-row group order and add order are the one-row kernel's.
var<workgroup> lad_q4t8: array<f32, 8>;
var<workgroup> p8a_q4t: array<f32, 256>;
var<workgroup> p8b_q4t: array<f32, 256>;

fn q4t_dot8v(w: u32, a: vec4<f32>, b: vec4<f32>) -> f32 {
    return (f32(w & 0xFu) - 8.0) * a.x
         + (f32((w >> 4u) & 0xFu) - 8.0) * a.y
         + (f32((w >> 8u) & 0xFu) - 8.0) * a.z
         + (f32((w >> 12u) & 0xFu) - 8.0) * a.w
         + (f32((w >> 16u) & 0xFu) - 8.0) * b.x
         + (f32((w >> 20u) & 0xFu) - 8.0) * b.y
         + (f32((w >> 24u) & 0xFu) - 8.0) * b.z
         + (f32((w >> 28u) & 0xFu) - 8.0) * b.w;
}

@compute @workgroup_size(256)
fn q4t_matvec8(@builtin(workgroup_id) wid: vec3<u32>,
               @builtin(num_workgroups) nwg: vec3<u32>,
               @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let sub = lid >> 6u;
    let l = lid & 63u;
    var base = wid.x * 8u;
    loop {
        if (base >= rows) { break; }
        let row_a = base + sub;
        let row_b = row_a + 4u;
        var acc_a = 0.0;
        var acc_b = 0.0;
        if (row_a < rows) {
            let live_b = row_b < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let xq = g * 8u;
                let x0 = q4v_x[xq];      let x1 = q4v_x[xq + 1u];
                let x2 = q4v_x[xq + 2u]; let x3 = q4v_x[xq + 3u];
                let x4 = q4v_x[xq + 4u]; let x5 = q4v_x[xq + 5u];
                let x6 = q4v_x[xq + 6u]; let x7 = q4v_x[xq + 7u];
                let ta = (row_a * gpr + g) * 9u;
                let sa = unpack2x16float(q4t_u16(ta)).x;
                let wa0 = q4t_u16(ta + 1u) | (q4t_u16(ta + 2u) << 16u);
                let wa1 = q4t_u16(ta + 3u) | (q4t_u16(ta + 4u) << 16u);
                let wa2 = q4t_u16(ta + 5u) | (q4t_u16(ta + 6u) << 16u);
                let wa3 = q4t_u16(ta + 7u) | (q4t_u16(ta + 8u) << 16u);
                acc_a = acc_a + sa
                    * (q4t_dot8v(wa0, x0, x1) + q4t_dot8v(wa1, x2, x3)
                     + q4t_dot8v(wa2, x4, x5) + q4t_dot8v(wa3, x6, x7));
                if (live_b) {
                    let tb = (row_b * gpr + g) * 9u;
                    let sb = unpack2x16float(q4t_u16(tb)).x;
                    let wb0 = q4t_u16(tb + 1u) | (q4t_u16(tb + 2u) << 16u);
                    let wb1 = q4t_u16(tb + 3u) | (q4t_u16(tb + 4u) << 16u);
                    let wb2 = q4t_u16(tb + 5u) | (q4t_u16(tb + 6u) << 16u);
                    let wb3 = q4t_u16(tb + 7u) | (q4t_u16(tb + 8u) << 16u);
                    acc_b = acc_b + sb
                        * (q4t_dot8v(wb0, x0, x1) + q4t_dot8v(wb1, x2, x3)
                         + q4t_dot8v(wb2, x4, x5) + q4t_dot8v(wb3, x6, x7));
                }
                g = g + 64u;
            }
        }
        p8a_q4t[lid] = acc_a;
        p8b_q4t[lid] = acc_b;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                p8a_q4t[lid] = p8a_q4t[lid] + p8a_q4t[lid + stride];
                p8b_q4t[lid] = p8b_q4t[lid] + p8b_q4t[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u) {
            if (row_a < rows) { q1y[row_a] = p8a_q4t[sub << 6u]; }
            if (row_b < rows) { q1y[row_b] = p8b_q4t[sub << 6u]; }
        }
        workgroupBarrier();
        base = base + nwg.x * 8u;
    }
}

// q4tp matvec: same nibble values as q4t, but the stride is a clean 16 B —
// so the words come straight off the u32 array instead of being assembled
// from u16 halves the way q4t's 2-aligned 18 B tiles force. The scale is a
// 5-bit rung on the row's ladder, kept in two planes after the nibbles.
//
// A workgroup owns one row at a time, so it expands that row's 32 rungs once
// into workgroup memory. Evaluating 2^(lo + code*step) per tile instead was
// measured on Metal to cost the model ~15% even though the kernel benchmarked
// faster standalone: the graph's dispatches serialize on each other, which
// exposes the dependent chain (code byte → exp2 → scale) that a free-running
// benchmark hides.
var<workgroup> lad_q4tp: array<f32, 32>;

fn q4tp_byte(off: u32) -> u32 {
    return (q1w[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
}

@compute @workgroup_size(64)
fn q4tp_matvec(@builtin(workgroup_id) wid: vec3<u32>,
               @builtin(num_workgroups) nwg: vec3<u32>,
               @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;                  // u32 index of row params
    let codes_b = rows * gpr * 16u + rows * 4u;      // byte offset of the codes
    let cstride = (gpr * 5u + 7u) / 8u;
    var row = wid.x;
    loop {
        if (row >= rows) { break; }
        if (lid < 32u) {
            let pr = unpack2x16float(q1w[params_w + row]);
            lad_q4tp[lid] = exp2(pr.x + f32(lid) * pr.y);
        }
        workgroupBarrier();
        var acc = 0.0;
        var g = lid;
        loop {
            if (g >= gpr) { break; }
            let bit = g * 5u;
            let cb = codes_b + row * cstride + (bit >> 3u);
            let sh = bit & 7u;
            // The 5-bit field spills into the next byte past bit 3; the row's
            // stride always holds that byte when it does.
            var cv = q4tp_byte(cb);
            if (sh > 3u) { cv = cv | (q4tp_byte(cb + 1u) << 8u); }
            let scale = lad_q4tp[(cv >> sh) & 31u];
            let base = (row * gpr + g) * 4u;
            let xb = g * 32u;
            var gsum = 0.0;
            for (var k = 0u; k < 4u; k = k + 1u) {
                gsum = gsum + q4b_dot8(q1w[base + k], xb + 8u * k);
            }
            acc = acc + scale * gsum;
            g = g + 64u;
        }
        partial_q1t[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { partial_q1t[lid] = partial_q1t[lid] + partial_q1t[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) { q1y[row] = partial_q1t[0]; }
        workgroupBarrier();
        row = row + nwg.x;
    }
}

// Narrow-matrix edition: 16 rows per workgroup for gpr <= 64 shapes
// (cols <= 2048: the GDN projections, o/qkv projections, lm_head), where
// the 8-row kernel gives each lane exactly ONE group and nothing to
// amortize. Four rows per 64-lane sub-block share every activation vec4
// four ways. Per-row lane layout and add order match the one-row kernel.
var<workgroup> lad_q16: array<f32, 512>;
var<workgroup> p16_a: array<f32, 256>;
var<workgroup> p16_b: array<f32, 256>;
var<workgroup> p16_c: array<f32, 256>;
var<workgroup> p16_d: array<f32, 256>;

@compute @workgroup_size(256)
fn q4tp_matvec16(@builtin(workgroup_id) wid: vec3<u32>,
                 @builtin(num_workgroups) nwg: vec3<u32>,
                 @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    // `_p0`: how many activation vectors share this weight. One is a matvec;
    // more is a batch. The BATCH is the fast axis of the dispatch, so the
    // workgroups that read the same weight rows are neighbours and meet in
    // L2; walking the whole output space instead put them `rows/16` apart,
    // which streams the weight once per batch element and defeats the point.
    // Reuse is still L2's to give — this is not a register-blocked B kernel —
    // so the win is a measurement, not a claim.
    let nb = max(q1p._p0, 1u);
    let blocks = (rows + 15u) / 16u;
    var wb = wid.x;
    loop {
        if (wb >= blocks * nb) { break; }
        let bi = wb % nb;
        let base = (wb / nb) * 16u;
        let bofs = bi * rows;
        // 16 rows x 32 rungs: each thread stages two.
        for (var q = lid; q < 512u; q = q + 256u) {
            let r = base + (q >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q16[q] = exp2(pr.x + f32(q & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        // w_* index the weight; r_* index the output, one batch apart.
        let w_a = base + sub * 4u;
        let w_b = w_a + 1u;
        let w_c = w_a + 2u;
        let w_d = w_a + 3u;
        let r_a = bofs + w_a;
        let r_b = bofs + w_b;
        let r_c = bofs + w_c;
        let r_d = bofs + w_d;
        // `_p1`: the low-rank group width, which slides the activation
        // window with the row. The FOUR rows this thread owns are
        // consecutive, so they share a window only when the width divides
        // the 16-row block — the caller checks that.
        var xblk = 0u;
        if (nb > 1u) { xblk = bi * gpr * 8u; }
        else if (q1p._p1 > 0u) { xblk = (w_a / q1p._p1) * gpr * 8u; }
        var aa = 0.0;
        var ab = 0.0;
        var ac = 0.0;
        var ad = 0.0;
        if (w_a < rows) {
            let all_live = w_d < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                let xq = xblk + g * 8u;
                let x0 = q4v_x[xq];      let x1 = q4v_x[xq + 1u];
                let x2 = q4v_x[xq + 2u]; let x3 = q4v_x[xq + 3u];
                let x4 = q4v_x[xq + 4u]; let x5 = q4v_x[xq + 5u];
                let x6 = q4v_x[xq + 6u]; let x7 = q4v_x[xq + 7u];
                let cra = codes_b + w_a * cstride + cbo;
                var cva = q4tp_byte(cra);
                if (sh > 3u) { cva = cva | (q4tp_byte(cra + 1u) << 8u); }
                let sa = lad_q16[(sub * 4u << 5u) + ((cva >> sh) & 31u)];
                let va = q4v_w[w_a * gpr + g];
                aa = aa + sa
                    * (q4v_dot8(va.x, x0, x1) + q4v_dot8(va.y, x2, x3)
                     + q4v_dot8(va.z, x4, x5) + q4v_dot8(va.w, x6, x7));
                if (all_live || w_b < rows) {
                    let crb = codes_b + w_b * cstride + cbo;
                    var cvb = q4tp_byte(crb);
                    if (sh > 3u) { cvb = cvb | (q4tp_byte(crb + 1u) << 8u); }
                    let sb = lad_q16[((sub * 4u + 1u) << 5u) + ((cvb >> sh) & 31u)];
                    let vb = q4v_w[w_b * gpr + g];
                    ab = ab + sb
                        * (q4v_dot8(vb.x, x0, x1) + q4v_dot8(vb.y, x2, x3)
                         + q4v_dot8(vb.z, x4, x5) + q4v_dot8(vb.w, x6, x7));
                }
                if (all_live || w_c < rows) {
                    let crc = codes_b + w_c * cstride + cbo;
                    var cvc = q4tp_byte(crc);
                    if (sh > 3u) { cvc = cvc | (q4tp_byte(crc + 1u) << 8u); }
                    let sc = lad_q16[((sub * 4u + 2u) << 5u) + ((cvc >> sh) & 31u)];
                    let vc = q4v_w[w_c * gpr + g];
                    ac = ac + sc
                        * (q4v_dot8(vc.x, x0, x1) + q4v_dot8(vc.y, x2, x3)
                         + q4v_dot8(vc.z, x4, x5) + q4v_dot8(vc.w, x6, x7));
                }
                if (all_live || w_d < rows) {
                    let crd = codes_b + w_d * cstride + cbo;
                    var cvd = q4tp_byte(crd);
                    if (sh > 3u) { cvd = cvd | (q4tp_byte(crd + 1u) << 8u); }
                    let sd = lad_q16[((sub * 4u + 3u) << 5u) + ((cvd >> sh) & 31u)];
                    let vd = q4v_w[w_d * gpr + g];
                    ad = ad + sd
                        * (q4v_dot8(vd.x, x0, x1) + q4v_dot8(vd.y, x2, x3)
                         + q4v_dot8(vd.z, x4, x5) + q4v_dot8(vd.w, x6, x7));
                }
                g = g + 64u;
            }
        }
        p16_a[lid] = aa;
        p16_b[lid] = ab;
        p16_c[lid] = ac;
        p16_d[lid] = ad;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                p16_a[lid] = p16_a[lid] + p16_a[lid + stride];
                p16_b[lid] = p16_b[lid] + p16_b[lid + stride];
                p16_c[lid] = p16_c[lid] + p16_c[lid + stride];
                p16_d[lid] = p16_d[lid] + p16_d[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u) {
            if (w_a < rows) { q1y[r_a] = p16_a[sub << 6u]; }
            if (w_b < rows) { q1y[r_b] = p16_b[sub << 6u]; }
            if (w_c < rows) { q1y[r_c] = p16_c[sub << 6u]; }
            if (w_d < rows) { q1y[r_d] = p16_d[sub << 6u]; }
        }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}

// q4tp matvec, tall edition: 4 rows per 256-thread workgroup, and the group's
// 16 B of nibbles arrive as ONE vec4<u32> load instead of four scalar loads.
// Written for dense-FFN shapes (17408x5120: the one-row kernel left a 27B
// dense model at ~5% of the card's bandwidth); the weight buffer is bound
// TWICE — the scalar u32 view for params and 5-bit codes (they live at
// unaligned offsets, and the buffer tail may not be 16 B-round, which a vec4
// view would silently clamp) and a vec4 view for the nibble tiles, whose
// region is 16 B-exact by construction. Each row's lane layout, add order and
// 64-slot reduction tree are byte-identical to q4tp_matvec, so the kernels
// are interchangeable under greedy parity.
@group(0) @binding(4) var<storage, read> q4v_w : array<vec4<u32>>;
// The activations again, as vec4: the scalar kernel issues 32 x-loads per
// 16 B of weights and is LSU-bound long before it is bandwidth-bound
// (measured 190 GB/s of 1.79 TB/s on the dense-FFN shapes). Components are
// consumed in the exact q4b_dot8 order.
@group(0) @binding(5) var<storage, read> q4v_x : array<vec4<f32>>;

var<workgroup> lad_q4v: array<f32, 256>;
var<workgroup> partial_q4v: array<f32, 256>;
var<workgroup> partial_q4vb: array<f32, 256>;

// A nibble as f32 WITHOUT an integer-to-float conversion: OR it into the
// mantissa of 2^23 and subtract 2^23 + 8. Exact — n − 8 for every n in
// 0..16 — so each kernel's arithmetic is unchanged to the bit; what
// changes is the instruction mix. I2F issues at 1/8 of the FMA rate on
// NVIDIA and there were eight of them per weight word, more than the
// eight FMAs the word is for. `CMF_MAGIC_UNPACK=0` restores the
// conversions for an A/B (the Rust side swaps the bodies).
fn q4v_nib(w: u32, sh: u32) -> f32 {
    return bitcast<f32>(((w >> sh) & 0xFu) | 0x4B000000u) - 8388616.0;
}
fn q4v_dot8(w: u32, a: vec4<f32>, b: vec4<f32>) -> f32 {
    return q4v_nib(w, 0u) * a.x
         + q4v_nib(w, 4u) * a.y
         + q4v_nib(w, 8u) * a.z
         + q4v_nib(w, 12u) * a.w
         + q4v_nib(w, 16u) * b.x
         + q4v_nib(w, 20u) * b.y
         + q4v_nib(w, 24u) * b.z
         + q4v_nib(w, 28u) * b.w;
}

// The 2-way group unroll of `q4tp_matvec4` (CMF_MV_U2=1): two
// independent load chains per lane per iteration, four accumulator
// registers as the price. MEASURED AND NULL on an RTX 5090 with
// Qwen3.8-27B: 50.1/50.6 tok/s against the base's 50.0/53.0 (noise
// band ±3), greedy output bit-identical. Occupancy already hides what
// a second in-flight vec4 would have hidden — kept opt-in so the idea
// is not re-opened, the same contract as CMF_MV_GRID's null.
@compute @workgroup_size(256)
fn q4tp_matvec4_u2(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(num_workgroups) nwg: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    // 8 rows per workgroup, register-blocked in pairs: sub-block `sub` owns
    // rows base+sub and base+sub+4, and every x vec4 fetched for a group
    // feeds BOTH rows' dot chains — the x side of the LSU load nearly
    // halves. Each row's group order and add order stay those of the
    // one-row kernel.
    // `_p0`: how many activation vectors share this weight. One is a matvec;
    // more is a batch. The BATCH is the fast axis of the dispatch, so the
    // workgroups that read the same weight rows are neighbours and meet in
    // L2; walking the whole output space instead put them `rows/16` apart,
    // which streams the weight once per batch element and defeats the point.
    // Reuse is still L2's to give — this is not a register-blocked B kernel —
    // so the win is a measurement, not a claim.
    let nb = max(q1p._p0, 1u);
    let blocks = (rows + 7u) / 8u;
    var wb = wid.x;
    loop {
        if (wb >= blocks * nb) { break; }
        let bi = wb % nb;
        let base = (wb / nb) * 8u;
        let bofs = bi * rows;
        {
            let r = base + (lid >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4v[lid] = exp2(pr.x + f32(lid & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let wrow_a = base + sub;
        let wrow_b = base + sub + 4u;
        let row_a = bofs + wrow_a;
        let row_b = bofs + wrow_b;
        let live_a = wrow_a < rows;
        let live_b = wrow_b < rows;
        // In vec4 units: (row / lora) * gpr * 32 floats.
        var xblk = 0u;
        if (nb > 1u) { xblk = bi * gpr * 8u; }
        else if (q1p._p1 > 0u) { xblk = (wrow_a / q1p._p1) * gpr * 8u; }
        var acc_a = 0.0;
        var acc_b = 0.0;
        var acc_a2 = 0.0;
        var acc_b2 = 0.0;
        if (live_a) {
            let crow_a = codes_b + wrow_a * cstride;
            let crow_b = codes_b + wrow_b * cstride;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let g2 = g + 64u;
                let live2 = g2 < gpr;
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv_a = q4tp_byte(crow_a + cbo);
                if (sh > 3u) { cv_a = cv_a | (q4tp_byte(crow_a + cbo + 1u) << 8u); }
                let v_a = q4v_w[wrow_a * gpr + g];
                // `_p1` is the low-rank group width. Set, it slides the
                // activation window with the row — which is the ONLY thing
                // the grouped output projection does differently, and the
                // reason it had a kernel of its own reading 3.82 ms against
                // this one's 1.24 on comparable weights. Rows base..base+7
                // share a window whenever the width is a multiple of 8, and
                // the caller only takes this path then.
                let xq = xblk + g * 8u;
                let x0 = q4v_x[xq];      let x1 = q4v_x[xq + 1u];
                let x2 = q4v_x[xq + 2u]; let x3 = q4v_x[xq + 3u];
                let x4 = q4v_x[xq + 4u]; let x5 = q4v_x[xq + 5u];
                let x6 = q4v_x[xq + 6u]; let x7 = q4v_x[xq + 7u];
                let sa = lad_q4v[(sub << 5u) + ((cv_a >> sh) & 31u)];
                acc_a = acc_a + sa
                    * (q4v_dot8(v_a.x, x0, x1) + q4v_dot8(v_a.y, x2, x3)
                     + q4v_dot8(v_a.z, x4, x5) + q4v_dot8(v_a.w, x6, x7));
                if (live_b) {
                    var cv_b = q4tp_byte(crow_b + cbo);
                    if (sh > 3u) { cv_b = cv_b | (q4tp_byte(crow_b + cbo + 1u) << 8u); }
                    let v_b = q4v_w[wrow_b * gpr + g];
                    let sb = lad_q4v[128u + (sub << 5u) + ((cv_b >> sh) & 31u)];
                    acc_b = acc_b + sb
                        * (q4v_dot8(v_b.x, x0, x1) + q4v_dot8(v_b.y, x2, x3)
                         + q4v_dot8(v_b.z, x4, x5) + q4v_dot8(v_b.w, x6, x7));
                }
                if (live2) {

                let bit_2 = g2 * 5u;
                let cbo_2 = bit_2 >> 3u;
                let sh_2 = bit_2 & 7u;
                var cv_a_2_2 = q4tp_byte(crow_a + cbo_2);
                if (sh_2 > 3u) { cv_a_2_2 = cv_a_2_2 | (q4tp_byte(crow_a + cbo_2 + 1u) << 8u); }
                let v_a_2 = q4v_w[wrow_a * gpr + g2];
                // `_p1` is the low-rank group width. Set, it slides the
                // activation window with the row — which is the ONLY thing
                // the grouped output projection does differently, and the
                // reason it had a kernel of its own reading 3.82 ms against
                // this one's 1.24 on comparable weights. Rows base..base+7
                // sh_2are a window whenever the width is a multiple of 8, and
                // the caller only takes this path then.
                let xq_2 = xblk + g2 * 8u;
                let x0_2 = q4v_x[xq_2];      let x1_2 = q4v_x[xq_2 + 1u];
                let x2_2 = q4v_x[xq_2 + 2u]; let x3_2 = q4v_x[xq_2 + 3u];
                let x4_2 = q4v_x[xq_2 + 4u]; let x5_2 = q4v_x[xq_2 + 5u];
                let x6_2 = q4v_x[xq_2 + 6u]; let x7_2 = q4v_x[xq_2 + 7u];
                let sa_2 = lad_q4v[(sub << 5u) + ((cv_a_2_2 >> sh_2) & 31u)];
                acc_a2 = acc_a2 + sa_2
                    * (q4v_dot8(v_a_2.x, x0_2, x1_2) + q4v_dot8(v_a_2.y, x2_2, x3_2)
                     + q4v_dot8(v_a_2.z, x4_2, x5_2) + q4v_dot8(v_a_2.w, x6_2, x7_2));
                if (live_b) {
                    var cv_b_2_2 = q4tp_byte(crow_b + cbo_2);
                    if (sh_2 > 3u) { cv_b_2_2 = cv_b_2_2 | (q4tp_byte(crow_b + cbo_2 + 1u) << 8u); }
                    let v_b_2 = q4v_w[wrow_b * gpr + g2];
                    let sb_2 = lad_q4v[128u + (sub << 5u) + ((cv_b_2_2 >> sh_2) & 31u)];
                    acc_b2 = acc_b2 + sb_2
                        * (q4v_dot8(v_b_2.x, x0_2, x1_2) + q4v_dot8(v_b_2.y, x2_2, x3_2)
                         + q4v_dot8(v_b_2.z, x4_2, x5_2) + q4v_dot8(v_b_2.w, x6_2, x7_2));
                }
                }
                g = g + 128u;
            }
        }
        partial_q4v[lid] = acc_a + acc_a2;
        partial_q4vb[lid] = acc_b + acc_b2;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q4v[lid] = partial_q4v[lid] + partial_q4v[lid + stride];
                partial_q4vb[lid] = partial_q4vb[lid] + partial_q4vb[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u && wrow_a < rows) { q1y[row_a] = partial_q4v[sub << 6u]; }
        if (l == 0u && wrow_b < rows) { q1y[row_b] = partial_q4vb[sub << 6u]; }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}

// Reduction-removed PROBE of `q4tp_matvec4` (CMF_MV_NORED=1):
// garbage answers, honest timing — prices the barrier tree alone.
@compute @workgroup_size(256)
fn q4tp_matvec4_nored(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(num_workgroups) nwg: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    // 8 rows per workgroup, register-blocked in pairs: sub-block `sub` owns
    // rows base+sub and base+sub+4, and every x vec4 fetched for a group
    // feeds BOTH rows' dot chains — the x side of the LSU load nearly
    // halves. Each row's group order and add order stay those of the
    // one-row kernel.
    // `_p0`: how many activation vectors share this weight. One is a matvec;
    // more is a batch. The BATCH is the fast axis of the dispatch, so the
    // workgroups that read the same weight rows are neighbours and meet in
    // L2; walking the whole output space instead put them `rows/16` apart,
    // which streams the weight once per batch element and defeats the point.
    // Reuse is still L2's to give — this is not a register-blocked B kernel —
    // so the win is a measurement, not a claim.
    let nb = max(q1p._p0, 1u);
    let blocks = (rows + 7u) / 8u;
    var wb = wid.x;
    loop {
        if (wb >= blocks * nb) { break; }
        let bi = wb % nb;
        let base = (wb / nb) * 8u;
        let bofs = bi * rows;
        {
            let r = base + (lid >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4v[lid] = exp2(pr.x + f32(lid & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let wrow_a = base + sub;
        let wrow_b = base + sub + 4u;
        let row_a = bofs + wrow_a;
        let row_b = bofs + wrow_b;
        let live_a = wrow_a < rows;
        let live_b = wrow_b < rows;
        // In vec4 units: (row / lora) * gpr * 32 floats.
        var xblk = 0u;
        if (nb > 1u) { xblk = bi * gpr * 8u; }
        else if (q1p._p1 > 0u) { xblk = (wrow_a / q1p._p1) * gpr * 8u; }
        var acc_a = 0.0;
        var acc_b = 0.0;
        if (live_a) {
            let crow_a = codes_b + wrow_a * cstride;
            let crow_b = codes_b + wrow_b * cstride;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv_a = q4tp_byte(crow_a + cbo);
                if (sh > 3u) { cv_a = cv_a | (q4tp_byte(crow_a + cbo + 1u) << 8u); }
                let v_a = q4v_w[wrow_a * gpr + g];
                // `_p1` is the low-rank group width. Set, it slides the
                // activation window with the row — which is the ONLY thing
                // the grouped output projection does differently, and the
                // reason it had a kernel of its own reading 3.82 ms against
                // this one's 1.24 on comparable weights. Rows base..base+7
                // share a window whenever the width is a multiple of 8, and
                // the caller only takes this path then.
                let xq = xblk + g * 8u;
                let x0 = q4v_x[xq];      let x1 = q4v_x[xq + 1u];
                let x2 = q4v_x[xq + 2u]; let x3 = q4v_x[xq + 3u];
                let x4 = q4v_x[xq + 4u]; let x5 = q4v_x[xq + 5u];
                let x6 = q4v_x[xq + 6u]; let x7 = q4v_x[xq + 7u];
                let sa = lad_q4v[(sub << 5u) + ((cv_a >> sh) & 31u)];
                acc_a = acc_a + sa
                    * (q4v_dot8(v_a.x, x0, x1) + q4v_dot8(v_a.y, x2, x3)
                     + q4v_dot8(v_a.z, x4, x5) + q4v_dot8(v_a.w, x6, x7));
                if (live_b) {
                    var cv_b = q4tp_byte(crow_b + cbo);
                    if (sh > 3u) { cv_b = cv_b | (q4tp_byte(crow_b + cbo + 1u) << 8u); }
                    let v_b = q4v_w[wrow_b * gpr + g];
                    let sb = lad_q4v[128u + (sub << 5u) + ((cv_b >> sh) & 31u)];
                    acc_b = acc_b + sb
                        * (q4v_dot8(v_b.x, x0, x1) + q4v_dot8(v_b.y, x2, x3)
                         + q4v_dot8(v_b.z, x4, x5) + q4v_dot8(v_b.w, x6, x7));
                }
                g = g + 64u;
            }
        }
        // PROBE: the tree is gone and THE ANSWER IS GARBAGE (lane 0's
        // partial only). The point is the time — whether eight barriers
        // per 8-row block are the missing third of the bus, priced
        // BEFORE building the two-blocks-per-reduction kernel.
        if (l == 0u && wrow_a < rows) { q1y[row_a] = acc_a; }
        if (l == 0u && wrow_b < rows) { q1y[row_b] = acc_b; }
        wb = wb + nwg.x;
    }
}

// Dual bindings: the second projection's weight (u32 + vec4 views)
// and output. Only the dual entry reads them, so every other entry's
// auto layout is untouched.
@group(0) @binding(6) var<storage, read> q1w2 : array<u32>;
@group(0) @binding(7) var<storage, read> q4v_w2 : array<vec4<u32>>;
@group(0) @binding(8) var<storage, read_write> q1y2 : array<f32>;
fn q4tp_byte2(off: u32) -> u32 {
    return (q1w2[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
}

// DUAL matvec (CMF_MV_DUAL=1): two independent projections of ONE
// input in ONE dispatch — gate+up, and q/k/v by pairs. Exists because
// the elimination table ended at the dispatch boundary: every in-kernel
// component measured null while the identical access pattern streams
// 1623 GB/s in isolation, and the pass serializes each 10-17 us wave
// against the next ("dispatches within a pass are serialized"). Side B
// mirrors side A over its own bindings; rows differ, so the layout
// arithmetic (params_w, codes_b) is per side.
@compute @workgroup_size(256)
fn q4tp_matvec4_dual(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(num_workgroups) nwg: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let rows2 = q1p._p1;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let params_w2 = rows2 * gpr * 4u;
    let codes_b2 = rows2 * gpr * 16u + rows2 * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let blocks1 = (rows + 7u) / 8u;
    let blocks2 = (rows2 + 7u) / 8u;
    var wb = wid.x;
    loop {
        if (wb >= blocks1 + blocks2) { break; }
        if (wb < blocks1) {
            let base = wb * 8u;
            let bofs = 0u;
        {
            let r = base + (lid >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4v[lid] = exp2(pr.x + f32(lid & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let wrow_a = base + sub;
        let wrow_b = base + sub + 4u;
        let row_a = bofs + wrow_a;
        let row_b = bofs + wrow_b;
        let live_a = wrow_a < rows;
        let live_b = wrow_b < rows;
        // In vec4 units: (row / lora) * gpr * 32 floats.
        let xblk = 0u;
        var acc_a = 0.0;
        var acc_b = 0.0;
        if (live_a) {
            let crow_a = codes_b + wrow_a * cstride;
            let crow_b = codes_b + wrow_b * cstride;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv_a = q4tp_byte(crow_a + cbo);
                if (sh > 3u) { cv_a = cv_a | (q4tp_byte(crow_a + cbo + 1u) << 8u); }
                let v_a = q4v_w[wrow_a * gpr + g];
                // `_p1` is the low-rank group width. Set, it slides the
                // activation window with the row — which is the ONLY thing
                // the grouped output projection does differently, and the
                // reason it had a kernel of its own reading 3.82 ms against
                // this one's 1.24 on comparable weights. Rows base..base+7
                // share a window whenever the width is a multiple of 8, and
                // the caller only takes this path then.
                let xq = xblk + g * 8u;
                let x0 = q4v_x[xq];      let x1 = q4v_x[xq + 1u];
                let x2 = q4v_x[xq + 2u]; let x3 = q4v_x[xq + 3u];
                let x4 = q4v_x[xq + 4u]; let x5 = q4v_x[xq + 5u];
                let x6 = q4v_x[xq + 6u]; let x7 = q4v_x[xq + 7u];
                let sa = lad_q4v[(sub << 5u) + ((cv_a >> sh) & 31u)];
                acc_a = acc_a + sa
                    * (q4v_dot8(v_a.x, x0, x1) + q4v_dot8(v_a.y, x2, x3)
                     + q4v_dot8(v_a.z, x4, x5) + q4v_dot8(v_a.w, x6, x7));
                if (live_b) {
                    var cv_b = q4tp_byte(crow_b + cbo);
                    if (sh > 3u) { cv_b = cv_b | (q4tp_byte(crow_b + cbo + 1u) << 8u); }
                    let v_b = q4v_w[wrow_b * gpr + g];
                    let sb = lad_q4v[128u + (sub << 5u) + ((cv_b >> sh) & 31u)];
                    acc_b = acc_b + sb
                        * (q4v_dot8(v_b.x, x0, x1) + q4v_dot8(v_b.y, x2, x3)
                         + q4v_dot8(v_b.z, x4, x5) + q4v_dot8(v_b.w, x6, x7));
                }
                g = g + 64u;
            }
        }
        partial_q4v[lid] = acc_a;
        partial_q4vb[lid] = acc_b;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q4v[lid] = partial_q4v[lid] + partial_q4v[lid + stride];
                partial_q4vb[lid] = partial_q4vb[lid] + partial_q4vb[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u && wrow_a < rows) { q1y[row_a] = partial_q4v[sub << 6u]; }
        if (l == 0u && wrow_b < rows) { q1y[row_b] = partial_q4vb[sub << 6u]; }
        } else {
            let base = (wb - blocks1) * 8u;
            let bofs = 0u;
        {
            let r = base + (lid >> 5u);
            if (r < rows2) {
                let pr = unpack2x16float(q1w2[params_w2 + r]);
                lad_q4v[lid] = exp2(pr.x + f32(lid & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let wrow_a = base + sub;
        let wrow_b = base + sub + 4u;
        let row_a = bofs + wrow_a;
        let row_b = bofs + wrow_b;
        let live_a = wrow_a < rows2;
        let live_b = wrow_b < rows2;
        // In vec4 units: (row / lora) * gpr * 32 floats.
        let xblk = 0u;
        var acc_a = 0.0;
        var acc_b = 0.0;
        if (live_a) {
            let crow_a = codes_b2 + wrow_a * cstride;
            let crow_b = codes_b2 + wrow_b * cstride;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv_a = q4tp_byte2(crow_a + cbo);
                if (sh > 3u) { cv_a = cv_a | (q4tp_byte2(crow_a + cbo + 1u) << 8u); }
                let v_a = q4v_w2[wrow_a * gpr + g];
                // `_p1` is the low-rank group width. Set, it slides the
                // activation window with the row — which is the ONLY thing
                // the grouped output projection does differently, and the
                // reason it had a kernel of its own reading 3.82 ms against
                // this one's 1.24 on comparable weights. Rows base..base+7
                // share a window whenever the width is a multiple of 8, and
                // the caller only takes this path then.
                let xq = xblk + g * 8u;
                let x0 = q4v_x[xq];      let x1 = q4v_x[xq + 1u];
                let x2 = q4v_x[xq + 2u]; let x3 = q4v_x[xq + 3u];
                let x4 = q4v_x[xq + 4u]; let x5 = q4v_x[xq + 5u];
                let x6 = q4v_x[xq + 6u]; let x7 = q4v_x[xq + 7u];
                let sa = lad_q4v[(sub << 5u) + ((cv_a >> sh) & 31u)];
                acc_a = acc_a + sa
                    * (q4v_dot8(v_a.x, x0, x1) + q4v_dot8(v_a.y, x2, x3)
                     + q4v_dot8(v_a.z, x4, x5) + q4v_dot8(v_a.w, x6, x7));
                if (live_b) {
                    var cv_b = q4tp_byte2(crow_b + cbo);
                    if (sh > 3u) { cv_b = cv_b | (q4tp_byte2(crow_b + cbo + 1u) << 8u); }
                    let v_b = q4v_w2[wrow_b * gpr + g];
                    let sb = lad_q4v[128u + (sub << 5u) + ((cv_b >> sh) & 31u)];
                    acc_b = acc_b + sb
                        * (q4v_dot8(v_b.x, x0, x1) + q4v_dot8(v_b.y, x2, x3)
                         + q4v_dot8(v_b.z, x4, x5) + q4v_dot8(v_b.w, x6, x7));
                }
                g = g + 64u;
            }
        }
        partial_q4v[lid] = acc_a;
        partial_q4vb[lid] = acc_b;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q4v[lid] = partial_q4v[lid] + partial_q4v[lid + stride];
                partial_q4vb[lid] = partial_q4vb[lid] + partial_q4vb[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u && wrow_a < rows2) { q1y2[row_a] = partial_q4v[sub << 6u]; }
        if (l == 0u && wrow_b < rows2) { q1y2[row_b] = partial_q4vb[sub << 6u]; }
        }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}


// Binding 9: the up-projection's output; binding 5 carries gate's.
@group(0) @binding(9) var<storage, read> q4v_u : array<vec4<f32>>;
fn fsilu(g: vec4<f32>, u: vec4<f32>) -> vec4<f32> {
    return u * g / (vec4<f32>(1.0) + exp(-g));
}

// The FFN's `down` with SiLU folded in (part of CMF_MV_DUAL=1): reads
// gate and up straight from their matvecs and mixes them per element,
// deleting the silu dispatch, its barrier, and the act buffer. The
// elimination table said the cost lives between dispatches; this and
// the dual cut the FFN from five waves to two.
@compute @workgroup_size(256)
fn q4tp_matvec4_dsilu(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(num_workgroups) nwg: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    // 8 rows per workgroup, register-blocked in pairs: sub-block `sub` owns
    // rows base+sub and base+sub+4, and every x vec4 fetched for a group
    // feeds BOTH rows' dot chains — the x side of the LSU load nearly
    // halves. Each row's group order and add order stay those of the
    // one-row kernel.
    // `_p0`: how many activation vectors share this weight. One is a matvec;
    // more is a batch. The BATCH is the fast axis of the dispatch, so the
    // workgroups that read the same weight rows are neighbours and meet in
    // L2; walking the whole output space instead put them `rows/16` apart,
    // which streams the weight once per batch element and defeats the point.
    // Reuse is still L2's to give — this is not a register-blocked B kernel —
    // so the win is a measurement, not a claim.
    let nb = max(q1p._p0, 1u);
    let blocks = (rows + 7u) / 8u;
    var wb = wid.x;
    loop {
        if (wb >= blocks * nb) { break; }
        let bi = wb % nb;
        let base = (wb / nb) * 8u;
        let bofs = bi * rows;
        {
            let r = base + (lid >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4v[lid] = exp2(pr.x + f32(lid & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let wrow_a = base + sub;
        let wrow_b = base + sub + 4u;
        let row_a = bofs + wrow_a;
        let row_b = bofs + wrow_b;
        let live_a = wrow_a < rows;
        let live_b = wrow_b < rows;
        // In vec4 units: (row / lora) * gpr * 32 floats.
        let xblk = 0u;
        var acc_a = 0.0;
        var acc_b = 0.0;
        if (live_a) {
            let crow_a = codes_b + wrow_a * cstride;
            let crow_b = codes_b + wrow_b * cstride;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv_a = q4tp_byte(crow_a + cbo);
                if (sh > 3u) { cv_a = cv_a | (q4tp_byte(crow_a + cbo + 1u) << 8u); }
                let v_a = q4v_w[wrow_a * gpr + g];
                // `_p1` is the low-rank group width. Set, it slides the
                // activation window with the row — which is the ONLY thing
                // the grouped output projection does differently, and the
                // reason it had a kernel of its own reading 3.82 ms against
                // this one's 1.24 on comparable weights. Rows base..base+7
                // share a window whenever the width is a multiple of 8, and
                // the caller only takes this path then.
                let xq = xblk + g * 8u;
                // Inline SiLU: this matvec CONSUMES the FFN's gate/up
                // directly — act = up * g * sigmoid(g) computed per
                // element read. Redundant across row-blocks, and that
                // redundancy buys out a whole dispatch, its barrier,
                // and the act buffer's round trip.
                let x0 = fsilu(q4v_x[xq],      q4v_u[xq]);
                let x1 = fsilu(q4v_x[xq + 1u], q4v_u[xq + 1u]);
                let x2 = fsilu(q4v_x[xq + 2u], q4v_u[xq + 2u]);
                let x3 = fsilu(q4v_x[xq + 3u], q4v_u[xq + 3u]);
                let x4 = fsilu(q4v_x[xq + 4u], q4v_u[xq + 4u]);
                let x5 = fsilu(q4v_x[xq + 5u], q4v_u[xq + 5u]);
                let x6 = fsilu(q4v_x[xq + 6u], q4v_u[xq + 6u]);
                let x7 = fsilu(q4v_x[xq + 7u], q4v_u[xq + 7u]);
                let sa = lad_q4v[(sub << 5u) + ((cv_a >> sh) & 31u)];
                acc_a = acc_a + sa
                    * (q4v_dot8(v_a.x, x0, x1) + q4v_dot8(v_a.y, x2, x3)
                     + q4v_dot8(v_a.z, x4, x5) + q4v_dot8(v_a.w, x6, x7));
                if (live_b) {
                    var cv_b = q4tp_byte(crow_b + cbo);
                    if (sh > 3u) { cv_b = cv_b | (q4tp_byte(crow_b + cbo + 1u) << 8u); }
                    let v_b = q4v_w[wrow_b * gpr + g];
                    let sb = lad_q4v[128u + (sub << 5u) + ((cv_b >> sh) & 31u)];
                    acc_b = acc_b + sb
                        * (q4v_dot8(v_b.x, x0, x1) + q4v_dot8(v_b.y, x2, x3)
                         + q4v_dot8(v_b.z, x4, x5) + q4v_dot8(v_b.w, x6, x7));
                }
                g = g + 64u;
            }
        }
        partial_q4v[lid] = acc_a;
        partial_q4vb[lid] = acc_b;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q4v[lid] = partial_q4v[lid] + partial_q4v[lid + stride];
                partial_q4vb[lid] = partial_q4vb[lid] + partial_q4vb[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u && wrow_a < rows) { q1y[row_a] = partial_q4v[sub << 6u]; }
        if (l == 0u && wrow_b < rows) { q1y[row_b] = partial_q4vb[sub << 6u]; }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}

@compute @workgroup_size(256)
fn q4tp_matvec4(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(num_workgroups) nwg: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    // 8 rows per workgroup, register-blocked in pairs: sub-block `sub` owns
    // rows base+sub and base+sub+4, and every x vec4 fetched for a group
    // feeds BOTH rows' dot chains — the x side of the LSU load nearly
    // halves. Each row's group order and add order stay those of the
    // one-row kernel.
    // `_p0`: how many activation vectors share this weight. One is a matvec;
    // more is a batch. The BATCH is the fast axis of the dispatch, so the
    // workgroups that read the same weight rows are neighbours and meet in
    // L2; walking the whole output space instead put them `rows/16` apart,
    // which streams the weight once per batch element and defeats the point.
    // Reuse is still L2's to give — this is not a register-blocked B kernel —
    // so the win is a measurement, not a claim.
    let nb = max(q1p._p0, 1u);
    let blocks = (rows + 7u) / 8u;
    var wb = wid.x;
    loop {
        if (wb >= blocks * nb) { break; }
        let bi = wb % nb;
        let base = (wb / nb) * 8u;
        let bofs = bi * rows;
        {
            let r = base + (lid >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4v[lid] = exp2(pr.x + f32(lid & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let wrow_a = base + sub;
        let wrow_b = base + sub + 4u;
        let row_a = bofs + wrow_a;
        let row_b = bofs + wrow_b;
        let live_a = wrow_a < rows;
        let live_b = wrow_b < rows;
        // In vec4 units: (row / lora) * gpr * 32 floats.
        var xblk = 0u;
        if (nb > 1u) { xblk = bi * gpr * 8u; }
        else if (q1p._p1 > 0u) { xblk = (wrow_a / q1p._p1) * gpr * 8u; }
        var acc_a = 0.0;
        var acc_b = 0.0;
        if (live_a) {
            let crow_a = codes_b + wrow_a * cstride;
            let crow_b = codes_b + wrow_b * cstride;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv_a = q4tp_byte(crow_a + cbo);
                if (sh > 3u) { cv_a = cv_a | (q4tp_byte(crow_a + cbo + 1u) << 8u); }
                let v_a = q4v_w[wrow_a * gpr + g];
                // `_p1` is the low-rank group width. Set, it slides the
                // activation window with the row — which is the ONLY thing
                // the grouped output projection does differently, and the
                // reason it had a kernel of its own reading 3.82 ms against
                // this one's 1.24 on comparable weights. Rows base..base+7
                // share a window whenever the width is a multiple of 8, and
                // the caller only takes this path then.
                let xq = xblk + g * 8u;
                let x0 = q4v_x[xq];      let x1 = q4v_x[xq + 1u];
                let x2 = q4v_x[xq + 2u]; let x3 = q4v_x[xq + 3u];
                let x4 = q4v_x[xq + 4u]; let x5 = q4v_x[xq + 5u];
                let x6 = q4v_x[xq + 6u]; let x7 = q4v_x[xq + 7u];
                let sa = lad_q4v[(sub << 5u) + ((cv_a >> sh) & 31u)];
                acc_a = acc_a + sa
                    * (q4v_dot8(v_a.x, x0, x1) + q4v_dot8(v_a.y, x2, x3)
                     + q4v_dot8(v_a.z, x4, x5) + q4v_dot8(v_a.w, x6, x7));
                if (live_b) {
                    var cv_b = q4tp_byte(crow_b + cbo);
                    if (sh > 3u) { cv_b = cv_b | (q4tp_byte(crow_b + cbo + 1u) << 8u); }
                    let v_b = q4v_w[wrow_b * gpr + g];
                    let sb = lad_q4v[128u + (sub << 5u) + ((cv_b >> sh) & 31u)];
                    acc_b = acc_b + sb
                        * (q4v_dot8(v_b.x, x0, x1) + q4v_dot8(v_b.y, x2, x3)
                         + q4v_dot8(v_b.z, x4, x5) + q4v_dot8(v_b.w, x6, x7));
                }
                g = g + 64u;
            }
        }
        partial_q4v[lid] = acc_a;
        partial_q4vb[lid] = acc_b;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q4v[lid] = partial_q4v[lid] + partial_q4v[lid + stride];
                partial_q4vb[lid] = partial_q4vb[lid] + partial_q4vb[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u && wrow_a < rows) { q1y[row_a] = partial_q4v[sub << 6u]; }
        if (l == 0u && wrow_b < rows) { q1y[row_b] = partial_q4vb[sub << 6u]; }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}

// ── q4tp matvec, batch blocked INSIDE the workgroup: sixteen rows a
// workgroup and four rows a lane, exactly the streaming shape of
// `q4tp_matvec16w`, with the batch as an inner unrolled loop instead of
// a dispatch axis. The alternative arm below dispatches
// (row-block × batch) workgroups and re-reads the weight for each
// element; here the four weight vec4 are loaded once and every batch
// element consumes them from registers.
//
// WHY THE SHAPE MATTERS MORE THAN THE TRAFFIC: this matvec is memory
// bound end to end — the arithmetic-free probe (`CMF_MV_PROBE=1`) runs
// the same 17.6 ms token — so what decides the number is how many
// weight loads a lane keeps in flight, i.e. the register budget. An
// earlier version of this kernel unpacked the nibbles into f32 before
// the batch loop; that is FEWER arithmetic ops and it lost, because 32
// registers of dequantized weight crowded out the occupancy that hides
// the load latency. Weights stay PACKED here for that reason.
//
// ADD ORDER is the one-row kernel's, per row and per element.
// Batch ≤ 4: accumulator components are indexed by CONSTANTS only — a
// dynamic index spills the vec4 to stack and the GEMV runs at a
// fraction of the card, the trap the 16-row kernel documents.
@compute @workgroup_size(256)
fn q4tp_matvec4_bk(@builtin(workgroup_id) wid: vec3<u32>,
                   @builtin(num_workgroups) nwg: vec3<u32>,
                   @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let nb = max(q1p._p0, 1u);
    let blocks = (rows + 15u) / 16u;
    var wb = wid.x;
    loop {
        if (wb >= blocks) { break; }
        let base = wb * 16u;
        for (var t = lid; t < 512u; t = t + 256u) {
            let r = base + (t >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4w[t] = exp2(pr.x + f32(t & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        // acc<row>[element]: four rows, up to four activation vectors.
        var acc0 = vec4<f32>(0.0);
        var acc1 = vec4<f32>(0.0);
        var acc2 = vec4<f32>(0.0);
        var acc3 = vec4<f32>(0.0);
        if (r0 < rows) {
            let c0 = codes_b + r0 * cstride;
            let c1 = codes_b + r1 * cstride;
            let c2 = codes_b + r2 * cstride;
            let c3 = codes_b + r3 * cstride;
            let l1 = r1 < rows;
            let l2 = r2 < rows;
            let l3 = r3 < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                // Four rows' weights and scales: read ONCE per group,
                // then reused by every batch element below.
                var cv = q4tp_byte(c0 + cbo);
                if (sh > 3u) { cv = cv | (q4tp_byte(c0 + cbo + 1u) << 8u); }
                let v0 = q4v_w[r0 * gpr + g];
                let s0 = lad_q4w[(sub << 5u) + ((cv >> sh) & 31u)];
                var v1 = vec4<u32>(0u, 0u, 0u, 0u);
                var s1 = 0.0;
                if (l1) {
                    cv = q4tp_byte(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c1 + cbo + 1u) << 8u); }
                    v1 = q4v_w[r1 * gpr + g];
                    s1 = lad_q4w[128u + (sub << 5u) + ((cv >> sh) & 31u)];
                }
                var v2 = vec4<u32>(0u, 0u, 0u, 0u);
                var s2 = 0.0;
                if (l2) {
                    cv = q4tp_byte(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c2 + cbo + 1u) << 8u); }
                    v2 = q4v_w[r2 * gpr + g];
                    s2 = lad_q4w[256u + (sub << 5u) + ((cv >> sh) & 31u)];
                }
                var v3 = vec4<u32>(0u, 0u, 0u, 0u);
                var s3 = 0.0;
                if (l3) {
                    cv = q4tp_byte(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c3 + cbo + 1u) << 8u); }
                    v3 = q4v_w[r3 * gpr + g];
                    s3 = lad_q4w[384u + (sub << 5u) + ((cv >> sh) & 31u)];
                }
                let xg = g * 8u;
                {
                    let xa = q4v_x[xg];      let xb = q4v_x[xg + 1u];
                    let xc = q4v_x[xg + 2u]; let xd = q4v_x[xg + 3u];
                    let xe = q4v_x[xg + 4u]; let xf = q4v_x[xg + 5u];
                    let xgg = q4v_x[xg + 6u]; let xh = q4v_x[xg + 7u];
                    acc0.x = acc0.x + s0 * (q4v_dot8(v0.x, xa, xb) + q4v_dot8(v0.y, xc, xd)
                                          + q4v_dot8(v0.z, xe, xf) + q4v_dot8(v0.w, xgg, xh));
                    acc1.x = acc1.x + s1 * (q4v_dot8(v1.x, xa, xb) + q4v_dot8(v1.y, xc, xd)
                                          + q4v_dot8(v1.z, xe, xf) + q4v_dot8(v1.w, xgg, xh));
                    acc2.x = acc2.x + s2 * (q4v_dot8(v2.x, xa, xb) + q4v_dot8(v2.y, xc, xd)
                                          + q4v_dot8(v2.z, xe, xf) + q4v_dot8(v2.w, xgg, xh));
                    acc3.x = acc3.x + s3 * (q4v_dot8(v3.x, xa, xb) + q4v_dot8(v3.y, xc, xd)
                                          + q4v_dot8(v3.z, xe, xf) + q4v_dot8(v3.w, xgg, xh));
                }
                if (nb > 1u) {
                    let xq = xg + gpr * 8u;
                    let xa = q4v_x[xq];      let xb = q4v_x[xq + 1u];
                    let xc = q4v_x[xq + 2u]; let xd = q4v_x[xq + 3u];
                    let xe = q4v_x[xq + 4u]; let xf = q4v_x[xq + 5u];
                    let xgg = q4v_x[xq + 6u]; let xh = q4v_x[xq + 7u];
                    acc0.y = acc0.y + s0 * (q4v_dot8(v0.x, xa, xb) + q4v_dot8(v0.y, xc, xd)
                                          + q4v_dot8(v0.z, xe, xf) + q4v_dot8(v0.w, xgg, xh));
                    acc1.y = acc1.y + s1 * (q4v_dot8(v1.x, xa, xb) + q4v_dot8(v1.y, xc, xd)
                                          + q4v_dot8(v1.z, xe, xf) + q4v_dot8(v1.w, xgg, xh));
                    acc2.y = acc2.y + s2 * (q4v_dot8(v2.x, xa, xb) + q4v_dot8(v2.y, xc, xd)
                                          + q4v_dot8(v2.z, xe, xf) + q4v_dot8(v2.w, xgg, xh));
                    acc3.y = acc3.y + s3 * (q4v_dot8(v3.x, xa, xb) + q4v_dot8(v3.y, xc, xd)
                                          + q4v_dot8(v3.z, xe, xf) + q4v_dot8(v3.w, xgg, xh));
                }
                if (nb > 2u) {
                    let xq = xg + gpr * 16u;
                    let xa = q4v_x[xq];      let xb = q4v_x[xq + 1u];
                    let xc = q4v_x[xq + 2u]; let xd = q4v_x[xq + 3u];
                    let xe = q4v_x[xq + 4u]; let xf = q4v_x[xq + 5u];
                    let xgg = q4v_x[xq + 6u]; let xh = q4v_x[xq + 7u];
                    acc0.z = acc0.z + s0 * (q4v_dot8(v0.x, xa, xb) + q4v_dot8(v0.y, xc, xd)
                                          + q4v_dot8(v0.z, xe, xf) + q4v_dot8(v0.w, xgg, xh));
                    acc1.z = acc1.z + s1 * (q4v_dot8(v1.x, xa, xb) + q4v_dot8(v1.y, xc, xd)
                                          + q4v_dot8(v1.z, xe, xf) + q4v_dot8(v1.w, xgg, xh));
                    acc2.z = acc2.z + s2 * (q4v_dot8(v2.x, xa, xb) + q4v_dot8(v2.y, xc, xd)
                                          + q4v_dot8(v2.z, xe, xf) + q4v_dot8(v2.w, xgg, xh));
                    acc3.z = acc3.z + s3 * (q4v_dot8(v3.x, xa, xb) + q4v_dot8(v3.y, xc, xd)
                                          + q4v_dot8(v3.z, xe, xf) + q4v_dot8(v3.w, xgg, xh));
                }
                if (nb > 3u) {
                    let xq = xg + gpr * 24u;
                    let xa = q4v_x[xq];      let xb = q4v_x[xq + 1u];
                    let xc = q4v_x[xq + 2u]; let xd = q4v_x[xq + 3u];
                    let xe = q4v_x[xq + 4u]; let xf = q4v_x[xq + 5u];
                    let xgg = q4v_x[xq + 6u]; let xh = q4v_x[xq + 7u];
                    acc0.w = acc0.w + s0 * (q4v_dot8(v0.x, xa, xb) + q4v_dot8(v0.y, xc, xd)
                                          + q4v_dot8(v0.z, xe, xf) + q4v_dot8(v0.w, xgg, xh));
                    acc1.w = acc1.w + s1 * (q4v_dot8(v1.x, xa, xb) + q4v_dot8(v1.y, xc, xd)
                                          + q4v_dot8(v1.z, xe, xf) + q4v_dot8(v1.w, xgg, xh));
                    acc2.w = acc2.w + s2 * (q4v_dot8(v2.x, xa, xb) + q4v_dot8(v2.y, xc, xd)
                                          + q4v_dot8(v2.z, xe, xf) + q4v_dot8(v2.w, xgg, xh));
                    acc3.w = acc3.w + s3 * (q4v_dot8(v3.x, xa, xb) + q4v_dot8(v3.y, xc, xd)
                                          + q4v_dot8(v3.z, xe, xf) + q4v_dot8(v3.w, xgg, xh));
                }
                g = g + 64u;
            }
        }
        // One reduction per batch element, reusing the SAME workgroup
        // array: four vec4 arrays of 256 would be 16 KB of workgroup
        // storage and price this kernel out of the smaller devices.
        var e = 0u;
        loop {
            if (e >= nb) { break; }
            var mine = vec4<f32>(acc0.x, acc1.x, acc2.x, acc3.x);
            if (e == 1u) { mine = vec4<f32>(acc0.y, acc1.y, acc2.y, acc3.y); }
            if (e == 2u) { mine = vec4<f32>(acc0.z, acc1.z, acc2.z, acc3.z); }
            if (e == 3u) { mine = vec4<f32>(acc0.w, acc1.w, acc2.w, acc3.w); }
            partial_q4k[lid] = mine;
            workgroupBarrier();
            var stride = 32u;
            loop {
                if (stride == 0u) { break; }
                if (l < stride) {
                    partial_q4k[lid] = partial_q4k[lid] + partial_q4k[lid + stride];
                }
                workgroupBarrier();
                stride = stride >> 1u;
            }
            if (l == 0u) {
                let r = partial_q4k[sub << 6u];
                let yo = e * rows;
                if (r0 < rows) { q1y[yo + r0] = r.x; }
                if (r1 < rows) { q1y[yo + r1] = r.y; }
                if (r2 < rows) { q1y[yo + r2] = r.z; }
                if (r3 < rows) { q1y[yo + r3] = r.w; }
            }
            workgroupBarrier();
            e = e + 1u;
        }
        wb = wb + nwg.x;
    }
}

// The two halves of a packed u32 as ready f32 weights. `q4v_dot8` fuses
// unpack and multiply, which is right for ONE activation vector and
// wrong for a batch: it redoes the shift/mask/convert/bias per nibble
// for every element. Split, the unpack is paid once.
// `q4v_dot8` over ALREADY-unpacked nibbles: the same eight products
// summed left to right, so a kernel that unpacks once for a batch can
// still land on the one-vector kernel's bits.
fn q4v_d8u(lo: vec4<f32>, hi: vec4<f32>, a: vec4<f32>, b: vec4<f32>) -> f32 {
    return lo.x * a.x + lo.y * a.y + lo.z * a.z + lo.w * a.w
         + hi.x * b.x + hi.y * b.y + hi.z * b.z + hi.w * b.w;
}

fn q4v_lo4(w: u32) -> vec4<f32> {
    return vec4<f32>(q4v_nib(w, 0u), q4v_nib(w, 4u), q4v_nib(w, 8u), q4v_nib(w, 12u));
}
fn q4v_hi4(w: u32) -> vec4<f32> {
    return vec4<f32>(q4v_nib(w, 16u), q4v_nib(w, 20u), q4v_nib(w, 24u), q4v_nib(w, 28u));
}

// ── q4tp matvec, batch blocked with the unpack SHARED (`CMF_MV_BK=2`).
// Same sixteen rows and four rows a lane as `q4tp_matvec4_bk`, but the
// nibbles of a u32 become f32 once and every batch element multiplies
// against them, where the sibling re-unpacks per element.
//
// The arithmetic is what this is for, and the budget is measured, not
// guessed. At b=1 a dense FFN layer needs 8.9 ms of weight stream and
// 6.3 ms of arithmetic, so the arithmetic hides and the token is
// bus-bound (`CMF_MV_PROBE` shows stripping it saves 4%). At b=3 the
// stream still needs 6.5 ms while the arithmetic needs 13.8 — measured
// 11.83/13.83/16.03 ms at b=2/3/4, dead linear, because each element
// unpacks again. Sharing it costs 3.25 ops a weight plus one FMA an
// element instead of five an element: at b=3 that is 6.25 against 15,
// which lands the batch FFN back on its memory floor.
//
// The unpack is held for ONE u32 at a time (eight f32 × four rows = 32
// registers). Holding a whole group's 32 weights for four rows was 128
// registers, and an earlier attempt that did something close to it lost
// outright: occupancy hides load latency, and this kernel still has to
// stream weights while it computes.
//
// WHAT BINDS IT NOW IS THE ACTIVATION LOADS, not registers — measured,
// and the opposite of what was assumed. At b=3 this arm leaves the FFN
// at 237 us a layer where the weight stream alone needs 139 and the
// arithmetic 123, and the obvious suspect was register pressure (peak
// ~68 against the one-vector kernel's ~40, so half the resident
// workgroups). A variant that unpacked ROW PAIRS to halve the live set
// — same weight traffic, same total unpack, twice the x re-reads — was
// 52% SLOWER (FFN 20.39 against 13.40 ms, verify 68.5 against 53.1,
// decode 42 against 51). Registers were not the wall; the 24 activation
// vec4 a group-iteration are. The move that would pay is FEWER x loads,
// never more.
//
// The next step, scoped and unblocked: f16 activations halve the load
// COUNT (32 f32 a group is eight vec4; as f16 it is four). It cannot go
// in this module — `enable f16` would make the whole of WGSL require
// SHADER_F16, and this module has to compile on adapters without it —
// so it wants its own source string next to COOP_MM_F16_SRC, built only
// when the adapter reports the feature, which is the pattern those
// kernels already follow. Worth ~0.8 ms per batch position, ~3.5% of a
// speculative round. More rows a lane is the other direction and is
// register-bound: 8 rows needs the packed weights (32), the
// accumulators (32) and the unpacked window at once.
@compute @workgroup_size(256)
fn q4tp_matvec4_bku(@builtin(workgroup_id) wid: vec3<u32>,
                    @builtin(num_workgroups) nwg: vec3<u32>,
                    @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let nb = max(q1p._p0, 1u);
    let blocks = (rows + 15u) / 16u;
    var wb = wid.x;
    loop {
        if (wb >= blocks) { break; }
        let base = wb * 16u;
        for (var t = lid; t < 512u; t = t + 256u) {
            let r = base + (t >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4w[t] = exp2(pr.x + f32(t & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        // Batches wider than the four accumulator components run in
        // chunks of four: the weight is re-read once per CHUNK, not once
        // per element. A k=4 draft is b=5, and before this it fell off
        // the kernel entirely and paid five reads (measured: k=4 decoded
        // SLOWER than k=3, 43.8 against 51.0, for exactly that reason).
        var cbase = 0u;
        loop {
        if (cbase >= nb) { break; }
        var acc0 = vec4<f32>(0.0);
        var acc1 = vec4<f32>(0.0);
        var acc2 = vec4<f32>(0.0);
        var acc3 = vec4<f32>(0.0);
        if (r0 < rows) {
            let c0 = codes_b + r0 * cstride;
            let c1 = codes_b + r1 * cstride;
            let c2 = codes_b + r2 * cstride;
            let c3 = codes_b + r3 * cstride;
            let l1 = r1 < rows;
            let l2 = r2 < rows;
            let l3 = r3 < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv = q4tp_byte(c0 + cbo);
                if (sh > 3u) { cv = cv | (q4tp_byte(c0 + cbo + 1u) << 8u); }
                let v0 = q4v_w[r0 * gpr + g];
                let s0 = lad_q4w[(sub << 5u) + ((cv >> sh) & 31u)];
                var v1 = vec4<u32>(0u, 0u, 0u, 0u);
                var s1 = 0.0;
                if (l1) {
                    cv = q4tp_byte(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c1 + cbo + 1u) << 8u); }
                    v1 = q4v_w[r1 * gpr + g];
                    s1 = lad_q4w[128u + (sub << 5u) + ((cv >> sh) & 31u)];
                }
                var v2 = vec4<u32>(0u, 0u, 0u, 0u);
                var s2 = 0.0;
                if (l2) {
                    cv = q4tp_byte(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c2 + cbo + 1u) << 8u); }
                    v2 = q4v_w[r2 * gpr + g];
                    s2 = lad_q4w[256u + (sub << 5u) + ((cv >> sh) & 31u)];
                }
                var v3 = vec4<u32>(0u, 0u, 0u, 0u);
                var s3 = 0.0;
                if (l3) {
                    cv = q4tp_byte(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c3 + cbo + 1u) << 8u); }
                    v3 = q4v_w[r3 * gpr + g];
                    s3 = lad_q4w[384u + (sub << 5u) + ((cv >> sh) & 31u)];
                }
                // Per (row, element) the group's four u32 words are summed
                // FIRST — each word through the one-vector kernel's scalar
                // chain (`q4v_d8u` == `q4v_dot8` term for term) and the four
                // words left to right — and only then scaled once into the
                // accumulator: `acc += s * (d0 + d1 + d2 + d3)`, exactly
                // `q4tp_matvec16w`'s expression. It used to scale and add
                // per word with vec4 `dot`, which is a different rounding
                // and made a speculative verify's row disagree with the
                // plain token's on near-ties. Bit-identical now (tested).
                var t0 = vec4<f32>(0.0); var t1 = vec4<f32>(0.0);
                var t2 = vec4<f32>(0.0); var t3 = vec4<f32>(0.0);
                var j = 0u;
                loop {
                    if (j >= 4u) { break; }
                    var w0 = v0.x; var w1 = v1.x; var w2 = v2.x; var w3 = v3.x;
                    if (j == 1u) { w0 = v0.y; w1 = v1.y; w2 = v2.y; w3 = v3.y; }
                    if (j == 2u) { w0 = v0.z; w1 = v1.z; w2 = v2.z; w3 = v3.z; }
                    if (j == 3u) { w0 = v0.w; w1 = v1.w; w2 = v2.w; w3 = v3.w; }
                    let a0 = q4v_lo4(w0); let b0 = q4v_hi4(w0);
                    let a1 = q4v_lo4(w1); let b1 = q4v_hi4(w1);
                    let a2 = q4v_lo4(w2); let b2 = q4v_hi4(w2);
                    let a3 = q4v_lo4(w3); let b3 = q4v_hi4(w3);
                    let xj = g * 8u + j * 2u;
                    {
                        let x0q = xj + gpr * 8u * cbase;
                        let xa = q4v_x[x0q]; let xb = q4v_x[x0q + 1u];
                        t0.x = t0.x + q4v_d8u(a0, b0, xa, xb);
                        t1.x = t1.x + q4v_d8u(a1, b1, xa, xb);
                        t2.x = t2.x + q4v_d8u(a2, b2, xa, xb);
                        t3.x = t3.x + q4v_d8u(a3, b3, xa, xb);
                    }
                    if (cbase + 1u < nb) {
                        let xq = xj + gpr * 8u * (cbase + 1u);
                        let xa = q4v_x[xq]; let xb = q4v_x[xq + 1u];
                        t0.y = t0.y + q4v_d8u(a0, b0, xa, xb);
                        t1.y = t1.y + q4v_d8u(a1, b1, xa, xb);
                        t2.y = t2.y + q4v_d8u(a2, b2, xa, xb);
                        t3.y = t3.y + q4v_d8u(a3, b3, xa, xb);
                    }
                    if (cbase + 2u < nb) {
                        let xq = xj + gpr * 8u * (cbase + 2u);
                        let xa = q4v_x[xq]; let xb = q4v_x[xq + 1u];
                        t0.z = t0.z + q4v_d8u(a0, b0, xa, xb);
                        t1.z = t1.z + q4v_d8u(a1, b1, xa, xb);
                        t2.z = t2.z + q4v_d8u(a2, b2, xa, xb);
                        t3.z = t3.z + q4v_d8u(a3, b3, xa, xb);
                    }
                    if (cbase + 3u < nb) {
                        let xq = xj + gpr * 8u * (cbase + 3u);
                        let xa = q4v_x[xq]; let xb = q4v_x[xq + 1u];
                        t0.w = t0.w + q4v_d8u(a0, b0, xa, xb);
                        t1.w = t1.w + q4v_d8u(a1, b1, xa, xb);
                        t2.w = t2.w + q4v_d8u(a2, b2, xa, xb);
                        t3.w = t3.w + q4v_d8u(a3, b3, xa, xb);
                    }
                    j = j + 1u;
                }
                acc0 = acc0 + s0 * t0;
                acc1 = acc1 + s1 * t1;
                acc2 = acc2 + s2 * t2;
                acc3 = acc3 + s3 * t3;
                g = g + 64u;
            }
        }
        var e = 0u;
        loop {
            if (cbase + e >= nb || e >= 4u) { break; }
            var mine = vec4<f32>(acc0.x, acc1.x, acc2.x, acc3.x);
            if (e == 1u) { mine = vec4<f32>(acc0.y, acc1.y, acc2.y, acc3.y); }
            if (e == 2u) { mine = vec4<f32>(acc0.z, acc1.z, acc2.z, acc3.z); }
            if (e == 3u) { mine = vec4<f32>(acc0.w, acc1.w, acc2.w, acc3.w); }
            partial_q4k[lid] = mine;
            workgroupBarrier();
            var stride = 32u;
            loop {
                if (stride == 0u) { break; }
                if (l < stride) {
                    partial_q4k[lid] = partial_q4k[lid] + partial_q4k[lid + stride];
                }
                workgroupBarrier();
                stride = stride >> 1u;
            }
            if (l == 0u) {
                let r = partial_q4k[sub << 6u];
                let yo = (cbase + e) * rows;
                if (r0 < rows) { q1y[yo + r0] = r.x; }
                if (r1 < rows) { q1y[yo + r1] = r.y; }
                if (r2 < rows) { q1y[yo + r2] = r.z; }
                if (r3 < rows) { q1y[yo + r3] = r.w; }
            }
            workgroupBarrier();
            e = e + 1u;
        }
        cbase = cbase + 4u;
        }
        wb = wb + nwg.x;
    }
}


// ── INT8-ACTIVATION batched matvec (default; `CMF_VERIFY_I8=0` for f32): the speculative
// verify's b ≤ 8 rows over q4tp weights with the activations quantized to
// int8 per 32-group and the inner product on `dot4I8Packed` (dp4a: four
// MACs an instruction). The f32 batched kernel above unpacks a weight
// word into eight f32 (I2F/magic + sub) and pays 32 FMAs per (row,
// element) — measured 3.2x a single token for 5 rows, +4.7 ms per extra
// row, i.e. bound by per-row work, not by the weight stream. Here a word's
// eight nibbles become two packed int8x4 (mask, shift, mask — 3 int ops,
// shared across the batch) and each element costs two dp4a per word: for
// a 32-group per row that is 8 dp4a + one I2F + 2 FMA an element against
// 32 FMA + a share of 32 unpacks. The bias: Σ (n − 8)·x = Σ n·x − 8·Σx,
// with Σx per group precomputed by the quantizer.
//
// NOT bit-exact against the one-vector kernel: x carries an int8 rounding
// (per-32 symmetric, |err| ≤ s/2 with s = max|x|/127 — the Q8_1
// activation grid every q4 matvec in llama.cpp runs on). The verify's
// argmax is the batched kernel's own; a near-tie can therefore resolve
// differently from the plain path. `CMF_VERIFY_I8=0` for the f32 verify;
// the parity test bounds the drift.
//
// x8 layout (per element e, group g): two vec4<u32> at (e·gpr + g)·2:
//   [w0.even, w0.odd, w1.even, w1.odd], [w2.even, w2.odd, w3.even, w3.odd]
// where word j covers x[8j..8j+8) of the group, `even` packs x[8j+0,2,4,6]
// and `odd` x[8j+1,3,5,7] as int8 — the order the nibbles come out of a
// weight word (lo nibbles = even positions, hi = odd).
// xs (per element, group): vec2(s_x, 8·s_x·Σq).
@group(0) @binding(9) var<storage, read> q4v_x8 : array<vec4<u32>>;
@group(0) @binding(10) var<storage, read> q4v_xs : array<vec2<f32>>;
@group(0) @binding(11) var<storage, read_write> q4v_x8w : array<vec4<u32>>;
@group(0) @binding(12) var<storage, read_write> q4v_xsw : array<vec2<f32>>;

fn pack_i8x4(a: i32, b: i32, c: i32, d: i32) -> u32 {
    return (u32(a) & 0xFFu) | ((u32(b) & 0xFFu) << 8u) | ((u32(c) & 0xFFu) << 16u) | ((u32(d) & 0xFFu) << 24u);
}

// One thread per (element, group): 32 f32 → int8, packed even/odd per
// word, plus (s, 8·s·Σq). Params: np = gpr, rows = batch.
@compute @workgroup_size(256)
fn x_quant_i8(@builtin(global_invocation_id) gid: vec3<u32>) {
    let gpr = q1p.np;
    let nb = q1p.rows;
    let i = gid.x;
    if (i >= gpr * nb) { return; }
    let x0 = i * 8u;
    var amax = 0.0;
    for (var j = 0u; j < 8u; j = j + 1u) {
        let v = abs(q4v_x[x0 + j]);
        amax = max(amax, max(max(v.x, v.y), max(v.z, v.w)));
    }
    let s = select(amax / 127.0, 1.0, amax == 0.0);
    let inv = 1.0 / s;
    var sum = 0;
    var out0 = vec4<u32>(0u);
    var out1 = vec4<u32>(0u);
    for (var j = 0u; j < 4u; j = j + 1u) {
        let a = q4v_x[x0 + j * 2u];
        let b = q4v_x[x0 + j * 2u + 1u];
        let qa = vec4<i32>(round(a * inv));
        let qb = vec4<i32>(round(b * inv));
        let qa2 = clamp(qa, vec4<i32>(-127), vec4<i32>(127));
        let qb2 = clamp(qb, vec4<i32>(-127), vec4<i32>(127));
        sum = sum + qa2.x + qa2.y + qa2.z + qa2.w + qb2.x + qb2.y + qb2.z + qb2.w;
        let ev = pack_i8x4(qa2.x, qa2.z, qb2.x, qb2.z);
        let od = pack_i8x4(qa2.y, qa2.w, qb2.y, qb2.w);
        if (j == 0u) { out0.x = ev; out0.y = od; }
        if (j == 1u) { out0.z = ev; out0.w = od; }
        if (j == 2u) { out1.x = ev; out1.y = od; }
        if (j == 3u) { out1.z = ev; out1.w = od; }
    }
    q4v_x8w[i * 2u] = out0;
    q4v_x8w[i * 2u + 1u] = out1;
    q4v_xsw[i] = vec2<f32>(s, 8.0 * s * f32(sum));
}

// The batch is a PIPELINE CONSTANT: one pipeline per batch size 2..8,
// so the element loop unrolls and the accumulator picks below become
// static register writes instead of eight predicated selects per row.
override NB8: u32 = 4u;
@compute @workgroup_size(256)
fn q4tp_matvec4_bk8(@builtin(workgroup_id) wid: vec3<u32>,
                    @builtin(num_workgroups) nwg: vec3<u32>,
                    @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let nb = NB8;
    let blocks = (rows + 15u) / 16u;
    var wb = wid.x;
    loop {
        if (wb >= blocks) { break; }
        let base = wb * 16u;
        for (var t = lid; t < 512u; t = t + 256u) {
            let r = base + (t >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4w[t] = exp2(pr.x + f32(t & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        // acc[row] = two vec4 (elements 0..3, 4..7)
        var a0 = vec4<f32>(0.0); var a0h = vec4<f32>(0.0);
        var a1 = vec4<f32>(0.0); var a1h = vec4<f32>(0.0);
        var a2 = vec4<f32>(0.0); var a2h = vec4<f32>(0.0);
        var a3 = vec4<f32>(0.0); var a3h = vec4<f32>(0.0);
        if (r0 < rows) {
            let c0 = codes_b + r0 * cstride;
            let c1 = codes_b + r1 * cstride;
            let c2 = codes_b + r2 * cstride;
            let c3 = codes_b + r3 * cstride;
            let l1 = r1 < rows;
            let l2 = r2 < rows;
            let l3 = r3 < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv = q4tp_byte(c0 + cbo);
                if (sh > 3u) { cv = cv | (q4tp_byte(c0 + cbo + 1u) << 8u); }
                let s0 = lad_q4w[(sub << 5u) + ((cv >> sh) & 31u)];
                let v0 = q4v_w[r0 * gpr + g];
                var s1 = 0.0; var v1 = vec4<u32>(0u);
                if (l1) {
                    cv = q4tp_byte(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c1 + cbo + 1u) << 8u); }
                    s1 = lad_q4w[128u + (sub << 5u) + ((cv >> sh) & 31u)];
                    v1 = q4v_w[r1 * gpr + g];
                }
                var s2 = 0.0; var v2 = vec4<u32>(0u);
                if (l2) {
                    cv = q4tp_byte(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c2 + cbo + 1u) << 8u); }
                    s2 = lad_q4w[256u + (sub << 5u) + ((cv >> sh) & 31u)];
                    v2 = q4v_w[r2 * gpr + g];
                }
                var s3 = 0.0; var v3 = vec4<u32>(0u);
                if (l3) {
                    cv = q4tp_byte(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c3 + cbo + 1u) << 8u); }
                    s3 = lad_q4w[384u + (sub << 5u) + ((cv >> sh) & 31u)];
                    v3 = q4v_w[r3 * gpr + g];
                }
                // packed nibbles: even positions (lo) / odd (hi) per word
                let m = 0x0F0F0F0Fu;
                let lo0 = v0 & vec4<u32>(m);       let hi0 = (v0 >> vec4<u32>(4u)) & vec4<u32>(m);
                let lo1 = v1 & vec4<u32>(m);       let hi1 = (v1 >> vec4<u32>(4u)) & vec4<u32>(m);
                let lo2 = v2 & vec4<u32>(m);       let hi2 = (v2 >> vec4<u32>(4u)) & vec4<u32>(m);
                let lo3 = v3 & vec4<u32>(m);       let hi3 = (v3 >> vec4<u32>(4u)) & vec4<u32>(m);
                for (var e = 0u; e < NB8; e = e + 1u) {
                    let xi = (e * gpr + g) * 2u;
                    let xa = q4v_x8[xi];
                    let xb = q4v_x8[xi + 1u];
                    let xs = q4v_xs[e * gpr + g];
                    // row 0
                    var d = dot4I8Packed(lo0.x, xa.x) + dot4I8Packed(hi0.x, xa.y)
                          + dot4I8Packed(lo0.y, xa.z) + dot4I8Packed(hi0.y, xa.w)
                          + dot4I8Packed(lo0.z, xb.x) + dot4I8Packed(hi0.z, xb.y)
                          + dot4I8Packed(lo0.w, xb.z) + dot4I8Packed(hi0.w, xb.w);
                    let t0 = s0 * fma(f32(d), xs.x, -xs.y);
                    d = dot4I8Packed(lo1.x, xa.x) + dot4I8Packed(hi1.x, xa.y)
                      + dot4I8Packed(lo1.y, xa.z) + dot4I8Packed(hi1.y, xa.w)
                      + dot4I8Packed(lo1.z, xb.x) + dot4I8Packed(hi1.z, xb.y)
                      + dot4I8Packed(lo1.w, xb.z) + dot4I8Packed(hi1.w, xb.w);
                    let t1 = s1 * fma(f32(d), xs.x, -xs.y);
                    d = dot4I8Packed(lo2.x, xa.x) + dot4I8Packed(hi2.x, xa.y)
                      + dot4I8Packed(lo2.y, xa.z) + dot4I8Packed(hi2.y, xa.w)
                      + dot4I8Packed(lo2.z, xb.x) + dot4I8Packed(hi2.z, xb.y)
                      + dot4I8Packed(lo2.w, xb.z) + dot4I8Packed(hi2.w, xb.w);
                    let t2 = s2 * fma(f32(d), xs.x, -xs.y);
                    d = dot4I8Packed(lo3.x, xa.x) + dot4I8Packed(hi3.x, xa.y)
                      + dot4I8Packed(lo3.y, xa.z) + dot4I8Packed(hi3.y, xa.w)
                      + dot4I8Packed(lo3.z, xb.x) + dot4I8Packed(hi3.z, xb.y)
                      + dot4I8Packed(lo3.w, xb.z) + dot4I8Packed(hi3.w, xb.w);
                    let t3 = s3 * fma(f32(d), xs.x, -xs.y);
                    if (e == 0u) { a0.x += t0; a1.x += t1; a2.x += t2; a3.x += t3; }
                    if (e == 1u) { a0.y += t0; a1.y += t1; a2.y += t2; a3.y += t3; }
                    if (e == 2u) { a0.z += t0; a1.z += t1; a2.z += t2; a3.z += t3; }
                    if (e == 3u) { a0.w += t0; a1.w += t1; a2.w += t2; a3.w += t3; }
                    if (e == 4u) { a0h.x += t0; a1h.x += t1; a2h.x += t2; a3h.x += t3; }
                    if (e == 5u) { a0h.y += t0; a1h.y += t1; a2h.y += t2; a3h.y += t3; }
                    if (e == 6u) { a0h.z += t0; a1h.z += t1; a2h.z += t2; a3h.z += t3; }
                    if (e == 7u) { a0h.w += t0; a1h.w += t1; a2h.w += t2; a3h.w += t3; }
                }
                g = g + 64u;
            }
        }
        // reduce across the 64 lanes, one element at a time
        var e = 0u;
        loop {
            if (e >= nb) { break; }
            var mine = vec4<f32>(0.0);
            if (e == 0u) { mine = vec4<f32>(a0.x, a1.x, a2.x, a3.x); }
            if (e == 1u) { mine = vec4<f32>(a0.y, a1.y, a2.y, a3.y); }
            if (e == 2u) { mine = vec4<f32>(a0.z, a1.z, a2.z, a3.z); }
            if (e == 3u) { mine = vec4<f32>(a0.w, a1.w, a2.w, a3.w); }
            if (e == 4u) { mine = vec4<f32>(a0h.x, a1h.x, a2h.x, a3h.x); }
            if (e == 5u) { mine = vec4<f32>(a0h.y, a1h.y, a2h.y, a3h.y); }
            if (e == 6u) { mine = vec4<f32>(a0h.z, a1h.z, a2h.z, a3h.z); }
            if (e == 7u) { mine = vec4<f32>(a0h.w, a1h.w, a2h.w, a3h.w); }
            partial_q4k[lid] = mine;
            workgroupBarrier();
            var stride = 32u;
            loop {
                if (stride == 0u) { break; }
                if (l < stride) {
                    partial_q4k[lid] = partial_q4k[lid] + partial_q4k[lid + stride];
                }
                workgroupBarrier();
                stride = stride >> 1u;
            }
            if (l == 0u) {
                let r = partial_q4k[sub << 6u];
                let yo = e * rows;
                if (r0 < rows) { q1y[yo + r0] = r.x; }
                if (r1 < rows) { q1y[yo + r1] = r.y; }
                if (r2 < rows) { q1y[yo + r2] = r.z; }
                if (r3 < rows) { q1y[yo + r3] = r.w; }
            }
            workgroupBarrier();
            e = e + 1u;
        }
        wb = wb + nwg.x;
    }
}

// ── The batched (bku) matvec for TWO weights of one input batch in one
// dispatch — the x2 fusion for the batch graph (verify / batched
// prefill): gate+up, GDN qkv+z, attention k+v. Body generated from
// `q4tp_matvec4_bku`, side B over its own bindings; per-row, per-element
// arithmetic identical. Same params as x2: rows2 in `_p1`, batch in `_p0`.
@compute @workgroup_size(256)
fn q4tp_matvec4_bku_x2(@builtin(workgroup_id) wid: vec3<u32>,
                       @builtin(num_workgroups) nwg: vec3<u32>,
                       @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let rows2 = q1p._p1;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let params_w2 = rows2 * gpr * 4u;
    let codes_b2 = rows2 * gpr * 16u + rows2 * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let nb = max(q1p._p0, 1u);
    let blocks1 = (rows + 15u) / 16u;
    let blocks2 = (rows2 + 15u) / 16u;
    var wb = wid.x;
    loop {
        if (wb >= blocks1 + blocks2) { break; }
        if (wb < blocks1) {
        let base = wb * 16u;
        for (var t = lid; t < 512u; t = t + 256u) {
            let r = base + (t >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4w[t] = exp2(pr.x + f32(t & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        // Batches wider than the four accumulator components run in
        // chunks of four: the weight is re-read once per CHUNK, not once
        // per element. A k=4 draft is b=5, and before this it fell off
        // the kernel entirely and paid five reads (measured: k=4 decoded
        // SLOWER than k=3, 43.8 against 51.0, for exactly that reason).
        var cbase = 0u;
        loop {
        if (cbase >= nb) { break; }
        var acc0 = vec4<f32>(0.0);
        var acc1 = vec4<f32>(0.0);
        var acc2 = vec4<f32>(0.0);
        var acc3 = vec4<f32>(0.0);
        if (r0 < rows) {
            let c0 = codes_b + r0 * cstride;
            let c1 = codes_b + r1 * cstride;
            let c2 = codes_b + r2 * cstride;
            let c3 = codes_b + r3 * cstride;
            let l1 = r1 < rows;
            let l2 = r2 < rows;
            let l3 = r3 < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv = q4tp_byte(c0 + cbo);
                if (sh > 3u) { cv = cv | (q4tp_byte(c0 + cbo + 1u) << 8u); }
                let v0 = q4v_w[r0 * gpr + g];
                let s0 = lad_q4w[(sub << 5u) + ((cv >> sh) & 31u)];
                var v1 = vec4<u32>(0u, 0u, 0u, 0u);
                var s1 = 0.0;
                if (l1) {
                    cv = q4tp_byte(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c1 + cbo + 1u) << 8u); }
                    v1 = q4v_w[r1 * gpr + g];
                    s1 = lad_q4w[128u + (sub << 5u) + ((cv >> sh) & 31u)];
                }
                var v2 = vec4<u32>(0u, 0u, 0u, 0u);
                var s2 = 0.0;
                if (l2) {
                    cv = q4tp_byte(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c2 + cbo + 1u) << 8u); }
                    v2 = q4v_w[r2 * gpr + g];
                    s2 = lad_q4w[256u + (sub << 5u) + ((cv >> sh) & 31u)];
                }
                var v3 = vec4<u32>(0u, 0u, 0u, 0u);
                var s3 = 0.0;
                if (l3) {
                    cv = q4tp_byte(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c3 + cbo + 1u) << 8u); }
                    v3 = q4v_w[r3 * gpr + g];
                    s3 = lad_q4w[384u + (sub << 5u) + ((cv >> sh) & 31u)];
                }
                // Per (row, element) the group's four u32 words are summed
                // FIRST — each word through the one-vector kernel's scalar
                // chain (`q4v_d8u` == `q4v_dot8` term for term) and the four
                // words left to right — and only then scaled once into the
                // accumulator: `acc += s * (d0 + d1 + d2 + d3)`, exactly
                // `q4tp_matvec16w`'s expression. It used to scale and add
                // per word with vec4 `dot`, which is a different rounding
                // and made a speculative verify's row disagree with the
                // plain token's on near-ties. Bit-identical now (tested).
                var t0 = vec4<f32>(0.0); var t1 = vec4<f32>(0.0);
                var t2 = vec4<f32>(0.0); var t3 = vec4<f32>(0.0);
                var j = 0u;
                loop {
                    if (j >= 4u) { break; }
                    var w0 = v0.x; var w1 = v1.x; var w2 = v2.x; var w3 = v3.x;
                    if (j == 1u) { w0 = v0.y; w1 = v1.y; w2 = v2.y; w3 = v3.y; }
                    if (j == 2u) { w0 = v0.z; w1 = v1.z; w2 = v2.z; w3 = v3.z; }
                    if (j == 3u) { w0 = v0.w; w1 = v1.w; w2 = v2.w; w3 = v3.w; }
                    let a0 = q4v_lo4(w0); let b0 = q4v_hi4(w0);
                    let a1 = q4v_lo4(w1); let b1 = q4v_hi4(w1);
                    let a2 = q4v_lo4(w2); let b2 = q4v_hi4(w2);
                    let a3 = q4v_lo4(w3); let b3 = q4v_hi4(w3);
                    let xj = g * 8u + j * 2u;
                    {
                        let x0q = xj + gpr * 8u * cbase;
                        let xa = q4v_x[x0q]; let xb = q4v_x[x0q + 1u];
                        t0.x = t0.x + q4v_d8u(a0, b0, xa, xb);
                        t1.x = t1.x + q4v_d8u(a1, b1, xa, xb);
                        t2.x = t2.x + q4v_d8u(a2, b2, xa, xb);
                        t3.x = t3.x + q4v_d8u(a3, b3, xa, xb);
                    }
                    if (cbase + 1u < nb) {
                        let xq = xj + gpr * 8u * (cbase + 1u);
                        let xa = q4v_x[xq]; let xb = q4v_x[xq + 1u];
                        t0.y = t0.y + q4v_d8u(a0, b0, xa, xb);
                        t1.y = t1.y + q4v_d8u(a1, b1, xa, xb);
                        t2.y = t2.y + q4v_d8u(a2, b2, xa, xb);
                        t3.y = t3.y + q4v_d8u(a3, b3, xa, xb);
                    }
                    if (cbase + 2u < nb) {
                        let xq = xj + gpr * 8u * (cbase + 2u);
                        let xa = q4v_x[xq]; let xb = q4v_x[xq + 1u];
                        t0.z = t0.z + q4v_d8u(a0, b0, xa, xb);
                        t1.z = t1.z + q4v_d8u(a1, b1, xa, xb);
                        t2.z = t2.z + q4v_d8u(a2, b2, xa, xb);
                        t3.z = t3.z + q4v_d8u(a3, b3, xa, xb);
                    }
                    if (cbase + 3u < nb) {
                        let xq = xj + gpr * 8u * (cbase + 3u);
                        let xa = q4v_x[xq]; let xb = q4v_x[xq + 1u];
                        t0.w = t0.w + q4v_d8u(a0, b0, xa, xb);
                        t1.w = t1.w + q4v_d8u(a1, b1, xa, xb);
                        t2.w = t2.w + q4v_d8u(a2, b2, xa, xb);
                        t3.w = t3.w + q4v_d8u(a3, b3, xa, xb);
                    }
                    j = j + 1u;
                }
                acc0 = acc0 + s0 * t0;
                acc1 = acc1 + s1 * t1;
                acc2 = acc2 + s2 * t2;
                acc3 = acc3 + s3 * t3;
                g = g + 64u;
            }
        }
        var e = 0u;
        loop {
            if (cbase + e >= nb || e >= 4u) { break; }
            var mine = vec4<f32>(acc0.x, acc1.x, acc2.x, acc3.x);
            if (e == 1u) { mine = vec4<f32>(acc0.y, acc1.y, acc2.y, acc3.y); }
            if (e == 2u) { mine = vec4<f32>(acc0.z, acc1.z, acc2.z, acc3.z); }
            if (e == 3u) { mine = vec4<f32>(acc0.w, acc1.w, acc2.w, acc3.w); }
            partial_q4k[lid] = mine;
            workgroupBarrier();
            var stride = 32u;
            loop {
                if (stride == 0u) { break; }
                if (l < stride) {
                    partial_q4k[lid] = partial_q4k[lid] + partial_q4k[lid + stride];
                }
                workgroupBarrier();
                stride = stride >> 1u;
            }
            if (l == 0u) {
                let r = partial_q4k[sub << 6u];
                let yo = (cbase + e) * rows;
                if (r0 < rows) { q1y[yo + r0] = r.x; }
                if (r1 < rows) { q1y[yo + r1] = r.y; }
                if (r2 < rows) { q1y[yo + r2] = r.z; }
                if (r3 < rows) { q1y[yo + r3] = r.w; }
            }
            workgroupBarrier();
            e = e + 1u;
        }
        cbase = cbase + 4u;
        }
        } else {
        let base = (wb - blocks1) * 16u;
        for (var t = lid; t < 512u; t = t + 256u) {
            let r = base + (t >> 5u);
            if (r < rows2) {
                let pr = unpack2x16float(q1w2[params_w2 + r]);
                lad_q4w[t] = exp2(pr.x + f32(t & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        // Batches wider than the four accumulator components run in
        // chunks of four: the weight is re-read once per CHUNK, not once
        // per element. A k=4 draft is b=5, and before this it fell off
        // the kernel entirely and paid five reads (measured: k=4 decoded
        // SLOWER than k=3, 43.8 against 51.0, for exactly that reason).
        var cbase = 0u;
        loop {
        if (cbase >= nb) { break; }
        var acc0 = vec4<f32>(0.0);
        var acc1 = vec4<f32>(0.0);
        var acc2 = vec4<f32>(0.0);
        var acc3 = vec4<f32>(0.0);
        if (r0 < rows2) {
            let c0 = codes_b2 + r0 * cstride;
            let c1 = codes_b2 + r1 * cstride;
            let c2 = codes_b2 + r2 * cstride;
            let c3 = codes_b2 + r3 * cstride;
            let l1 = r1 < rows2;
            let l2 = r2 < rows2;
            let l3 = r3 < rows2;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv = q4tp_byte2(c0 + cbo);
                if (sh > 3u) { cv = cv | (q4tp_byte2(c0 + cbo + 1u) << 8u); }
                let v0 = q4v_w2[r0 * gpr + g];
                let s0 = lad_q4w[(sub << 5u) + ((cv >> sh) & 31u)];
                var v1 = vec4<u32>(0u, 0u, 0u, 0u);
                var s1 = 0.0;
                if (l1) {
                    cv = q4tp_byte2(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte2(c1 + cbo + 1u) << 8u); }
                    v1 = q4v_w2[r1 * gpr + g];
                    s1 = lad_q4w[128u + (sub << 5u) + ((cv >> sh) & 31u)];
                }
                var v2 = vec4<u32>(0u, 0u, 0u, 0u);
                var s2 = 0.0;
                if (l2) {
                    cv = q4tp_byte2(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte2(c2 + cbo + 1u) << 8u); }
                    v2 = q4v_w2[r2 * gpr + g];
                    s2 = lad_q4w[256u + (sub << 5u) + ((cv >> sh) & 31u)];
                }
                var v3 = vec4<u32>(0u, 0u, 0u, 0u);
                var s3 = 0.0;
                if (l3) {
                    cv = q4tp_byte2(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte2(c3 + cbo + 1u) << 8u); }
                    v3 = q4v_w2[r3 * gpr + g];
                    s3 = lad_q4w[384u + (sub << 5u) + ((cv >> sh) & 31u)];
                }
                // Per (row, element) the group's four u32 words are summed
                // FIRST — each word through the one-vector kernel's scalar
                // chain (`q4v_d8u` == `q4v_dot8` term for term) and the four
                // words left to right — and only then scaled once into the
                // accumulator: `acc += s * (d0 + d1 + d2 + d3)`, exactly
                // `q4tp_matvec16w`'s expression. It used to scale and add
                // per word with vec4 `dot`, which is a different rounding
                // and made a speculative verify's row disagree with the
                // plain token's on near-ties. Bit-identical now (tested).
                var t0 = vec4<f32>(0.0); var t1 = vec4<f32>(0.0);
                var t2 = vec4<f32>(0.0); var t3 = vec4<f32>(0.0);
                var j = 0u;
                loop {
                    if (j >= 4u) { break; }
                    var w0 = v0.x; var w1 = v1.x; var w2 = v2.x; var w3 = v3.x;
                    if (j == 1u) { w0 = v0.y; w1 = v1.y; w2 = v2.y; w3 = v3.y; }
                    if (j == 2u) { w0 = v0.z; w1 = v1.z; w2 = v2.z; w3 = v3.z; }
                    if (j == 3u) { w0 = v0.w; w1 = v1.w; w2 = v2.w; w3 = v3.w; }
                    let a0 = q4v_lo4(w0); let b0 = q4v_hi4(w0);
                    let a1 = q4v_lo4(w1); let b1 = q4v_hi4(w1);
                    let a2 = q4v_lo4(w2); let b2 = q4v_hi4(w2);
                    let a3 = q4v_lo4(w3); let b3 = q4v_hi4(w3);
                    let xj = g * 8u + j * 2u;
                    {
                        let x0q = xj + gpr * 8u * cbase;
                        let xa = q4v_x[x0q]; let xb = q4v_x[x0q + 1u];
                        t0.x = t0.x + q4v_d8u(a0, b0, xa, xb);
                        t1.x = t1.x + q4v_d8u(a1, b1, xa, xb);
                        t2.x = t2.x + q4v_d8u(a2, b2, xa, xb);
                        t3.x = t3.x + q4v_d8u(a3, b3, xa, xb);
                    }
                    if (cbase + 1u < nb) {
                        let xq = xj + gpr * 8u * (cbase + 1u);
                        let xa = q4v_x[xq]; let xb = q4v_x[xq + 1u];
                        t0.y = t0.y + q4v_d8u(a0, b0, xa, xb);
                        t1.y = t1.y + q4v_d8u(a1, b1, xa, xb);
                        t2.y = t2.y + q4v_d8u(a2, b2, xa, xb);
                        t3.y = t3.y + q4v_d8u(a3, b3, xa, xb);
                    }
                    if (cbase + 2u < nb) {
                        let xq = xj + gpr * 8u * (cbase + 2u);
                        let xa = q4v_x[xq]; let xb = q4v_x[xq + 1u];
                        t0.z = t0.z + q4v_d8u(a0, b0, xa, xb);
                        t1.z = t1.z + q4v_d8u(a1, b1, xa, xb);
                        t2.z = t2.z + q4v_d8u(a2, b2, xa, xb);
                        t3.z = t3.z + q4v_d8u(a3, b3, xa, xb);
                    }
                    if (cbase + 3u < nb) {
                        let xq = xj + gpr * 8u * (cbase + 3u);
                        let xa = q4v_x[xq]; let xb = q4v_x[xq + 1u];
                        t0.w = t0.w + q4v_d8u(a0, b0, xa, xb);
                        t1.w = t1.w + q4v_d8u(a1, b1, xa, xb);
                        t2.w = t2.w + q4v_d8u(a2, b2, xa, xb);
                        t3.w = t3.w + q4v_d8u(a3, b3, xa, xb);
                    }
                    j = j + 1u;
                }
                acc0 = acc0 + s0 * t0;
                acc1 = acc1 + s1 * t1;
                acc2 = acc2 + s2 * t2;
                acc3 = acc3 + s3 * t3;
                g = g + 64u;
            }
        }
        var e = 0u;
        loop {
            if (cbase + e >= nb || e >= 4u) { break; }
            var mine = vec4<f32>(acc0.x, acc1.x, acc2.x, acc3.x);
            if (e == 1u) { mine = vec4<f32>(acc0.y, acc1.y, acc2.y, acc3.y); }
            if (e == 2u) { mine = vec4<f32>(acc0.z, acc1.z, acc2.z, acc3.z); }
            if (e == 3u) { mine = vec4<f32>(acc0.w, acc1.w, acc2.w, acc3.w); }
            partial_q4k[lid] = mine;
            workgroupBarrier();
            var stride = 32u;
            loop {
                if (stride == 0u) { break; }
                if (l < stride) {
                    partial_q4k[lid] = partial_q4k[lid] + partial_q4k[lid + stride];
                }
                workgroupBarrier();
                stride = stride >> 1u;
            }
            if (l == 0u) {
                let r = partial_q4k[sub << 6u];
                let yo = (cbase + e) * rows2;
                if (r0 < rows2) { q1y2[yo + r0] = r.x; }
                if (r1 < rows2) { q1y2[yo + r1] = r.y; }
                if (r2 < rows2) { q1y2[yo + r2] = r.z; }
                if (r3 < rows2) { q1y2[yo + r3] = r.w; }
            }
            workgroupBarrier();
            e = e + 1u;
        }
        cbase = cbase + 4u;
        }
        }
        wb = wb + nwg.x;
    }
}

// ── TIMING PROBE, answers are GARBAGE (`CMF_MV_PROBE`). The quad-row
// kernel with its arithmetic removed and every LOAD kept: same grid,
// same weight/code/x traffic, same loop trip count, but no nibble
// unpack and no per-weight FMA. What it measures is the floor the
// memory system alone imposes on this access pattern — the number that
// decides whether the matvec is worth optimizing at all.
// `_p1` is a BITMASK: 1 = the base probe (no unpack, no FMA), +2 drops
// the activation loads, +4 drops the code-plane loads. The code plane
// is a SECOND stream, far from the weights, touched once per 16 bytes
// of weight; whether that costs anything is not derivable from the
// byte counts, only from running without it.
@compute @workgroup_size(256)
fn q4tp_matvec16w_probe(@builtin(workgroup_id) wid: vec3<u32>,
                        @builtin(num_workgroups) nwg: vec3<u32>,
                        @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let drop_x = (q1p._p1 & 2u) != 0u;
    let drop_c = (q1p._p1 & 4u) != 0u;
    let drop_r = (q1p._p1 & 8u) != 0u;
    let blocks = (rows + 15u) / 16u;
    var wb = wid.x;
    loop {
        if (wb >= blocks) { break; }
        let base = wb * 16u;
        // drop_l (&16): the ladder was the last piece without a bit —
        // every other suspect measured null. Garbage scales, honest time.
        if ((q1p._p1 & 16u) == 0u) {
            for (var t = lid; t < 512u; t = t + 256u) {
                let r = base + (t >> 5u);
                if (r < rows) {
                    let pr = unpack2x16float(q1w[params_w + r]);
                    lad_q4w[t] = exp2(pr.x + f32(t & 31u) * pr.y);
                }
            }
            workgroupBarrier();
        }
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        var acc = vec4<f32>(0.0);
        if (r0 < rows) {
            let c0 = codes_b + r0 * cstride;
            let c1 = codes_b + r1 * cstride;
            let c2 = codes_b + r2 * cstride;
            let c3 = codes_b + r3 * cstride;
            let l1 = r1 < rows;
            let l2 = r2 < rows;
            let l3 = r3 < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var xs = 1.0;
                if (!drop_x) {
                    let x0 = g * 8u;
                    xs = q4v_x[x0].x + q4v_x[x0 + 1u].x + q4v_x[x0 + 2u].x
                       + q4v_x[x0 + 3u].x + q4v_x[x0 + 4u].x + q4v_x[x0 + 5u].x
                       + q4v_x[x0 + 6u].x + q4v_x[x0 + 7u].x;
                }
                var cv = 0u;
                if (!drop_c) {
                    cv = q4tp_byte(c0 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c0 + cbo + 1u) << 8u); }
                }
                var v = q4v_w[r0 * gpr + g];
                acc.x = acc.x + f32(v.x ^ v.y ^ v.z ^ v.w ^ cv) * xs;
                if (l1) {
                    if (!drop_c) {
                        cv = q4tp_byte(c1 + cbo);
                        if (sh > 3u) { cv = cv | (q4tp_byte(c1 + cbo + 1u) << 8u); }
                    }
                    v = q4v_w[r1 * gpr + g];
                    acc.y = acc.y + f32(v.x ^ v.y ^ v.z ^ v.w ^ cv) * xs;
                }
                if (l2) {
                    if (!drop_c) {
                        cv = q4tp_byte(c2 + cbo);
                        if (sh > 3u) { cv = cv | (q4tp_byte(c2 + cbo + 1u) << 8u); }
                    }
                    v = q4v_w[r2 * gpr + g];
                    acc.z = acc.z + f32(v.x ^ v.y ^ v.z ^ v.w ^ cv) * xs;
                }
                if (l3) {
                    if (!drop_c) {
                        cv = q4tp_byte(c3 + cbo);
                        if (sh > 3u) { cv = cv | (q4tp_byte(c3 + cbo + 1u) << 8u); }
                    }
                    v = q4v_w[r3 * gpr + g];
                    acc.w = acc.w + f32(v.x ^ v.y ^ v.z ^ v.w ^ cv) * xs;
                }
                g = g + 64u;
            }
        }
        // +8 skips the cross-lane reduction: six workgroup barriers over
        // 256 threads for what, at gpr 160, is two and a half iterations
        // of actual streaming. Whether that fixed cost or the bus sets
        // the floor is exactly what the pair of numbers answers.
        if (drop_r) {
            if (l == 0u) {
                if (r0 < rows) { q1y[r0] = acc.x; }
                if (r1 < rows) { q1y[r1] = acc.y; }
                if (r2 < rows) { q1y[r2] = acc.z; }
                if (r3 < rows) { q1y[r3] = acc.w; }
            }
        } else {
            partial_q4k[lid] = acc;
            workgroupBarrier();
            var stride = 32u;
            loop {
                if (stride == 0u) { break; }
                if (l < stride) {
                    partial_q4k[lid] = partial_q4k[lid] + partial_q4k[lid + stride];
                }
                workgroupBarrier();
                stride = stride >> 1u;
            }
            if (l == 0u) {
                let r = partial_q4k[sub << 6u];
                if (r0 < rows) { q1y[r0] = r.x; }
                if (r1 < rows) { q1y[r1] = r.y; }
                if (r2 < rows) { q1y[r2] = r.z; }
                if (r3 < rows) { q1y[r3] = r.w; }
            }
            workgroupBarrier();
        }
        wb = wb + nwg.x;
    }
}

// ── GDN step, parallel edition, k-looped: a workgroup per (head,
// column-quad) with the position loop INSIDE — the occupancy of the
// parallel kernel (thousands of workgroups where the serial one raised
// nv) and none of the per-position dispatch drains. Every state slice
// is workgroup-local across positions, so the loop needs no cross-
// workgroup sync; the raw o lands per position and gdn_step_norm_k
// applies the gated RMSNorm after. Snapshots ride the update itself.
@group(0) @binding(7) var<storage, read_write> gdk_S4 : array<vec4<f32>>;
@group(0) @binding(8) var<storage, read_write> gdk_o4 : array<vec4<f32>>;
@group(0) @binding(10) var<storage, read_write> gdk_snap4 : array<vec4<f32>>;
var<workgroup> gdk_r2: array<f32, 256>;
var<workgroup> gdk_r3: array<f32, 256>;
var<workgroup> gdk_r4: array<f32, 256>;
@compute @workgroup_size(128)
fn gdn_step_par_k(@builtin(workgroup_id) wid: vec3<u32>,
                  @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let dj4 = wid.y;
    let t = lid.x;
    let dk = gdk_p.dk;
    let dv = gdk_p.dv;
    if (h >= gdk_p.nv || dj4 * 4u >= dv) { return; }
    let ko = h / gdk_p.rep;
    let dv4 = dv >> 2u;
    let s4base = (h * dk * dv) >> 2u;
    for (var i = 0u; i < gdk_p.kb; i = i + 1u) {
        let cq0 = i * gdk_p.cdim;
        let qs = cq0 + ko * dk;
        let ks = cq0 + gdk_p.kd + ko * dk;
        gdk_red[t] = select(0.0, gdk_cq[qs + t] * gdk_cq[qs + t], t < dk);
        workgroupBarrier();
        var stride = 64u;
        loop {
            if (stride == 0u) { break; }
            if (t < stride) { gdk_red[t] = gdk_red[t] + gdk_red[t + stride]; }
            workgroupBarrier();
            stride = stride / 2u;
        }
        let nq = gdk_red[0];
        workgroupBarrier();
        gdk_red[t] = select(0.0, gdk_cq[ks + t] * gdk_cq[ks + t], t < dk);
        workgroupBarrier();
        stride = 64u;
        loop {
            if (stride == 0u) { break; }
            if (t < stride) { gdk_red[t] = gdk_red[t] + gdk_red[t + stride]; }
            workgroupBarrier();
            stride = stride / 2u;
        }
        let nkn = gdk_red[0];
        workgroupBarrier();
        let invq = 1.0 / (sqrt(nq + 1e-6) * sqrt(f32(dk)));
        let invk = 1.0 / sqrt(nkn + 1e-6);
        let abo = i * gdk_p.nv;
        let g = exp(-exp(gdk_alog[h]) * gd_softplus(gdk_a[abo + h] + gdk_dtb[h]));
        let beta = 1.0 / (1.0 + exp(-gdk_b[abo + h]));
        let vto = cq0 + 2u * gdk_p.kd + h * dv + dj4 * 4u;
        let vt = vec4<f32>(gdk_cq[vto], gdk_cq[vto + 1u], gdk_cq[vto + 2u], gdk_cq[vto + 3u]);
        let kf_t = select(0.0, gdk_cq[ks + t] * invk, t < dk);
        let qf_t = select(0.0, gdk_cq[qs + t] * invq, t < dk);
        var kv4 = vec4<f32>(0.0);
        if (t < dk) {
            kv4 = gdk_S4[s4base + t * dv4 + dj4] * kf_t;
        }
        gdk_red[t] = kv4.x;
        gdk_r2[t] = kv4.y;
        gdk_r3[t] = kv4.z;
        gdk_r4[t] = kv4.w;
        workgroupBarrier();
        stride = 64u;
        loop {
            if (stride == 0u) { break; }
            if (t < stride) {
                gdk_red[t] = gdk_red[t] + gdk_red[t + stride];
                gdk_r2[t] = gdk_r2[t] + gdk_r2[t + stride];
                gdk_r3[t] = gdk_r3[t] + gdk_r3[t + stride];
                gdk_r4[t] = gdk_r4[t] + gdk_r4[t + stride];
            }
            workgroupBarrier();
            stride = stride / 2u;
        }
        let kv = vec4<f32>(gdk_red[0], gdk_r2[0], gdk_r3[0], gdk_r4[0]);
        workgroupBarrier();
        let delta = (vt - g * kv) * beta;
        var contrib = vec4<f32>(0.0);
        if (t < dk) {
            let idx = s4base + t * dv4 + dj4;
            let cell = g * gdk_S4[idx] + kf_t * delta;
            gdk_S4[idx] = cell;
            if (gdk_p.stride != 0u) {
                gdk_snap4[(i * gdk_p.stride + gdk_p.ring_els) / 4u + idx] = cell;
            }
            contrib = qf_t * cell;
        }
        gdk_red[t] = contrib.x;
        gdk_r2[t] = contrib.y;
        gdk_r3[t] = contrib.z;
        gdk_r4[t] = contrib.w;
        workgroupBarrier();
        stride = 64u;
        loop {
            if (stride == 0u) { break; }
            if (t < stride) {
                gdk_red[t] = gdk_red[t] + gdk_red[t + stride];
                gdk_r2[t] = gdk_r2[t] + gdk_r2[t + stride];
                gdk_r3[t] = gdk_r3[t] + gdk_r3[t + stride];
                gdk_r4[t] = gdk_r4[t] + gdk_r4[t + stride];
            }
            workgroupBarrier();
            stride = stride / 2u;
        }
        if (t == 0u) {
            gdk_o4[(i * gdk_p.nv * dv) / 4u + h * dv4 + dj4] =
                vec4<f32>(gdk_red[0], gdk_r2[0], gdk_r3[0], gdk_r4[0]);
        }
        workgroupBarrier();
    }
}

// Gated RMSNorm over the k raw o rows the parallel step left.
@compute @workgroup_size(256)
fn gdn_step_norm_k(@builtin(workgroup_id) wid: vec3<u32>,
                   @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let t = lid.x;
    let dv = gdk_p.dv;
    if (h >= gdk_p.nv) { return; }
    for (var i = 0u; i < gdk_p.kb; i = i + 1u) {
        let zo = i * gdk_p.nv * dv;
        gdk_red[t] = select(0.0, gdk_o[zo + h * dv + t] * gdk_o[zo + h * dv + t], t < dv);
        workgroupBarrier();
        let ss = gdk_reduce(t);
        workgroupBarrier();
        let inv = 1.0 / sqrt(ss / f32(dv) + gdk_p.eps);
        if (t < dv) {
            let zz = gdk_z[zo + h * dv + t];
            gdk_o[zo + h * dv + t] =
                gdk_o[zo + h * dv + t] * inv * gdk_norm[t] * (zz / (1.0 + exp(-zz)));
        }
        workgroupBarrier();
    }
}

// ── q4tp matvec, QUAD row blocking: sixteen rows a workgroup, each
// 64-lane sub-block owning FOUR rows 4 apart, so the eight x vec4 loads
// of a group feed four dot chains instead of two. The pair kernel's LSU
// is x-bound at large widths — 128 bytes of activations per group
// against ~20 of weights per row — and halving-again the x side is the
// same medicine the DSV4 down projection took (its 4-row twin). The
// accumulator is a NAMED vec4: constant component indexing only, or the
// registers spill to stack and the GEMV runs at a fraction of the card.
var<workgroup> lad_q4w: array<f32, 512>;
@compute @workgroup_size(256)
fn q4tp_matvec16w(@builtin(workgroup_id) wid: vec3<u32>,
                  @builtin(num_workgroups) nwg: vec3<u32>,
                  @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let blocks = (rows + 15u) / 16u;
    var wb = wid.x;
    loop {
        if (wb >= blocks) { break; }
        let base = wb * 16u;
        for (var t = lid; t < 512u; t = t + 256u) {
            let r = base + (t >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4w[t] = exp2(pr.x + f32(t & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        var acc = vec4<f32>(0.0);
        if (r0 < rows) {
            let c0 = codes_b + r0 * cstride;
            let c1 = codes_b + r1 * cstride;
            let c2 = codes_b + r2 * cstride;
            let c3 = codes_b + r3 * cstride;
            let l1 = r1 < rows;
            let l2 = r2 < rows;
            let l3 = r3 < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                let x0 = g * 8u;
                let xa = q4v_x[x0];      let xb = q4v_x[x0 + 1u];
                let xc = q4v_x[x0 + 2u]; let xd = q4v_x[x0 + 3u];
                let xe = q4v_x[x0 + 4u]; let xf = q4v_x[x0 + 5u];
                let xg = q4v_x[x0 + 6u]; let xh = q4v_x[x0 + 7u];
                var cv = q4tp_byte(c0 + cbo);
                if (sh > 3u) { cv = cv | (q4tp_byte(c0 + cbo + 1u) << 8u); }
                var v = q4v_w[r0 * gpr + g];
                acc.x = acc.x + lad_q4w[(sub << 5u) + ((cv >> sh) & 31u)]
                    * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                     + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                if (l1) {
                    cv = q4tp_byte(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c1 + cbo + 1u) << 8u); }
                    v = q4v_w[r1 * gpr + g];
                    acc.y = acc.y + lad_q4w[128u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                         + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                }
                if (l2) {
                    cv = q4tp_byte(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c2 + cbo + 1u) << 8u); }
                    v = q4v_w[r2 * gpr + g];
                    acc.z = acc.z + lad_q4w[256u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                         + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                }
                if (l3) {
                    cv = q4tp_byte(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c3 + cbo + 1u) << 8u); }
                    v = q4v_w[r3 * gpr + g];
                    acc.w = acc.w + lad_q4w[384u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                         + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                }
                g = g + 64u;
            }
        }
        partial_q4k[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q4k[lid] = partial_q4k[lid] + partial_q4k[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u) {
            let r = partial_q4k[sub << 6u];
            if (r0 < rows) { q1y[r0] = r.x; }
            if (r1 < rows) { q1y[r1] = r.y; }
            if (r2 < rows) { q1y[r2] = r.z; }
            if (r3 < rows) { q1y[r3] = r.w; }
        }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}



// ── TWO projections of one input in ONE dispatch, on the quad-row 16w
// body: blocks [0, blocks1) stream weight A into y, blocks [blocks1,
// blocks1+blocks2) stream weight B into y2 (rows2 in `_p1`). Per-row
// arithmetic and add order are `q4tp_matvec16w`'s exactly — the body is
// generated from it — so the outputs are bit-identical; only the
// dispatch boundary between the two projections is gone. Measured
// motive (RTX 5090 pod, 700 tiny dispatches in one submit): a dispatch
// costs ~10-20 us of launch + drain in this serialized chain, and a
// token issues ~700 of them — a third of the token. gate+up, the GDN's
// qkv+z and attention's k+v are pairs over the SAME x. WGSL cannot
// alias bindings, hence side B mirrors side A over its own slots.
@compute @workgroup_size(256)
fn q4tp_matvec16w_x2(@builtin(workgroup_id) wid: vec3<u32>,
                     @builtin(num_workgroups) nwg: vec3<u32>,
                     @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let rows2 = q1p._p1;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let params_w2 = rows2 * gpr * 4u;
    let codes_b2 = rows2 * gpr * 16u + rows2 * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let blocks1 = (rows + 15u) / 16u;
    let blocks2 = (rows2 + 15u) / 16u;
    var wb = wid.x;
    loop {
        if (wb >= blocks1 + blocks2) { break; }
        if (wb < blocks1) {
        let base = wb * 16u;
        for (var t = lid; t < 512u; t = t + 256u) {
            let r = base + (t >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4w[t] = exp2(pr.x + f32(t & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        var acc = vec4<f32>(0.0);
        if (r0 < rows) {
            let c0 = codes_b + r0 * cstride;
            let c1 = codes_b + r1 * cstride;
            let c2 = codes_b + r2 * cstride;
            let c3 = codes_b + r3 * cstride;
            let l1 = r1 < rows;
            let l2 = r2 < rows;
            let l3 = r3 < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                let x0 = g * 8u;
                let xa = q4v_x[x0];      let xb = q4v_x[x0 + 1u];
                let xc = q4v_x[x0 + 2u]; let xd = q4v_x[x0 + 3u];
                let xe = q4v_x[x0 + 4u]; let xf = q4v_x[x0 + 5u];
                let xg = q4v_x[x0 + 6u]; let xh = q4v_x[x0 + 7u];
                var cv = q4tp_byte(c0 + cbo);
                if (sh > 3u) { cv = cv | (q4tp_byte(c0 + cbo + 1u) << 8u); }
                var v = q4v_w[r0 * gpr + g];
                acc.x = acc.x + lad_q4w[(sub << 5u) + ((cv >> sh) & 31u)]
                    * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                     + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                if (l1) {
                    cv = q4tp_byte(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c1 + cbo + 1u) << 8u); }
                    v = q4v_w[r1 * gpr + g];
                    acc.y = acc.y + lad_q4w[128u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                         + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                }
                if (l2) {
                    cv = q4tp_byte(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c2 + cbo + 1u) << 8u); }
                    v = q4v_w[r2 * gpr + g];
                    acc.z = acc.z + lad_q4w[256u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                         + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                }
                if (l3) {
                    cv = q4tp_byte(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c3 + cbo + 1u) << 8u); }
                    v = q4v_w[r3 * gpr + g];
                    acc.w = acc.w + lad_q4w[384u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                         + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                }
                g = g + 64u;
            }
        }
        partial_q4k[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q4k[lid] = partial_q4k[lid] + partial_q4k[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u) {
            let r = partial_q4k[sub << 6u];
            if (r0 < rows) { q1y[r0] = r.x; }
            if (r1 < rows) { q1y[r1] = r.y; }
            if (r2 < rows) { q1y[r2] = r.z; }
            if (r3 < rows) { q1y[r3] = r.w; }
        }
        workgroupBarrier();
        } else {
        let base = (wb - blocks1) * 16u;
        for (var t = lid; t < 512u; t = t + 256u) {
            let r = base + (t >> 5u);
            if (r < rows2) {
                let pr = unpack2x16float(q1w2[params_w2 + r]);
                lad_q4w[t] = exp2(pr.x + f32(t & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        var acc = vec4<f32>(0.0);
        if (r0 < rows2) {
            let c0 = codes_b2 + r0 * cstride;
            let c1 = codes_b2 + r1 * cstride;
            let c2 = codes_b2 + r2 * cstride;
            let c3 = codes_b2 + r3 * cstride;
            let l1 = r1 < rows2;
            let l2 = r2 < rows2;
            let l3 = r3 < rows2;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                let x0 = g * 8u;
                let xa = q4v_x[x0];      let xb = q4v_x[x0 + 1u];
                let xc = q4v_x[x0 + 2u]; let xd = q4v_x[x0 + 3u];
                let xe = q4v_x[x0 + 4u]; let xf = q4v_x[x0 + 5u];
                let xg = q4v_x[x0 + 6u]; let xh = q4v_x[x0 + 7u];
                var cv = q4tp_byte2(c0 + cbo);
                if (sh > 3u) { cv = cv | (q4tp_byte2(c0 + cbo + 1u) << 8u); }
                var v = q4v_w2[r0 * gpr + g];
                acc.x = acc.x + lad_q4w[(sub << 5u) + ((cv >> sh) & 31u)]
                    * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                     + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                if (l1) {
                    cv = q4tp_byte2(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte2(c1 + cbo + 1u) << 8u); }
                    v = q4v_w2[r1 * gpr + g];
                    acc.y = acc.y + lad_q4w[128u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                         + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                }
                if (l2) {
                    cv = q4tp_byte2(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte2(c2 + cbo + 1u) << 8u); }
                    v = q4v_w2[r2 * gpr + g];
                    acc.z = acc.z + lad_q4w[256u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                         + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                }
                if (l3) {
                    cv = q4tp_byte2(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte2(c3 + cbo + 1u) << 8u); }
                    v = q4v_w2[r3 * gpr + g];
                    acc.w = acc.w + lad_q4w[384u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                         + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                }
                g = g + 64u;
            }
        }
        partial_q4k[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q4k[lid] = partial_q4k[lid] + partial_q4k[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u) {
            let r = partial_q4k[sub << 6u];
            if (r0 < rows2) { q1y2[r0] = r.x; }
            if (r1 < rows2) { q1y2[r1] = r.y; }
            if (r2 < rows2) { q1y2[r2] = r.z; }
            if (r3 < rows2) { q1y2[r3] = r.w; }
        }
        workgroupBarrier();
        }
        wb = wb + nwg.x;
    }
}

// ── FFN gate+up+SiLU in ONE dispatch on the quad-row 16w body: a
// workgroup owns eight rows of BOTH weights (r0/r1 gate, r2/r3 the same
// rows of up), reduces all four chains, and its lane 0 writes
// act[r] = silu(gate[r]) * up[r] straight into the down projection's
// input. Per-row arithmetic and add order are `q4tp_matvec16w`'s (body
// generated from it) and the SiLU is `silu_mul_pre`'s expression, so
// the activations are bit-identical to gate-matvec, up-matvec, silu —
// three dispatches and two 70 KB round trips become one dispatch.
// Weight A = gate (slots 0/4), weight B = up (slots 6/7), `_p0` = the
// swiglu limit as f32 bits (0 = none), rows = the FFN width.
@compute @workgroup_size(256)
fn q4tp_matvec16w_gu(@builtin(workgroup_id) wid: vec3<u32>,
                     @builtin(num_workgroups) nwg: vec3<u32>,
                     @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let params_w2 = params_w;
    let codes_b2 = codes_b;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let blocks = (rows + 7u) / 8u;
    var wb = wid.x;
    loop {
        if (wb >= blocks) { break; }
        let base = wb * 8u;
        for (var t = lid; t < 512u; t = t + 256u) {
            let slot = t >> 5u;
            if (slot < 8u) {
                let r = base + slot;
                if (r < rows) {
                    let pr = unpack2x16float(q1w[params_w + r]);
                    lad_q4w[t] = exp2(pr.x + f32(t & 31u) * pr.y);
                }
            } else {
                let r = base + slot - 8u;
                if (r < rows) {
                    let pr = unpack2x16float(q1w2[params_w2 + r]);
                    lad_q4w[t] = exp2(pr.x + f32(t & 31u) * pr.y);
                }
            }
        }
        workgroupBarrier();
        // r0/r1: gate rows base+sub, base+sub+4; r2/r3: the SAME up rows.
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub;
        let r3 = base + sub + 4u;
        var acc = vec4<f32>(0.0);
        if (r0 < rows) {
            let c0 = codes_b + r0 * cstride;
            let c1 = codes_b + r1 * cstride;
            let c2 = codes_b2 + r2 * cstride;
            let c3 = codes_b2 + r3 * cstride;
            let l1 = r1 < rows;
            let l2 = r2 < rows;
            let l3 = r3 < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                let x0 = g * 8u;
                let xa = q4v_x[x0];      let xb = q4v_x[x0 + 1u];
                let xc = q4v_x[x0 + 2u]; let xd = q4v_x[x0 + 3u];
                let xe = q4v_x[x0 + 4u]; let xf = q4v_x[x0 + 5u];
                let xg = q4v_x[x0 + 6u]; let xh = q4v_x[x0 + 7u];
                var cv = q4tp_byte(c0 + cbo);
                if (sh > 3u) { cv = cv | (q4tp_byte(c0 + cbo + 1u) << 8u); }
                var v = q4v_w[r0 * gpr + g];
                acc.x = acc.x + lad_q4w[(sub << 5u) + ((cv >> sh) & 31u)]
                    * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                     + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                if (l1) {
                    cv = q4tp_byte(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c1 + cbo + 1u) << 8u); }
                    v = q4v_w[r1 * gpr + g];
                    acc.y = acc.y + lad_q4w[128u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                         + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                }
                if (l2) {
                    cv = q4tp_byte2(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte2(c2 + cbo + 1u) << 8u); }
                    v = q4v_w2[r2 * gpr + g];
                    acc.z = acc.z + lad_q4w[256u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                         + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                }
                if (l3) {
                    cv = q4tp_byte2(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte2(c3 + cbo + 1u) << 8u); }
                    v = q4v_w2[r3 * gpr + g];
                    acc.w = acc.w + lad_q4w[384u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q4v_dot8(v.x, xa, xb) + q4v_dot8(v.y, xc, xd)
                         + q4v_dot8(v.z, xe, xf) + q4v_dot8(v.w, xg, xh));
                }
                g = g + 64u;
            }
        }
        partial_q4k[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q4k[lid] = partial_q4k[lid] + partial_q4k[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u) {
            let r = partial_q4k[sub << 6u];
            // act = (g / (1 + exp(-g))) * u — the silu_mul_pre expression,
            // term for term, with the reference's asymmetric swiglu limit.
            var ga = r.x; var gb = r.y; var ua = r.z; var ub = r.w;
            let lim = bitcast<f32>(q1p._p0);
            if (lim > 0.0) {
                ua = clamp(ua, -lim, lim); ub = clamp(ub, -lim, lim);
                ga = min(ga, lim); gb = min(gb, lim);
            }
            if (r0 < rows) { q1y[r0] = (ga / (1.0 + exp(-ga))) * ua; }
            if (r1 < rows) { q1y[r1] = (gb / (1.0 + exp(-gb))) * ub; }
        }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}

// ── q2tp decode matvec, quad-row 16w shape: the 2-bit plane of the q2tp
// profile (dense FFN gate/up when the model has no experts). A group is
// 32 codes in TWO u32 words (8 bytes, half of q4tp's 16), value =
// (code − 1.5) · scale; the scale is the q4tp ladder with rung 0 the
// format's exact zero (the ±0.5/±1.5 grid has none of its own), live
// rungs shifted down one — the same convention as the MoE q2tp kernels.
// Layout in u32 words: nibbles (row·gpr + g)·2, params rows·gpr·2 + row,
// codes byte plane at rows·gpr·8 + rows·4 + row·cstride.
// The 2-bit code's (c − 1.5) by the same magic, doubled: 2c into the
// mantissa of 2^23, minus 2^23 + 3, is 2c − 3 exactly, and the sum of
// exactly-doubled terms halved at the end is the undoubled sum to the
// bit (scaling by two commutes with every rounding here).
fn q2v_c2(w: u32, sh: u32, affine: u32) -> f32 {
    // The raw q2tp center is 1.5: 2c-3, while an explicit q2tp_affine
    // descriptor requests 1.0: 2c-2. Keep the center in the shader rather
    // than retagging the payload; `_p1` is set only after the descriptor
    // target has been validated by the Rust caller.
    return bitcast<f32>((((w >> sh) & 3u) << 1u) | 0x4B000000u)
        - select(8388611.0, 8388610.0, affine != 0u);
}
fn q2v_d16(w: u32, a: vec4<f32>, b: vec4<f32>, c: vec4<f32>, d: vec4<f32>, affine: u32) -> f32 {
    return (q2v_c2(w, 0u, affine) * a.x
         + q2v_c2(w, 2u, affine) * a.y
         + q2v_c2(w, 4u, affine) * a.z
         + q2v_c2(w, 6u, affine) * a.w
         + q2v_c2(w, 8u, affine) * b.x
         + q2v_c2(w, 10u, affine) * b.y
         + q2v_c2(w, 12u, affine) * b.z
         + q2v_c2(w, 14u, affine) * b.w
         + q2v_c2(w, 16u, affine) * c.x
         + q2v_c2(w, 18u, affine) * c.y
         + q2v_c2(w, 20u, affine) * c.z
         + q2v_c2(w, 22u, affine) * c.w
         + q2v_c2(w, 24u, affine) * d.x
         + q2v_c2(w, 26u, affine) * d.y
         + q2v_c2(w, 28u, affine) * d.z
         + q2v_c2(w, 30u, affine) * d.w) * 0.5;
}
@compute @workgroup_size(256)
fn q2tp_matvec16w(@builtin(workgroup_id) wid: vec3<u32>,
                  @builtin(num_workgroups) nwg: vec3<u32>,
                  @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 2u;
    let codes_b = rows * gpr * 8u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let blocks = (rows + 15u) / 16u;
    var wb = wid.x;
    loop {
        if (wb >= blocks) { break; }
        let base = wb * 16u;
        // ladder: 32 rungs a row, rung 0 = 0.0, rung r>0 = exp2(lo + (r-1)·step)
        for (var t = lid; t < 512u; t = t + 256u) {
            let r = base + (t >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                let rung = t & 31u;
                lad_q4w[t] = select(exp2(pr.x + f32(max(rung, 1u) - 1u) * pr.y), 0.0, rung == 0u);
            }
        }
        workgroupBarrier();
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        var acc = vec4<f32>(0.0);
        if (r0 < rows) {
            let c0 = codes_b + r0 * cstride;
            let c1 = codes_b + r1 * cstride;
            let c2 = codes_b + r2 * cstride;
            let c3 = codes_b + r3 * cstride;
            let l1 = r1 < rows;
            let l2 = r2 < rows;
            let l3 = r3 < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                let x0 = g * 8u;
                let xa = q4v_x[x0];      let xb = q4v_x[x0 + 1u];
                let xc = q4v_x[x0 + 2u]; let xd = q4v_x[x0 + 3u];
                let xe = q4v_x[x0 + 4u]; let xf = q4v_x[x0 + 5u];
                let xg = q4v_x[x0 + 6u]; let xh = q4v_x[x0 + 7u];
                var cv = q4tp_byte(c0 + cbo);
                if (sh > 3u) { cv = cv | (q4tp_byte(c0 + cbo + 1u) << 8u); }
                var wi = (r0 * gpr + g) * 2u;
                acc.x = acc.x + lad_q4w[(sub << 5u) + ((cv >> sh) & 31u)]
                    * (q2v_d16(q1w[wi], xa, xb, xc, xd, q1p._p1) + q2v_d16(q1w[wi + 1u], xe, xf, xg, xh, q1p._p1));
                if (l1) {
                    cv = q4tp_byte(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c1 + cbo + 1u) << 8u); }
                    wi = (r1 * gpr + g) * 2u;
                    acc.y = acc.y + lad_q4w[128u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q2v_d16(q1w[wi], xa, xb, xc, xd, q1p._p1) + q2v_d16(q1w[wi + 1u], xe, xf, xg, xh, q1p._p1));
                }
                if (l2) {
                    cv = q4tp_byte(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c2 + cbo + 1u) << 8u); }
                    wi = (r2 * gpr + g) * 2u;
                    acc.z = acc.z + lad_q4w[256u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q2v_d16(q1w[wi], xa, xb, xc, xd, q1p._p1) + q2v_d16(q1w[wi + 1u], xe, xf, xg, xh, q1p._p1));
                }
                if (l3) {
                    cv = q4tp_byte(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q4tp_byte(c3 + cbo + 1u) << 8u); }
                    wi = (r3 * gpr + g) * 2u;
                    acc.w = acc.w + lad_q4w[384u + (sub << 5u) + ((cv >> sh) & 31u)]
                        * (q2v_d16(q1w[wi], xa, xb, xc, xd, q1p._p1) + q2v_d16(q1w[wi + 1u], xe, xf, xg, xh, q1p._p1));
                }
                g = g + 64u;
            }
        }
        partial_q4k[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q4k[lid] = partial_q4k[lid] + partial_q4k[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u) {
            let r = partial_q4k[sub << 6u];
            if (r0 < rows) { q1y[r0] = r.x; }
            if (r1 < rows) { q1y[r1] = r.y; }
            if (r2 < rows) { q1y[r2] = r.z; }
            if (r3 < rows) { q1y[r3] = r.w; }
        }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}

// ── q2tp affine decode, one-vector INT8/DP4A shape.  This is deliberately
// separate from q4tp_matvec4_bk8: q2tp's affine symbols are exact signed
// ternaries (code − 1), so the q4tp `−8·sum(q)` correction is both unnecessary
// and wrong here.  The activation is the existing per-32-group symmetric Q8
// grid produced by x_quant_i8.  The x8 layout is already interleaved as
// even/odd lanes for each 8-value word; two q2 bytes therefore become the
// two packed int8 words needed by the same eight DP4A dots.
//
// This entry point is opt-in from Rust (`CMF_Q2_DP4A=1`) and is only admitted
// for the explicitly affine Prism descriptor.  A caller that cannot satisfy
// those guards stays on q2tp_matvec16w, preserving the validated scalar path.
fn q2i8_byte(off: u32) -> u32 {
    return (q1w[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
}

fn q2i8_pack_even(a: u32, b: u32) -> u32 {
    let x0 = i32(a & 3u) - 1;
    let x1 = i32((a >> 4u) & 3u) - 1;
    let x2 = i32(b & 3u) - 1;
    let x3 = i32((b >> 4u) & 3u) - 1;
    return (u32(x0) & 0xFFu) | ((u32(x1) & 0xFFu) << 8u)
        | ((u32(x2) & 0xFFu) << 16u) | ((u32(x3) & 0xFFu) << 24u);
}

fn q2i8_pack_odd(a: u32, b: u32) -> u32 {
    let x0 = i32((a >> 2u) & 3u) - 1;
    let x1 = i32((a >> 6u) & 3u) - 1;
    let x2 = i32((b >> 2u) & 3u) - 1;
    let x3 = i32((b >> 6u) & 3u) - 1;
    return (u32(x0) & 0xFFu) | ((u32(x1) & 0xFFu) << 8u)
        | ((u32(x2) & 0xFFu) << 16u) | ((u32(x3) & 0xFFu) << 24u);
}

fn q2i8_dot8(a: u32, b: u32, even: u32, odd: u32) -> i32 {
    return dot4I8Packed(q2i8_pack_even(a, b), even)
        + dot4I8Packed(q2i8_pack_odd(a, b), odd);
}

fn q2i8_dot_group(base: u32, x: vec4<u32>, y: vec4<u32>) -> i32 {
    let b0 = q2i8_byte(base);
    let b1 = q2i8_byte(base + 1u);
    let b2 = q2i8_byte(base + 2u);
    let b3 = q2i8_byte(base + 3u);
    let b4 = q2i8_byte(base + 4u);
    let b5 = q2i8_byte(base + 5u);
    let b6 = q2i8_byte(base + 6u);
    let b7 = q2i8_byte(base + 7u);
    return q2i8_dot8(b0, b1, x.x, x.y)
        + q2i8_dot8(b2, b3, x.z, x.w)
        + q2i8_dot8(b4, b5, y.x, y.y)
        + q2i8_dot8(b6, b7, y.z, y.w);
}

var<workgroup> partial_q2i8: array<vec4<f32>, 256>;

@compute @workgroup_size(256)
fn q2tp_matvec1_i8(@builtin(workgroup_id) wid: vec3<u32>,
                   @builtin(num_workgroups) nwg: vec3<u32>,
                   @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 2u;
    let codes_b = rows * gpr * 8u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let blocks = (rows + 15u) / 16u;
    var wb = wid.x;
    loop {
        if (wb >= blocks) { break; }
        let base = wb * 16u;
        for (var t = lid; t < 512u; t = t + 256u) {
            let r = base + (t >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                let rung = t & 31u;
                lad_q4w[t] = select(
                    exp2(pr.x + f32(max(rung, 1u) - 1u) * pr.y),
                    0.0,
                    rung == 0u,
                );
            }
        }
        workgroupBarrier();
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        var acc = vec4<f32>(0.0);
        if (r0 < rows) {
            let c0 = codes_b + r0 * cstride;
            let c1 = codes_b + r1 * cstride;
            let c2 = codes_b + r2 * cstride;
            let c3 = codes_b + r3 * cstride;
            let l1 = r1 < rows;
            let l2 = r2 < rows;
            let l3 = r3 < rows;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv = q2i8_byte(c0 + cbo);
                if (sh > 3u) { cv = cv | (q2i8_byte(c0 + cbo + 1u) << 8u); }
                let rung0 = (cv >> sh) & 31u;
                let x0 = q4v_x8[g * 2u];
                let x1 = q4v_x8[g * 2u + 1u];
                let sx = q4v_xs[g].x;
                let d0 = q2i8_dot_group((r0 * gpr + g) * 8u, x0, x1);
                acc.x = acc.x + lad_q4w[(sub << 5u) + rung0] * sx * f32(d0);
                if (l1) {
                    cv = q2i8_byte(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q2i8_byte(c1 + cbo + 1u) << 8u); }
                    let rung1 = (cv >> sh) & 31u;
                    let d1 = q2i8_dot_group((r1 * gpr + g) * 8u, x0, x1);
                    acc.y = acc.y + lad_q4w[128u + (sub << 5u) + rung1] * sx * f32(d1);
                }
                if (l2) {
                    cv = q2i8_byte(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q2i8_byte(c2 + cbo + 1u) << 8u); }
                    let rung2 = (cv >> sh) & 31u;
                    let d2 = q2i8_dot_group((r2 * gpr + g) * 8u, x0, x1);
                    acc.z = acc.z + lad_q4w[256u + (sub << 5u) + rung2] * sx * f32(d2);
                }
                if (l3) {
                    cv = q2i8_byte(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q2i8_byte(c3 + cbo + 1u) << 8u); }
                    let rung3 = (cv >> sh) & 31u;
                    let d3 = q2i8_dot_group((r3 * gpr + g) * 8u, x0, x1);
                    acc.w = acc.w + lad_q4w[384u + (sub << 5u) + rung3] * sx * f32(d3);
                }
                g = g + 64u;
            }
        }
        partial_q2i8[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q2i8[lid] = partial_q2i8[lid] + partial_q2i8[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u) {
            let r = partial_q2i8[sub << 6u];
            if (r0 < rows) { q1y[r0] = r.x; }
            if (r1 < rows) { q1y[r1] = r.y; }
            if (r2 < rows) { q1y[r2] = r.z; }
            if (r3 < rows) { q1y[r3] = r.w; }
        }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}


// ── q4tp matvec, k-in-registers batch: ONE weight decode serves up to
// FOUR activation vectors. The nb-dispatch batch streams the whole weight
// once per element and hopes for L2; at 2-8 MB a layer there is nothing
// left to hope with, and a k=3 verify paid the weight bandwidth three
// times. Here the batch lives in a vec4 accumulator (named, constant
// -indexed — a dynamically indexed array would spill to stack and run at
// a fraction of the card, the register-spill lesson). Rows past kb read
// clamped garbage and are zeroed by the mask.
var<workgroup> partial_q4k: array<vec4<f32>, 256>;
var<workgroup> partial_q4kb: array<vec4<f32>, 256>;
@compute @workgroup_size(256)
fn q4tp_matvec4_k(@builtin(workgroup_id) wid: vec3<u32>,
                  @builtin(num_workgroups) nwg: vec3<u32>,
                  @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let kb = max(q1p._p0, 1u);
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let xr = gpr * 8u;
    let mask = vec4<f32>(
        1.0,
        select(0.0, 1.0, kb > 1u),
        select(0.0, 1.0, kb > 2u),
        select(0.0, 1.0, kb > 3u),
    );
    let blocks = (rows + 7u) / 8u;
    var wb = wid.x;
    loop {
        if (wb >= blocks) { break; }
        let base = wb * 8u;
        {
            let r = base + (lid >> 5u);
            if (r < rows) {
                let pr = unpack2x16float(q1w[params_w + r]);
                lad_q4v[lid] = exp2(pr.x + f32(lid & 31u) * pr.y);
            }
        }
        workgroupBarrier();
        let wrow_a = base + sub;
        let wrow_b = base + sub + 4u;
        let live_b = wrow_b < rows;
        var acc_a = vec4<f32>(0.0);
        var acc_b = vec4<f32>(0.0);
        if (wrow_a < rows) {
            let crow_a = codes_b + wrow_a * cstride;
            let crow_b = codes_b + wrow_b * cstride;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                var cv_a = q4tp_byte(crow_a + cbo);
                if (sh > 3u) { cv_a = cv_a | (q4tp_byte(crow_a + cbo + 1u) << 8u); }
                let v_a = q4v_w[wrow_a * gpr + g];
                let sa = lad_q4v[(sub << 5u) + ((cv_a >> sh) & 31u)];
                let x0 = g * 8u;
                let d_a = vec4<f32>(
                    q4v_dot8(v_a.x, q4v_x[x0], q4v_x[x0 + 1u])
                        + q4v_dot8(v_a.y, q4v_x[x0 + 2u], q4v_x[x0 + 3u])
                        + q4v_dot8(v_a.z, q4v_x[x0 + 4u], q4v_x[x0 + 5u])
                        + q4v_dot8(v_a.w, q4v_x[x0 + 6u], q4v_x[x0 + 7u]),
                    q4v_dot8(v_a.x, q4v_x[xr + x0], q4v_x[xr + x0 + 1u])
                        + q4v_dot8(v_a.y, q4v_x[xr + x0 + 2u], q4v_x[xr + x0 + 3u])
                        + q4v_dot8(v_a.z, q4v_x[xr + x0 + 4u], q4v_x[xr + x0 + 5u])
                        + q4v_dot8(v_a.w, q4v_x[xr + x0 + 6u], q4v_x[xr + x0 + 7u]),
                    q4v_dot8(v_a.x, q4v_x[2u * xr + x0], q4v_x[2u * xr + x0 + 1u])
                        + q4v_dot8(v_a.y, q4v_x[2u * xr + x0 + 2u], q4v_x[2u * xr + x0 + 3u])
                        + q4v_dot8(v_a.z, q4v_x[2u * xr + x0 + 4u], q4v_x[2u * xr + x0 + 5u])
                        + q4v_dot8(v_a.w, q4v_x[2u * xr + x0 + 6u], q4v_x[2u * xr + x0 + 7u]),
                    q4v_dot8(v_a.x, q4v_x[3u * xr + x0], q4v_x[3u * xr + x0 + 1u])
                        + q4v_dot8(v_a.y, q4v_x[3u * xr + x0 + 2u], q4v_x[3u * xr + x0 + 3u])
                        + q4v_dot8(v_a.z, q4v_x[3u * xr + x0 + 4u], q4v_x[3u * xr + x0 + 5u])
                        + q4v_dot8(v_a.w, q4v_x[3u * xr + x0 + 6u], q4v_x[3u * xr + x0 + 7u]),
                );
                acc_a = acc_a + sa * d_a * mask;
                if (live_b) {
                    var cv_b = q4tp_byte(crow_b + cbo);
                    if (sh > 3u) { cv_b = cv_b | (q4tp_byte(crow_b + cbo + 1u) << 8u); }
                    let v_b = q4v_w[wrow_b * gpr + g];
                    let sb = lad_q4v[128u + (sub << 5u) + ((cv_b >> sh) & 31u)];
                    let d_b = vec4<f32>(
                        q4v_dot8(v_b.x, q4v_x[x0], q4v_x[x0 + 1u])
                            + q4v_dot8(v_b.y, q4v_x[x0 + 2u], q4v_x[x0 + 3u])
                            + q4v_dot8(v_b.z, q4v_x[x0 + 4u], q4v_x[x0 + 5u])
                            + q4v_dot8(v_b.w, q4v_x[x0 + 6u], q4v_x[x0 + 7u]),
                        q4v_dot8(v_b.x, q4v_x[xr + x0], q4v_x[xr + x0 + 1u])
                            + q4v_dot8(v_b.y, q4v_x[xr + x0 + 2u], q4v_x[xr + x0 + 3u])
                            + q4v_dot8(v_b.z, q4v_x[xr + x0 + 4u], q4v_x[xr + x0 + 5u])
                            + q4v_dot8(v_b.w, q4v_x[xr + x0 + 6u], q4v_x[xr + x0 + 7u]),
                        q4v_dot8(v_b.x, q4v_x[2u * xr + x0], q4v_x[2u * xr + x0 + 1u])
                            + q4v_dot8(v_b.y, q4v_x[2u * xr + x0 + 2u], q4v_x[2u * xr + x0 + 3u])
                            + q4v_dot8(v_b.z, q4v_x[2u * xr + x0 + 4u], q4v_x[2u * xr + x0 + 5u])
                            + q4v_dot8(v_b.w, q4v_x[2u * xr + x0 + 6u], q4v_x[2u * xr + x0 + 7u]),
                        q4v_dot8(v_b.x, q4v_x[3u * xr + x0], q4v_x[3u * xr + x0 + 1u])
                            + q4v_dot8(v_b.y, q4v_x[3u * xr + x0 + 2u], q4v_x[3u * xr + x0 + 3u])
                            + q4v_dot8(v_b.z, q4v_x[3u * xr + x0 + 4u], q4v_x[3u * xr + x0 + 5u])
                            + q4v_dot8(v_b.w, q4v_x[3u * xr + x0 + 6u], q4v_x[3u * xr + x0 + 7u]),
                    );
                    acc_b = acc_b + sb * d_b * mask;
                }
                g = g + 64u;
            }
        }
        partial_q4k[lid] = acc_a;
        partial_q4kb[lid] = acc_b;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q4k[lid] = partial_q4k[lid] + partial_q4k[lid + stride];
                partial_q4kb[lid] = partial_q4kb[lid] + partial_q4kb[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u && wrow_a < rows) {
            let ra = partial_q4k[sub << 6u];
            q1y[wrow_a] = ra.x;
            if (kb > 1u) { q1y[rows + wrow_a] = ra.y; }
            if (kb > 2u) { q1y[2u * rows + wrow_a] = ra.z; }
            if (kb > 3u) { q1y[3u * rows + wrow_a] = ra.w; }
        }
        if (l == 0u && live_b) {
            let rb = partial_q4kb[sub << 6u];
            q1y[wrow_b] = rb.x;
            if (kb > 1u) { q1y[rows + wrow_b] = rb.y; }
            if (kb > 2u) { q1y[2u * rows + wrow_b] = rb.z; }
            if (kb > 3u) { q1y[3u * rows + wrow_b] = rb.w; }
        }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}

// ── Fold-select MoE twins: gu/down recompute the top-k FROM THE ROUTER
// LOGITS inside every workgroup — redundant arithmetic, but the serial
// select hop disappears and the layer chain loses one ~25 us dispatch
// latency. The comparator (max, lowest index on ties) is order-free, so
// every workgroup lands on the same experts; softmax summation order
// differs from the retired select kernel only in reduction shape.
// Slot 3 carries the LOGITS where the plain twins carry the selection.
@group(0) @binding(3) var<storage, read> mgf_logit : array<f32>;
struct MgfP { n_exp: u32, _a: u32, _b: u32, _c: u32 };
@group(0) @binding(7) var<uniform> mgf_p : MgfP;

var<workgroup> mgf_lg: array<f32, 256>;
var<workgroup> mgf_v:  array<f32, 64>;
var<workgroup> mgf_i:  array<u32, 64>;

// top-(slot+1) of n logits with 64 lanes; returns the slot'th expert id.
fn mgf_pick(slot: u32, n: u32, lid: u32) -> u32 {
    var chosen = 0u;
    for (var s = 0u; s <= slot; s = s + 1u) {
        var best = -3.0e38;
        var bi = 0xFFFFu;
        var i = lid;
        loop {
            if (i >= n) { break; }
            let v = mgf_lg[i];
            if (v > best || (v == best && i < bi)) { best = v; bi = i; }
            i = i + 64u;
        }
        mgf_v[lid] = best;
        mgf_i[lid] = bi;
        workgroupBarrier();
        var st = 32u;
        loop {
            if (st == 0u) { break; }
            if (lid < st) {
                let b = mgf_v[lid + st];
                let ib = mgf_i[lid + st];
                if (b > mgf_v[lid] || (b == mgf_v[lid] && ib < mgf_i[lid])) {
                    mgf_v[lid] = b;
                    mgf_i[lid] = ib;
                }
            }
            workgroupBarrier();
            st = st >> 1u;
        }
        chosen = mgf_i[0];
        workgroupBarrier();
        if (lid == 0u && s < slot) { mgf_lg[chosen] = -3.0e38; }
        workgroupBarrier();
    }
    return chosen;
}

@compute @workgroup_size(64)
fn moe_gate_up_q2tp_f(@builtin(workgroup_id) wid: vec3<u32>,
                      @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let slot = wid.y;
    let gpr = mg_p.gpr;
    let rows = mg_p.inter;
    let n = mgf_p.n_exp;
    let mat16 = mg_p.mat16;
    // Stage logits once (shared expert = last slot, id n).
    var i = lid;
    loop {
        if (i >= n) { break; }
        mgf_lg[i] = mgf_logit[i];
        i = i + 64u;
    }
    workgroupBarrier();
    var id = n;
    if (slot < mg_p.slots - 1u) {
        id = mgf_pick(slot, n, lid);
    }
    let base16 = id * mat16;
    let nib16 = base16 + row * gpr * 4u;
    let par16 = base16 + rows * gpr * 4u + row * 2u;
    let cst = (gpr * 5u + 7u) / 8u;
    let cod8 = (base16 + rows * gpr * 4u + rows * 2u) * 2u + row * cst;

    let gl = unpack2x16float(mg_g16(par16) | (mg_g16(par16 + 1u) << 16u));
    let ul = unpack2x16float(mg_u16f(par16) | (mg_u16f(par16 + 1u) << 16u));
    var ag = 0.0;
    var au = 0.0;
    for (var g = lid; g < gpr; g = g + 64u) {
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cg = mgp_gu8(cod8 + cb);
        var cu = mgp_uu8(cod8 + cb);
        if (shf > 3u) {
            cg = cg | (mgp_gu8(cod8 + cb + 1u) << 8u);
            cu = cu | (mgp_uu8(cod8 + cb + 1u) << 8u);
        }
        let cgv = (cg >> shf) & 31u;
        let cuv = (cu >> shf) & 31u;
        let sg = select(exp2(gl.x + f32(max(cgv, 1u) - 1u) * gl.y), 0.0, cgv == 0u);
        let su = select(exp2(ul.x + f32(max(cuv, 1u) - 1u) * ul.y), 0.0, cuv == 0u);
        let w32 = (nib16 + g * 4u) >> 1u;
        let xq = g * 8u;
        let x0 = mg_xv[xq];      let x1 = mg_xv[xq + 1u];
        let x2 = mg_xv[xq + 2u]; let x3 = mg_xv[xq + 3u];
        let x4 = mg_xv[xq + 4u]; let x5 = mg_xv[xq + 5u];
        let x6 = mg_xv[xq + 6u]; let x7 = mg_xv[xq + 7u];
        let dg = mg_dot16v(mg_gw[w32], x0, x1, x2, x3)
               + mg_dot16v(mg_gw[w32 + 1u], x4, x5, x6, x7);
        let du = mg_dot16v(mg_uw[w32], x0, x1, x2, x3)
               + mg_dot16v(mg_uw[w32 + 1u], x4, x5, x6, x7);
        ag = ag + sg * dg;
        au = au + su * du;
    }
    mg_pg[lid] = ag;
    mg_pu[lid] = au;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) {
            mg_pg[lid] = mg_pg[lid] + mg_pg[lid + stride];
            mg_pu[lid] = mg_pu[lid] + mg_pu[lid + stride];
        }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        let g = mg_pg[0];
        var gg = g;
        var uu = mg_pu[0];
        if (mg_p.lim > 0.0) {
            uu = clamp(uu, -mg_p.lim, mg_p.lim);
            gg = min(gg, mg_p.lim);
        }
        mg_act[slot * mg_p.inter + row] = (gg / (1.0 + exp(-gg))) * uu;
    }
}

// down twin: recomputes ids AND weights (softmax over picked + shared
// sigmoid). Slot 2 = logits, slot 3 = shared-gate weight (f32 bits),
// slot 6 = the token's activations for the shared-gate dot.
@group(0) @binding(2) var<storage, read> mdf_logit : array<f32>;
@group(0) @binding(3) var<storage, read> mdf_sgw   : array<u32>;
@group(0) @binding(6) var<storage, read> mdf_x     : array<f32>;
struct MdfP { n_exp: u32, top_k: u32, norm: u32, pk: u32 };
@group(0) @binding(7) var<uniform> mdf_p : MdfP;

var<workgroup> mdf_lg: array<f32, 256>;
var<workgroup> mdf_v:  array<f32, 64>;
var<workgroup> mdf_i:  array<u32, 64>;
var<workgroup> mdf_sel: array<u32, 16>;
var<workgroup> mdf_wt:  array<f32, 16>;

@compute @workgroup_size(64)
fn moe_down_q4tp_f(@builtin(workgroup_id) wid: vec3<u32>,
                   @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let gpr = md_p.gpr;
    let rows = md_p.hidden;
    let n = mdf_p.n_exp;
    let kk = mdf_p.top_k;
    let sg_kind = mdf_p.pk & 0xFFu;
    let sg_hidden = mdf_p.pk >> 8u;
    // shared gate dot (same strided shape as the retired select kernel,
    // 64 lanes instead of 256)
    var sgv = 0.0;
    if (sg_kind == 4u) {
        var d = 0.0;
        var i = lid;
        loop {
            if (i >= sg_hidden) { break; }
            d = d + bitcast<f32>(mdf_sgw[i]) * mdf_x[i];
            i = i + 64u;
        }
        mdf_v[lid] = d;
        workgroupBarrier();
        var st = 32u;
        loop {
            if (st == 0u) { break; }
            if (lid < st) { mdf_v[lid] = mdf_v[lid] + mdf_v[lid + st]; }
            workgroupBarrier();
            st = st >> 1u;
        }
        sgv = mdf_v[0];
        workgroupBarrier();
    }
    // logits + max + denom with 64-lane reductions
    var i2 = lid;
    loop {
        if (i2 >= n) { break; }
        mdf_lg[i2] = mdf_logit[i2];
        i2 = i2 + 64u;
    }
    workgroupBarrier();
    var mbest = -3.0e38;
    var i3 = lid;
    loop {
        if (i3 >= n) { break; }
        mbest = max(mbest, mdf_lg[i3]);
        i3 = i3 + 64u;
    }
    mdf_v[lid] = mbest;
    workgroupBarrier();
    var st2 = 32u;
    loop {
        if (st2 == 0u) { break; }
        if (lid < st2) { mdf_v[lid] = max(mdf_v[lid], mdf_v[lid + st2]); }
        workgroupBarrier();
        st2 = st2 >> 1u;
    }
    let mx = mdf_v[0];
    workgroupBarrier();
    var dsum = 0.0;
    var i4 = lid;
    loop {
        if (i4 >= n) { break; }
        dsum = dsum + exp(mdf_lg[i4] - mx);
        i4 = i4 + 64u;
    }
    mdf_v[lid] = dsum;
    workgroupBarrier();
    st2 = 32u;
    loop {
        if (st2 == 0u) { break; }
        if (lid < st2) { mdf_v[lid] = mdf_v[lid] + mdf_v[lid + st2]; }
        workgroupBarrier();
        st2 = st2 >> 1u;
    }
    let denom = mdf_v[0];
    workgroupBarrier();
    // top-k, weights, optional renorm; shared expert last
    var wsum = 0.0;
    for (var s = 0u; s < kk; s = s + 1u) {
        var best = -3.0e38;
        var bi = 0xFFFFu;
        var i5 = lid;
        loop {
            if (i5 >= n) { break; }
            let v = mdf_lg[i5];
            if (v > best || (v == best && i5 < bi)) { best = v; bi = i5; }
            i5 = i5 + 64u;
        }
        mdf_v[lid] = best;
        mdf_i[lid] = bi;
        workgroupBarrier();
        var st3 = 32u;
        loop {
            if (st3 == 0u) { break; }
            if (lid < st3) {
                let b = mdf_v[lid + st3];
                let ib = mdf_i[lid + st3];
                if (b > mdf_v[lid] || (b == mdf_v[lid] && ib < mdf_i[lid])) {
                    mdf_v[lid] = b;
                    mdf_i[lid] = ib;
                }
            }
            workgroupBarrier();
            st3 = st3 >> 1u;
        }
        if (lid == 0u) {
            mdf_sel[s] = mdf_i[0];
            mdf_wt[s] = exp(mdf_v[0] - mx) / denom;
        }
        workgroupBarrier();
        wsum = wsum + exp(mdf_v[0] - mx) / denom;
        if (lid == 0u) { mdf_lg[mdf_i[0]] = -3.0e38; }
        workgroupBarrier();
    }
    if (lid == 0u) {
        if (mdf_p.norm != 0u) {
            for (var s = 0u; s < kk; s = s + 1u) { mdf_wt[s] = mdf_wt[s] / wsum; }
        }
        mdf_sel[kk] = n;
        mdf_wt[kk] = 1.0 / (1.0 + exp(-sgv));
    }
    workgroupBarrier();
    let cst = (gpr * 5u + 7u) / 8u;
    let total = md_p.slots * gpr;
    var acc = 0.0;
    for (var i6 = lid; i6 < total; i6 = i6 + 64u) {
        let slot = i6 / gpr;
        let g = i6 % gpr;
        let base16 = mdf_sel[slot] * md_p.mat16;
        let par16 = base16 + rows * gpr * 8u + row * 2u;
        let cod8 = (base16 + rows * gpr * 8u + rows * 2u) * 2u + row * cst;
        let pl = unpack2x16float(md_u16(par16) | (md_u16(par16 + 1u) << 16u));
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cv = mdp_u8(cod8 + cb);
        if (shf > 3u) { cv = cv | (mdp_u8(cod8 + cb + 1u) << 8u); }
        let scale = exp2(pl.x + f32((cv >> shf) & 31u) * pl.y);
        let t16 = base16 + (row * gpr + g) * 8u;
        let xb = (slot * gpr + g) * 32u;
        var d = 0.0;
        for (var k2 = 0u; k2 < 4u; k2 = k2 + 1u) {
            let w = md_u16(t16 + 2u * k2) | (md_u16(t16 + 1u + 2u * k2) << 16u);
            d = d + md_dot8(w, xb + 8u * k2);
        }
        acc = acc + mdf_wt[slot] * scale * d;
    }
    md_pt[lid] = acc;
    workgroupBarrier();
    var st4 = 32u;
    loop {
        if (st4 == 0u) { break; }
        if (lid < st4) { md_pt[lid] = md_pt[lid] + md_pt[lid + st4]; }
        workgroupBarrier();
        st4 = st4 >> 1u;
    }
    if (lid == 0u) { md_y[row] = md_pt[0]; }
}

// ── Multi-step greedy tail: argmax over the logits on the device, then
// re-embed the winner — k decode steps ride ONE submit and the CPU sees
// k token ids instead of k megabytes of logits. Ties pick an arbitrary
// maximal index (same class as the documented GPU float-order ties).
struct AmP { n: u32, parts: u32, st: u32, _p: u32 };
@group(0) @binding(0) var<storage, read>       am_x   : array<f32>;
@group(0) @binding(1) var<storage, read_write> am_pv  : array<f32>;
@group(0) @binding(2) var<storage, read_write> am_pi  : array<u32>;
@group(0) @binding(3) var<uniform>             am_p   : AmP;
var<workgroup> am_wv: array<f32, 256>;
var<workgroup> am_wi: array<u32, 256>;

@compute @workgroup_size(256)
fn argmax_part(@builtin(workgroup_id) wid: vec3<u32>,
               @builtin(local_invocation_index) lid: u32) {
    var best = -3.0e38;
    var bi = 0u;
    var i = wid.x * 256u + lid;
    let stride = am_p.parts * 256u;
    loop {
        if (i >= am_p.n) { break; }
        let v = am_x[i];
        if (v > best) { best = v; bi = i; }
        i = i + stride;
    }
    am_wv[lid] = best;
    am_wi[lid] = bi;
    workgroupBarrier();
    var s = 128u;
    loop {
        if (s == 0u) { break; }
        if (lid < s && am_wv[lid + s] > am_wv[lid]) {
            am_wv[lid] = am_wv[lid + s];
            am_wi[lid] = am_wi[lid + s];
        }
        workgroupBarrier();
        s = s >> 1u;
    }
    if (lid == 0u) {
        am_pv[wid.x] = am_wv[0];
        am_pi[wid.x] = am_wi[0];
    }
}

@group(0) @binding(0) var<storage, read>       af_pv : array<f32>;
@group(0) @binding(1) var<storage, read>       af_pi : array<u32>;
@group(0) @binding(2) var<storage, read_write> af_ids: array<u32>;
@group(0) @binding(3) var<uniform>             af_p  : AmP;
var<workgroup> af_wv: array<f32, 256>;
var<workgroup> af_wi: array<u32, 256>;

@compute @workgroup_size(256)
fn argmax_final(@builtin(local_invocation_index) lid: u32) {
    var best = -3.0e38;
    var bi = 0u;
    var i = lid;
    loop {
        if (i >= af_p.parts) { break; }
        if (af_pv[i] > best) { best = af_pv[i]; bi = af_pi[i]; }
        i = i + 256u;
    }
    af_wv[lid] = best;
    af_wi[lid] = bi;
    workgroupBarrier();
    var s = 128u;
    loop {
        if (s == 0u) { break; }
        if (lid < s && af_wv[lid + s] > af_wv[lid]) {
            af_wv[lid] = af_wv[lid + s];
            af_wi[lid] = af_wi[lid + s];
        }
        workgroupBarrier();
        s = s >> 1u;
    }
    if (lid == 0u) { af_ids[af_p.st] = af_wi[0]; }
}

// One thread = one hidden element of the winner's q4tp embedding row.
// `mult` carries the model's embed multiplier as f32 bits.
struct EgP { hidden: u32, gpr: u32, rows: u32, st: u32, mult: u32, _a: u32, _b: u32, _c: u32 };
@group(0) @binding(0) var<storage, read>       eg_w  : array<u32>;
@group(0) @binding(1) var<storage, read>       eg_ids: array<u32>;
@group(0) @binding(2) var<storage, read_write> eg_h  : array<f32>;
@group(0) @binding(3) var<uniform>             eg_p  : EgP;

fn eg_byte(off: u32) -> u32 {
    return (eg_w[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
}

@compute @workgroup_size(256)
fn embed_gather_q4tp(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= eg_p.hidden) { return; }
    let r = eg_ids[eg_p.st];
    let gpr = eg_p.gpr;
    let g = i / 32u;
    let k = i % 32u;
    let params_w = eg_p.rows * gpr * 4u;
    let codes_b = eg_p.rows * gpr * 16u + eg_p.rows * 4u;
    let cst = (gpr * 5u + 7u) / 8u;
    let pr = unpack2x16float(eg_w[params_w + r]);
    let bit = g * 5u;
    let cb = codes_b + r * cst + (bit >> 3u);
    let sh = bit & 7u;
    var cv = eg_byte(cb);
    if (sh > 3u) { cv = cv | (eg_byte(cb + 1u) << 8u); }
    let sc = exp2(pr.x + f32((cv >> sh) & 31u) * pr.y);
    let nb = (r * gpr + g) * 16u + (k >> 1u);
    let byte = eg_byte(nb);
    let q = select(byte & 0xFu, (byte >> 4u) & 0xFu, (k & 1u) == 1u);
    eg_h[i] = (f32(q) - 8.0) * sc * bitcast<f32>(eg_p.mult);
}

// Fused SiLU(gate)·up → Q4Block down-proj matvec: eliminates the standalone
// silu dispatch (saves one inter-pass pipeline flush per layer).
@group(0) @binding(0) var<storage, read>       sd_w : array<u32>;
@group(0) @binding(1) var<storage, read>       sd_gate : array<f32>;
@group(0) @binding(2) var<storage, read>       sd_up : array<f32>;
@group(0) @binding(3) var<storage, read_write> sd_y : array<f32>;
@group(0) @binding(4) var<uniform>             sd_p : Q1Params;

var<workgroup> partial_sd: array<f32, 64>;

fn sd_dot8(w: u32, xi: u32) -> f32 {
    let g0 = sd_gate[xi];     let g1 = sd_gate[xi + 1u];
    let g2 = sd_gate[xi + 2u]; let g3 = sd_gate[xi + 3u];
    let g4 = sd_gate[xi + 4u]; let g5 = sd_gate[xi + 5u];
    let g6 = sd_gate[xi + 6u]; let g7 = sd_gate[xi + 7u];
    return (f32(w & 0xFu) - 8.0) * (g0 / (1.0 + exp(-g0)) * sd_up[xi])
         + (f32((w >> 4u) & 0xFu) - 8.0) * (g1 / (1.0 + exp(-g1)) * sd_up[xi + 1u])
         + (f32((w >> 8u) & 0xFu) - 8.0) * (g2 / (1.0 + exp(-g2)) * sd_up[xi + 2u])
         + (f32((w >> 12u) & 0xFu) - 8.0) * (g3 / (1.0 + exp(-g3)) * sd_up[xi + 3u])
         + (f32((w >> 16u) & 0xFu) - 8.0) * (g4 / (1.0 + exp(-g4)) * sd_up[xi + 4u])
         + (f32((w >> 20u) & 0xFu) - 8.0) * (g5 / (1.0 + exp(-g5)) * sd_up[xi + 5u])
         + (f32((w >> 24u) & 0xFu) - 8.0) * (g6 / (1.0 + exp(-g6)) * sd_up[xi + 6u])
         + (f32((w >> 28u) & 0xFu) - 8.0) * (g7 / (1.0 + exp(-g7)) * sd_up[xi + 7u]);
}

@compute @workgroup_size(64)
fn silu_down_matvec(@builtin(workgroup_id) wid: vec3<u32>,
                    @builtin(num_workgroups) nwg: vec3<u32>,
                    @builtin(local_invocation_index) lid: u32) {
    let gpr = sd_p.np;
    let rows = sd_p.rows;
    let scales_off = rows * gpr * 16u;
    var row = wid.x;
    loop {
        if (row >= rows) { break; }
        var acc = 0.0;
        var g = lid;
        loop {
            if (g >= gpr) { break; }
            let gi = row * gpr + g;
            let sc_byte = scales_off + gi * 2u;
            let sc16 = (sd_w[sc_byte >> 2u] >> ((sc_byte & 3u) * 8u)) & 0xFFFFu;
            let scale = unpack2x16float(sc16).x;
            let pk4 = gi * 4u;
            let xb = g * 32u;
            let gsum = sd_dot8(sd_w[pk4], xb)
                     + sd_dot8(sd_w[pk4 + 1u], xb + 8u)
                     + sd_dot8(sd_w[pk4 + 2u], xb + 16u)
                     + sd_dot8(sd_w[pk4 + 3u], xb + 24u);
            acc = acc + scale * gsum;
            g = g + 64u;
        }
        partial_sd[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { partial_sd[lid] = partial_sd[lid] + partial_sd[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) { sd_y[row] = partial_sd[0]; }
        workgroupBarrier();
        row = row + nwg.x;
    }
}

// q1t register-blocked GEMM (prefill) — the WGSL cousin of the Metal q1t_mul_mm
// and structurally identical to q8_mul_mm here; only the W staging decodes
// base-3 ternary × per-group f16 scale (no row_scale; scale folds into the
// staged weight). Own 4-slot bindings. The overlay is a second pass.
struct Q1tMmP { cols4: u32, rows: u32, nb: u32, _p: u32 };
@group(0) @binding(0) var<storage, read>       qmm : array<u32>;
@group(0) @binding(1) var<storage, read>       xmm : array<f32>;
// The same activations as vec4 (same-slot rule): the GEMM stages four
// consecutive floats per thread per K-step, and col0 is a multiple of 4,
// so that is one 16-byte load instead of four scalar ones.
@group(0) @binding(1) var<storage, read>       xmm4 : array<vec4<f32>>;
@group(0) @binding(2) var<storage, read_write> ymm : array<f32>;
@group(0) @binding(3) var<uniform>             pmm : Q1tMmP;

fn qmm_byte(off: u32) -> u32 {
    return (qmm[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
}
var<workgroup> q1t_at: array<f32, 64 * 16>;
var<workgroup> q1t_wt: array<f32, 64 * 16>;

@compute @workgroup_size(16, 16)
fn q1t_mul_mm(@builtin(workgroup_id) wid: vec3<u32>,
              @builtin(local_invocation_id) lid: vec3<u32>) {
    let cols = pmm.cols4 * 4u;
    let gpr = cols >> 5u;
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;
    let tid = lid.y * 16u + lid.x;
    // Sixteen named scalars, not array<array<f32,4>,4> — see q4t_mul_mm.
    var a00 = 0.0; var a01 = 0.0; var a02 = 0.0; var a03 = 0.0;
    var a10 = 0.0; var a11 = 0.0; var a12 = 0.0; var a13 = 0.0;
    var a20 = 0.0; var a21 = 0.0; var a22 = 0.0; var a23 = 0.0;
    var a30 = 0.0; var a31 = 0.0; var a32 = 0.0; var a33 = 0.0;
    var k0 = 0u;
    loop {
        if (k0 >= cols) { break; }
        for (var t = tid; t < 64u * 4u; t = t + 256u) {
            let m = t / 4u;
            let k4 = t % 4u;
            var xv = vec4<f32>(0.0);
            let col0 = k0 + k4 * 4u;
            if (m0 + m < pmm.nb && col0 < cols) {
                // cols is a multiple of 32 and col0 of 4 — vec4-aligned.
                xv = xmm4[((m0 + m) * cols + col0) >> 2u];
            }
            let dst = m * 16u + k4 * 4u;
            q1t_at[dst] = xv.x; q1t_at[dst + 1u] = xv.y;
            q1t_at[dst + 2u] = xv.z; q1t_at[dst + 3u] = xv.w;
        }
        for (var t = tid; t < 64u * 4u; t = t + 256u) {
            let n = t / 4u;
            let k4 = t % 4u;
            var wv = vec4<f32>(0.0);
            let col0 = k0 + k4 * 4u;
            if (n0 + n < pmm.rows && col0 < cols) {
                let g = col0 >> 5u;
                let toff = ((n0 + n) * gpr + g) * 9u;
                let sc16 = qmm_byte(toff) | (qmm_byte(toff + 1u) << 8u);
                let scale = unpack2x16float(sc16).x;
                let codes = toff + 2u;
                for (var d = 0u; d < 4u; d = d + 1u) {
                    let p = (col0 + d) - g * 32u;
                    let b = qmm_byte(codes + p / 5u);
                    let code = (Q1T_LUT[b] >> ((p % 5u) * 2u)) & 3u;
                    var sgn = 0.0;
                    if (code == 1u) { sgn = 1.0; } else if (code == 2u) { sgn = -1.0; }
                    wv[d] = sgn * scale;
                }
            }
            let dst = n * 16u + k4 * 4u;
            q1t_wt[dst] = wv.x; q1t_wt[dst + 1u] = wv.y;
            q1t_wt[dst + 2u] = wv.z; q1t_wt[dst + 3u] = wv.w;
        }
        workgroupBarrier();
        let ab = lid.y * 64u;
        let wb = lid.x * 64u;
        for (var k = 0u; k < 16u; k = k + 1u) {
            let x0 = q1t_at[ab + k];
            let x1 = q1t_at[ab + 16u + k];
            let x2 = q1t_at[ab + 32u + k];
            let x3 = q1t_at[ab + 48u + k];
            let y0 = q1t_wt[wb + k];
            let y1 = q1t_wt[wb + 16u + k];
            let y2 = q1t_wt[wb + 32u + k];
            let y3 = q1t_wt[wb + 48u + k];
            a00 = a00 + x0 * y0; a01 = a01 + x0 * y1;
            a02 = a02 + x0 * y2; a03 = a03 + x0 * y3;
            a10 = a10 + x1 * y0; a11 = a11 + x1 * y1;
            a12 = a12 + x1 * y2; a13 = a13 + x1 * y3;
            a20 = a20 + x2 * y0; a21 = a21 + x2 * y1;
            a22 = a22 + x2 * y2; a23 = a23 + x2 * y3;
            a30 = a30 + x3 * y0; a31 = a31 + x3 * y1;
            a32 = a32 + x3 * y2; a33 = a33 + x3 * y3;
        }
        workgroupBarrier();
        k0 = k0 + 16u;
    }
    let mb = m0 + lid.y * 4u;
    let nb2 = n0 + lid.x * 4u;
    q4t_store4(mb, nb2, a00, a01, a02, a03);
    q4t_store4(mb + 1u, nb2, a10, a11, a12, a13);
    q4t_store4(mb + 2u, nb2, a20, a21, a22, a23);
    q4t_store4(mb + 3u, nb2, a30, a31, a32, a33);
}

// q4t register-blocked GEMM (imagegen DiT prefill / any wide q4t
// batch) — the WGSL cousin of the Metal q4t_mul_mm and structurally
// identical to q1t_mul_mm above; only the W staging decodes 18-byte
// q4t tiles (f16 scale + 16 nibble bytes per 32-weight group).
// Shares the 4-slot qmm/xmm/ymm/pmm bindings.
// Staged k-major and read as vec4. The obvious layout — 16 k values per
// row, read with `tile[lid.x * 64u + k]` — gives every one of the sixteen
// threads in a row the SAME shared-memory bank (stride 64 floats, 32
// banks), so each read serialises sixteen ways. k-major makes one row of
// the tile 64 contiguous floats: sixteen threads take sixteen vec4s that
// cover all 32 banks twice, conflict-free, and one load now feeds four
// FMAs instead of one.
// Row-major with a 17-float stride, not 16. Sixteen threads of a row read
// one column of the tile; at stride 16 every one of them lands in the same
// bank of 32 and each read serialises sixteen ways. The odd pad walks them
// across the banks instead. (The alternative — k-major vec4 tiles — is
// faster still, but it needs four threads to write four lanes of one
// shared vec4, which a backend that lowers a dynamic component write to
// read-modify-write turns into a data race. Metal does.)
const TSTRIDE: u32 = 17u;
// Activations go k-major as whole vec4s — one thread owns one vec4, so no
// lane is written by two threads and the inner read is conflict-free.
// Weights stay row-major with the odd pad: staging them k-major would make
// each thread decode four rows, and four scale ladders with an exp2 apiece
// cost more than the two-way conflict the pad leaves behind.
var<workgroup> q4t_at4: array<vec4<f32>, 16 * 16>;
var<workgroup> q4t_wt: array<f32, 64 * 17>;

fn q4t_store4(m: u32, n0: u32, v0: f32, v1: f32, v2: f32, v3: f32) {
    if (m >= pmm.nb) { return; }
    let base = m * pmm.rows + n0;
    if (n0 < pmm.rows) { ymm[base] = v0; }
    if (n0 + 1u < pmm.rows) { ymm[base + 1u] = v1; }
    if (n0 + 2u < pmm.rows) { ymm[base + 2u] = v2; }
    if (n0 + 3u < pmm.rows) { ymm[base + 3u] = v3; }
}

// q4tp register-blocked GEMM — the q4t kernel above with one block swapped:
// a 16 B nibble stride instead of the 18 B tile, and the scale off the row's
// ladder. Shares q4t_store4 and the q4t_at/q4t_wt staging arrays; only one
// entry point runs per dispatch, so the workgroup allocation is not doubled.
@compute @workgroup_size(16, 16)
fn q4tp_mul_mm(@builtin(workgroup_id) wid: vec3<u32>,
              @builtin(local_invocation_id) lid: vec3<u32>) {
    let cols = pmm.cols4 * 4u;
    let gpr = cols >> 5u;
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;
    let tid = lid.y * 16u + lid.x;
    // The 4x4 register block is SIXTEEN NAMED SCALARS, not
    // array<array<f32,4>,4>: indexed by loop variables the array is a
    // private array, which this backend puts in stack memory — the
    // accumulators leave registers and the GEMM runs at a fraction of
    // the card (measured 373 GFLOP/s of an RTX 3090's ~35 TFLOP/s).
    var a00 = 0.0; var a01 = 0.0; var a02 = 0.0; var a03 = 0.0;
    var a10 = 0.0; var a11 = 0.0; var a12 = 0.0; var a13 = 0.0;
    var a20 = 0.0; var a21 = 0.0; var a22 = 0.0; var a23 = 0.0;
    var a30 = 0.0; var a31 = 0.0; var a32 = 0.0; var a33 = 0.0;
    var k0 = 0u;
    loop {
        if (k0 >= cols) { break; }
        // Each thread owns ONE whole vec4 of the tile. Writing single
        // lanes of a shared vec4 from four threads is a data race on any
        // backend that lowers a dynamic component write to read-modify-write
        // the whole vector — Metal does, and three lanes in four came back
        // zero while Vulkan was fine.
        {
            // `lid.x` selects the input-column group and `lid.y` selects
            // the four input rows.  The old transposed assignment made
            // every output block consume the wrong batch rows; k=1 hid it
            // behind a plausible-looking dot product, while a real prefill
            // returned corrupted hidden rows.
            let kk = tid % 16u;
            let slot = tid / 16u;
            let col = k0 + kk;
            let m = m0 + slot * 4u;
            var xv = vec4<f32>(0.0);
            if (col < cols) {
                let i0 = m * cols + col;
                if (m < pmm.nb) { xv.x = xmm4[i0 >> 2u][i0 & 3u]; }
                let i1 = i0 + cols;
                if (m + 1u < pmm.nb) { xv.y = xmm4[i1 >> 2u][i1 & 3u]; }
                let i2 = i1 + cols;
                if (m + 2u < pmm.nb) { xv.z = xmm4[i2 >> 2u][i2 & 3u]; }
                let i3 = i2 + cols;
                if (m + 3u < pmm.nb) { xv.w = xmm4[i3 >> 2u][i3 & 3u]; }
            }
            q4t_at4[kk * 16u + slot] = xv;
        }
        for (var t = tid; t < 64u * 4u; t = t + 256u) {
            let n = t / 4u;
            let k4 = t % 4u;
            var wv = vec4<f32>(0.0);
            let col0 = k0 + k4 * 4u;
            if (n0 + n < pmm.rows && col0 < cols) {
                let g = col0 >> 5u;
                let wrow = n0 + n;
                let params_b = pmm.rows * gpr * 16u;
                let codes_b = params_b + pmm.rows * 4u;
                let cstride = (gpr * 5u + 7u) / 8u;
                let bit = g * 5u;
                let cb = codes_b + wrow * cstride + (bit >> 3u);
                let sh = bit & 7u;
                var cv = qmm_byte(cb);
                if (sh > 3u) { cv = cv | (qmm_byte(cb + 1u) << 8u); }
                let pr = unpack2x16float(qmm[(params_b >> 2u) + wrow]);
                let scale = exp2(pr.x + f32((cv >> sh) & 31u) * pr.y);
                let toff = (wrow * gpr + g) * 16u;
                let p = col0 - g * 32u;
                let bo = toff + p / 2u;
                let w32 = qmm[bo >> 2u];
                let sh0 = (bo & 3u) * 8u;
                let b0 = (w32 >> sh0) & 0xFFu;
                var b1 = 0u;
                if ((bo & 3u) == 3u) {
                    b1 = qmm[(bo >> 2u) + 1u] & 0xFFu;
                } else {
                    b1 = (w32 >> (sh0 + 8u)) & 0xFFu;
                }
                wv[0u] = (f32(b0 & 0xFu) - 8.0) * scale;
                wv[1u] = (f32(b0 >> 4u) - 8.0) * scale;
                wv[2u] = (f32(b1 & 0xFu) - 8.0) * scale;
                wv[3u] = (f32(b1 >> 4u) - 8.0) * scale;
            }
            let dst = n * TSTRIDE + k4 * 4u;
            q4t_wt[dst] = wv.x; q4t_wt[dst + 1u] = wv.y;
            q4t_wt[dst + 2u] = wv.z; q4t_wt[dst + 3u] = wv.w;
        }
        workgroupBarrier();
        let wb = lid.x * 4u * TSTRIDE;
        for (var k = 0u; k < 16u; k = k + 1u) {
            let xv = q4t_at4[k * 16u + lid.y];
            let x0 = xv.x; let x1 = xv.y; let x2 = xv.z; let x3 = xv.w;
            let y0 = q4t_wt[wb + k];
            let y1 = q4t_wt[wb + TSTRIDE + k];
            let y2 = q4t_wt[wb + 2u * TSTRIDE + k];
            let y3 = q4t_wt[wb + 3u * TSTRIDE + k];
            a00 = a00 + x0 * y0; a01 = a01 + x0 * y1;
            a02 = a02 + x0 * y2; a03 = a03 + x0 * y3;
            a10 = a10 + x1 * y0; a11 = a11 + x1 * y1;
            a12 = a12 + x1 * y2; a13 = a13 + x1 * y3;
            a20 = a20 + x2 * y0; a21 = a21 + x2 * y1;
            a22 = a22 + x2 * y2; a23 = a23 + x2 * y3;
            a30 = a30 + x3 * y0; a31 = a31 + x3 * y1;
            a32 = a32 + x3 * y2; a33 = a33 + x3 * y3;
        }
        workgroupBarrier();
        k0 = k0 + 16u;
    }
    let mb = m0 + lid.y * 4u;
    let nb2 = n0 + lid.x * 4u;
    q4t_store4(mb, nb2, a00, a01, a02, a03);
    q4t_store4(mb + 1u, nb2, a10, a11, a12, a13);
    q4t_store4(mb + 2u, nb2, a20, a21, a22, a23);
    q4t_store4(mb + 3u, nb2, a30, a31, a32, a33);
}

// The same GEMM over a `q2tp` weight plane. Three things differ and
// nothing else does: the weight plane is 8 bytes a group rather than
// 16, a byte holds four 2-bit codes instead of two nibbles, and the
// scale ladder is shifted down a rung because rung 0 is spent naming
// the exact zero the +-0.5/+-1.5 grid cannot reach. The params and
// codes planes are byte-identical to q4tp's.
@compute @workgroup_size(16, 16)
fn q2tp_mul_mm(@builtin(workgroup_id) wid: vec3<u32>,
              @builtin(local_invocation_id) lid: vec3<u32>) {
    let cols = pmm.cols4 * 4u;
    let gpr = cols >> 5u;
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;
    let tid = lid.y * 16u + lid.x;
    // The 4x4 register block is SIXTEEN NAMED SCALARS, not
    // array<array<f32,4>,4>: indexed by loop variables the array is a
    // private array, which this backend puts in stack memory — the
    // accumulators leave registers and the GEMM runs at a fraction of
    // the card (measured 373 GFLOP/s of an RTX 3090's ~35 TFLOP/s).
    var a00 = 0.0; var a01 = 0.0; var a02 = 0.0; var a03 = 0.0;
    var a10 = 0.0; var a11 = 0.0; var a12 = 0.0; var a13 = 0.0;
    var a20 = 0.0; var a21 = 0.0; var a22 = 0.0; var a23 = 0.0;
    var a30 = 0.0; var a31 = 0.0; var a32 = 0.0; var a33 = 0.0;
    var k0 = 0u;
    loop {
        if (k0 >= cols) { break; }
        // Each thread owns ONE whole vec4 of the tile. Writing single
        // lanes of a shared vec4 from four threads is a data race on any
        // backend that lowers a dynamic component write to read-modify-write
        // the whole vector — Metal does, and three lanes in four came back
        // zero while Vulkan was fine.
        {
            // Keep the shared activation tile indexed [column_group][row_group].
            let kk = tid % 16u;
            let slot = tid / 16u;
            let col = k0 + kk;
            let m = m0 + slot * 4u;
            var xv = vec4<f32>(0.0);
            if (col < cols) {
                let i0 = m * cols + col;
                if (m < pmm.nb) { xv.x = xmm4[i0 >> 2u][i0 & 3u]; }
                let i1 = i0 + cols;
                if (m + 1u < pmm.nb) { xv.y = xmm4[i1 >> 2u][i1 & 3u]; }
                let i2 = i1 + cols;
                if (m + 2u < pmm.nb) { xv.z = xmm4[i2 >> 2u][i2 & 3u]; }
                let i3 = i2 + cols;
                if (m + 3u < pmm.nb) { xv.w = xmm4[i3 >> 2u][i3 & 3u]; }
            }
            q4t_at4[kk * 16u + slot] = xv;
        }
        for (var t = tid; t < 64u * 4u; t = t + 256u) {
            let n = t / 4u;
            let k4 = t % 4u;
            var wv = vec4<f32>(0.0);
            let col0 = k0 + k4 * 4u;
            if (n0 + n < pmm.rows && col0 < cols) {
                let g = col0 >> 5u;
                let wrow = n0 + n;
                let params_b = pmm.rows * gpr * 8u;
                let codes_b = params_b + pmm.rows * 4u;
                let cstride = (gpr * 5u + 7u) / 8u;
                let bit = g * 5u;
                let cb = codes_b + wrow * cstride + (bit >> 3u);
                let sh = bit & 7u;
                var cv = qmm_byte(cb);
                if (sh > 3u) { cv = cv | (qmm_byte(cb + 1u) << 8u); }
                let pr = unpack2x16float(qmm[(params_b >> 2u) + wrow]);
                // Rung 0 means the group is exactly zero — the 4-level
                // grid cannot spell it — and 1..=31 are the q4tp ladder
                // shifted down one.
                let code = (cv >> sh) & 31u;
                var scale = 0.0;
                if (code > 0u) { scale = exp2(pr.x + f32(code - 1u) * pr.y); }
                let toff = (wrow * gpr + g) * 8u;
                let p = col0 - g * 32u;
                // Four 2-bit weights to a byte, LSB first, so one byte
                // is exactly the vec4 this thread owns.
                let bo = toff + p / 4u;
                let by = (qmm[bo >> 2u] >> ((bo & 3u) * 8u)) & 0xFFu;
                let center = select(1.5, 1.0, pmm._p != 0u);
                wv[0u] = (f32(by & 3u) - center) * scale;
                wv[1u] = (f32((by >> 2u) & 3u) - center) * scale;
                wv[2u] = (f32((by >> 4u) & 3u) - center) * scale;
                wv[3u] = (f32((by >> 6u) & 3u) - center) * scale;
            }
            let dst = n * TSTRIDE + k4 * 4u;
            q4t_wt[dst] = wv.x; q4t_wt[dst + 1u] = wv.y;
            q4t_wt[dst + 2u] = wv.z; q4t_wt[dst + 3u] = wv.w;
        }
        workgroupBarrier();
        let wb = lid.x * 4u * TSTRIDE;
        for (var k = 0u; k < 16u; k = k + 1u) {
            let xv = q4t_at4[k * 16u + lid.y];
            let x0 = xv.x; let x1 = xv.y; let x2 = xv.z; let x3 = xv.w;
            let y0 = q4t_wt[wb + k];
            let y1 = q4t_wt[wb + TSTRIDE + k];
            let y2 = q4t_wt[wb + 2u * TSTRIDE + k];
            let y3 = q4t_wt[wb + 3u * TSTRIDE + k];
            a00 = a00 + x0 * y0; a01 = a01 + x0 * y1;
            a02 = a02 + x0 * y2; a03 = a03 + x0 * y3;
            a10 = a10 + x1 * y0; a11 = a11 + x1 * y1;
            a12 = a12 + x1 * y2; a13 = a13 + x1 * y3;
            a20 = a20 + x2 * y0; a21 = a21 + x2 * y1;
            a22 = a22 + x2 * y2; a23 = a23 + x2 * y3;
            a30 = a30 + x3 * y0; a31 = a31 + x3 * y1;
            a32 = a32 + x3 * y2; a33 = a33 + x3 * y3;
        }
        workgroupBarrier();
        k0 = k0 + 16u;
    }
    let mb = m0 + lid.y * 4u;
    let nb2 = n0 + lid.x * 4u;
    q4t_store4(mb, nb2, a00, a01, a02, a03);
    q4t_store4(mb + 1u, nb2, a10, a11, a12, a13);
    q4t_store4(mb + 2u, nb2, a20, a21, a22, a23);
    q4t_store4(mb + 3u, nb2, a30, a31, a32, a33);
}

@compute @workgroup_size(16, 16)
fn q4t_mul_mm(@builtin(workgroup_id) wid: vec3<u32>,
              @builtin(local_invocation_id) lid: vec3<u32>) {
    let cols = pmm.cols4 * 4u;
    let gpr = cols >> 5u;
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;
    let tid = lid.y * 16u + lid.x;
    // The 4x4 register block is SIXTEEN NAMED SCALARS, not
    // array<array<f32,4>,4>: indexed by loop variables the array is a
    // private array, which this backend puts in stack memory — the
    // accumulators leave registers and the GEMM runs at a fraction of
    // the card (measured 373 GFLOP/s of an RTX 3090's ~35 TFLOP/s).
    var a00 = 0.0; var a01 = 0.0; var a02 = 0.0; var a03 = 0.0;
    var a10 = 0.0; var a11 = 0.0; var a12 = 0.0; var a13 = 0.0;
    var a20 = 0.0; var a21 = 0.0; var a22 = 0.0; var a23 = 0.0;
    var a30 = 0.0; var a31 = 0.0; var a32 = 0.0; var a33 = 0.0;
    var k0 = 0u;
    loop {
        if (k0 >= cols) { break; }
        // One whole vec4 per thread; see `q4tp_mul_mm` for why lanes of a
        // shared vec4 must not be written from four threads.
        {
            // Keep the shared activation tile indexed [column_group][row_group].
            let kk = tid % 16u;
            let slot = tid / 16u;
            let col = k0 + kk;
            let m = m0 + slot * 4u;
            var xv = vec4<f32>(0.0);
            if (col < cols) {
                let i0 = m * cols + col;
                if (m < pmm.nb) { xv.x = xmm4[i0 >> 2u][i0 & 3u]; }
                let i1 = i0 + cols;
                if (m + 1u < pmm.nb) { xv.y = xmm4[i1 >> 2u][i1 & 3u]; }
                let i2 = i1 + cols;
                if (m + 2u < pmm.nb) { xv.z = xmm4[i2 >> 2u][i2 & 3u]; }
                let i3 = i2 + cols;
                if (m + 3u < pmm.nb) { xv.w = xmm4[i3 >> 2u][i3 & 3u]; }
            }
            q4t_at4[kk * 16u + slot] = xv;
        }
        for (var t = tid; t < 64u * 4u; t = t + 256u) {
            let n = t / 4u;
            let k4 = t % 4u;
            var wv = vec4<f32>(0.0);
            let col0 = k0 + k4 * 4u;
            if (n0 + n < pmm.rows && col0 < cols) {
                let g = col0 >> 5u;
                let toff = ((n0 + n) * gpr + g) * 18u;
                let sc16 = qmm_byte(toff) | (qmm_byte(toff + 1u) << 8u);
                let scale = unpack2x16float(sc16).x;
                let p = col0 - g * 32u;
                let b0 = qmm_byte(toff + 2u + p / 2u);
                let b1 = qmm_byte(toff + 3u + p / 2u);
                wv[0u] = (f32(b0 & 0xFu) - 8.0) * scale;
                wv[1u] = (f32(b0 >> 4u) - 8.0) * scale;
                wv[2u] = (f32(b1 & 0xFu) - 8.0) * scale;
                wv[3u] = (f32(b1 >> 4u) - 8.0) * scale;
            }
            let dst = n * TSTRIDE + k4 * 4u;
            q4t_wt[dst] = wv.x; q4t_wt[dst + 1u] = wv.y;
            q4t_wt[dst + 2u] = wv.z; q4t_wt[dst + 3u] = wv.w;
        }
        workgroupBarrier();
        let wb = lid.x * 4u * TSTRIDE;
        for (var k = 0u; k < 16u; k = k + 1u) {
            let xv = q4t_at4[k * 16u + lid.y];
            let x0 = xv.x; let x1 = xv.y; let x2 = xv.z; let x3 = xv.w;
            let y0 = q4t_wt[wb + k];
            let y1 = q4t_wt[wb + TSTRIDE + k];
            let y2 = q4t_wt[wb + 2u * TSTRIDE + k];
            let y3 = q4t_wt[wb + 3u * TSTRIDE + k];
            a00 = a00 + x0 * y0; a01 = a01 + x0 * y1;
            a02 = a02 + x0 * y2; a03 = a03 + x0 * y3;
            a10 = a10 + x1 * y0; a11 = a11 + x1 * y1;
            a12 = a12 + x1 * y2; a13 = a13 + x1 * y3;
            a20 = a20 + x2 * y0; a21 = a21 + x2 * y1;
            a22 = a22 + x2 * y2; a23 = a23 + x2 * y3;
            a30 = a30 + x3 * y0; a31 = a31 + x3 * y1;
            a32 = a32 + x3 * y2; a33 = a33 + x3 * y3;
        }
        workgroupBarrier();
        k0 = k0 + 16u;
    }
    let mb = m0 + lid.y * 4u;
    let nb2 = n0 + lid.x * 4u;
    q4t_store4(mb, nb2, a00, a01, a02, a03);
    q4t_store4(mb + 1u, nb2, a10, a11, a12, a13);
    q4t_store4(mb + 2u, nb2, a20, a21, a22, a23);
    q4t_store4(mb + 3u, nb2, a30, a31, a32, a33);
}

@compute @workgroup_size(64)
fn q1t_overlay_mm(@builtin(global_invocation_id) gid: vec3<u32>) {
    let row = gid.x;
    if (row >= pmm.rows) { return; }
    let cols = pmm.cols4 * 4u;
    let gpr = cols >> 5u;
    let base_len = pmm.rows * gpr * 9u;
    let ent = base_len + (pmm.rows + 1u) * 4u;
    let rp0 = base_len + row * 4u;
    let c0 = qmm_byte(rp0) | (qmm_byte(rp0 + 1u) << 8u) | (qmm_byte(rp0 + 2u) << 16u) | (qmm_byte(rp0 + 3u) << 24u);
    let rp1 = base_len + (row + 1u) * 4u;
    let c1 = qmm_byte(rp1) | (qmm_byte(rp1 + 1u) << 8u) | (qmm_byte(rp1 + 2u) << 16u) | (qmm_byte(rp1 + 3u) << 24u);
    for (var p = c0; p < c1; p = p + 1u) {
        let e = ent + p * 4u;
        let col = qmm_byte(e) | (qmm_byte(e + 1u) << 8u);
        let val = unpack2x16float(qmm_byte(e + 2u) | (qmm_byte(e + 3u) << 8u)).x;
        for (var bi = 0u; bi < pmm.nb; bi = bi + 1u) {
            ymm[bi * pmm.rows + row] = ymm[bi * pmm.rows + row] + val * xmm[bi * cols + col];
        }
    }
}

// ── MoE inside the whole-token graph ────────────────────────────────────────
// Router logits/shared-gate logit arrive from ordinary matvecs; these three
// kernels keep the routing DECISION and every selected expert on-device, so
// a MoE layer costs one extra pass over a dense one instead of a CPU sync.
// Expert weights live in three per-layer concat buffers (q4t tiles, expert e
// at u16 offset e·mat16); the SHARED expert is the last block, pinned by the
// select kernel at slot top_k with a sigmoid weight.

// `sg_kind` = 4 means this kernel computes the shared-expert gate itself
// from `ms_sgw` · `ms_x` and the host skips that matvec entirely. It is a
// ONE-ROW projection: 2048 multiply-adds for a whole dispatch, and a
// dispatch costs ~23 us on this stack against ~5 us for a pass — measured
// by sweeping the layer count. Folding it into a kernel that already runs
// one workgroup is free. Any other dtype keeps the separate matvec and
// this kernel reads its result from `ms_slog`.
struct MoeSelP { n_exp: u32, top_k: u32, norm: u32, pk: u32, scale: f32, _s0: u32, _s1: u32, _s2: u32 };
@group(0) @binding(0) var<storage, read>       ms_logit : array<f32>;
@group(0) @binding(1) var<storage, read>       ms_slog  : array<f32>;
@group(0) @binding(2) var<storage, read_write> ms_sel   : array<u32>;
@group(0) @binding(3) var<storage, read_write> ms_w     : array<f32>;
@group(0) @binding(4) var<uniform>             ms_p     : MoeSelP;
@group(0) @binding(5) var<storage, read>       ms_sgw   : array<u32>;
@group(0) @binding(6) var<storage, read>       ms_x     : array<f32>;
// Per-expert SELECTION bias (noaux_tc): ranks the top-k choice, never the
// mixing weights. A 4-byte dummy rides here when the model has none.
@group(0) @binding(7) var<storage, read>       ms_bias  : array<f32>;
var<workgroup> ms_lg:  array<f32, 256>;
// The per-expert mixing score: softmax prob or sigmoid, depending on
// `norm` bit 1. Kept apart from ms_lg because with a bias the RANKING
// key and the WEIGHT differ, and conflating them is exactly the noaux_tc
// mistake.
var<workgroup> ms_sc:  array<f32, 256>;
var<workgroup> ms_red: array<f32, 256>;
var<workgroup> ms_ri:  array<u32, 256>;
var<workgroup> ms_pick: u32;
var<workgroup> ms_sg:  f32;

// One workgroup, ALL-parallel: softmax reductions, then k rounds of an
// argmax reduce (the selected logit is neutralized between rounds). A
// serial one-thread top-k here measured ~270 ns per L2-latency-bound
// probe — 22 ms/token across 40 layers at k=8, the whole decode wall.
// Ties pick the LOWEST index, matching the CPU scan. n_exp ≤ 256.
@compute @workgroup_size(256)
fn moe_select(@builtin(local_invocation_index) lid: u32) {
    // Shared-expert gate first: the reduction scratch below is reused, so
    // this has to land in ms_sg before the router work starts.
    let sg_kind = ms_p.pk & 0xFFu;
    let sg_hidden = ms_p.pk >> 8u;
    if (sg_kind == 4u) {
        var d = 0.0;
        var i = lid;
        loop {
            if (i >= sg_hidden) { break; }
            d = d + bitcast<f32>(ms_sgw[i]) * ms_x[i];
            i = i + 256u;
        }
        ms_red[lid] = d;
        workgroupBarrier();
        var st = 128u;
        loop {
            if (st == 0u) { break; }
            if (lid < st) { ms_red[lid] = ms_red[lid] + ms_red[lid + st]; }
            workgroupBarrier();
            st = st >> 1u;
        }
        if (lid == 0u) { ms_sg = ms_red[0]; }
    } else {
        if (lid == 0u) { ms_sg = ms_slog[0]; }
    }
    workgroupBarrier();
    let n = ms_p.n_exp;
    var v = -3.0e38;
    if (lid < n) { v = ms_logit[lid]; }
    ms_lg[lid] = v;
    ms_red[lid] = v;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { ms_red[lid] = max(ms_red[lid], ms_red[lid + stride]); }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    let mx = ms_red[0];
    workgroupBarrier();
    ms_red[lid] = select(0.0, exp(v - mx), lid < n);
    workgroupBarrier();
    stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { ms_red[lid] = ms_red[lid] + ms_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    let denom = ms_red[0];
    workgroupBarrier();
    // norm word: bit0 renorm, bit1 sigmoid scores, bit2 selection bias,
    // bit3 shared expert present. Softmax models pass 0/1, so every bit
    // pattern the old contract could send decodes to the old behaviour.
    let mflags = ms_p.norm;
    let msig = (mflags & 2u) != 0u;
    var msc = 0.0;
    if (lid < n) {
        if (msig) {
            msc = 1.0 / (1.0 + exp(-v));
        } else {
            msc = exp(v - mx) / denom;
        }
    }
    ms_sc[lid] = msc;
    if (msig) {
        var mkey = msc;
        if ((mflags & 4u) != 0u && lid < n) { mkey = mkey + ms_bias[lid]; }
        ms_lg[lid] = select(-3.0e38, mkey, lid < n);
    }
    workgroupBarrier();
    let k = ms_p.top_k;
    var wsum = 0.0;
    for (var slot = 0u; slot < k; slot = slot + 1u) {
        ms_red[lid] = ms_lg[lid];
        ms_ri[lid] = lid;
        workgroupBarrier();
        stride = 128u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) {
                let a = ms_red[lid];
                let b = ms_red[lid + stride];
                let ia = ms_ri[lid];
                let ib = ms_ri[lid + stride];
                if (b > a || (b == a && ib < ia)) {
                    ms_red[lid] = b;
                    ms_ri[lid] = ib;
                }
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) {
            let bi = ms_ri[0];
            ms_sel[slot] = bi;
            ms_w[slot] = ms_sc[bi];
            ms_pick = bi;
        }
        workgroupBarrier();
        wsum = wsum + ms_sc[ms_ri[0]];
        if (lid == ms_pick) { ms_lg[lid] = -3.0e38; }
        workgroupBarrier();
    }
    if (lid == 0u) {
        if ((mflags & 1u) != 0u) {
            // LFM2 floors the denominator (HF's + 1e-6); softmax probs
            // already sum near 1 and keep the bare sum, as before.
            let dn = select(wsum, wsum + 1e-6, msig);
            for (var slot = 0u; slot < k; slot = slot + 1u) { ms_w[slot] = ms_w[slot] / dn; }
        }
        // routed_scaling_factor (DeepSeek-V3 lineage, HunYuan hy_v3 2.826):
        // on the ROUTED mix only; the shared expert below is not scaled.
        // Every other family passes 1.0, which changes no bit.
        for (var slot = 0u; slot < k; slot = slot + 1u) { ms_w[slot] = ms_w[slot] * ms_p.scale; }
        if ((mflags & 8u) != 0u) {
            ms_sel[k] = n;
            // bit4: the shared expert has NO gate (hy_v3) — weight 1.
            if ((mflags & 16u) != 0u) {
                ms_w[k] = 1.0;
            } else {
                ms_w[k] = 1.0 / (1.0 + exp(-ms_sg));
            }
        }
    }
}

// gate+up+SiLU for every selected expert: workgroup (row, slot) does BOTH q4t
// row dots (they share the activation reads) and writes act = silu(g)·u.
struct MoeGuP { gpr: u32, inter: u32, slots: u32, mat16: u32 , lim: f32, _p0: u32, _p1: u32, _p2: u32 };
// Three scalars, not a vec3: a vec3<u32> aligns to 16 in uniform layout and
// pushes the struct to 48 bytes, while the buffer handed in is 32.
// `lim` is DeepSeek-V4's swiglu_limit: the up projection is clamped both
// ways, the gate only from above. Zero means no clamp, which is every other
// architecture that reaches these kernels.
@group(0) @binding(0) var<storage, read>       mg_gw  : array<u32>;
@group(0) @binding(1) var<storage, read>       mg_uw  : array<u32>;
@group(0) @binding(2) var<storage, read>       mg_x   : array<f32>;
@group(0) @binding(3) var<storage, read>       mg_sel : array<u32>;
@group(0) @binding(4) var<storage, read_write> mg_act : array<f32>;
@group(0) @binding(5) var<uniform>             mg_p   : MoeGuP;
var<workgroup> mg_pg: array<f32, 64>;
var<workgroup> mg_pu: array<f32, 64>;

// The same activations as `mg_x`, at the SAME SLOT, seen as vec4. The
// scalar view costs one load per weight, which left the dense FFN kernel
// at ~11% of the card's bandwidth until the vec4 rewrite (+2.7x there).
//
// A second global on slot 2 rather than a new slot 6, because an auto
// layout lists only the bindings its entry point actually USES: the q2tp
// kernel stopped touching `mg_x`, naga dropped slot 2, and the 7-entry
// bind group met a 6-entry layout. Same slot = the bind group is
// unchanged for every kernel here.
@group(0) @binding(2) var<storage, read> mg_xv : array<vec4<f32>>;

fn mg_g16(o: u32) -> u32 { return (mg_gw[o >> 1u] >> ((o & 1u) * 16u)) & 0xFFFFu; }
fn mg_u16f(o: u32) -> u32 { return (mg_uw[o >> 1u] >> ((o & 1u) * 16u)) & 0xFFFFu; }
fn mg_dot8(w: u32, xi: u32) -> f32 {
    return (f32(w & 0xFu) - 8.0) * mg_x[xi]
         + (f32((w >> 4u) & 0xFu) - 8.0) * mg_x[xi + 1u]
         + (f32((w >> 8u) & 0xFu) - 8.0) * mg_x[xi + 2u]
         + (f32((w >> 12u) & 0xFu) - 8.0) * mg_x[xi + 3u]
         + (f32((w >> 16u) & 0xFu) - 8.0) * mg_x[xi + 4u]
         + (f32((w >> 20u) & 0xFu) - 8.0) * mg_x[xi + 5u]
         + (f32((w >> 24u) & 0xFu) - 8.0) * mg_x[xi + 6u]
         + (f32((w >> 28u) & 0xFu) - 8.0) * mg_x[xi + 7u];
}

@compute @workgroup_size(64)
fn moe_gate_up(@builtin(workgroup_id) wid: vec3<u32>,
               @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let slot = wid.y;
    let gpr = mg_p.gpr;
    let base = mg_sel[slot] * mg_p.mat16 + row * gpr * 9u;
    var ag = 0.0;
    var au = 0.0;
    for (var g = lid; g < gpr; g = g + 64u) {
        let t16 = base + g * 9u;
        let sg = unpack2x16float(mg_g16(t16)).x;
        let su = unpack2x16float(mg_u16f(t16)).x;
        let xb = g * 32u;
        var dg = 0.0;
        var du = 0.0;
        for (var k = 0u; k < 4u; k = k + 1u) {
            let wg = mg_g16(t16 + 1u + 2u * k) | (mg_g16(t16 + 2u + 2u * k) << 16u);
            let wu = mg_u16f(t16 + 1u + 2u * k) | (mg_u16f(t16 + 2u + 2u * k) << 16u);
            dg = dg + mg_dot8(wg, xb + 8u * k);
            du = du + mg_dot8(wu, xb + 8u * k);
        }
        ag = ag + sg * dg;
        au = au + su * du;
    }
    mg_pg[lid] = ag;
    mg_pu[lid] = au;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) {
            mg_pg[lid] = mg_pg[lid] + mg_pg[lid + stride];
            mg_pu[lid] = mg_pu[lid] + mg_pu[lid + stride];
        }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        let g = mg_pg[0];
        var gg = g;
        var uu = mg_pu[0];
        if (mg_p.lim > 0.0) {
            uu = clamp(uu, -mg_p.lim, mg_p.lim);
            gg = min(gg, mg_p.lim);
        }
        mg_act[slot * mg_p.inter + row] = (gg / (1.0 + exp(-gg))) * uu;
    }
}

// Weighted down-projection: one workgroup per hidden row accumulates
// Σ_slot w[slot]·(down[sel[slot]] row · act[slot]) over the flattened
// (slot, group) space, then overwrites y[row] (the graph's usual FFN
// output slot — the existing fused residual add consumes it).
struct MoeDnP { gpr: u32, hidden: u32, slots: u32, mat16: u32 };
@group(0) @binding(0) var<storage, read>       md_w   : array<u32>;
@group(0) @binding(1) var<storage, read>       md_act : array<f32>;
@group(0) @binding(2) var<storage, read>       md_sel : array<u32>;
@group(0) @binding(3) var<storage, read>       md_wt  : array<f32>;
@group(0) @binding(4) var<storage, read_write> md_y   : array<f32>;
@group(0) @binding(5) var<uniform>             md_p   : MoeDnP;
var<workgroup> md_pt: array<f32, 64>;

// The activations again as vec4, on the SAME slot — the scalar view costs
// one load per weight and put the down kernel at 33 us for ~5 MB of reads.
@group(0) @binding(1) var<storage, read> md_actv : array<vec4<f32>>;

fn md_u16(o: u32) -> u32 { return (md_w[o >> 1u] >> ((o & 1u) * 16u)) & 0xFFFFu; }
fn md_dot8v(w: u32, a: vec4<f32>, b: vec4<f32>) -> f32 {
    return (f32(w & 0xFu) - 8.0) * a.x
         + (f32((w >> 4u) & 0xFu) - 8.0) * a.y
         + (f32((w >> 8u) & 0xFu) - 8.0) * a.z
         + (f32((w >> 12u) & 0xFu) - 8.0) * a.w
         + (f32((w >> 16u) & 0xFu) - 8.0) * b.x
         + (f32((w >> 20u) & 0xFu) - 8.0) * b.y
         + (f32((w >> 24u) & 0xFu) - 8.0) * b.z
         + (f32((w >> 28u) & 0xFu) - 8.0) * b.w;
}
fn md_dot8(w: u32, xi: u32) -> f32 {
    return (f32(w & 0xFu) - 8.0) * md_act[xi]
         + (f32((w >> 4u) & 0xFu) - 8.0) * md_act[xi + 1u]
         + (f32((w >> 8u) & 0xFu) - 8.0) * md_act[xi + 2u]
         + (f32((w >> 12u) & 0xFu) - 8.0) * md_act[xi + 3u]
         + (f32((w >> 16u) & 0xFu) - 8.0) * md_act[xi + 4u]
         + (f32((w >> 20u) & 0xFu) - 8.0) * md_act[xi + 5u]
         + (f32((w >> 24u) & 0xFu) - 8.0) * md_act[xi + 6u]
         + (f32((w >> 28u) & 0xFu) - 8.0) * md_act[xi + 7u];
}

@compute @workgroup_size(64)
fn moe_down(@builtin(workgroup_id) wid: vec3<u32>,
            @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let gpr = md_p.gpr;
    let total = md_p.slots * gpr;
    var acc = 0.0;
    for (var i = lid; i < total; i = i + 64u) {
        let slot = i / gpr;
        let g = i % gpr;
        let t16 = md_sel[slot] * md_p.mat16 + (row * gpr + g) * 9u;
        let scale = unpack2x16float(md_u16(t16)).x;
        let xb = (slot * gpr + g) * 32u;
        var d = 0.0;
        for (var k = 0u; k < 4u; k = k + 1u) {
            let w = md_u16(t16 + 1u + 2u * k) | (md_u16(t16 + 2u + 2u * k) << 16u);
            d = d + md_dot8(w, xb + 8u * k);
        }
        acc = acc + md_wt[slot] * scale * d;
    }
    md_pt[lid] = acc;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { md_pt[lid] = md_pt[lid] + md_pt[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { md_y[row] = md_pt[0]; }
}

// ── q4tp twins of the two MoE kernels. Identical nibble math and
// identical bindings; only where the scale comes from differs. q4t
// carries an f16 scale inside each 18-byte tile, q4tp packs the nibbles
// 16-byte tight and puts a 5-bit rung index into a side plane, read
// against the row's geometric ladder `2^(lo + code·step)`. Per expert
// the blob is [nibbles | row params (f16 lo, f16 step) | 5-bit codes],
// so the two extra plane offsets fall out of rows/gpr.
fn mgp_gu8(o: u32) -> u32 { return (mg_gw[o >> 2u] >> ((o & 3u) * 8u)) & 0xFFu; }
fn mgp_uu8(o: u32) -> u32 { return (mg_uw[o >> 2u] >> ((o & 3u) * 8u)) & 0xFFu; }

@compute @workgroup_size(64)
fn moe_gate_up_q4tp(@builtin(workgroup_id) wid: vec3<u32>,
                    @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let slot = wid.y;
    let gpr = mg_p.gpr;
    let rows = mg_p.inter;
    let base16 = mg_sel[slot] * mg_p.mat16;
    let nib16 = base16 + row * gpr * 8u;
    let par16 = base16 + rows * gpr * 8u + row * 2u;
    let cst = (gpr * 5u + 7u) / 8u;
    let cod8 = (base16 + rows * gpr * 8u + rows * 2u) * 2u + row * cst;

    let gl = unpack2x16float(mg_g16(par16) | (mg_g16(par16 + 1u) << 16u));
    let ul = unpack2x16float(mg_u16f(par16) | (mg_u16f(par16 + 1u) << 16u));
    var ag = 0.0;
    var au = 0.0;
    for (var g = lid; g < gpr; g = g + 64u) {
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cg = mgp_gu8(cod8 + cb);
        var cu = mgp_uu8(cod8 + cb);
        // A 5-bit field starting past bit 3 spills into the next byte.
        if (shf > 3u) {
            cg = cg | (mgp_gu8(cod8 + cb + 1u) << 8u);
            cu = cu | (mgp_uu8(cod8 + cb + 1u) << 8u);
        }
        let sg = exp2(gl.x + f32((cg >> shf) & 31u) * gl.y);
        let su = exp2(ul.x + f32((cu >> shf) & 31u) * ul.y);
        let t16 = nib16 + g * 8u;
        let xb = g * 32u;
        var dg = 0.0;
        var du = 0.0;
        for (var k = 0u; k < 4u; k = k + 1u) {
            let wg = mg_g16(t16 + 2u * k) | (mg_g16(t16 + 1u + 2u * k) << 16u);
            let wu = mg_u16f(t16 + 2u * k) | (mg_u16f(t16 + 1u + 2u * k) << 16u);
            dg = dg + mg_dot8(wg, xb + 8u * k);
            du = du + mg_dot8(wu, xb + 8u * k);
        }
        ag = ag + sg * dg;
        au = au + su * du;
    }
    mg_pg[lid] = ag;
    mg_pu[lid] = au;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) {
            mg_pg[lid] = mg_pg[lid] + mg_pg[lid + stride];
            mg_pu[lid] = mg_pu[lid] + mg_pu[lid + stride];
        }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        let g = mg_pg[0];
        var gg = g;
        var uu = mg_pu[0];
        if (mg_p.lim > 0.0) {
            uu = clamp(uu, -mg_p.lim, mg_p.lim);
            gg = min(gg, mg_p.lim);
        }
        mg_act[slot * mg_p.inter + row] = (gg / (1.0 + exp(-gg))) * uu;
    }
}

// q2tp gate/up: the q4tp kernel with a 2-bit weight plane. A group is
// 32 weights in 8 bytes (4 u16 units) instead of 16, and one u32 carries
// SIXTEEN weights, so the group is two words and two dot16s. The params
// and 5-bit code planes are byte-identical to q4tp — only the plane
// offsets move, since they sit behind a half-size weight plane.
// 16 two-bit weights against four staged vec4s — same add order as the
// scalar mg_dot16 below, which greedy parity depends on.
fn mg_dot16v(w: u32, a: vec4<f32>, b: vec4<f32>, c: vec4<f32>, d: vec4<f32>) -> f32 {
    return (f32(w & 3u) - 1.5) * a.x
         + (f32((w >> 2u) & 3u) - 1.5) * a.y
         + (f32((w >> 4u) & 3u) - 1.5) * a.z
         + (f32((w >> 6u) & 3u) - 1.5) * a.w
         + (f32((w >> 8u) & 3u) - 1.5) * b.x
         + (f32((w >> 10u) & 3u) - 1.5) * b.y
         + (f32((w >> 12u) & 3u) - 1.5) * b.z
         + (f32((w >> 14u) & 3u) - 1.5) * b.w
         + (f32((w >> 16u) & 3u) - 1.5) * c.x
         + (f32((w >> 18u) & 3u) - 1.5) * c.y
         + (f32((w >> 20u) & 3u) - 1.5) * c.z
         + (f32((w >> 22u) & 3u) - 1.5) * c.w
         + (f32((w >> 24u) & 3u) - 1.5) * d.x
         + (f32((w >> 26u) & 3u) - 1.5) * d.y
         + (f32((w >> 28u) & 3u) - 1.5) * d.z
         + (f32((w >> 30u) & 3u) - 1.5) * d.w;
}

fn mg_dot16(w: u32, xi: u32) -> f32 {
    return (f32(w & 3u) - 1.5) * mg_x[xi]
         + (f32((w >> 2u) & 3u) - 1.5) * mg_x[xi + 1u]
         + (f32((w >> 4u) & 3u) - 1.5) * mg_x[xi + 2u]
         + (f32((w >> 6u) & 3u) - 1.5) * mg_x[xi + 3u]
         + (f32((w >> 8u) & 3u) - 1.5) * mg_x[xi + 4u]
         + (f32((w >> 10u) & 3u) - 1.5) * mg_x[xi + 5u]
         + (f32((w >> 12u) & 3u) - 1.5) * mg_x[xi + 6u]
         + (f32((w >> 14u) & 3u) - 1.5) * mg_x[xi + 7u]
         + (f32((w >> 16u) & 3u) - 1.5) * mg_x[xi + 8u]
         + (f32((w >> 18u) & 3u) - 1.5) * mg_x[xi + 9u]
         + (f32((w >> 20u) & 3u) - 1.5) * mg_x[xi + 10u]
         + (f32((w >> 22u) & 3u) - 1.5) * mg_x[xi + 11u]
         + (f32((w >> 24u) & 3u) - 1.5) * mg_x[xi + 12u]
         + (f32((w >> 26u) & 3u) - 1.5) * mg_x[xi + 13u]
         + (f32((w >> 28u) & 3u) - 1.5) * mg_x[xi + 14u]
         + (f32((w >> 30u) & 3u) - 1.5) * mg_x[xi + 15u];
}

@compute @workgroup_size(64)
fn moe_gate_up_q2tp(@builtin(workgroup_id) wid: vec3<u32>,
                    @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let slot = wid.y;
    let gpr = mg_p.gpr;
    let rows = mg_p.inter;
    let base16 = mg_sel[slot] * mg_p.mat16;
    let nib16 = base16 + row * gpr * 4u;
    let par16 = base16 + rows * gpr * 4u + row * 2u;
    let cst = (gpr * 5u + 7u) / 8u;
    let cod8 = (base16 + rows * gpr * 4u + rows * 2u) * 2u + row * cst;

    let gl = unpack2x16float(mg_g16(par16) | (mg_g16(par16 + 1u) << 16u));
    let ul = unpack2x16float(mg_u16f(par16) | (mg_u16f(par16 + 1u) << 16u));
    var ag = 0.0;
    var au = 0.0;
    for (var g = lid; g < gpr; g = g + 64u) {
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cg = mgp_gu8(cod8 + cb);
        var cu = mgp_uu8(cod8 + cb);
        if (shf > 3u) {
            cg = cg | (mgp_gu8(cod8 + cb + 1u) << 8u);
            cu = cu | (mgp_uu8(cod8 + cb + 1u) << 8u);
        }
        // Rung 0 is the format's exact zero (the ±0.5/±1.5 grid has no
        // zero of its own); live rungs are the ladder shifted down one.
        let cgv = (cg >> shf) & 31u;
        let cuv = (cu >> shf) & 31u;
        let sg = select(exp2(gl.x + f32(max(cgv, 1u) - 1u) * gl.y), 0.0, cgv == 0u);
        let su = select(exp2(ul.x + f32(max(cuv, 1u) - 1u) * ul.y), 0.0, cuv == 0u);
        // Group base in u16 units is a multiple of 4, so the two 32-bit
        // words land on u32 lanes (nib16 >> 1) and (nib16 >> 1) + 1.
        let w32 = (nib16 + g * 4u) >> 1u;
        // One group = 32 activations = 8 vec4s, shared by gate and up.
        let xq = g * 8u;
        let x0 = mg_xv[xq];      let x1 = mg_xv[xq + 1u];
        let x2 = mg_xv[xq + 2u]; let x3 = mg_xv[xq + 3u];
        let x4 = mg_xv[xq + 4u]; let x5 = mg_xv[xq + 5u];
        let x6 = mg_xv[xq + 6u]; let x7 = mg_xv[xq + 7u];
        let dg = mg_dot16v(mg_gw[w32], x0, x1, x2, x3)
               + mg_dot16v(mg_gw[w32 + 1u], x4, x5, x6, x7);
        let du = mg_dot16v(mg_uw[w32], x0, x1, x2, x3)
               + mg_dot16v(mg_uw[w32 + 1u], x4, x5, x6, x7);
        ag = ag + sg * dg;
        au = au + su * du;
    }
    mg_pg[lid] = ag;
    mg_pu[lid] = au;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) {
            mg_pg[lid] = mg_pg[lid] + mg_pg[lid + stride];
            mg_pu[lid] = mg_pu[lid] + mg_pu[lid + stride];
        }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        let g = mg_pg[0];
        var gg = g;
        var uu = mg_pu[0];
        if (mg_p.lim > 0.0) {
            uu = clamp(uu, -mg_p.lim, mg_p.lim);
            gg = min(gg, mg_p.lim);
        }
        mg_act[slot * mg_p.inter + row] = (gg / (1.0 + exp(-gg))) * uu;
    }
}

// FOUR output rows to a workgroup, 64 lanes each. `gpr` is 128 on the
// release, so a row has no use for more than 64 lanes — but four rows give
// the memory system four more independent streams to overlap, and the MoE
// is 5.6 ms of a 33 ms chain against a bandwidth floor near 1.8.
var<workgroup> mgm_pg: array<f32, 256>;
var<workgroup> mgm_pu: array<f32, 256>;
@compute @workgroup_size(256)
fn moe_gate_up_q2tp_m(@builtin(workgroup_id) wid: vec3<u32>,
                    @builtin(local_invocation_index) lid: u32) {
    let sub = lid / 64u;
    let lane = lid % 64u;
    let row = wid.x * 4u + sub;
    let slot = wid.y;
    if (row >= mg_p.inter) { return; }
    let gpr = mg_p.gpr;
    let rows = mg_p.inter;
    let base16 = mg_sel[slot] * mg_p.mat16;
    let nib16 = base16 + row * gpr * 4u;
    let par16 = base16 + rows * gpr * 4u + row * 2u;
    let cst = (gpr * 5u + 7u) / 8u;
    let cod8 = (base16 + rows * gpr * 4u + rows * 2u) * 2u + row * cst;

    let gl = unpack2x16float(mg_g16(par16) | (mg_g16(par16 + 1u) << 16u));
    let ul = unpack2x16float(mg_u16f(par16) | (mg_u16f(par16 + 1u) << 16u));
    var ag = 0.0;
    var au = 0.0;
    for (var g = lane; g < gpr; g = g + 64u) {
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cg = mgp_gu8(cod8 + cb);
        var cu = mgp_uu8(cod8 + cb);
        if (shf > 3u) {
            cg = cg | (mgp_gu8(cod8 + cb + 1u) << 8u);
            cu = cu | (mgp_uu8(cod8 + cb + 1u) << 8u);
        }
        // Rung 0 is the format's exact zero (the ±0.5/±1.5 grid has no
        // zero of its own); live rungs are the ladder shifted down one.
        let cgv = (cg >> shf) & 31u;
        let cuv = (cu >> shf) & 31u;
        let sg = select(exp2(gl.x + f32(max(cgv, 1u) - 1u) * gl.y), 0.0, cgv == 0u);
        let su = select(exp2(ul.x + f32(max(cuv, 1u) - 1u) * ul.y), 0.0, cuv == 0u);
        // Group base in u16 units is a multiple of 4, so the two 32-bit
        // words land on u32 lanes (nib16 >> 1) and (nib16 >> 1) + 1.
        let w32 = (nib16 + g * 4u) >> 1u;
        // One group = 32 activations = 8 vec4s, shared by gate and up.
        let xq = g * 8u;
        let x0 = mg_xv[xq];      let x1 = mg_xv[xq + 1u];
        let x2 = mg_xv[xq + 2u]; let x3 = mg_xv[xq + 3u];
        let x4 = mg_xv[xq + 4u]; let x5 = mg_xv[xq + 5u];
        let x6 = mg_xv[xq + 6u]; let x7 = mg_xv[xq + 7u];
        let dg = mg_dot16v(mg_gw[w32], x0, x1, x2, x3)
               + mg_dot16v(mg_gw[w32 + 1u], x4, x5, x6, x7);
        let du = mg_dot16v(mg_uw[w32], x0, x1, x2, x3)
               + mg_dot16v(mg_uw[w32 + 1u], x4, x5, x6, x7);
        ag = ag + sg * dg;
        au = au + su * du;
    }
    mgm_pg[lid] = ag;
    mgm_pu[lid] = au;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lane < stride) {
            mgm_pg[lid] = mgm_pg[lid] + mgm_pg[lid + stride];
            mgm_pu[lid] = mgm_pu[lid] + mgm_pu[lid + stride];
        }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lane == 0u) {
        let g = mgm_pg[sub * 64u];
        var gg = g;
        var uu = mgm_pu[sub * 64u];
        if (mg_p.lim > 0.0) {
            uu = clamp(uu, -mg_p.lim, mg_p.lim);
            gg = min(gg, mg_p.lim);
        }
        mg_act[slot * mg_p.inter + row] = (gg / (1.0 + exp(-gg))) * uu;
    }
}

fn mdp_u8(o: u32) -> u32 { return (md_w[o >> 2u] >> ((o & 3u) * 8u)) & 0xFFu; }

@compute @workgroup_size(64)
fn moe_down_q4tp(@builtin(workgroup_id) wid: vec3<u32>,
                 @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let gpr = md_p.gpr;
    let rows = md_p.hidden;
    let cst = (gpr * 5u + 7u) / 8u;
    let total = md_p.slots * gpr;
    var acc = 0.0;
    for (var i = lid; i < total; i = i + 64u) {
        let slot = i / gpr;
        let g = i % gpr;
        let base16 = md_sel[slot] * md_p.mat16;
        let par16 = base16 + rows * gpr * 8u + row * 2u;
        let cod8 = (base16 + rows * gpr * 8u + rows * 2u) * 2u + row * cst;
        let pl = unpack2x16float(md_u16(par16) | (md_u16(par16 + 1u) << 16u));
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cv = mdp_u8(cod8 + cb);
        if (shf > 3u) { cv = cv | (mdp_u8(cod8 + cb + 1u) << 8u); }
        let scale = exp2(pl.x + f32((cv >> shf) & 31u) * pl.y);
        let t16 = base16 + (row * gpr + g) * 8u;
        let xq = (slot * gpr + g) * 8u;
        let x0 = md_actv[xq];      let x1 = md_actv[xq + 1u];
        let x2 = md_actv[xq + 2u]; let x3 = md_actv[xq + 3u];
        let x4 = md_actv[xq + 4u]; let x5 = md_actv[xq + 5u];
        let x6 = md_actv[xq + 6u]; let x7 = md_actv[xq + 7u];
        let w0 = md_u16(t16) | (md_u16(t16 + 1u) << 16u);
        let w1 = md_u16(t16 + 2u) | (md_u16(t16 + 3u) << 16u);
        let w2 = md_u16(t16 + 4u) | (md_u16(t16 + 5u) << 16u);
        let w3 = md_u16(t16 + 6u) | (md_u16(t16 + 7u) << 16u);
        let d = md_dot8v(w0, x0, x1) + md_dot8v(w1, x2, x3)
              + md_dot8v(w2, x4, x5) + md_dot8v(w3, x6, x7);
        acc = acc + md_wt[slot] * scale * d;
    }
    md_pt[lid] = acc;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { md_pt[lid] = md_pt[lid] + md_pt[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { md_y[row] = md_pt[0]; }
}

// 256 threads. The loop walks slots·gpr = 576 terms on the release, which
// at 64 threads is nine iterations of dependent loads per thread; at 256 it
// is two and a bit, and the workgroup count (one per hidden row, 4096) was
// never the problem.
var<workgroup> mdm_pt: array<f32, 256>;
@compute @workgroup_size(256)
fn moe_down_q4tp_m(@builtin(workgroup_id) wid: vec3<u32>,
                 @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let gpr = md_p.gpr;
    let rows = md_p.hidden;
    let cst = (gpr * 5u + 7u) / 8u;
    let total = md_p.slots * gpr;
    var acc = 0.0;
    for (var i = lid; i < total; i = i + 256u) {
        let slot = i / gpr;
        let g = i % gpr;
        let base16 = md_sel[slot] * md_p.mat16;
        let par16 = base16 + rows * gpr * 8u + row * 2u;
        let cod8 = (base16 + rows * gpr * 8u + rows * 2u) * 2u + row * cst;
        let pl = unpack2x16float(md_u16(par16) | (md_u16(par16 + 1u) << 16u));
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cv = mdp_u8(cod8 + cb);
        if (shf > 3u) { cv = cv | (mdp_u8(cod8 + cb + 1u) << 8u); }
        let scale = exp2(pl.x + f32((cv >> shf) & 31u) * pl.y);
        let t16 = base16 + (row * gpr + g) * 8u;
        let xq = (slot * gpr + g) * 8u;
        let x0 = md_actv[xq];      let x1 = md_actv[xq + 1u];
        let x2 = md_actv[xq + 2u]; let x3 = md_actv[xq + 3u];
        let x4 = md_actv[xq + 4u]; let x5 = md_actv[xq + 5u];
        let x6 = md_actv[xq + 6u]; let x7 = md_actv[xq + 7u];
        let w0 = md_u16(t16) | (md_u16(t16 + 1u) << 16u);
        let w1 = md_u16(t16 + 2u) | (md_u16(t16 + 3u) << 16u);
        let w2 = md_u16(t16 + 4u) | (md_u16(t16 + 5u) << 16u);
        let w3 = md_u16(t16 + 6u) | (md_u16(t16 + 7u) << 16u);
        let d = md_dot8v(w0, x0, x1) + md_dot8v(w1, x2, x3)
              + md_dot8v(w2, x4, x5) + md_dot8v(w3, x6, x7);
        acc = acc + md_wt[slot] * scale * d;
    }
    mdm_pt[lid] = acc;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { mdm_pt[lid] = mdm_pt[lid] + mdm_pt[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { md_y[row] = mdm_pt[0]; }
}

// ── Batched MoE: the whole layer's routing and experts with a TOKEN
// dimension, for the batch (prefill) graph. The per-token encoding of
// this block cost ~7 commands per token per layer — at k=32 over 40
// layers that is ~9000 encoder commands a chunk, and the chunk clocked
// at the same 16 ms/token as the per-position path it was meant to
// beat. These three kernels replace all of it with THREE dispatches per
// layer and zero buffer-to-buffer copies.
//
// The router matvec lives INSIDE the select kernel: one thread = one
// expert row (n_exp <= 256 = workgroup size), reading x straight from
// the batch hidden at its token offset. No logits buffer, no row
// staging. f32 router weights only — the converter leaves the router
// unquantized; anything else falls back to the per-token path.
struct MoeSelBP { n_exp: u32, top_k: u32, norm: u32, pk: u32, scale: f32, _s0: u32, _s1: u32, _s2: u32 };
@group(0) @binding(0) var<storage, read>       sb_lgin: array<f32>;
@group(0) @binding(1) var<storage, read>       sb_x   : array<f32>;
@group(0) @binding(2) var<storage, read_write> sb_sel : array<u32>;
@group(0) @binding(3) var<storage, read_write> sb_w   : array<f32>;
@group(0) @binding(4) var<uniform>             sb_p   : MoeSelBP;
@group(0) @binding(5) var<storage, read>       sb_sgw : array<u32>;
// Per-expert SELECTION bias (noaux_tc), a 4-byte dummy when absent.
@group(0) @binding(6) var<storage, read>       sb_bias: array<f32>;
var<workgroup> sb_lg:  array<f32, 256>;
var<workgroup> sb_red: array<f32, 256>;
var<workgroup> sb_ri:  array<u32, 256>;
var<workgroup> sb_sg:  f32;
// The mixing score per expert (softmax prob or sigmoid); the RANKING key
// in sb_lg may carry the selection bias, the weight never does.
var<workgroup> sb_sc:  array<f32, 256>;

@compute @workgroup_size(256)
fn moe_select_b(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let t = wid.x;
    let n = sb_p.n_exp;
    let sg_kind = sb_p.pk & 0xFFu;
    let hidden = sb_p.pk >> 8u;
    let xb = t * hidden;
    // Logits arrive from the same f32 matvec kernel the parity-proven
    // path used, one slice per token — computing them here with a
    // different summation order shifted near-tied experts and broke
    // token parity with the CPU.
    var v = -3.0e38;
    if (lid < n) { v = sb_lgin[t * n + lid]; }
    // Shared-expert gate on this token's x (f32 weights, same fold as
    // the single-token kernel).
    if (sg_kind == 4u) {
        var d = 0.0;
        var i = lid;
        loop {
            if (i >= hidden) { break; }
            d = d + bitcast<f32>(sb_sgw[i]) * sb_x[xb + i];
            i = i + 256u;
        }
        sb_red[lid] = d;
        workgroupBarrier();
        var st = 128u;
        loop {
            if (st == 0u) { break; }
            if (lid < st) { sb_red[lid] = sb_red[lid] + sb_red[lid + st]; }
            workgroupBarrier();
            st = st >> 1u;
        }
        if (lid == 0u) { sb_sg = sb_red[0]; }
        workgroupBarrier();
    }
    sb_lg[lid] = v;
    sb_red[lid] = v;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { sb_red[lid] = max(sb_red[lid], sb_red[lid + stride]); }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    let mx = sb_red[0];
    workgroupBarrier();
    sb_red[lid] = select(0.0, exp(v - mx), lid < n);
    workgroupBarrier();
    stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { sb_red[lid] = sb_red[lid] + sb_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    let denom = sb_red[0];
    workgroupBarrier();
    // norm word: bit0 renorm, bit1 sigmoid scores, bit2 selection bias,
    // bit4 ungated shared expert — the single-token kernel's contract.
    let mflags = sb_p.norm;
    let msig = (mflags & 2u) != 0u;
    var msc = 0.0;
    if (lid < n) {
        if (msig) {
            msc = 1.0 / (1.0 + exp(-v));
        } else {
            msc = exp(v - mx) / denom;
        }
    }
    sb_sc[lid] = msc;
    if (msig) {
        var mkey = msc;
        if ((mflags & 4u) != 0u && lid < n) { mkey = mkey + sb_bias[lid]; }
        sb_lg[lid] = select(-3.0e38, mkey, lid < n);
    }
    workgroupBarrier();
    let kk = sb_p.top_k;
    let ob = t * (kk + 1u);
    var wsum = 0.0;
    for (var slot = 0u; slot < kk; slot = slot + 1u) {
        sb_red[lid] = sb_lg[lid];
        sb_ri[lid] = lid;
        workgroupBarrier();
        stride = 128u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) {
                let a = sb_red[lid];
                let b = sb_red[lid + stride];
                let ia = sb_ri[lid];
                let ib = sb_ri[lid + stride];
                if (b > a || (b == a && ib < ia)) {
                    sb_red[lid] = b;
                    sb_ri[lid] = ib;
                }
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) {
            let bi = sb_ri[0];
            sb_sel[ob + slot] = bi;
            sb_w[ob + slot] = sb_sc[bi];
        }
        workgroupBarrier();
        wsum = wsum + sb_sc[sb_ri[0]];
        if (lid == sb_ri[0]) { sb_lg[lid] = -3.0e38; }
        workgroupBarrier();
    }
    if (lid == 0u) {
        if ((mflags & 1u) != 0u) {
            let dn = select(wsum, wsum + 1e-6, msig);
            for (var slot = 0u; slot < kk; slot = slot + 1u) {
                sb_w[ob + slot] = sb_w[ob + slot] / dn;
            }
        }
        for (var slot = 0u; slot < kk; slot = slot + 1u) {
            sb_w[ob + slot] = sb_w[ob + slot] * sb_p.scale;
        }
        sb_sel[ob + kk] = n;
        if ((mflags & 16u) != 0u) {
            sb_w[ob + kk] = 1.0;
        } else {
            sb_w[ob + kk] = 1.0 / (1.0 + exp(-sb_sg));
        }
    }
}

// gate+up+SiLU with a token axis: workgroup (row, slot, token).
struct MoeGuBP { gpr: u32, inter: u32, slots: u32, mat16: u32 , lim: f32, _p0: u32, _p1: u32, _p2: u32 };
// Three scalars, not a vec3: a vec3<u32> aligns to 16 in uniform layout and
// pushes the struct to 48 bytes, while the buffer handed in is 32.
// `lim` is DeepSeek-V4's swiglu_limit: the up projection is clamped both
// ways, the gate only from above. Zero means no clamp, which is every other
// architecture that reaches these kernels.
@group(0) @binding(0) var<storage, read>       gb_gw  : array<u32>;
@group(0) @binding(1) var<storage, read>       gb_uw  : array<u32>;
@group(0) @binding(2) var<storage, read>       gb_x   : array<f32>;
@group(0) @binding(3) var<storage, read>       gb_sel : array<u32>;
@group(0) @binding(4) var<storage, read_write> gb_act : array<f32>;
@group(0) @binding(5) var<uniform>             gb_p   : MoeGuBP;
var<workgroup> gb_pg: array<f32, 64>;
var<workgroup> gb_pu: array<f32, 64>;

fn gb_g16(o: u32) -> u32 { return (gb_gw[o >> 1u] >> ((o & 1u) * 16u)) & 0xFFFFu; }
fn gb_u16(o: u32) -> u32 { return (gb_uw[o >> 1u] >> ((o & 1u) * 16u)) & 0xFFFFu; }
fn gb_gu8(o: u32) -> u32 { return (gb_gw[o >> 2u] >> ((o & 3u) * 8u)) & 0xFFu; }
fn gb_uu8(o: u32) -> u32 { return (gb_uw[o >> 2u] >> ((o & 3u) * 8u)) & 0xFFu; }
fn gb_dot8(w: u32, xi: u32) -> f32 {
    return (f32(w & 0xFu) - 8.0) * gb_x[xi]
         + (f32((w >> 4u) & 0xFu) - 8.0) * gb_x[xi + 1u]
         + (f32((w >> 8u) & 0xFu) - 8.0) * gb_x[xi + 2u]
         + (f32((w >> 12u) & 0xFu) - 8.0) * gb_x[xi + 3u]
         + (f32((w >> 16u) & 0xFu) - 8.0) * gb_x[xi + 4u]
         + (f32((w >> 20u) & 0xFu) - 8.0) * gb_x[xi + 5u]
         + (f32((w >> 24u) & 0xFu) - 8.0) * gb_x[xi + 6u]
         + (f32((w >> 28u) & 0xFu) - 8.0) * gb_x[xi + 7u];
}

@compute @workgroup_size(64)
fn moe_gate_up_q4tp_b(@builtin(workgroup_id) wid: vec3<u32>,
                      @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let slot = wid.y;
    let t = wid.z;
    let gpr = gb_p.gpr;
    let rows = gb_p.inter;
    let hidden = gpr * 32u;
    let xoff = t * hidden;
    let base16 = gb_sel[t * gb_p.slots + slot] * gb_p.mat16;
    let nib16 = base16 + row * gpr * 8u;
    let par16 = base16 + rows * gpr * 8u + row * 2u;
    let cst = (gpr * 5u + 7u) / 8u;
    let cod8 = (base16 + rows * gpr * 8u + rows * 2u) * 2u + row * cst;
    let gl = unpack2x16float(gb_g16(par16) | (gb_g16(par16 + 1u) << 16u));
    let ul = unpack2x16float(gb_u16(par16) | (gb_u16(par16 + 1u) << 16u));
    var ag = 0.0;
    var au = 0.0;
    for (var g = lid; g < gpr; g = g + 64u) {
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cg = gb_gu8(cod8 + cb);
        var cu = gb_uu8(cod8 + cb);
        if (shf > 3u) {
            cg = cg | (gb_gu8(cod8 + cb + 1u) << 8u);
            cu = cu | (gb_uu8(cod8 + cb + 1u) << 8u);
        }
        let sg = exp2(gl.x + f32((cg >> shf) & 31u) * gl.y);
        let su = exp2(ul.x + f32((cu >> shf) & 31u) * ul.y);
        let t16 = nib16 + g * 8u;
        let xb = xoff + g * 32u;
        var dg = 0.0;
        var du = 0.0;
        for (var k = 0u; k < 4u; k = k + 1u) {
            let wg = gb_g16(t16 + 2u * k) | (gb_g16(t16 + 1u + 2u * k) << 16u);
            let wu = gb_u16(t16 + 2u * k) | (gb_u16(t16 + 1u + 2u * k) << 16u);
            dg = dg + gb_dot8(wg, xb + 8u * k);
            du = du + gb_dot8(wu, xb + 8u * k);
        }
        ag = ag + sg * dg;
        au = au + su * du;
    }
    gb_pg[lid] = ag;
    gb_pu[lid] = au;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) {
            gb_pg[lid] = gb_pg[lid] + gb_pg[lid + stride];
            gb_pu[lid] = gb_pu[lid] + gb_pu[lid + stride];
        }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        let g = gb_pg[0];
        var gg = g;
        var uu = gb_pu[0];
        if (gb_p.lim > 0.0) {
            uu = clamp(uu, -gb_p.lim, gb_p.lim);
            gg = min(gg, gb_p.lim);
        }
        gb_act[(t * gb_p.slots + slot) * gb_p.inter + row] =
            (gg / (1.0 + exp(-gg))) * uu;
    }
}

// Four output rows share one activation load. Q4TP gate/up is unusually
// activation-traffic heavy: every row used to reread the same 4096-float
// vector for both matrices. The per-row accumulation and reduction order
// stay identical to `moe_gate_up_q4tp_b`; only the loads are hoisted.
var<workgroup> gb4_pg: array<f32, 64>;
var<workgroup> gb4_pu: array<f32, 64>;

fn gb_dot8v(w: u32, a: vec4<f32>, b: vec4<f32>) -> f32 {
    return (f32(w & 0xFu) - 8.0) * a.x
         + (f32((w >> 4u) & 0xFu) - 8.0) * a.y
         + (f32((w >> 8u) & 0xFu) - 8.0) * a.z
         + (f32((w >> 12u) & 0xFu) - 8.0) * a.w
         + (f32((w >> 16u) & 0xFu) - 8.0) * b.x
         + (f32((w >> 20u) & 0xFu) - 8.0) * b.y
         + (f32((w >> 24u) & 0xFu) - 8.0) * b.z
         + (f32((w >> 28u) & 0xFu) - 8.0) * b.w;
}

@compute @workgroup_size(64)
fn moe_gate_up_q4tp_b_r4(@builtin(workgroup_id) wid: vec3<u32>,
                         @builtin(local_invocation_index) lid: u32) {
    let row0 = wid.x * 4u;
    let slot = wid.y;
    let t = wid.z;
    let gpr = gb_p.gpr;
    let rows = gb_p.inter;
    let hidden = gpr * 32u;
    let xoff = t * hidden;
    let base16 = gb_sel[t * gb_p.slots + slot] * gb_p.mat16;
    let cst = (gpr * 5u + 7u) / 8u;
    let cod0 = (base16 + rows * gpr * 8u + rows * 2u) * 2u;
    var ag0 = 0.0; var au0 = 0.0;
    var ag1 = 0.0; var au1 = 0.0;
    var ag2 = 0.0; var au2 = 0.0;
    var ag3 = 0.0; var au3 = 0.0;
    for (var g = lid; g < gpr; g = g + 64u) {
        let xb = xoff + g * 32u;
        let x0a = vec4<f32>(gb_x[xb],      gb_x[xb + 1u],  gb_x[xb + 2u],  gb_x[xb + 3u]);
        let x0b = vec4<f32>(gb_x[xb + 4u], gb_x[xb + 5u],  gb_x[xb + 6u],  gb_x[xb + 7u]);
        let x1a = vec4<f32>(gb_x[xb + 8u], gb_x[xb + 9u],  gb_x[xb + 10u], gb_x[xb + 11u]);
        let x1b = vec4<f32>(gb_x[xb + 12u],gb_x[xb + 13u], gb_x[xb + 14u], gb_x[xb + 15u]);
        let x2a = vec4<f32>(gb_x[xb + 16u],gb_x[xb + 17u], gb_x[xb + 18u], gb_x[xb + 19u]);
        let x2b = vec4<f32>(gb_x[xb + 20u],gb_x[xb + 21u], gb_x[xb + 22u], gb_x[xb + 23u]);
        let x3a = vec4<f32>(gb_x[xb + 24u],gb_x[xb + 25u], gb_x[xb + 26u], gb_x[xb + 27u]);
        let x3b = vec4<f32>(gb_x[xb + 28u],gb_x[xb + 29u], gb_x[xb + 30u], gb_x[xb + 31u]);
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        for (var r = 0u; r < 4u; r = r + 1u) {
            let row = row0 + r;
            if (row >= rows) { break; }
            let par16 = base16 + rows * gpr * 8u + row * 2u;
            // All Q4TP row/group offsets are even u16 indices. Read each
            // packed word once; the scalar helper otherwise issues two
            // storage lookups for its low and high halves.
            let gl = unpack2x16float(gb_gw[par16 >> 1u]);
            let ul = unpack2x16float(gb_uw[par16 >> 1u]);
            let cod8 = cod0 + row * cst;
            var cg = gb_gu8(cod8 + cb);
            var cu = gb_uu8(cod8 + cb);
            if (shf > 3u) {
                cg = cg | (gb_gu8(cod8 + cb + 1u) << 8u);
                cu = cu | (gb_uu8(cod8 + cb + 1u) << 8u);
            }
            let sg = exp2(gl.x + f32((cg >> shf) & 31u) * gl.y);
            let su = exp2(ul.x + f32((cu >> shf) & 31u) * ul.y);
            let t16 = base16 + row * gpr * 8u + g * 8u;
            let w32 = t16 >> 1u;
            let wg0 = gb_gw[w32];      let wu0 = gb_uw[w32];
            let wg1 = gb_gw[w32 + 1u]; let wu1 = gb_uw[w32 + 1u];
            let wg2 = gb_gw[w32 + 2u]; let wu2 = gb_uw[w32 + 2u];
            let wg3 = gb_gw[w32 + 3u]; let wu3 = gb_uw[w32 + 3u];
            let dg = gb_dot8v(wg0, x0a, x0b) + gb_dot8v(wg1, x1a, x1b)
                   + gb_dot8v(wg2, x2a, x2b) + gb_dot8v(wg3, x3a, x3b);
            let du = gb_dot8v(wu0, x0a, x0b) + gb_dot8v(wu1, x1a, x1b)
                   + gb_dot8v(wu2, x2a, x2b) + gb_dot8v(wu3, x3a, x3b);
            if (r == 0u) { ag0 = ag0 + sg * dg; au0 = au0 + su * du; }
            if (r == 1u) { ag1 = ag1 + sg * dg; au1 = au1 + su * du; }
            if (r == 2u) { ag2 = ag2 + sg * dg; au2 = au2 + su * du; }
            if (r == 3u) { ag3 = ag3 + sg * dg; au3 = au3 + su * du; }
        }
    }
    for (var r = 0u; r < 4u; r = r + 1u) {
        var ag = ag0; var au = au0;
        if (r == 1u) { ag = ag1; au = au1; }
        if (r == 2u) { ag = ag2; au = au2; }
        if (r == 3u) { ag = ag3; au = au3; }
        gb4_pg[lid] = ag;
        gb4_pu[lid] = au;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) {
                gb4_pg[lid] = gb4_pg[lid] + gb4_pg[lid + stride];
                gb4_pu[lid] = gb4_pu[lid] + gb4_pu[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u && row0 + r < rows) {
            var gg = gb4_pg[0];
            var uu = gb4_pu[0];
            if (gb_p.lim > 0.0) {
                uu = clamp(uu, -gb_p.lim, gb_p.lim);
                gg = min(gg, gb_p.lim);
            }
            gb_act[(t * gb_p.slots + slot) * gb_p.inter + row0 + r] =
                (gg / (1.0 + exp(-gg))) * uu;
        }
        workgroupBarrier();
    }
}

// Weighted down-projection with a token axis; writes STRAIGHT into the
// batch FFN output at the token's row — no staging row, no copy back.
struct MoeDnBP { gpr: u32, hidden: u32, slots: u32, mat16: u32 };
@group(0) @binding(0) var<storage, read>       db_w   : array<u32>;
@group(0) @binding(1) var<storage, read>       db_act : array<f32>;
@group(0) @binding(2) var<storage, read>       db_sel : array<u32>;
@group(0) @binding(3) var<storage, read>       db_wt  : array<f32>;
@group(0) @binding(4) var<storage, read_write> db_y   : array<f32>;
@group(0) @binding(5) var<uniform>             db_p   : MoeDnBP;
// The same weight buffer, seen 16 bytes at a time: one q4tp group tile is
// exactly one vec4<u32>, and the scalar view above pays four transactions
// for it. Only the b4 kernel binds this view.
@group(0) @binding(6) var<storage, read>       db_wv  : array<vec4<u32>>;
// The activations too: the four-row kernel once staged them through a
// function-scope array, and dynamic indexing sent that array to local
// memory — every dot read paid a spill. Eight named vec4 registers do
// what the array was meant to.
@group(0) @binding(7) var<storage, read>       db_x4  : array<vec4<f32>>;

// Takes and returns the running sum so the fold order is EXACTLY the
// scalar kernel's left-to-right chain — grouping the eight terms first
// would round differently.
fn db_dotv(acc: f32, w: u32, a: vec4<f32>, b: vec4<f32>) -> f32 {
    return acc
         + (f32(w & 0xFu) - 8.0) * a.x
         + (f32((w >> 4u) & 0xFu) - 8.0) * a.y
         + (f32((w >> 8u) & 0xFu) - 8.0) * a.z
         + (f32((w >> 12u) & 0xFu) - 8.0) * a.w
         + (f32((w >> 16u) & 0xFu) - 8.0) * b.x
         + (f32((w >> 20u) & 0xFu) - 8.0) * b.y
         + (f32((w >> 24u) & 0xFu) - 8.0) * b.z
         + (f32((w >> 28u) & 0xFu) - 8.0) * b.w;
}
var<workgroup> db_pt: array<f32, 64>;

fn db_u16(o: u32) -> u32 { return (db_w[o >> 1u] >> ((o & 1u) * 16u)) & 0xFFFFu; }
fn db_u8(o: u32) -> u32 { return (db_w[o >> 2u] >> ((o & 3u) * 8u)) & 0xFFu; }
fn db_dot8(w: u32, xi: u32) -> f32 {
    return (f32(w & 0xFu) - 8.0) * db_act[xi]
         + (f32((w >> 4u) & 0xFu) - 8.0) * db_act[xi + 1u]
         + (f32((w >> 8u) & 0xFu) - 8.0) * db_act[xi + 2u]
         + (f32((w >> 12u) & 0xFu) - 8.0) * db_act[xi + 3u]
         + (f32((w >> 16u) & 0xFu) - 8.0) * db_act[xi + 4u]
         + (f32((w >> 20u) & 0xFu) - 8.0) * db_act[xi + 5u]
         + (f32((w >> 24u) & 0xFu) - 8.0) * db_act[xi + 6u]
         + (f32((w >> 28u) & 0xFu) - 8.0) * db_act[xi + 7u];
}

@compute @workgroup_size(64)
fn moe_down_q4tp_b(@builtin(workgroup_id) wid: vec3<u32>,
                   @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let t = wid.y;
    let gpr = db_p.gpr;
    let rows = db_p.hidden;
    let inter = gpr * 32u;
    let cst = (gpr * 5u + 7u) / 8u;
    let sb = t * db_p.slots;
    let ab = t * db_p.slots * inter;
    let total = db_p.slots * gpr;
    var acc = 0.0;
    for (var i = lid; i < total; i = i + 64u) {
        let slot = i / gpr;
        let g = i % gpr;
        let base16 = db_sel[sb + slot] * db_p.mat16;
        let par16 = base16 + rows * gpr * 8u + row * 2u;
        let cod8 = (base16 + rows * gpr * 8u + rows * 2u) * 2u + row * cst;
        let pl = unpack2x16float(db_u16(par16) | (db_u16(par16 + 1u) << 16u));
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cv = db_u8(cod8 + cb);
        if (shf > 3u) { cv = cv | (db_u8(cod8 + cb + 1u) << 8u); }
        let scale = exp2(pl.x + f32((cv >> shf) & 31u) * pl.y);
        let t16 = base16 + (row * gpr + g) * 8u;
        let xb = ab + (slot * gpr + g) * 32u;
        var d = 0.0;
        for (var k = 0u; k < 4u; k = k + 1u) {
            let w = db_u16(t16 + 2u * k) | (db_u16(t16 + 1u + 2u * k) << 16u);
            d = d + db_dot8(w, xb + 8u * k);
        }
        acc = acc + db_wt[sb + slot] * scale * d;
    }
    db_pt[lid] = acc;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { db_pt[lid] = db_pt[lid] + db_pt[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { db_y[t * rows + row] = db_pt[0]; }
}

// The down-projection split for parallelism: one workgroup per
// (row, slot, token) writes its weighted partial, a second dispatch sums
// the slots in ascending order. The single-dispatch kernel above gives
// each (row, token) all nine slots and stalls on latency — measured 0.71
// ms of the layer against gate/up's 0.22 for the same bytes. Round-off
// class: the partial's reduction tree differs from the fused loop's.
@compute @workgroup_size(64)
fn moe_down_q4tp_part(@builtin(workgroup_id) wid: vec3<u32>,
                      @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let slot = wid.y;
    let t = wid.z;
    let gpr = db_p.gpr;
    let rows = db_p.hidden;
    let cst = (gpr * 5u + 7u) / 8u;
    let sb = t * db_p.slots;
    let ab = t * db_p.slots * gpr * 32u;
    let base16 = db_sel[sb + slot] * db_p.mat16;
    let par16 = base16 + rows * gpr * 8u + row * 2u;
    let pl = unpack2x16float(db_u16(par16) | (db_u16(par16 + 1u) << 16u));
    let cod8 = (base16 + rows * gpr * 8u + rows * 2u) * 2u + row * cst;
    var acc = 0.0;
    for (var g = lid; g < gpr; g = g + 64u) {
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cv = db_u8(cod8 + cb);
        if (shf > 3u) { cv = cv | (db_u8(cod8 + cb + 1u) << 8u); }
        let scale = exp2(pl.x + f32((cv >> shf) & 31u) * pl.y);
        let t16 = base16 + (row * gpr + g) * 8u;
        let xb = ab + (slot * gpr + g) * 32u;
        var d = 0.0;
        for (var k = 0u; k < 4u; k = k + 1u) {
            let w = db_u16(t16 + 2u * k) | (db_u16(t16 + 1u + 2u * k) << 16u);
            d = d + db_dot8(w, xb + 8u * k);
        }
        acc = acc + scale * d;
    }
    db_pt[lid] = acc;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { db_pt[lid] = db_pt[lid] + db_pt[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        db_y[(t * db_p.slots + slot) * rows + row] = db_wt[sb + slot] * db_pt[0];
    }
}

// …and the slot sum, one thread per (row, token), slots ascending.
@group(0) @binding(0) var<storage, read>       dr_part : array<f32>;
@group(0) @binding(1) var<storage, read_write> dr_y    : array<f32>;
@group(0) @binding(2) var<uniform>             dr_p    : MoeDnBP;
@compute @workgroup_size(256)
fn moe_down_q4tp_red(@builtin(workgroup_id) wid: vec3<u32>,
                     @builtin(local_invocation_index) lid: u32) {
    let row = wid.x * 256u + lid;
    let t = wid.y;
    if (row >= dr_p.hidden) { return; }
    var acc = 0.0;
    for (var s = 0u; s < dr_p.slots; s = s + 1u) {
        acc = acc + dr_part[(t * dr_p.slots + s) * dr_p.hidden + row];
    }
    dr_y[t * dr_p.hidden + row] = acc;
}

// The 2-bit down-projection for the DRAFT's experts: the same ladder
// and packing as gate/up (five-bit rung codes, sixteen weights a word),
// pointed the other way. Draft-only fidelity: nothing downstream needs
// this to match a walk bit for bit — a trunk pass verifies every token.
@compute @workgroup_size(64)
fn moe_down_q2tp_b(@builtin(workgroup_id) wid: vec3<u32>,
                   @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let t = wid.y;
    let gpr = db_p.gpr;
    let rows = db_p.hidden;
    let cst = (gpr * 5u + 7u) / 8u;
    let sb = t * db_p.slots;
    let ab4 = t * db_p.slots * gpr * 8u;
    let total = db_p.slots * gpr;
    var cur = 0xFFFFFFFFu;
    var base16 = 0u;
    var cod8 = 0u;
    var sw = 0.0;
    var pl = vec2<f32>(0.0, 0.0);
    var acc = 0.0;
    for (var i = lid; i < total; i = i + 64u) {
        let slot = i / gpr;
        let g = i % gpr;
        if (slot != cur) {
            cur = slot;
            base16 = db_sel[sb + slot] * db_p.mat16;
            let par16 = base16 + rows * gpr * 4u + row * 2u;
            pl = unpack2x16float(db_u16(par16) | (db_u16(par16 + 1u) << 16u));
            cod8 = (base16 + rows * gpr * 4u + rows * 2u) * 2u + row * cst;
            sw = db_wt[sb + slot];
        }
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cv = db_u8(cod8 + cb);
        if (shf > 3u) { cv = cv | (db_u8(cod8 + cb + 1u) << 8u); }
        let code = (cv >> shf) & 31u;
        let scale = select(exp2(pl.x + f32(max(code, 1u) - 1u) * pl.y), 0.0, code == 0u);
        let w32 = (base16 + (row * gpr + g) * 4u) >> 1u;
        let x4 = ab4 + (slot * gpr + g) * 8u;
        let d = mg_dot16v(db_w[w32], db_x4[x4], db_x4[x4 + 1u], db_x4[x4 + 2u], db_x4[x4 + 3u])
              + mg_dot16v(db_w[w32 + 1u], db_x4[x4 + 4u], db_x4[x4 + 5u], db_x4[x4 + 6u],
                          db_x4[x4 + 7u]);
        acc = acc + sw * scale * d;
    }
    db_pt[lid] = acc;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { db_pt[lid] = db_pt[lid] + db_pt[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { db_y[t * rows + row] = db_pt[0]; }
}

// Four rows per workgroup: the x span is loaded once per thread and
// feeds four weight tiles, cutting the L2 activation traffic that limits
// the one-row kernel four-fold. Each row's accumulation order equals the
// one-row kernel's, so per-row sums are bit-identical.
@compute @workgroup_size(64)
fn moe_down_q4tp_b4(@builtin(workgroup_id) wid: vec3<u32>,
                    @builtin(local_invocation_index) lid: u32) {
    let row0 = wid.x * 4u;
    let t = wid.y;
    let gpr = db_p.gpr;
    let rows = db_p.hidden;
    let cst = (gpr * 5u + 7u) / 8u;
    let sb = t * db_p.slots;
    let ab = t * db_p.slots * gpr * 32u;
    let total = db_p.slots * gpr;
    var cur = 0xFFFFFFFFu;
    var base16 = 0u;
    var codb = 0u;
    var sw = 0.0;
    var pl = vec2<f32>(0.0, 0.0);
    var pl1 = vec2<f32>(0.0, 0.0);
    var pl2 = vec2<f32>(0.0, 0.0);
    var pl3 = vec2<f32>(0.0, 0.0);
    var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
    for (var i = lid; i < total; i = i + 64u) {
        let slot = i / gpr;
        let g = i % gpr;
        if (slot != cur) {
            cur = slot;
            base16 = db_sel[sb + slot] * db_p.mat16;
            let par16 = base16 + rows * gpr * 8u + row0 * 2u;
            pl = unpack2x16float(db_u16(par16) | (db_u16(par16 + 1u) << 16u));
            pl1 = unpack2x16float(db_u16(par16 + 2u) | (db_u16(par16 + 3u) << 16u));
            pl2 = unpack2x16float(db_u16(par16 + 4u) | (db_u16(par16 + 5u) << 16u));
            pl3 = unpack2x16float(db_u16(par16 + 6u) | (db_u16(par16 + 7u) << 16u));
            codb = (base16 + rows * gpr * 8u + rows * 2u) * 2u;
            sw = db_wt[sb + slot];
        }
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        let x4 = (ab + (slot * gpr + g) * 32u) >> 2u;
        // The x span once, into eight NAMED vec4s — registers, not a
        // spilled array.
        let v0 = db_x4[x4];      let v1 = db_x4[x4 + 1u];
        let v2 = db_x4[x4 + 2u]; let v3 = db_x4[x4 + 3u];
        let v4 = db_x4[x4 + 4u]; let v5 = db_x4[x4 + 5u];
        let v6 = db_x4[x4 + 6u]; let v7 = db_x4[x4 + 7u];
        for (var r = 0u; r < 4u; r = r + 1u) {
            let row = row0 + r;
            if (row >= rows) { break; }
            let cod8 = codb + row * cst;
            var cv = db_u8(cod8 + cb);
            if (shf > 3u) { cv = cv | (db_u8(cod8 + cb + 1u) << 8u); }
            var plr = pl;
            if (r == 1u) { plr = pl1; }
            if (r == 2u) { plr = pl2; }
            if (r == 3u) { plr = pl3; }
            let scale = exp2(plr.x + f32((cv >> shf) & 31u) * plr.y);
            let wv = db_wv[(base16 + (row * gpr + g) * 8u) >> 3u];
            let d = db_dotv(
                db_dotv(db_dotv(db_dotv(0.0, wv.x, v0, v1), wv.y, v2, v3), wv.z, v4, v5),
                wv.w, v6, v7,
            );
            let v = sw * scale * d;
            if (r == 0u) { a0 = a0 + v; }
            if (r == 1u) { a1 = a1 + v; }
            if (r == 2u) { a2 = a2 + v; }
            if (r == 3u) { a3 = a3 + v; }
        }
    }
    // Four reductions through the same shared tree, one row at a time —
    // the tree per row equals the one-row kernel's.
    for (var r = 0u; r < 4u; r = r + 1u) {
        var acc = a0;
        if (r == 1u) { acc = a1; }
        if (r == 2u) { acc = a2; }
        if (r == 3u) { acc = a3; }
        db_pt[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { db_pt[lid] = db_pt[lid] + db_pt[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u && row0 + r < rows) { db_y[t * rows + row0 + r] = db_pt[0]; }
        workgroupBarrier();
    }
}

// The same weighted down-projection with the loads it always meant: the
// tile's four u32 words read directly (the u16 pair OR-ed together above
// is two loads of ONE word — legal only because the tile base is even,
// which the host checks), and the slot's params fetched once per slot
// instead of once per group. The per-thread accumulation order is the
// original loop's, so the sum is bit-identical.
@compute @workgroup_size(64)
fn moe_down_q4tp_b2(@builtin(workgroup_id) wid: vec3<u32>,
                    @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let t = wid.y;
    let gpr = db_p.gpr;
    let rows = db_p.hidden;
    let cst = (gpr * 5u + 7u) / 8u;
    let sb = t * db_p.slots;
    let ab = t * db_p.slots * gpr * 32u;
    let total = db_p.slots * gpr;
    var cur = 0xFFFFFFFFu;
    var base16 = 0u;
    var cod8 = 0u;
    var sw = 0.0;
    var pl = vec2<f32>(0.0, 0.0);
    var acc = 0.0;
    for (var i = lid; i < total; i = i + 64u) {
        let slot = i / gpr;
        let g = i % gpr;
        if (slot != cur) {
            cur = slot;
            base16 = db_sel[sb + slot] * db_p.mat16;
            let par16 = base16 + rows * gpr * 8u + row * 2u;
            pl = unpack2x16float(db_u16(par16) | (db_u16(par16 + 1u) << 16u));
            cod8 = (base16 + rows * gpr * 8u + rows * 2u) * 2u + row * cst;
            sw = db_wt[sb + slot];
        }
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cv = db_u8(cod8 + cb);
        if (shf > 3u) { cv = cv | (db_u8(cod8 + cb + 1u) << 8u); }
        let scale = exp2(pl.x + f32((cv >> shf) & 31u) * pl.y);
        let w32 = (base16 + (row * gpr + g) * 8u) >> 1u;
        let xb = ab + (slot * gpr + g) * 32u;
        var d = 0.0;
        for (var k = 0u; k < 4u; k = k + 1u) {
            d = d + db_dot8(db_w[w32 + k], xb + 8u * k);
        }
        acc = acc + sw * scale * d;
    }
    db_pt[lid] = acc;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { db_pt[lid] = db_pt[lid] + db_pt[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { db_y[t * rows + row] = db_pt[0]; }
}


// ── O(1) Nystrom attention on the graph (spec: nystrom.rs step/far_insert).
// State lives on the device after one upload per seal epoch: ring window,
// sinks, landmarks, and each head's flash-scaled far skeleton. Per o1
// layer per token: THREE dispatches replacing kv_append+attend, and the
// work is O(m + w) instead of O(ctx) — the graph's exact attend was
// 54 -> 37.8 tok/s from 4K to 16K while o1 holds flat by construction.
struct O1P { hpg: u32, m: u32, w: u32, nsrect: u32, d: u32, dv: u32, scale: f32, goff: u32 };

@group(0) @binding(0) var<storage, read_write> of_meta : array<u32>;
@group(0) @binding(1) var<storage, read>       of_rk   : array<f32>;
@group(0) @binding(2) var<storage, read>       of_rv   : array<f32>;
@group(0) @binding(3) var<storage, read>       of_qt   : array<f32>;
@group(0) @binding(4) var<storage, read_write> of_mz   : array<f32>;
@group(0) @binding(5) var<storage, read_write> of_th   : array<f32>;
@group(0) @binding(6) var<uniform>             of_p    : O1P;
var<workgroup> of_part: array<f32, 64>;
var<workgroup> of_rs: f32;
var<workgroup> of_e: f32;

// One workgroup per (group, head, landmark): absorb the evicted window
// slot into this head's far accumulators (nystrom.rs far_insert).
@compute @workgroup_size(64)
fn o1_far(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_index) lid: u32) {
    let hm = of_p.hpg * of_p.m;
    let g = wid.x / hm;
    let rr = wid.x % hm;
    let h = rr / of_p.m;
    let i = rr % of_p.m;
    let len = of_meta[g * 4u];
    if (len < of_p.w) { return; }
    let slot = of_meta[g * 4u + 1u];
    let d = of_p.d;
    let qb = ((g * of_p.hpg + h) * of_p.m + i) * d;
    let kb = (g * of_p.w + slot) * d;
    var acc = 0.0;
    var t = lid;
    loop {
        if (t >= d) { break; }
        acc = acc + of_qt[qb + t] * of_rk[kb + t];
        t = t + 64u;
    }
    of_part[lid] = acc;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { of_part[lid] = of_part[lid] + of_part[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    let mzb = (g * of_p.hpg + h) * 2u * of_p.m;
    if (lid == 0u) {
        let l = of_part[0] * of_p.scale;
        var mm = of_mz[mzb + i];
        var rs = 1.0;
        if (l > mm) {
            rs = exp(mm - l);
            of_mz[mzb + of_p.m + i] = of_mz[mzb + of_p.m + i] * rs;
            mm = l;
            of_mz[mzb + i] = l;
        }
        let e = exp(l - mm);
        of_mz[mzb + of_p.m + i] = of_mz[mzb + of_p.m + i] + e;
        of_rs = rs;
        of_e = e;
    }
    workgroupBarrier();
    let rs = of_rs;
    let e = of_e;
    let thb = ((g * of_p.hpg + h) * of_p.m + i) * of_p.dv;
    let vb = (g * of_p.w + slot) * of_p.dv;
    var u = lid;
    loop {
        if (u >= of_p.dv) { break; }
        of_th[thb + u] = of_th[thb + u] * rs + e * of_rv[vb + u];
        u = u + 64u;
    }
}

@group(0) @binding(0) var<storage, read_write> op_meta : array<u32>;
@group(0) @binding(1) var<storage, read>       op_k    : array<f32>;
@group(0) @binding(2) var<storage, read>       op_v    : array<f32>;
@group(0) @binding(3) var<storage, read_write> op_rk   : array<f32>;
@group(0) @binding(4) var<storage, read_write> op_rv   : array<f32>;
@group(0) @binding(5) var<uniform>             op_p    : O1P;

// One workgroup per group: push this token's rotated K and V into the
// window ring (after o1_far has read the slot being overwritten).
@compute @workgroup_size(256)
fn o1_push(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_index) lid: u32) {
    let g = wid.x;
    let len = op_meta[g * 4u];
    let head = op_meta[g * 4u + 1u];
    var slot = len;
    if (len == op_p.w) { slot = head; }
    let d = op_p.d;
    var t = lid;
    loop {
        if (t >= d) { break; }
        // `goff` is the batch row's Q element offset.  K/V rows are shorter
        // by `hpg`, so divide once here instead of carrying a second offset
        // through the fixed 32-byte uniform.  The token graph keeps goff=0.
        op_rk[(g * op_p.w + slot) * d + t] =
            op_k[(op_p.goff / op_p.hpg) + g * d + t];
        t = t + 256u;
    }
    t = lid;
    loop {
        if (t >= op_p.dv) { break; }
        op_rv[(g * op_p.w + slot) * op_p.dv + t] =
            op_v[(op_p.goff / op_p.hpg) + g * op_p.dv + t];
        t = t + 256u;
    }
    workgroupBarrier();
    if (lid == 0u) {
        if (len == op_p.w) {
            op_meta[g * 4u + 1u] = (head + 1u) % op_p.w;
            op_meta[g * 4u + 2u] = op_meta[g * 4u + 2u] + 1u;
        } else {
            op_meta[g * 4u] = len + 1u;
        }
    }
}

@group(0) @binding(0)  var<storage, read>       oa_meta : array<u32>;
@group(0) @binding(1)  var<storage, read>       oa_q    : array<f32>;
@group(0) @binding(2)  var<storage, read>       oa_rk   : array<f32>;
@group(0) @binding(3)  var<storage, read>       oa_rv   : array<f32>;
@group(0) @binding(4)  var<storage, read>       oa_sk   : array<f32>;
@group(0) @binding(5)  var<storage, read>       oa_sv   : array<f32>;
@group(0) @binding(6)  var<storage, read>       oa_kt   : array<f32>;
@group(0) @binding(7)  var<storage, read>       oa_mu   : array<f32>;
@group(0) @binding(8)  var<storage, read>       oa_mz   : array<f32>;
@group(0) @binding(9)  var<storage, read>       oa_th   : array<f32>;
@group(0) @binding(10) var<storage, read_write> oa_out  : array<f32>;
@group(0) @binding(11) var<uniform>             oa_p    : O1P;
var<workgroup> oa_qs:  array<f32, 256>;
// The admission contract permits sink + sliding window up to 2052 rows
// (w2048/sink4 is the first extended profile).  Keep the workgroup softmax
// scratch at that same bound.  Near scores are distributed across the 256
// lanes below, so the array is a bounded ring-sized workspace rather than a
// requirement that one lane own every key.
var<workgroup> oa_scr: array<f32, 2052>;
var<workgroup> oa_f:   array<f32, 32>;
var<workgroup> oa_u:   array<f32, 32>;
var<workgroup> oa_red: array<f32, 256>;
var<workgroup> oa_sc:  array<f32, 4>; // [c/f staging, far_den, den, have_far]

// One workgroup per (group, head): the whole Nystrom step output.
// Per-score dots run one THREAD per key/landmark (serial over d) — no
// barriers in the hot part, and the same product order as the CPU's
// scalar loop.
@compute @workgroup_size(256)
fn o1_attend(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_index) lid: u32) {
    let g = wid.x / oa_p.hpg;
    let h = wid.x % oa_p.hpg;
    let d = oa_p.d;
    let dv = oa_p.dv;
    let m = oa_p.m;
    let ns = oa_p.nsrect & 0xFFu;
    let rect_fm = (oa_p.nsrect >> 8u) != 0u;
    let len = oa_meta[g * 4u];
    let farl = oa_meta[g * 4u + 2u];
    let n = ns + len;
    let gh = g * oa_p.hpg + h;
    // q into shared
    var t = lid;
    loop {
        if (t >= d) { break; }
        // `oa_q` is the one-row rope scratch (the batch input is consumed by
        // attn_rope before this kernel), so its base stays zero just like the
        // token graph. `goff` belongs only to the batch K/V source in push.
        oa_qs[t] = oa_q[gh * d + t];
        t = t + 256u;
    }
    workgroupBarrier();
    // Near scores: each lane walks a bounded grid-stride slice.  The old
    // one-key-per-lane mapping was correct only while sink+window <= 196;
    // keeping the same mapping for the extended window left most scores
    // uninitialised and also collided with the landmark lanes.
    var s = lid;
    loop {
        if (s >= n) { break; }
        var acc = 0.0;
        if (s < ns) {
            let kb = (g * ns + s) * d;
            for (var j = 0u; j < d; j = j + 1u) { acc = acc + oa_qs[j] * oa_sk[kb + j]; }
        } else {
            let kb = (g * oa_p.w + (s - ns)) * d;
            for (var j = 0u; j < d; j = j + 1u) { acc = acc + oa_qs[j] * oa_rk[kb + j]; }
        }
        oa_scr[s] = acc * oa_p.scale;
        s = s + 256u;
    }
    // Landmark scores use their own small array, so every lane can also
    // participate in the near grid above without a lane-number reservation.
    if (lid < m && farl > 0u) {
        let a = lid;
        var acc = 0.0;
        let ktb = (g * m + a) * d;
        for (var j = 0u; j < d; j = j + 1u) {
            acc = acc + oa_qs[j] * oa_kt[ktb + j];
        }
        oa_f[a] = acc * oa_p.scale;
    }
    workgroupBarrier();
    // c = max near score (single thread — n <= 2052, bounded).  The
    // landmark inverse readout below has m independent output columns.  Keep
    // the score/exp order on lane 0, publish its two scalars, and let one lane
    // own each column's ascending-a accumulation.  This removes the old
    // 32-column serial loop without changing any per-column arithmetic.
    if (lid == 0u) {
        var c = -3.0e38;
        for (var sidx = 0u; sidx < n; sidx = sidx + 1u) { c = max(c, oa_scr[sidx]); }
        oa_sc[0] = c;
        oa_sc[1] = 0.0;
        if (farl > 0u) {
            var f = -3.0e38;
            for (var a = 0u; a < m; a = a + 1u) { f = max(f, oa_f[a]); }
            for (var a = 0u; a < m; a = a + 1u) { oa_f[a] = exp(oa_f[a] - f); }
            oa_sc[1] = f;
        }
    }
    workgroupBarrier();
    // Each output column keeps the exact old ascending-a accumulation order.
    // The guard is uniform in farl for a given workgroup; the barrier after
    // it is unconditional so far-empty admissions and m < 32 remain safe.
    if (lid < m && farl > 0u) {
        let b = lid;
        var uacc = 0.0;
        for (var a = 0u; a < m; a = a + 1u) {
            uacc = uacc + oa_f[a] * oa_mu[(gh * m + a) * m + b];
        }
        if (rect_fm) { uacc = max(uacc, 0.0); }
        oa_u[b] = uacc;
    }
    workgroupBarrier();
    if (lid == 0u) {
        let c = oa_sc[0];
        let f = oa_sc[1];
        var c_all = c;
        var far_den = 0.0;
        var have_far = 0.0;
        if (farl > 0u) {
            let mzb = gh * 2u * m;
            for (var b = 0u; b < m; b = b + 1u) {
                c_all = max(c_all, f + oa_mz[mzb + b]);
            }
            for (var b = 0u; b < m; b = b + 1u) {
                let gain = oa_u[b] * exp(f + oa_mz[mzb + b] - c_all);
                oa_u[b] = gain;
                far_den = far_den + gain * oa_mz[mzb + m + b];
            }
            if (far_den >= 0.0) { have_far = 1.0; } else { far_den = 0.0; }
        }
        var den = far_den;
        for (var sidx = 0u; sidx < n; sidx = sidx + 1u) {
            let pv = exp(oa_scr[sidx] - c_all);
            oa_scr[sidx] = pv;
            den = den + pv;
        }
        oa_sc[0] = c_all;
        oa_sc[1] = far_den;
        oa_sc[2] = max(den, 1e-30);
        oa_sc[3] = have_far;
    }
    workgroupBarrier();
    let den = oa_sc[2];
    let have_far = oa_sc[3] > 0.5;
    t = lid;
    loop {
        if (t >= dv) { break; }
        var acc = 0.0;
        if (have_far) {
            for (var b = 0u; b < m; b = b + 1u) {
                acc = acc + oa_u[b] * oa_th[(gh * m + b) * dv + t];
            }
        }
        for (var sidx = 0u; sidx < ns; sidx = sidx + 1u) {
            acc = acc + oa_scr[sidx] * oa_sv[(g * ns + sidx) * dv + t];
        }
        for (var sidx = ns; sidx < n; sidx = sidx + 1u) {
            acc = acc + oa_scr[sidx] * oa_rv[(g * oa_p.w + (sidx - ns)) * dv + t];
        }
        oa_out[gh * dv + t] = acc / den;
        t = t + 256u;
    }
}

// ── DiT attention (imagegen): scores GEMM -> row softmax -> P·V.
// Same 64x64 tile / 4x4 named-scalar register block as the quantized
// GEMMs; the operands are plain f32 here, so the staging is a copy.
// `hpk` is query heads per KV head (GQA) and `ntok` the block's full token
// count: with the heads batched into one dispatch, a shader has to find its
// own head's slice, and the two strides it needs are not derivable from m,
// k and n alone.
struct DitP { m: u32, k: u32, n: u32, scale: f32, s0: u32, causal: u32, hpk: u32, ntok: u32, };
@group(0) @binding(0) var<storage, read> da: array<f32>;
@group(0) @binding(1) var<storage, read> db: array<f32>;
@group(0) @binding(2) var<storage, read_write> dc: array<f32>;
@group(0) @binding(3) var<uniform> dp: DitP;
// Padded rows, as in `q4tp_mul_mm`: at a stride of 16 the sixteen threads
// of a row shared one bank and every read serialised.
var<workgroup> dit_at: array<f32, 64 * 17>;
var<workgroup> dit_bt: array<f32, 64 * 17>;

fn dit_gemm(wid: vec3<u32>, lid: vec3<u32>, bt: bool, ab: u32, bb: u32, cb: u32) {
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;
    let tid = lid.y * 16u + lid.x;
    var a00 = 0.0; var a01 = 0.0; var a02 = 0.0; var a03 = 0.0;
    var a10 = 0.0; var a11 = 0.0; var a12 = 0.0; var a13 = 0.0;
    var a20 = 0.0; var a21 = 0.0; var a22 = 0.0; var a23 = 0.0;
    var a30 = 0.0; var a31 = 0.0; var a32 = 0.0; var a33 = 0.0;
    var k0 = 0u;
    loop {
        if (k0 >= dp.k) { break; }
        // One whole vec4 per thread; four threads writing lanes of the
        // same shared vec4 is a race wherever a dynamic component write
        // becomes read-modify-write.
        for (var q = 0u; q < 4u; q = q + 1u) {
            let r = tid / 4u + q * 64u;
            if (r < 64u) {
                let c4 = (tid % 4u) * 4u;
                for (var e = 0u; e < 4u; e = e + 1u) {
                    let kk = k0 + c4 + e;
                    var va = 0.0;
                    if (m0 + r < dp.m && kk < dp.k) { va = da[ab + (m0 + r) * dp.k + kk]; }
                    dit_at[r * TSTRIDE + c4 + e] = va;
                    var vb = 0.0;
                    if (n0 + r < dp.n && kk < dp.k) {
                        if (bt) { vb = db[bb + (n0 + r) * dp.k + kk]; }
                        else { vb = db[bb + kk * dp.n + n0 + r]; }
                    }
                    dit_bt[r * TSTRIDE + c4 + e] = vb;
                }
            }
        }
        workgroupBarrier();
        let ab2 = lid.y * 4u * TSTRIDE;
        let wb2 = lid.x * 4u * TSTRIDE;
        for (var k = 0u; k < 16u; k = k + 1u) {
            let x0 = dit_at[ab2 + k];
            let x1 = dit_at[ab2 + TSTRIDE + k];
            let x2 = dit_at[ab2 + 2u * TSTRIDE + k];
            let x3 = dit_at[ab2 + 3u * TSTRIDE + k];
            let y0 = dit_bt[wb2 + k];
            let y1 = dit_bt[wb2 + TSTRIDE + k];
            let y2 = dit_bt[wb2 + 2u * TSTRIDE + k];
            let y3 = dit_bt[wb2 + 3u * TSTRIDE + k];
            a00 = a00 + x0 * y0; a01 = a01 + x0 * y1;
            a02 = a02 + x0 * y2; a03 = a03 + x0 * y3;
            a10 = a10 + x1 * y0; a11 = a11 + x1 * y1;
            a12 = a12 + x1 * y2; a13 = a13 + x1 * y3;
            a20 = a20 + x2 * y0; a21 = a21 + x2 * y1;
            a22 = a22 + x2 * y2; a23 = a23 + x2 * y3;
            a30 = a30 + x3 * y0; a31 = a31 + x3 * y1;
            a32 = a32 + x3 * y2; a33 = a33 + x3 * y3;
        }
        workgroupBarrier();
        k0 = k0 + 16u;
    }
    let mb = m0 + lid.y * 4u;
    let nb2 = n0 + lid.x * 4u;
    dit_store4(mb, nb2, a00, a01, a02, a03, cb);
    dit_store4(mb + 1u, nb2, a10, a11, a12, a13, cb);
    dit_store4(mb + 2u, nb2, a20, a21, a22, a23, cb);
    dit_store4(mb + 3u, nb2, a30, a31, a32, a33, cb);
}

fn dit_store4(m: u32, n0: u32, v0: f32, v1: f32, v2: f32, v3: f32, cb: u32) {
    if (m >= dp.m) { return; }
    let base = cb + m * dp.n + n0;
    if (n0 < dp.n) { dc[base] = v0 * dp.scale; }
    if (n0 + 1u < dp.n) { dc[base + 1u] = v1 * dp.scale; }
    if (n0 + 2u < dp.n) { dc[base + 2u] = v2 * dp.scale; }
    if (n0 + 3u < dp.n) { dc[base + 3u] = v3 * dp.scale; }
}

@compute @workgroup_size(16, 16)
fn dit_qk(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    // q and k are head-major over the block's whole token span; the scores
    // are m x n per head, packed back to back.
    let h = wid.z;
    dit_gemm(wid, lid, true,
             h * dp.ntok * dp.k,
             (h / dp.hpk) * dp.ntok * dp.k,
             h * dp.m * dp.n);
}

@compute @workgroup_size(16, 16)
fn dit_pv(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.z;
    dit_gemm(wid, lid, false,
             h * dp.m * dp.k,
             (h / dp.hpk) * dp.ntok * dp.n,
             h * dp.ntok * dp.n);
}

var<workgroup> dit_red: array<f32, 256>;

// Row softmax over dc, one workgroup per row of dp.n columns.
@compute @workgroup_size(256)
fn dit_softmax(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let row = wid.y * dp.m * dp.n + wid.x * dp.n;
    let t = lid.x;
    // Causal bound: query `wid.x` may see keys 0..=s0+wid.x. Masked
    // entries are zeroed rather than set to -inf so the P·V GEMM that
    // follows reads a clean matrix.
    var lim = dp.n;
    if (dp.causal != 0u) { lim = min(dp.n, dp.s0 + wid.x + 1u); }
    var mx = -3.4e38;
    for (var j = t; j < lim; j = j + 256u) { mx = max(mx, dc[row + j]); }
    dit_red[t] = mx;
    workgroupBarrier();
    for (var s = 128u; s > 0u; s = s >> 1u) {
        if (t < s) { dit_red[t] = max(dit_red[t], dit_red[t + s]); }
        workgroupBarrier();
    }
    let m = dit_red[0];
    workgroupBarrier();
    var sum = 0.0;
    for (var j = t; j < lim; j = j + 256u) {
        let e = exp(dc[row + j] - m);
        dc[row + j] = e;
        sum = sum + e;
    }
    for (var j = lim + t; j < dp.n; j = j + 256u) { dc[row + j] = 0.0; }
    dit_red[t] = sum;
    workgroupBarrier();
    for (var s = 128u; s > 0u; s = s >> 1u) {
        if (t < s) { dit_red[t] = dit_red[t] + dit_red[t + s]; }
        workgroupBarrier();
    }
    let inv = 1.0 / dit_red[0];
    for (var j = t; j < lim; j = j + 256u) { dc[row + j] = dc[row + j] * inv; }
}

// [nh][n][hd] panel -> [n][nh*hd].
@compute @workgroup_size(256)
fn dit_unstack(@builtin(global_invocation_id) gid: vec3<u32>,
               @builtin(num_workgroups) nwg: vec3<u32>) {
    let i = gid.x + gid.y * nwg.x * 256u;
    let total = dp.m * dp.k * dp.n;   // nh * n * hd
    if (i >= total) { return; }
    let hd = dp.n;
    let n = dp.k;
    let h = i / (n * hd);
    let rest = i % (n * hd);
    let tok = rest / hd;
    let d = rest % hd;
    dc[tok * dp.m * hd + h * hd + d] = da[i];
}


// ── Per-head RMS and the rope tail (DeepSeek-V4) ────────────────────────────
//
// Three places need this and they differ only in flags: the queries take a
// second RMS on each head after wq_b and then a forward rotation; the shared
// KV vector takes a forward rotation alone; and attention's output takes the
// INVERSE rotation, a detail no naming convention would suggest.
//
// The rotation pairs ADJACENT coordinates — the reference builds its complex
// numbers with unflatten(-1, (2)) — and pairing halves instead agrees with it
// exactly at position 0 and nowhere else. That cost a night once.

struct RpP { nh: u32, hd: u32, rd: u32, flags: u32 };   // flags: 1 = rms, 2 = inverse
@group(0) @binding(0) var<storage, read_write> rp_x    : array<f32>;   // nh*hd
@group(0) @binding(1) var<storage, read>       rp_freq : array<f32>;   // rd/2
@group(0) @binding(2) var<uniform>             rp_p    : RpP;
@group(0) @binding(3) var<storage, read>       rp_pos  : array<f32>;   // [position, eps]

var<workgroup> rp_red: array<f32, 256>;

@compute @workgroup_size(256)
fn rope_heads(@builtin(workgroup_id) wid: vec3<u32>,
              @builtin(local_invocation_index) lid: u32) {
    let h = wid.x;
    if (h >= rp_p.nh) { return; }
    let hd = rp_p.hd;
    let rd = rp_p.rd;
    let b = h * hd;

    if ((rp_p.flags & 1u) != 0u) {
        var acc = 0.0;
        var i = lid;
        loop {
            if (i >= hd) { break; }
            let v = rp_x[b + i];
            acc = acc + v * v;
            i = i + 256u;
        }
        rp_red[lid] = acc;
        workgroupBarrier();
        var stride = 128u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { rp_red[lid] = rp_red[lid] + rp_red[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        let inv = inverseSqrt(rp_red[0] / f32(hd) + rp_pos[1]);
        workgroupBarrier();
        var j = lid;
        loop {
            if (j >= hd) { break; }
            rp_x[b + j] = rp_x[b + j] * inv;
            j = j + 256u;
        }
        workgroupBarrier();
    }

    // The tail only, adjacent pairs.
    let base = b + hd - rd;
    let pos = rp_pos[0];
    var t = lid;
    loop {
        if (t >= rd / 2u) { break; }
        let th = pos * rp_freq[t];
        var sn = sin(th);
        let cs = cos(th);
        if ((rp_p.flags & 2u) != 0u) { sn = -sn; }
        let a = rp_x[base + 2u * t];
        let c = rp_x[base + 2u * t + 1u];
        rp_x[base + 2u * t] = a * cs - c * sn;
        rp_x[base + 2u * t + 1u] = a * sn + c * cs;
        t = t + 256u;
    }
}

// ── Grouped low-rank output projection, stage A (DeepSeek-V4) ───────────────
//
// wo_a is block-diagonal wearing a dense disguise. It is stored as one
// [groups*lora, per_group] matrix, but row i multiplies ONLY the slice of the
// attention output that group i/lora owns — a matvec whose activation window
// slides with the row. One added term in the index buys the whole operator.
//
// It earns a kernel of its own by size: on the release checkpoint wo_a is 33M
// weights read once per layer per token, the largest single thing still on
// the CPU once the experts are away.
//
// q4tp layout; the per-row math is copied from q4tp_matvec unchanged, so the
// two agree bit for bit wherever they overlap — that is, at lora = rows,
// where the window stops sliding.

@compute @workgroup_size(64)
fn o_lora_a(@builtin(workgroup_id) wid: vec3<u32>,
            @builtin(num_workgroups) nwg: vec3<u32>,
            @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let lora = q1p._p0;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    var row = wid.x;
    loop {
        if (row >= rows) { break; }
        if (lid < 32u) {
            let pr = unpack2x16float(q1w[params_w + row]);
            lad_q4tp[lid] = exp2(pr.x + f32(lid) * pr.y);
        }
        workgroupBarrier();
        // A row is exactly one group's width, so the slice offset needs no
        // parameter of its own: per_group = gpr * 32.
        let xoff = (row / lora) * gpr * 32u;
        var acc = 0.0;
        var g = lid;
        loop {
            if (g >= gpr) { break; }
            let bit = g * 5u;
            let cb = codes_b + row * cstride + (bit >> 3u);
            let sh = bit & 7u;
            var cv = q4tp_byte(cb);
            if (sh > 3u) { cv = cv | (q4tp_byte(cb + 1u) << 8u); }
            let scale = lad_q4tp[(cv >> sh) & 31u];
            let base = (row * gpr + g) * 4u;
            let xb = xoff + g * 32u;
            var gsum = 0.0;
            for (var k = 0u; k < 4u; k = k + 1u) {
                gsum = gsum + q4b_dot8(q1w[base + k], xb + 8u * k);
            }
            acc = acc + scale * gsum;
            g = g + 64u;
        }
        partial_q1t[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { partial_q1t[lid] = partial_q1t[lid] + partial_q1t[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) { q1y[row] = partial_q1t[0]; }
        workgroupBarrier();
        row = row + nwg.x;
    }
}

// The same grouped projection with 256 threads a row. `f32_matvec_w` showed
// what a 64-thread workgroup costs on this card; this kernel is the other
// one the chain leans on, and the output projection was 5.0 ms of a 45 ms
// chain. Its own partial array, because the 64-wide one is exactly 64 long.
var<workgroup> olw_part: array<f32, 256>;

@compute @workgroup_size(256)
fn o_lora_a_w(@builtin(workgroup_id) wid: vec3<u32>,
            @builtin(num_workgroups) nwg: vec3<u32>,
            @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let lora = q1p._p0;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    var row = wid.x;
    loop {
        if (row >= rows) { break; }
        if (lid < 32u) {
            let pr = unpack2x16float(q1w[params_w + row]);
            lad_q4tp[lid] = exp2(pr.x + f32(lid) * pr.y);
        }
        workgroupBarrier();
        // A row is exactly one group's width, so the slice offset needs no
        // parameter of its own: per_group = gpr * 32.
        let xoff = (row / lora) * gpr * 32u;
        var acc = 0.0;
        var g = lid;
        loop {
            if (g >= gpr) { break; }
            let bit = g * 5u;
            let cb = codes_b + row * cstride + (bit >> 3u);
            let sh = bit & 7u;
            var cv = q4tp_byte(cb);
            if (sh > 3u) { cv = cv | (q4tp_byte(cb + 1u) << 8u); }
            let scale = lad_q4tp[(cv >> sh) & 31u];
            let base = (row * gpr + g) * 4u;
            let xb = xoff + g * 32u;
            var gsum = 0.0;
            for (var k = 0u; k < 4u; k = k + 1u) {
                gsum = gsum + q4b_dot8(q1w[base + k], xb + 8u * k);
            }
            acc = acc + scale * gsum;
            g = g + 256u;
        }
        olw_part[lid] = acc;
        workgroupBarrier();
        var stride = 128u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { olw_part[lid] = olw_part[lid] + olw_part[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u) { q1y[row] = olw_part[0]; }
        workgroupBarrier();
        row = row + nwg.x;
    }
}

var<workgroup> olm_part: array<f32, 256>;
var<workgroup> olm_lad: array<f32, 128>;

// FOUR rows at once, 64 lanes each. `gpr` is 128 on the release, so a row
// cannot use more than 64 lanes without half of them striding past the end —
// but four rows in one workgroup give the memory system four independent
// load streams to overlap, which is the same trick the 8-row q4tp kernel
// uses and the reason it beat the one-row one by 6.8 ms a token.
@compute @workgroup_size(256)
fn o_lora_a_m(@builtin(workgroup_id) wid: vec3<u32>,
              @builtin(num_workgroups) nwg: vec3<u32>,
              @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let lora = q1p._p0;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid / 64u;          // which of the four rows
    let lane = lid % 64u;
    var row = wid.x * 4u + sub;
    loop {
        if (row >= rows) { break; }
        if (lane < 32u) {
            let pr = unpack2x16float(q1w[params_w + row]);
            olm_lad[sub * 32u + lane] = exp2(pr.x + f32(lane) * pr.y);
        }
        workgroupBarrier();
        let xoff = (row / lora) * gpr * 32u;
        var acc = 0.0;
        var g = lane;
        loop {
            if (g >= gpr) { break; }
            let bit = g * 5u;
            let cb = codes_b + row * cstride + (bit >> 3u);
            let sh = bit & 7u;
            var cv = q4tp_byte(cb);
            if (sh > 3u) { cv = cv | (q4tp_byte(cb + 1u) << 8u); }
            let scale = olm_lad[sub * 32u + ((cv >> sh) & 31u)];
            let base = (row * gpr + g) * 4u;
            let xb = xoff + g * 32u;
            var gsum = 0.0;
            for (var k = 0u; k < 4u; k = k + 1u) {
                gsum = gsum + q4b_dot8(q1w[base + k], xb + 8u * k);
            }
            acc = acc + scale * gsum;
            g = g + 64u;
        }
        olm_part[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lane < stride) {
                olm_part[lid] = olm_part[lid] + olm_part[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lane == 0u) { q1y[row] = olm_part[sub * 64u]; }
        workgroupBarrier();
        row = row + nwg.x * 4u;
    }
}

// ── Sparse attention, split in two (DeepSeek-V4) ────────────────────────────
//
// The one-workgroup-per-head version leaves 64 workgroups on a card with
// 150-odd multiprocessors, and measured 0.54 ms a layer — the whole cost of
// the attention block once its encoding was cached away. Scores are cheap and
// genuinely per-head; the weighted sum is nh*hd independent outputs. So:
// scores in one dispatch of nh groups, the sum in another of nh*hd/256.

struct Sa2P { nh: u32, hd: u32, m: u32, scale: f32 };

@group(0) @binding(0) var<storage, read>       s2_q    : array<f32>;   // nh*hd
@group(0) @binding(1) var<storage, read>       s2_kv   : array<f32>;
@group(0) @binding(2) var<storage, read>       s2_idx  : array<u32>;   // m
@group(0) @binding(3) var<storage, read>       s2_sink : array<f32>;   // nh
@group(0) @binding(4) var<storage, read_write> s2_w    : array<f32>;   // nh*m
@group(0) @binding(5) var<uniform>             s2_p    : Sa2P;

var<workgroup> s2_red: array<f32, 256>;
var<workgroup> s2_sc:  array<f32, 1024>;
var<workgroup> s2_max: f32;
var<workgroup> s2_den: f32;

@compute @workgroup_size(256)
fn sa_scores(@builtin(workgroup_id) wid: vec3<u32>,
             @builtin(local_invocation_index) lid: u32) {
    let h = wid.x;
    if (h >= s2_p.nh) { return; }
    let hd = s2_p.hd;
    let m = s2_p.m;
    let qb = h * hd;

    var mx = s2_sink[h];
    var t = lid;
    loop {
        if (t >= m) { break; }
        let kb = s2_idx[t] * hd;
        var d = 0.0;
        for (var i = 0u; i < hd; i = i + 1u) { d = d + s2_q[qb + i] * s2_kv[kb + i]; }
        d = d * s2_p.scale;
        s2_sc[t] = d;
        mx = max(mx, d);
        t = t + 256u;
    }
    s2_red[lid] = mx;
    workgroupBarrier();
    var st = 128u;
    loop {
        if (st == 0u) { break; }
        if (lid < st) { s2_red[lid] = max(s2_red[lid], s2_red[lid + st]); }
        workgroupBarrier();
        st = st >> 1u;
    }
    if (lid == 0u) { s2_max = s2_red[0]; }
    workgroupBarrier();

    // The learned sink enters the denominator and NOT the numerator: that is
    // what lets a head attend to nothing at all.
    var acc = 0.0;
    var u = lid;
    loop {
        if (u >= m) { break; }
        let e = exp(s2_sc[u] - s2_max);
        s2_sc[u] = e;
        acc = acc + e;
        u = u + 256u;
    }
    s2_red[lid] = acc;
    workgroupBarrier();
    st = 128u;
    loop {
        if (st == 0u) { break; }
        if (lid < st) { s2_red[lid] = s2_red[lid] + s2_red[lid + st]; }
        workgroupBarrier();
        st = st >> 1u;
    }
    if (lid == 0u) { s2_den = s2_red[0] + exp(s2_sink[h] - s2_max); }
    workgroupBarrier();
    let inv = 1.0 / s2_den;
    var v = lid;
    loop {
        if (v >= m) { break; }
        s2_w[h * m + v] = s2_sc[v] * inv;
        v = v + 256u;
    }
}

@group(0) @binding(0) var<storage, read>       sy_w   : array<f32>;   // nh*m
@group(0) @binding(1) var<storage, read>       sy_kv  : array<f32>;
@group(0) @binding(2) var<storage, read>       sy_idx : array<u32>;
@group(0) @binding(3) var<storage, read_write> sy_out : array<f32>;   // nh*hd
@group(0) @binding(4) var<uniform>             sy_p   : Sa2P;

@compute @workgroup_size(256)
fn sa_apply(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    let hd = sy_p.hd;
    if (i >= sy_p.nh * hd) { return; }
    let h = i / hd;
    let d = i % hd;
    let m = sy_p.m;
    let wb = h * m;
    var acc = 0.0;
    for (var t = 0u; t < m; t = t + 1u) {
        acc = acc + sy_w[wb + t] * sy_kv[sy_idx[t] * hd + d];
    }
    sy_out[i] = acc;
}

// ── The KV compressor's pooling step (DeepSeek-V4) ──────────────────────────
//
// A softmax over the slot axis taken PER DIMENSION — not per token — then the
// weighted sum. Both compressors end here; they differ only in how the slots
// are gathered, so the gather lives inside the kernel and the graph never has
// to materialise the interleaved copy the CPU builds.
//
// Overlapping (ratio 4 in the release): each token contributes 2*width
// values, the first half belonging to the window that began half a stride
// earlier. Fold time pools 2*ratio slots — the previous window's taking their
// first half, the current window's taking their second. A missing previous
// window votes with -inf, which is also how a whole column of absent slots
// leaves the output at zero instead of dividing by nothing.

struct KpP { slots: u32, width: u32, ratio: u32, flags: u32 };
// flags: 1 = overlapping, 2 = a previous window exists, 4 = add the APE bias

@group(0) @binding(0) var<storage, read>       kp_pkv : array<f32>;
@group(0) @binding(1) var<storage, read>       kp_psc : array<f32>;
@group(0) @binding(2) var<storage, read>       kp_ckv : array<f32>;
@group(0) @binding(3) var<storage, read>       kp_csc : array<f32>;
@group(0) @binding(4) var<storage, read>       kp_ape : array<f32>;
@group(0) @binding(5) var<storage, read_write> kp_out : array<f32>;
@group(0) @binding(6) var<uniform>             kp_p   : KpP;

const KP_NINF: f32 = -3.0e38;

@compute @workgroup_size(256)
fn kv_pool(@builtin(global_invocation_id) gid: vec3<u32>) {
    let d = gid.x;
    let w = kp_p.width;
    if (d >= w) { return; }
    let slots = kp_p.slots;
    let r = kp_p.ratio;
    let overlap = (kp_p.flags & 1u) != 0u;
    let have_prev = (kp_p.flags & 2u) != 0u;
    let use_ape = (kp_p.flags & 4u) != 0u;

    // Pass one: the maximum, so the exponentials cannot overflow. A column
    // that is entirely absent stays at -inf and the slot is left at zero.
    var mx = KP_NINF;
    for (var t = 0u; t < slots; t = t + 1u) {
        var sc = KP_NINF;
        if (overlap) {
            if (t < r) {
                if (have_prev) { sc = kp_psc[t * 2u * w + d]; }
            } else {
                sc = kp_csc[(t - r) * 2u * w + w + d];
            }
        } else {
            sc = kp_csc[t * w + d];
            if (use_ape) { sc = sc + kp_ape[t * w + d]; }
        }
        mx = max(mx, sc);
    }
    if (mx <= KP_NINF) { kp_out[d] = 0.0; return; }

    var den = 0.0;
    var acc = 0.0;
    for (var t = 0u; t < slots; t = t + 1u) {
        var sc = KP_NINF;
        var kv = 0.0;
        if (overlap) {
            if (t < r) {
                if (have_prev) {
                    sc = kp_psc[t * 2u * w + d];
                    kv = kp_pkv[t * 2u * w + d];
                }
            } else {
                sc = kp_csc[(t - r) * 2u * w + w + d];
                kv = kp_ckv[(t - r) * 2u * w + w + d];
            }
        } else {
            sc = kp_csc[t * w + d];
            if (use_ape) { sc = sc + kp_ape[t * w + d]; }
            kv = kp_ckv[t * w + d];
        }
        if (sc > KP_NINF) {
            let e = exp(sc - mx);
            den = den + e;
            acc = acc + e * kv;
        }
    }
    if (den <= 0.0) { kp_out[d] = 0.0; return; }
    kp_out[d] = acc / den;
}

// ── The sparse indexer: scores, then the top-k (DeepSeek-V4) ────────────────
//
// The relu comes BEFORE the per-head weighting, so a head can vote for a
// position or abstain but never against it. Getting that order wrong produces
// scores that look reasonable and a top-k that is quietly different.

struct IxP { nh: u32, hd: u32, n_pos: u32, limit: u32 };

@group(0) @binding(0) var<storage, read>       ix_q   : array<f32>;   // nh*hd
@group(0) @binding(1) var<storage, read>       ix_kv  : array<f32>;   // n_pos*hd
@group(0) @binding(2) var<storage, read>       ix_w   : array<f32>;   // nh
@group(0) @binding(3) var<storage, read_write> ix_out : array<f32>;   // n_pos
@group(0) @binding(4) var<uniform>             ix_p   : IxP;

var<workgroup> ix_red: array<f32, 256>;

@compute @workgroup_size(256)
fn index_scores(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let t = wid.x;
    if (t >= ix_p.n_pos) { return; }
    if (t >= ix_p.limit) {
        if (lid == 0u) { ix_out[t] = KP_NINF; }
        return;
    }
    let hd = ix_p.hd;
    let kb = t * hd;
    // One lane per head: the relu makes the heads non-additive before their
    // weights, so a head's dot has to be finished by whoever owns it.
    var acc = 0.0;
    var h = lid;
    loop {
        if (h >= ix_p.nh) { break; }
        var dot = 0.0;
        let qb = h * hd;
        for (var i = 0u; i < hd; i = i + 1u) {
            dot = dot + ix_q[qb + i] * ix_kv[kb + i];
        }
        acc = acc + max(dot, 0.0) * ix_w[h];
        h = h + 256u;
    }
    ix_red[lid] = acc;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { ix_red[lid] = ix_red[lid] + ix_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { ix_out[t] = ix_red[0]; }
}

// Sparse attention, split over the ATTENDED POSITIONS.
//
// The one-workgroup-per-head kernel launches 64 workgroups on the release —
// 64 of the ~200 the card can hold, so it sits at a few percent occupancy and
// waits on memory latency rather than running out of arithmetic: the skip
// probe puts it at 18.1 ms of a 51.3 ms chain while doing 42M MACs a layer,
// which is about 1% of the machine.
//
// So each head's positions are cut into chunks, one workgroup a chunk, and
// each returns its own softmax frame — the running max, the running
// denominator, and an accumulator weighted against ITS max. The merge
// rescales the frames into a common max. Same flash-decoding shape as
// `gqa_attend_part`/`gqa_attend_merge` for the other architecture, and the
// same caveat: the sum happens in a different order, so this is a contract
// change, not an identity.

struct SapP { nh: u32, hd: u32, m: u32, nc: u32, ck: u32, _a: u32, _b: u32, scale: f32 };

@group(0) @binding(0) var<storage, read>       sp_q    : array<f32>;   // nh*hd
@group(0) @binding(1) var<storage, read>       sp_kv   : array<f32>;   // n*hd
@group(0) @binding(2) var<storage, read>       sp_idx  : array<u32>;   // m
@group(0) @binding(3) var<storage, read_write> sp_acc  : array<f32>;   // nh*nc*hd
@group(0) @binding(4) var<storage, read_write> sp_mx   : array<f32>;   // nh*nc
@group(0) @binding(5) var<storage, read_write> sp_ln   : array<f32>;   // nh*nc
@group(0) @binding(6) var<uniform>             sp_p    : SapP;

var<workgroup> sp_red: array<f32, 256>;
var<workgroup> sp_w:   array<f32, 256>;
var<workgroup> sp_mval: f32;
var<workgroup> sp_den: f32;

@compute @workgroup_size(256)
fn sparse_attend_part(@builtin(workgroup_id) wid: vec3<u32>,
                      @builtin(local_invocation_index) lid: u32) {
    let h = wid.x;
    let c = wid.y;
    if (h >= sp_p.nh || c >= sp_p.nc) { return; }
    let hd = sp_p.hd;
    let qb = h * hd;
    let lo = c * sp_p.ck;
    var hi = lo + sp_p.ck;
    if (hi > sp_p.m) { hi = sp_p.m; }
    let slot = h * sp_p.nc + c;
    if (lo >= hi) {
        if (lid == 0u) { sp_mx[slot] = -3.0e38; sp_ln[slot] = 0.0; }
        var z = lid;
        loop { if (z >= hd) { break; } sp_acc[slot * hd + z] = 0.0; z = z + 256u; }
        return;
    }
    let len = hi - lo;

    // 1. this chunk's scores.
    //
    // THIRTY-TWO THREADS TO A POSITION, not one. A thread that owns a whole
    // position reads its 512-float row by itself while its neighbour reads a
    // row two kilobytes away: every load is its own cache line and the
    // workgroup coalesces nothing. Splitting across k instead means the 32
    // threads of a group ask for 32 CONSECUTIVE floats at a time, which is
    // one transaction — and the head's positions are then done eight at a
    // time instead of 256, which costs a barrier per eight and is worth it.
    let lane = lid % 32u;
    let grp = lid / 32u;
    var t = grp;
    loop {
        if (t >= len) { break; }
        let p = sp_idx[lo + t];
        var d = 0.0;
        var k = lane;
        loop {
            if (k >= hd) { break; }
            d = d + sp_q[qb + k] * sp_kv[p * hd + k];
            k = k + 32u;
        }
        sp_red[lid] = d;
        workgroupBarrier();
        if (lane == 0u) {
            var sd = 0.0;
            for (var j = 0u; j < 32u; j = j + 1u) { sd = sd + sp_red[grp * 32u + j]; }
            sp_w[t] = sd * sp_p.scale;
        }
        workgroupBarrier();
        t = t + 8u;
    }
    var mx = -3.0e38;
    var t2 = lid;
    loop {
        if (t2 >= len) { break; }
        mx = max(mx, sp_w[t2]);
        t2 = t2 + 256u;
    }
    sp_red[lid] = mx;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { sp_red[lid] = max(sp_red[lid], sp_red[lid + stride]); }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { sp_mval = sp_red[0]; }
    workgroupBarrier();
    let mval = sp_mval;

    // 2. weights against THIS chunk's max, and its denominator
    var den = 0.0;
    t = lid;
    loop {
        if (t >= len) { break; }
        let e = exp(sp_w[t] - mval);
        sp_w[t] = e;
        den = den + e;
        t = t + 256u;
    }
    sp_red[lid] = den;
    workgroupBarrier();
    stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { sp_red[lid] = sp_red[lid] + sp_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        sp_den = sp_red[0];
        sp_mx[slot] = mval;
        sp_ln[slot] = sp_red[0];
    }
    workgroupBarrier();

    // 3. the chunk's unnormalised accumulator, parallel over the output dim
    var k2 = lid;
    loop {
        if (k2 >= hd) { break; }
        var acc = 0.0;
        for (var i = 0u; i < len; i = i + 1u) {
            acc = acc + sp_w[i] * sp_kv[sp_idx[lo + i] * hd + k2];
        }
        sp_acc[slot * hd + k2] = acc;
        k2 = k2 + 256u;
    }
}

struct SamP { nh: u32, hd: u32, nc: u32, _a: u32 };

@group(0) @binding(0) var<storage, read>       sm_acc  : array<f32>;   // nh*nc*hd
@group(0) @binding(1) var<storage, read>       sm_mx   : array<f32>;   // nh*nc
@group(0) @binding(2) var<storage, read>       sm_ln   : array<f32>;   // nh*nc
@group(0) @binding(3) var<storage, read>       sm_sink : array<f32>;   // nh
@group(0) @binding(4) var<storage, read_write> sm_out  : array<f32>;   // nh*hd
@group(0) @binding(5) var<uniform>             sm_p    : SamP;

var<workgroup> sm_m: f32;
var<workgroup> sm_d: f32;

@compute @workgroup_size(256)
fn sparse_attend_merge(@builtin(workgroup_id) wid: vec3<u32>,
                       @builtin(local_invocation_index) lid: u32) {
    let h = wid.x;
    if (h >= sm_p.nh) { return; }
    let hd = sm_p.hd;
    let nc = sm_p.nc;
    // The sink enters the denominator and nothing else — which is what lets a
    // head attend to nothing at all.
    if (lid == 0u) {
        var mm = sm_sink[h];
        for (var c = 0u; c < nc; c = c + 1u) { mm = max(mm, sm_mx[h * nc + c]); }
        var dd = exp(sm_sink[h] - mm);
        for (var c = 0u; c < nc; c = c + 1u) {
            dd = dd + exp(sm_mx[h * nc + c] - mm) * sm_ln[h * nc + c];
        }
        sm_m = mm;
        sm_d = dd;
    }
    workgroupBarrier();
    let mm = sm_m;
    let inv = 1.0 / sm_d;
    var k = lid;
    loop {
        if (k >= hd) { break; }
        var y = 0.0;
        for (var c = 0u; c < nc; c = c + 1u) {
            y = y + exp(sm_mx[h * nc + c] - mm) * sm_acc[(h * nc + c) * hd + k];
        }
        sm_out[h * hd + k] = y * inv;
        k = k + 256u;
    }
}

// A copy, as a dispatch. `copy_buffer_to_buffer` cannot be recorded inside a
// compute pass, so every small copy in the prep path — this token's slot in
// the compressor window, the window shift, the closed window becoming the
// previous one — used to end a pass and start another. At ~30 passes a layer
// that bookkeeping was most of what a decode step cost, so the copies become
// dispatches and the whole layer becomes a handful of passes.

struct BlP { n: u32, soff: u32, doff: u32, _a: u32 };

@group(0) @binding(0) var<storage, read>       bl_src : array<f32>;
@group(0) @binding(1) var<storage, read_write> bl_dst : array<f32>;
@group(0) @binding(2) var<uniform>             bl_p   : BlP;

@compute @workgroup_size(256)
fn blit(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= bl_p.n) { return; }
    bl_dst[bl_p.doff + i] = bl_src[bl_p.soff + i];
}

// Top-k without a sort. The CPU picks by repeated argmax, first maximum wins,
// and returns the winners in index order; the same set falls out of a rank —
// how many positions beat me, counting an equal score as a win only if it
// sits at a lower index — kept when it is below k. Two O(n^2) passes over one
// workgroup, which for the compressed axis is a few thousand comparisons and
// needs neither a scan nor an atomic to stay deterministic.

struct TkP { n: u32, k: u32, _a: u32, _b: u32 };

@group(0) @binding(0) var<storage, read>       tk_s   : array<f32>;   // n
@group(0) @binding(1) var<storage, read_write> tk_idx : array<u32>;   // k
@group(0) @binding(2) var<storage, read_write> tk_cnt : array<u32>;   // 1
@group(0) @binding(3) var<uniform>             tk_p   : TkP;

var<workgroup> tk_keep: array<u32, 4096>;

// A thousand threads. Both passes are O(n²) over the attended list and
// they run in ONE workgroup — at the release's 640 positions that is
// 410 000 comparisons on a single multiprocessor.
@compute @workgroup_size(1024)
fn top_k_index(@builtin(local_invocation_index) lid: u32) {
    let n = tk_p.n;
    var i = lid;
    loop {
        if (i >= n) { break; }
        let si = tk_s[i];
        var keep = 0u;
        if (si > KP_NINF) {
            var rank = 0u;
            for (var j = 0u; j < n; j = j + 1u) {
                let sj = tk_s[j];
                if (sj > KP_NINF) {
                    if (sj > si || (sj == si && j < i)) { rank = rank + 1u; }
                }
            }
            if (rank < tk_p.k) { keep = 1u; }
        }
        tk_keep[i] = keep;
        i = i + 1024u;
    }
    workgroupBarrier();
    // Position among the kept, by index — counted rather than scanned, which
    // costs one more pass and removes every ordering question.
    var m = lid;
    loop {
        if (m >= n) { break; }
        if (tk_keep[m] == 1u) {
            var before = 0u;
            for (var j = 0u; j < m; j = j + 1u) { before = before + tk_keep[j]; }
            tk_idx[before] = m;
        }
        m = m + 1024u;
    }
    workgroupBarrier();
    if (lid == 0u) {
        var total = 0u;
        for (var j = 0u; j < n; j = j + 1u) { total = total + tk_keep[j]; }
        tk_cnt[0] = total;
    }
}

// ── The attended-position list, assembled on the device ─────────────────────
//
// The sliding window first, in cache order, then the compressed positions the
// indexer picked — shifted by the window's CAPACITY, not by how much of it is
// in use, because that is where the compressed region starts in the layer's
// cache buffer.
//
// The length never has to come back from the card: it is `win_len + min(topk,
// finite positions)`, and the host knows both. Only the CONTENTS are a device
// secret, which is what lets the whole token stay in one submission.

struct IbP { win_len: u32, window: u32, k: u32, _p: u32 };
@group(0) @binding(0) var<storage, read>       ib_pick : array<u32>;
@group(0) @binding(1) var<storage, read_write> ib_out  : array<u32>;
@group(0) @binding(2) var<uniform>             ib_p    : IbP;

@compute @workgroup_size(256)
fn idx_build(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i < ib_p.win_len) { ib_out[i] = i; return; }
    let j = i - ib_p.win_len;
    if (j < ib_p.k) { ib_out[i] = ib_p.window + ib_pick[j]; }
}

// ── MoE routing: sqrt-softplus, noaux_tc bias, top-k (DeepSeek-V4) ──────────
//
// The bias shifts the CHOICE and never the weight: the weight of a chosen
// expert is its pre-bias score. Swapping those two — an easy thing to do when
// the bias is right there — leaves a model that still speaks and routes
// slightly wrong forever.
//
// Ranking replaces the repeated argmax and gives selection order for free:
// rank i = how many experts beat it, an equal score counting only from a
// lower index, which is precisely what "first maximum wins" means. Ranks of
// the finite entries are dense from zero, so the count is just how many
// slots got filled.

struct RtP { n: u32, top_k: u32, flags: u32, scale: f32 };
// flags: 1 = bias present, 2 = mask present, 4 = indices forced (hash layers),
//        8 = pin the shared expert in slot top_k with weight 1,
//       16 = the packed set is a SUBSET: rt_map turns a global expert id into
//            a slot, or 0xFFFFFFFF when that expert did not fit on the card,
//       32 = Qwen routing: rank raw logits and softmax over the selected top-k.
//
// A cold pick is not dropped and not substituted — it is handed back. The
// slot gets weight zero so the device contributes nothing for it, and the
// expert's global id and its real weight go into rt_cold for the host to
// finish. Routing therefore still ranges over every expert, which is the
// whole difference between this and a mask.
//
// With bit 8 the output is the `msel`/`mwt` pair the batched expert kernels
// read: top_k routed slots then the shared one, every slot written. Slots the
// router could not fill (a mask that closes too much) get weight ZERO rather
// than being left short — the kernels downstream take a fixed slot count, and
// a stale index with a live weight is how a token gets an expert nobody chose.

@group(0) @binding(0) var<storage, read>       rt_s      : array<f32>;   // n
@group(0) @binding(1) var<storage, read>       rt_bias   : array<f32>;   // n
@group(0) @binding(2) var<storage, read>       rt_mask   : array<u32>;   // n
@group(0) @binding(3) var<storage, read>       rt_forced : array<u32>;   // top_k + optional shared weight bits
@group(0) @binding(4) var<storage, read_write> rt_idx    : array<u32>;   // top_k
@group(0) @binding(5) var<storage, read_write> rt_w      : array<f32>;   // top_k
@group(0) @binding(6) var<storage, read_write> rt_cnt    : array<u32>;   // 1
@group(0) @binding(7) var<uniform>             rt_p      : RtP;
@group(0) @binding(8) var<storage, read>       rt_map    : array<u32>;   // n
@group(0) @binding(9) var<storage, read_write> rt_cold   : array<u32>;   // 2*top_k

var<workgroup> rt_sc:   array<f32, 1024>;   // sqrt(softplus(score))
var<workgroup> rt_sh:   array<f32, 1024>;   // the same, biased and masked
var<workgroup> rt_used: array<u32, 64>;

// A THOUSAND threads. The ranking is O(n²) — for each expert it counts how
// many beat it — and with 256 routed experts that is 65 536 comparisons in
// ONE workgroup, which measured 1.65 ms of a 30.6 ms chain. Counting ranks
// is order-independent, so the tie-break (equal score, lower index wins) is
// untouched by how many threads do the counting.
@compute @workgroup_size(1024)
fn moe_route(@builtin(local_invocation_index) lid: u32) {
    let n = rt_p.n;
    let k = rt_p.top_k;
    let has_bias = (rt_p.flags & 1u) != 0u;
    let has_mask = (rt_p.flags & 2u) != 0u;
    let forced = (rt_p.flags & 4u) != 0u;

    // `shared` is a WGSL reserved word. Naming it that compiled here and
    // failed at pipeline creation, which took the whole context down and made
    // every GPU test pass by skipping.
    let pin_shared = (rt_p.flags & 8u) != 0u;
    let subset = (rt_p.flags & 16u) != 0u;
    let qwen = (rt_p.flags & 32u) != 0u;
    let shared_gated = (rt_p.flags & 64u) != 0u;
    let preweighted = (rt_p.flags & 128u) != 0u;
    if (lid < k) {
        rt_used[lid] = 0u;
        rt_idx[lid] = 0u;
        rt_w[lid] = 0.0;
        rt_cold[2u * lid] = 0xFFFFFFFFu;
        rt_cold[2u * lid + 1u] = 0u;
        // Route-statistics mirror, second half: every global winner and its
        // normalized PRE-bias route weight, regardless of where it lives.
        rt_cold[2u * k + 2u * lid] = 0xFFFFFFFFu;
        rt_cold[2u * k + 2u * lid + 1u] = 0u;
    }
    // The shared expert sits LAST in the packing, which is `n` only when
    // every expert was packed. With a subset it is n_pack, carried in the
    // flags' upper bits — writing `n` there pointed the kernel past the end
    // of the buffer at whatever followed.
    let shared_slot = rt_p.flags >> 8u;
    if (pin_shared && lid == 0u) {
        rt_idx[k] = shared_slot;
        if (shared_gated) {
            rt_w[k] = bitcast<f32>(rt_forced[k]);
        } else {
            rt_w[k] = 1.0;
        }
    }
    // Same reason: the zero-fill is a storage write that the ranking lanes
    // must not race with.
    storageBarrier();
    var i = lid;
    loop {
        if (i >= n) { break; }
        let v = rt_s[i];
        // DeepSeek-V4 ranks sqrt(softplus(logit)); Qwen ranks the raw logit.
        // Qwen's softmax is applied after top-k below.  With top-k
        // renormalisation this is algebraically identical to softmax over all
        // 512 experts followed by division by the selected mass, without an
        // overflow-prone full softmax in workgroup memory.
        var sc = v;
        if (!qwen) {
            var sp = v;
            if (v <= 20.0) { sp = log(1.0 + exp(v)); }
            sc = sqrt(sp);
        }
        rt_sc[i] = sc;
        var sh = sc;
        if (has_bias) { sh = sh + rt_bias[i]; }
        if (has_mask && rt_mask[i] == 0u) { sh = KP_NINF; }
        rt_sh[i] = sh;
        i = i + 1024u;
    }
    workgroupBarrier();

    if (forced) {
        if (lid < k) {
            let e = rt_forced[lid];
            var w = 0.0;
            if (e < n) {
                if (preweighted) { w = rt_bias[e]; }
                else { w = rt_sc[e]; }
            }
            rt_cold[2u * k + 2u * lid] = e;
            rt_cold[2u * k + 2u * lid + 1u] = bitcast<u32>(w);
            if (subset && e < n) {
                let slot = rt_map[e];
                if (slot == 0xFFFFFFFFu) {
                    rt_idx[lid] = 0u;
                    rt_w[lid] = 0.0;
                    rt_cold[2u * lid] = e;
                    rt_cold[2u * lid + 1u] = bitcast<u32>(w);
                } else {
                    rt_idx[lid] = slot;
                    rt_w[lid] = w;
                }
            } else {
                rt_idx[lid] = e;
                rt_w[lid] = w;
            }
            rt_used[lid] = 1u;
        }
    } else {
        var m = lid;
        loop {
            if (m >= n) { break; }
            let si = rt_sh[m];
            if (si > KP_NINF) {
                var rank = 0u;
                for (var j = 0u; j < n; j = j + 1u) {
                    let sj = rt_sh[j];
                    if (sj > KP_NINF) {
                        if (sj > si || (sj == si && j < m)) { rank = rank + 1u; }
                    }
                }
                if (rank < k) {
                    rt_used[rank] = 1u;
                    rt_cold[2u * k + 2u * rank] = m;
                    rt_cold[2u * k + 2u * rank + 1u] = bitcast<u32>(rt_sc[m]);
                    // What the kernel believes it was handed, in the slots the
                    // winners do not use.
                    if (subset) {
                        let slot = rt_map[m];
                        if (slot == 0xFFFFFFFFu) {
                            // Cold: the device computes nothing for it and the
                            // host is told which expert and with what weight.
                            rt_idx[rank] = 0u;
                            rt_w[rank] = 0.0;
                            rt_cold[2u * rank] = m;
                            rt_cold[2u * rank + 1u] = bitcast<u32>(rt_sc[m]);
                        } else {
                            rt_idx[rank] = slot;
                            rt_w[rank] = rt_sc[m];
                        }
                    } else {
                        rt_idx[rank] = m;
                        rt_w[rank] = rt_sc[m];
                    }
                }
            }
            m = m + 1024u;
        }
    }
    // BOTH barriers. The ranking above writes rt_idx/rt_w, which are STORAGE
    // buffers, and the lane that normalises them below reads what every other
    // lane wrote. workgroupBarrier orders workgroup memory only; without the
    // storage barrier those writes need not be visible yet. With 8 experts it
    // happened to work, with 256 it did not — and the failure is a routing
    // weight quietly attached to the wrong expert.
    workgroupBarrier();
    storageBarrier();

    // Normalisation is a handful of terms; one lane keeps the add order fixed.
    if (lid == 0u) {

        var cnt = 0u;
        for (var j = 0u; j < k; j = j + 1u) {
            if (rt_used[j] == 1u) { cnt = cnt + 1u; }
        }
        rt_cnt[0] = cnt;
        // The sum runs over the chosen experts INCLUDING the cold ones — the
        // reference normalises across the whole top-k, and leaving them out
        // would inflate every surviving weight.
        // Qwen weights are exp(logit-max_selected).  Cold and resident picks
        // carry the same raw score at this point, so transform both before
        // the common normalisation.  DeepSeek keeps its sqrt-softplus score.
        var qmx = -3.0e38;
        if (qwen && !preweighted) {
            for (var j = 0u; j < cnt; j = j + 1u) {
                var v = rt_w[j];
                if (rt_cold[2u * j] != 0xFFFFFFFFu) {
                    v = bitcast<f32>(rt_cold[2u * j + 1u]);
                }
                qmx = max(qmx, v);
            }
            for (var j = 0u; j < cnt; j = j + 1u) {
                if (rt_cold[2u * j] != 0xFFFFFFFFu) {
                    rt_cold[2u * j + 1u] =
                        bitcast<u32>(exp(bitcast<f32>(rt_cold[2u * j + 1u]) - qmx));
                } else {
                    rt_w[j] = exp(rt_w[j] - qmx);
                }
                rt_cold[2u * k + 2u * j + 1u] =
                    bitcast<u32>(exp(bitcast<f32>(rt_cold[2u * k + 2u * j + 1u]) - qmx));
            }
        }
        var sum = 0.0;
        for (var j = 0u; j < cnt; j = j + 1u) {
            sum = sum + rt_w[j];
            if (rt_cold[2u * j] != 0xFFFFFFFFu) {
                sum = sum + bitcast<f32>(rt_cold[2u * j + 1u]);
            }
        }
        if (sum > 0.0 && !preweighted) {
            let inv = rt_p.scale / sum;
            for (var j = 0u; j < cnt; j = j + 1u) {
                rt_w[j] = rt_w[j] * inv;
                if (rt_cold[2u * j] != 0xFFFFFFFFu) {
                    rt_cold[2u * j + 1u] =
                        bitcast<u32>(bitcast<f32>(rt_cold[2u * j + 1u]) * inv);
                }
                rt_cold[2u * k + 2u * j + 1u] =
                    bitcast<u32>(bitcast<f32>(rt_cold[2u * k + 2u * j + 1u]) * inv);
            }
        }
        // The shared expert is not part of that normalisation — it rides at
        // weight 1 whatever the router decided.
    }
}

// ── Sparse attention over an index list (DeepSeek-V4) ───────────────────────
//
// Not the sliding-window attention the canonical graph encodes: the keys are
// named by an INDEX LIST — the window's positions followed by whichever
// compressed ones the indexer chose — and one KV vector of a single head's
// width serves all query heads. The learned sink enters the denominator and
// contributes nothing to the numerator, which is what lets a head attend to
// nothing at all.
//
// One workgroup per head. `hd` is 512 in the release, so the accumulator
// fits in workgroup storage and the whole head is one pass.

struct SaP { nh: u32, hd: u32, m: u32, scale: f32 };

@group(0) @binding(0) var<storage, read>       sa_q    : array<f32>;   // nh*hd
@group(0) @binding(1) var<storage, read>       sa_kv   : array<f32>;   // n*hd
@group(0) @binding(2) var<storage, read>       sa_idx  : array<u32>;   // m
@group(0) @binding(3) var<storage, read>       sa_sink : array<f32>;   // nh
@group(0) @binding(4) var<storage, read_write> sa_out  : array<f32>;   // nh*hd
@group(0) @binding(5) var<uniform>             sa_p    : SaP;

var<workgroup> sa_red: array<f32, 256>;
// Scores, then weights, for every attended position. Bounding it here bounds
// the index list: window (128) + index_topk (512) fits with room.
var<workgroup> sa_w: array<f32, 1024>;
var<workgroup> sa_max: f32;
var<workgroup> sa_den: f32;

@compute @workgroup_size(256)
fn sparse_attend(@builtin(workgroup_id) wid: vec3<u32>,
                 @builtin(local_invocation_index) lid: u32) {
    let h = wid.x;
    if (h >= sa_p.nh) { return; }
    let hd = sa_p.hd;
    let m = sa_p.m;
    let qb = h * hd;

    // 1. scores, kept — recomputing them per output dimension would cost
    //    hd times more, and computing them twice was the first draft's waste.
    var mx = sa_sink[h];
    var t = lid;
    loop {
        if (t >= m) { break; }
        let p = sa_idx[t];
        var d = 0.0;
        for (var k = 0u; k < hd; k = k + 1u) {
            d = d + sa_q[qb + k] * sa_kv[p * hd + k];
        }
        let sc = d * sa_p.scale;
        sa_w[t] = sc;
        mx = max(mx, sc);
        t = t + 256u;
    }
    sa_red[lid] = mx;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { sa_red[lid] = max(sa_red[lid], sa_red[lid + stride]); }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    let mval = sa_red[0];
    workgroupBarrier();

    // 2. weights and the denominator, the sink taking its share of the latter
    var den = 0.0;
    t = lid;
    loop {
        if (t >= m) { break; }
        let w = exp(sa_w[t] - mval);
        sa_w[t] = w;
        den = den + w;
        t = t + 256u;
    }
    sa_red[lid] = den;
    workgroupBarrier();
    stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { sa_red[lid] = sa_red[lid] + sa_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { sa_den = sa_red[0] + exp(sa_sink[h] - mval); }
    workgroupBarrier();
    let inv = 1.0 / sa_den;

    // 3. the weighted sum, parallel over the OUTPUT dimension. Splitting by
    //    position instead makes every thread accumulate into the same
    //    sa_acc[k] — a data race that the first draft called "serialised".
    var k = lid;
    loop {
        if (k >= hd) { break; }
        var acc = 0.0;
        for (var i = 0u; i < m; i = i + 1u) {
            acc = acc + sa_w[i] * sa_kv[sa_idx[i] * hd + k];
        }
        sa_out[qb + k] = acc * inv;
        k = k + 256u;
    }
}

// ── Hyper-connections on the device (DeepSeek-V4) ───────────────────────────
//
// The hidden state is `hc` copies of a `dim` vector, and a block folds them
// to one, runs, then expands back through a Sinkhorn-normalized mixing
// matrix. There is no ordinary residual, so this is not an add — it is the
// join between every pair of blocks, and leaving it on the CPU is what
// forces a round trip per layer.
//
// Sizes are small where it matters: hc is 4, mix_hc is 24, and only the fold
// runs over dim. One workgroup owns the whole thing.

struct HcP { hc: u32, dim: u32, iters: u32, eps: f32, nrm: u32, _a: u32, _b: u32, _c: u32 };

@group(0) @binding(0) var<storage, read>       hc_state : array<f32>;   // hc*dim
@group(0) @binding(1) var<storage, read>       hc_mix   : array<f32>;   // mix_hc, raw
@group(0) @binding(2) var<storage, read>       hc_sc    : array<f32>;   // 3
@group(0) @binding(3) var<storage, read>       hc_base  : array<f32>;   // mix_hc
@group(0) @binding(4) var<storage, read_write> hc_fold  : array<f32>;   // dim
@group(0) @binding(5) var<storage, read_write> hc_post  : array<f32>;   // hc
@group(0) @binding(6) var<storage, read_write> hc_comb  : array<f32>;   // hc*hc
@group(0) @binding(7) var<uniform>             hc_p     : HcP;
// The norm that ALWAYS follows the fold. Its own dispatch cost as much as
// the fold did and it reduces over the same vector this workgroup just
// wrote, so it is a phase, not a kernel. `hc_nrm != 0` turns it on.
@group(0) @binding(8) var<storage, read>       hc_nw    : array<f32>;   // dim
@group(0) @binding(9) var<storage, read_write> hc_out   : array<f32>;   // dim

// A thousand threads: this kernel reduces, so it is one workgroup by
// construction, and on one SM only more threads can hide the latency.
var<workgroup> hc_red: array<f32, 1024>;
var<workgroup> hc_pre_w: array<f32, 8>;
var<workgroup> hc_cmb_w: array<f32, 64>;
var<workgroup> hc_rsq: f32;

@compute @workgroup_size(1024)
fn hc_pre_fold(@builtin(local_invocation_index) lid: u32) {
    let hc = hc_p.hc;
    let dim = hc_p.dim;
    let n = hc * dim;

    // rsqrt(mean(state^2) + eps) — the reference scales the mixes by this,
    // and it is a mean over ALL copies, not per copy.
    var acc = 0.0;
    var i = lid;
    loop {
        if (i >= n) { break; }
        let v = hc_state[i];
        acc = acc + v * v;
        i = i + 1024u;
    }
    hc_red[lid] = acc;
    workgroupBarrier();
    var stride = 512u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { hc_red[lid] = hc_red[lid] + hc_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        hc_rsq = inverseSqrt(hc_red[0] / f32(n) + hc_p.eps);
    }
    workgroupBarrier();
    let rsq = hc_rsq;

    // pre / post / comb. Thread 0 does it: hc is 4, and the Sinkhorn is a
    // sequential fixed point over a 4x4 — parallelising it would cost more
    // in barriers than it saves.
    if (lid == 0u) {
        for (var j = 0u; j < hc; j = j + 1u) {
            let m = hc_mix[j] * rsq * hc_sc[0] + hc_base[j];
            hc_pre_w[j] = 1.0 / (1.0 + exp(-m)) + hc_p.eps;
            let m2 = hc_mix[hc + j] * rsq * hc_sc[1] + hc_base[hc + j];
            hc_post[j] = 2.0 / (1.0 + exp(-m2));
        }
        // row softmax, then the alternating normalisation
        for (var j = 0u; j < hc; j = j + 1u) {
            var mx = -1e30;
            for (var k = 0u; k < hc; k = k + 1u) {
                let v = hc_mix[2u * hc + j * hc + k] * rsq * hc_sc[2]
                      + hc_base[2u * hc + j * hc + k];
                hc_cmb_w[j * hc + k] = v;
                mx = max(mx, v);
            }
            var sum = 0.0;
            for (var k = 0u; k < hc; k = k + 1u) {
                let e = exp(hc_cmb_w[j * hc + k] - mx);
                hc_cmb_w[j * hc + k] = e;
                sum = sum + e;
            }
            for (var k = 0u; k < hc; k = k + 1u) {
                hc_cmb_w[j * hc + k] = hc_cmb_w[j * hc + k] / sum + hc_p.eps;
            }
        }
        for (var k = 0u; k < hc; k = k + 1u) {          // first column pass
            var sum = 0.0;
            for (var j = 0u; j < hc; j = j + 1u) { sum = sum + hc_cmb_w[j * hc + k]; }
            for (var j = 0u; j < hc; j = j + 1u) {
                hc_cmb_w[j * hc + k] = hc_cmb_w[j * hc + k] / (sum + hc_p.eps);
            }
        }
        for (var it = 1u; it < hc_p.iters; it = it + 1u) {
            for (var j = 0u; j < hc; j = j + 1u) {
                var sum = 0.0;
                for (var k = 0u; k < hc; k = k + 1u) { sum = sum + hc_cmb_w[j * hc + k]; }
                for (var k = 0u; k < hc; k = k + 1u) {
                    hc_cmb_w[j * hc + k] = hc_cmb_w[j * hc + k] / (sum + hc_p.eps);
                }
            }
            for (var k = 0u; k < hc; k = k + 1u) {
                var sum = 0.0;
                for (var j = 0u; j < hc; j = j + 1u) { sum = sum + hc_cmb_w[j * hc + k]; }
                for (var j = 0u; j < hc; j = j + 1u) {
                    hc_cmb_w[j * hc + k] = hc_cmb_w[j * hc + k] / (sum + hc_p.eps);
                }
            }
        }
        for (var j = 0u; j < hc * hc; j = j + 1u) { hc_comb[j] = hc_cmb_w[j]; }
    }
    workgroupBarrier();

    // fold: y[d] = sum_j pre[j] * state[j*dim + d], and its sum of squares
    var acc3 = 0.0;
    var d = lid;
    loop {
        if (d >= dim) { break; }
        var y = 0.0;
        for (var j = 0u; j < hc; j = j + 1u) {
            y = y + hc_pre_w[j] * hc_state[j * dim + d];
        }
        hc_fold[d] = y;
        acc3 = acc3 + y * y;
        d = d + 1024u;
    }
    if (hc_p.nrm == 0u) { return; }
    hc_red[lid] = acc3;
    workgroupBarrier();
    var st2 = 512u;
    loop {
        if (st2 == 0u) { break; }
        if (lid < st2) { hc_red[lid] = hc_red[lid] + hc_red[lid + st2]; }
        workgroupBarrier();
        st2 = st2 >> 1u;
    }
    let inv = inverseSqrt(hc_red[0] / f32(dim) + hc_p.eps);
    var d2 = lid;
    loop {
        if (d2 >= dim) { break; }
        hc_out[d2] = hc_fold[d2] * inv * hc_nw[d2];
        d2 = d2 + 1024u;
    }
}

// The whole hyper-connection join in ONE dispatch.
//
// Expand, the mix projection, the fold with its Sinkhorn, and the norm are
// four dependent steps over a state of hc·dim floats — 16 384 on the
// release. As four dispatches they cost four kernel launches a piece and
// this happens TWICE a layer, which on 42 layers is 336 launches a token:
// the skip probe put "everything that is neither attention nor MoE" at
// 27.9 ms of a 51.3 ms chain, and this is the largest single item in it.
//
// One workgroup owns all four, with barriers where a dispatch boundary used
// to be. That is a real trade — 256 threads instead of the whole card — but
// the arithmetic is ~400k MACs and a launch is ~30 µs, so the trade is not
// close. `post`/`comb` are read by the expand and WRITTEN by the fold, in
// that order, which the barrier between them makes safe.

struct HbP { hc: u32, dim: u32, iters: u32, eps: f32, mix_hc: u32, _a: u32, _b: u32, _c: u32 };

@group(0) @binding(0)  var<storage, read>       hb_x     : array<f32>;   // dim
@group(0) @binding(1)  var<storage, read>       hb_res   : array<f32>;   // hc*dim
@group(0) @binding(2)  var<storage, read_write> hb_post  : array<f32>;   // hc
@group(0) @binding(3)  var<storage, read_write> hb_comb  : array<f32>;   // hc*hc
@group(0) @binding(4)  var<storage, read>       hb_mixw  : array<f32>;   // mix_hc*hc*dim
@group(0) @binding(5)  var<storage, read>       hb_sc    : array<f32>;   // 3
@group(0) @binding(6)  var<storage, read>       hb_base  : array<f32>;   // mix_hc
@group(0) @binding(7)  var<storage, read>       hb_nw    : array<f32>;   // dim
@group(0) @binding(8)  var<storage, read_write> hb_state : array<f32>;   // hc*dim
@group(0) @binding(9)  var<storage, read_write> hb_fold  : array<f32>;   // dim
@group(0) @binding(10) var<storage, read_write> hb_norm  : array<f32>;   // dim
@group(0) @binding(11) var<uniform>             hb_p     : HbP;

var<workgroup> hb_red: array<f32, 256>;
var<workgroup> hb_mix: array<f32, 64>;
var<workgroup> hb_pre: array<f32, 8>;
var<workgroup> hb_cmb: array<f32, 64>;
var<workgroup> hb_rsq: f32;

const HB_LANES: u32 = 8u;   // threads per mix output; 256/8 = 32 outputs max

@compute @workgroup_size(256)
fn hc_block(@builtin(local_invocation_index) lid: u32) {
    let hc = hb_p.hc;
    let dim = hb_p.dim;
    let n = hc * dim;
    let mh = hb_p.mix_hc;

    // ── 1. expand ─────────────────────────────────────────────────────────
    var i = lid;
    loop {
        if (i >= n) { break; }
        let j = i / dim;
        let d = i % dim;
        var y = hb_post[j] * hb_x[d];
        for (var k = 0u; k < hc; k = k + 1u) {
            y = y + hb_comb[k * hc + j] * hb_res[k * dim + d];
        }
        hb_state[i] = y;
        i = i + 256u;
    }
    workgroupBarrier();

    // ── 2. the mix projection, all outputs at once ────────────────────────
    // Eight threads to an output, striding the shared axis: one barrier for
    // the lot instead of one tree reduction per output.
    let o = lid / HB_LANES;
    let sub = lid % HB_LANES;
    var acc = 0.0;
    if (o < mh) {
        let base = o * n;
        var t = sub;
        loop {
            if (t >= n) { break; }
            acc = acc + hb_mixw[base + t] * hb_state[t];
            t = t + HB_LANES;
        }
    }
    hb_red[lid] = acc;
    workgroupBarrier();
    if (sub == 0u && o < mh) {
        var sm = 0.0;
        for (var t = 0u; t < HB_LANES; t = t + 1u) { sm = sm + hb_red[o * HB_LANES + t]; }
        hb_mix[o] = sm;
    }

    // ── 3. rsqrt(mean(state^2) + eps), over ALL copies ────────────────────
    var acc2 = 0.0;
    var i2 = lid;
    loop {
        if (i2 >= n) { break; }
        let v = hb_state[i2];
        acc2 = acc2 + v * v;
        i2 = i2 + 256u;
    }
    workgroupBarrier();           // hb_red is being reused
    hb_red[lid] = acc2;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { hb_red[lid] = hb_red[lid] + hb_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { hb_rsq = inverseSqrt(hb_red[0] / f32(n) + hb_p.eps); }
    workgroupBarrier();
    let rsq = hb_rsq;

    // ── 4. pre / post / comb, thread 0: hc is 4 and the Sinkhorn is a
    //       sequential fixed point over a 4x4.
    if (lid == 0u) {
        for (var j = 0u; j < hc; j = j + 1u) {
            let m = hb_mix[j] * rsq * hb_sc[0] + hb_base[j];
            hb_pre[j] = 1.0 / (1.0 + exp(-m)) + hb_p.eps;
            let m2 = hb_mix[hc + j] * rsq * hb_sc[1] + hb_base[hc + j];
            hb_post[j] = 2.0 / (1.0 + exp(-m2));
        }
        for (var j = 0u; j < hc; j = j + 1u) {
            var mx = -1e30;
            for (var k = 0u; k < hc; k = k + 1u) {
                let v = hb_mix[2u * hc + j * hc + k] * rsq * hb_sc[2]
                      + hb_base[2u * hc + j * hc + k];
                hb_cmb[j * hc + k] = v;
                mx = max(mx, v);
            }
            var sum = 0.0;
            for (var k = 0u; k < hc; k = k + 1u) {
                let e = exp(hb_cmb[j * hc + k] - mx);
                hb_cmb[j * hc + k] = e;
                sum = sum + e;
            }
            for (var k = 0u; k < hc; k = k + 1u) {
                hb_cmb[j * hc + k] = hb_cmb[j * hc + k] / sum + hb_p.eps;
            }
        }
        for (var k = 0u; k < hc; k = k + 1u) {
            var sum = 0.0;
            for (var j = 0u; j < hc; j = j + 1u) { sum = sum + hb_cmb[j * hc + k]; }
            for (var j = 0u; j < hc; j = j + 1u) {
                hb_cmb[j * hc + k] = hb_cmb[j * hc + k] / (sum + hb_p.eps);
            }
        }
        for (var it = 1u; it < hb_p.iters; it = it + 1u) {
            for (var j = 0u; j < hc; j = j + 1u) {
                var sum = 0.0;
                for (var k = 0u; k < hc; k = k + 1u) { sum = sum + hb_cmb[j * hc + k]; }
                for (var k = 0u; k < hc; k = k + 1u) {
                    hb_cmb[j * hc + k] = hb_cmb[j * hc + k] / (sum + hb_p.eps);
                }
            }
            for (var k = 0u; k < hc; k = k + 1u) {
                var sum = 0.0;
                for (var j = 0u; j < hc; j = j + 1u) { sum = sum + hb_cmb[j * hc + k]; }
                for (var j = 0u; j < hc; j = j + 1u) {
                    hb_cmb[j * hc + k] = hb_cmb[j * hc + k] / (sum + hb_p.eps);
                }
            }
        }
        for (var j = 0u; j < hc * hc; j = j + 1u) { hb_comb[j] = hb_cmb[j]; }
    }
    workgroupBarrier();

    // ── 5. fold, then the norm over it ────────────────────────────────────
    var acc3 = 0.0;
    var d = lid;
    loop {
        if (d >= dim) { break; }
        var y = 0.0;
        for (var j = 0u; j < hc; j = j + 1u) {
            y = y + hb_pre[j] * hb_state[j * dim + d];
        }
        hb_fold[d] = y;
        acc3 = acc3 + y * y;
        d = d + 256u;
    }
    hb_red[lid] = acc3;
    workgroupBarrier();
    stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { hb_red[lid] = hb_red[lid] + hb_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    let inv = inverseSqrt(hb_red[0] / f32(dim) + hb_p.eps);
    var d2 = lid;
    loop {
        if (d2 >= dim) { break; }
        hb_norm[d2] = hb_fold[d2] * inv * hb_nw[d2];
        d2 = d2 + 256u;
    }
}

// expand: state[j*dim+d] = post[j]*x[d] + sum_k comb[k*hc+j]*residual[k*dim+d]
// Summing over the FIRST index of comb, as the reference does — reading it
// the other way transposes the mixing and is not detectable by eye.
@group(0) @binding(0) var<storage, read>       he_x    : array<f32>;   // dim
@group(0) @binding(1) var<storage, read>       he_res  : array<f32>;   // hc*dim
@group(0) @binding(2) var<storage, read>       he_post : array<f32>;   // hc
@group(0) @binding(3) var<storage, read>       he_comb : array<f32>;   // hc*hc
@group(0) @binding(4) var<storage, read_write> he_out  : array<f32>;   // hc*dim
@group(0) @binding(5) var<uniform>             he_p    : HcP;
@compute @workgroup_size(256)
fn hc_post_expand(@builtin(global_invocation_id) gid: vec3<u32>) {
    let hc = he_p.hc;
    let dim = he_p.dim;
    let i = gid.x;
    if (i >= hc * dim) { return; }
    let j = i / dim;
    let d = i % dim;
    var y = he_post[j] * he_x[d];
    for (var k = 0u; k < hc; k = k + 1u) {
        y = y + he_comb[k * hc + j] * he_res[k * dim + d];
    }
    he_out[i] = y;
}

// ── Token-axis twins for the DeepSeek-V4 batched frame ─────────────────────
//
// The batch used to encode one FRAME per token: at B=5 that is five times the
// dispatches of a single token, and the chain is launch-latency-bound, so the
// batch was slower per token than the walk. These twins carry the token on a
// grid axis instead: per-token operands live in ONE buffer strided by the
// token, weights are read once per dispatch. Every kernel is an arithmetic
// copy of its single-token original — same add order, same tie-breaks — so
// greedy parity carries over term for term.

// rope_heads with grid (head, token): x is [b, nh*hd], the position comes
// from a per-token table instead of a uniform.
@group(0) @binding(0) var<storage, read_write> br_x    : array<f32>;
@group(0) @binding(1) var<storage, read>       br_freq : array<f32>;
@group(0) @binding(2) var<uniform>             br_p    : RpP;
@group(0) @binding(3) var<storage, read>       br_meta : array<vec2<f32>>;  // [pos, eps] per token
var<workgroup> br_red: array<f32, 256>;
@compute @workgroup_size(256)
fn bt_rope_heads(@builtin(workgroup_id) wid: vec3<u32>,
                 @builtin(local_invocation_index) lid: u32) {
    let h = wid.x;
    let t = wid.y;
    if (h >= br_p.nh) { return; }
    let hd = br_p.hd;
    let rd = br_p.rd;
    let b = t * br_p.nh * hd + h * hd;
    if ((br_p.flags & 1u) != 0u) {
        var acc = 0.0;
        var i = lid;
        loop {
            if (i >= hd) { break; }
            let v = br_x[b + i];
            acc = acc + v * v;
            i = i + 256u;
        }
        br_red[lid] = acc;
        workgroupBarrier();
        var stride = 128u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { br_red[lid] = br_red[lid] + br_red[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        let inv = inverseSqrt(br_red[0] / f32(hd) + br_meta[t].y);
        workgroupBarrier();
        var j = lid;
        loop {
            if (j >= hd) { break; }
            br_x[b + j] = br_x[b + j] * inv;
            j = j + 256u;
        }
        workgroupBarrier();
    }
    let base = b + hd - rd;
    let pos = br_meta[t].x;
    var tt = lid;
    loop {
        if (tt >= rd / 2u) { break; }
        let th = pos * br_freq[tt];
        var sn = sin(th);
        let cs = cos(th);
        if ((br_p.flags & 2u) != 0u) { sn = -sn; }
        let a = br_x[base + 2u * tt];
        let c = br_x[base + 2u * tt + 1u];
        br_x[base + 2u * tt] = a * cs - c * sn;
        br_x[base + 2u * tt + 1u] = a * sn + c * cs;
        tt = tt + 256u;
    }
}

// hc_pre_fold with grid (token): all per-token operands strided.
// `hc_p._a` carries mix_hc (the mixes stride).
@group(0) @binding(0) var<storage, read>       bhf_state : array<f32>;   // b, hc*dim
@group(0) @binding(1) var<storage, read>       bhf_mix   : array<f32>;   // b, mix_hc
@group(0) @binding(2) var<storage, read>       bhf_sc    : array<f32>;   // 3
@group(0) @binding(3) var<storage, read>       bhf_base  : array<f32>;   // mix_hc
@group(0) @binding(4) var<storage, read_write> bhf_fold  : array<f32>;   // b, dim
@group(0) @binding(5) var<storage, read_write> bhf_post  : array<f32>;   // b, hc
@group(0) @binding(6) var<storage, read_write> bhf_comb  : array<f32>;   // b, hc*hc
@group(0) @binding(7) var<uniform>             bhf_p     : HcP;
@group(0) @binding(8) var<storage, read>       bhf_nw    : array<f32>;   // dim
@group(0) @binding(9) var<storage, read_write> bhf_out   : array<f32>;   // b, dim
var<workgroup> bhf_red: array<f32, 1024>;
var<workgroup> bhf_pre_w: array<f32, 8>;
var<workgroup> bhf_cmb_w: array<f32, 64>;
var<workgroup> bhf_rsq: f32;
@compute @workgroup_size(1024)
fn bt_hc_pre_fold(@builtin(workgroup_id) wid: vec3<u32>,
                  @builtin(local_invocation_index) lid: u32) {
    let hc = bhf_p.hc;
    let dim = bhf_p.dim;
    let n = hc * dim;
    let t = wid.x;
    let sb = t * n;
    let mb = t * bhf_p._a;
    var acc = 0.0;
    var i = lid;
    loop {
        if (i >= n) { break; }
        let v = bhf_state[sb + i];
        acc = acc + v * v;
        i = i + 1024u;
    }
    bhf_red[lid] = acc;
    workgroupBarrier();
    var stride = 512u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { bhf_red[lid] = bhf_red[lid] + bhf_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        bhf_rsq = inverseSqrt(bhf_red[0] / f32(n) + bhf_p.eps);
    }
    workgroupBarrier();
    let rsq = bhf_rsq;
    if (lid == 0u) {
        for (var j = 0u; j < hc; j = j + 1u) {
            let m = bhf_mix[mb + j] * rsq * bhf_sc[0] + bhf_base[j];
            bhf_pre_w[j] = 1.0 / (1.0 + exp(-m)) + bhf_p.eps;
            let m2 = bhf_mix[mb + hc + j] * rsq * bhf_sc[1] + bhf_base[hc + j];
            bhf_post[t * hc + j] = 2.0 / (1.0 + exp(-m2));
        }
        for (var j = 0u; j < hc; j = j + 1u) {
            var mx = -1e30;
            for (var k = 0u; k < hc; k = k + 1u) {
                let v = bhf_mix[mb + 2u * hc + j * hc + k] * rsq * bhf_sc[2]
                      + bhf_base[2u * hc + j * hc + k];
                bhf_cmb_w[j * hc + k] = v;
                mx = max(mx, v);
            }
            var sum = 0.0;
            for (var k = 0u; k < hc; k = k + 1u) {
                let e = exp(bhf_cmb_w[j * hc + k] - mx);
                bhf_cmb_w[j * hc + k] = e;
                sum = sum + e;
            }
            for (var k = 0u; k < hc; k = k + 1u) {
                bhf_cmb_w[j * hc + k] = bhf_cmb_w[j * hc + k] / sum + bhf_p.eps;
            }
        }
        for (var k = 0u; k < hc; k = k + 1u) {
            var sum = 0.0;
            for (var j = 0u; j < hc; j = j + 1u) { sum = sum + bhf_cmb_w[j * hc + k]; }
            for (var j = 0u; j < hc; j = j + 1u) {
                bhf_cmb_w[j * hc + k] = bhf_cmb_w[j * hc + k] / (sum + bhf_p.eps);
            }
        }
        for (var it = 1u; it < bhf_p.iters; it = it + 1u) {
            for (var j = 0u; j < hc; j = j + 1u) {
                var sum = 0.0;
                for (var k = 0u; k < hc; k = k + 1u) { sum = sum + bhf_cmb_w[j * hc + k]; }
                for (var k = 0u; k < hc; k = k + 1u) {
                    bhf_cmb_w[j * hc + k] = bhf_cmb_w[j * hc + k] / (sum + bhf_p.eps);
                }
            }
            for (var k = 0u; k < hc; k = k + 1u) {
                var sum = 0.0;
                for (var j = 0u; j < hc; j = j + 1u) { sum = sum + bhf_cmb_w[j * hc + k]; }
                for (var j = 0u; j < hc; j = j + 1u) {
                    bhf_cmb_w[j * hc + k] = bhf_cmb_w[j * hc + k] / (sum + bhf_p.eps);
                }
            }
        }
        for (var j = 0u; j < hc * hc; j = j + 1u) { bhf_comb[t * hc * hc + j] = bhf_cmb_w[j]; }
    }
    workgroupBarrier();
    var acc3 = 0.0;
    var d = lid;
    loop {
        if (d >= dim) { break; }
        var y = 0.0;
        for (var j = 0u; j < hc; j = j + 1u) {
            y = y + bhf_pre_w[j] * bhf_state[sb + j * dim + d];
        }
        bhf_fold[t * dim + d] = y;
        acc3 = acc3 + y * y;
        d = d + 1024u;
    }
    if (bhf_p.nrm == 0u) { return; }
    bhf_red[lid] = acc3;
    workgroupBarrier();
    var st2 = 512u;
    loop {
        if (st2 == 0u) { break; }
        if (lid < st2) { bhf_red[lid] = bhf_red[lid] + bhf_red[lid + st2]; }
        workgroupBarrier();
        st2 = st2 >> 1u;
    }
    let inv = inverseSqrt(bhf_red[0] / f32(dim) + bhf_p.eps);
    var d2 = lid;
    loop {
        if (d2 >= dim) { break; }
        bhf_out[t * dim + d2] = bhf_fold[t * dim + d2] * inv * bhf_nw[d2];
        d2 = d2 + 1024u;
    }
}

// hc_post_expand with grid (ceil(hc*dim/256), token).
@group(0) @binding(0) var<storage, read>       bhe_x    : array<f32>;   // b, dim
@group(0) @binding(1) var<storage, read>       bhe_res  : array<f32>;   // b, hc*dim
@group(0) @binding(2) var<storage, read>       bhe_post : array<f32>;   // b, hc
@group(0) @binding(3) var<storage, read>       bhe_comb : array<f32>;   // b, hc*hc
@group(0) @binding(4) var<storage, read_write> bhe_out  : array<f32>;   // b, hc*dim
@group(0) @binding(5) var<uniform>             bhe_p    : HcP;
@compute @workgroup_size(256)
fn bt_hc_post_expand(@builtin(workgroup_id) wid: vec3<u32>,
                     @builtin(local_invocation_id) lid: vec3<u32>) {
    let hc = bhe_p.hc;
    let dim = bhe_p.dim;
    let i = wid.x * 256u + lid.x;
    if (i >= hc * dim) { return; }
    let t = wid.y;
    let j = i / dim;
    let d = i % dim;
    var y = bhe_post[t * hc + j] * bhe_x[t * dim + d];
    for (var k = 0u; k < hc; k = k + 1u) {
        y = y + bhe_comb[t * hc * hc + k * hc + j] * bhe_res[t * hc * dim + k * dim + d];
    }
    bhe_out[t * hc * dim + i] = y;
}

// f32_matvec_w with grid (row, token): the weight is read for the whole
// batch, x and y stride by the token.
@group(0) @binding(0) var<storage, read>       bfw_w : array<f32>;
@group(0) @binding(1) var<storage, read>       bfw_x : array<f32>;   // b, cols
@group(0) @binding(2) var<storage, read_write> bfw_y : array<f32>;   // b, rows
@group(0) @binding(3) var<uniform>             bfw_p : F32WP;
var<workgroup> bfw_part: array<f32, 256>;
@compute @workgroup_size(256)
fn bt_f32_matvec_w(@builtin(workgroup_id) wid: vec3<u32>,
                   @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    if (row >= bfw_p.rows) { return; }
    let t = wid.y;
    let base = row * bfw_p.cols;
    let xb = t * bfw_p.cols;
    var acc = 0.0;
    var i = lid;
    loop {
        if (i >= bfw_p.cols) { break; }
        acc = acc + bfw_w[base + i] * bfw_x[xb + i];
        i = i + 256u;
    }
    bfw_part[lid] = acc;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { bfw_part[lid] = bfw_part[lid] + bfw_part[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    if (lid == 0u) { bfw_y[t * bfw_p.rows + row] = bfw_part[0]; }
}

// f32_matvec_x with grid (row, token): the 1024-thread tree of the walk's
// few-row path, unchanged — the batch must reduce in the SAME order or a
// near-tied indexer top-k flips and the verified text is not the walked one.
@group(0) @binding(0) var<storage, read>       bfx_w : array<f32>;
@group(0) @binding(1) var<storage, read>       bfx_x : array<f32>;   // b, cols
@group(0) @binding(2) var<storage, read_write> bfx_y : array<f32>;   // b, rows
@group(0) @binding(3) var<uniform>             bfx_p : F32WP;
var<workgroup> bfx_part: array<f32, 1024>;
@compute @workgroup_size(1024)
fn bt_f32_matvec_x(@builtin(workgroup_id) wid: vec3<u32>,
                   @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    if (row >= bfx_p.rows) { return; }
    let t = wid.y;
    let base = row * bfx_p.cols;
    let xb = t * bfx_p.cols;
    var acc = 0.0;
    var i = lid;
    loop {
        if (i >= bfx_p.cols) { break; }
        acc = acc + bfx_w[base + i] * bfx_x[xb + i];
        i = i + 1024u;
    }
    bfx_part[lid] = acc;
    workgroupBarrier();
    var stride = 512u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { bfx_part[lid] = bfx_part[lid] + bfx_part[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    if (lid == 0u) { bfx_y[t * bfx_p.rows + row] = bfx_part[0]; }
}

// moe_route with grid (token): logits, winners, weights, counts, forced rows
// and cold mirrors all stride by the token; the bias, mask and slot map are
// the layer's and shared.
@group(0) @binding(0) var<storage, read>       bt_s      : array<f32>;   // b, n
@group(0) @binding(1) var<storage, read>       bt_bias   : array<f32>;   // n
@group(0) @binding(2) var<storage, read>       bt_mask   : array<u32>;   // n
@group(0) @binding(3) var<storage, read>       bt_forced : array<u32>;   // b, top_k
@group(0) @binding(4) var<storage, read_write> bt_idx    : array<u32>;   // b, top_k+1
@group(0) @binding(5) var<storage, read_write> bt_w      : array<f32>;   // b, top_k+1
@group(0) @binding(6) var<storage, read_write> bt_cnt    : array<u32>;   // b
@group(0) @binding(7) var<uniform>             bt_p      : RtP;
@group(0) @binding(8) var<storage, read>       bt_map    : array<u32>;   // n
@group(0) @binding(9) var<storage, read_write> bt_cold   : array<u32>;   // b, 4*top_k
var<workgroup> btr_sc:   array<f32, 1024>;
var<workgroup> btr_sh:   array<f32, 1024>;
var<workgroup> btr_used: array<u32, 64>;
@compute @workgroup_size(1024)
fn bt_moe_route(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let n = bt_p.n;
    let k = bt_p.top_k;
    let t = wid.x;
    let lb = t * n;
    let ob = t * (k + 1u);
    let cb = t * 4u * k;
    let has_bias = (bt_p.flags & 1u) != 0u;
    let has_mask = (bt_p.flags & 2u) != 0u;
    let forced = (bt_p.flags & 4u) != 0u;
    let pin_shared = (bt_p.flags & 8u) != 0u;
    let subset = (bt_p.flags & 16u) != 0u;
    if (lid < k) {
        btr_used[lid] = 0u;
        bt_idx[ob + lid] = 0u;
        bt_w[ob + lid] = 0.0;
        bt_cold[cb + 2u * lid] = 0xFFFFFFFFu;
        bt_cold[cb + 2u * lid + 1u] = 0u;
        bt_cold[cb + 2u * k + 2u * lid] = 0xFFFFFFFFu;
        bt_cold[cb + 2u * k + 2u * lid + 1u] = 0u;
    }
    let shared_slot = bt_p.flags >> 8u;
    if (pin_shared && lid == 0u) {
        bt_idx[ob + k] = shared_slot;
        bt_w[ob + k] = 1.0;
    }
    storageBarrier();
    var i = lid;
    loop {
        if (i >= n) { break; }
        let v = bt_s[lb + i];
        var sp = v;
        if (v <= 20.0) { sp = log(1.0 + exp(v)); }
        let sc = sqrt(sp);
        btr_sc[i] = sc;
        var sh = sc;
        if (has_bias) { sh = sh + bt_bias[i]; }
        if (has_mask && bt_mask[i] == 0u) { sh = KP_NINF; }
        btr_sh[i] = sh;
        i = i + 1024u;
    }
    workgroupBarrier();
    if (forced) {
        if (lid < k) {
            let e = bt_forced[t * k + lid];
            var w = 0.0;
            if (e < n) { w = btr_sc[e]; }
            bt_cold[cb + 2u * k + 2u * lid] = e;
            bt_cold[cb + 2u * k + 2u * lid + 1u] = bitcast<u32>(w);
            if (subset && e < n) {
                let slot = bt_map[e];
                if (slot == 0xFFFFFFFFu) {
                    bt_idx[ob + lid] = 0u;
                    bt_w[ob + lid] = 0.0;
                    bt_cold[cb + 2u * lid] = e;
                    bt_cold[cb + 2u * lid + 1u] = bitcast<u32>(w);
                } else {
                    bt_idx[ob + lid] = slot;
                    bt_w[ob + lid] = w;
                }
            } else {
                bt_idx[ob + lid] = e;
                bt_w[ob + lid] = w;
            }
            btr_used[lid] = 1u;
        }
    } else {
        var m = lid;
        loop {
            if (m >= n) { break; }
            let si = btr_sh[m];
            if (si > KP_NINF) {
                var rank = 0u;
                for (var j = 0u; j < n; j = j + 1u) {
                    let sj = btr_sh[j];
                    if (sj > KP_NINF) {
                        if (sj > si || (sj == si && j < m)) { rank = rank + 1u; }
                    }
                }
                if (rank < k) {
                    btr_used[rank] = 1u;
                    bt_cold[cb + 2u * k + 2u * rank] = m;
                    bt_cold[cb + 2u * k + 2u * rank + 1u] = bitcast<u32>(btr_sc[m]);
                    if (subset) {
                        let slot = bt_map[m];
                        if (slot == 0xFFFFFFFFu) {
                            bt_idx[ob + rank] = 0u;
                            bt_w[ob + rank] = 0.0;
                            bt_cold[cb + 2u * rank] = m;
                            bt_cold[cb + 2u * rank + 1u] = bitcast<u32>(btr_sc[m]);
                        } else {
                            bt_idx[ob + rank] = slot;
                            bt_w[ob + rank] = btr_sc[m];
                        }
                    } else {
                        bt_idx[ob + rank] = m;
                        bt_w[ob + rank] = btr_sc[m];
                    }
                }
            }
            m = m + 1024u;
        }
    }
    workgroupBarrier();
    storageBarrier();
    if (lid == 0u) {
        var cnt = 0u;
        for (var j = 0u; j < k; j = j + 1u) {
            if (btr_used[j] == 1u) { cnt = cnt + 1u; }
        }
        bt_cnt[t] = cnt;
        var sum = 0.0;
        for (var j = 0u; j < cnt; j = j + 1u) {
            sum = sum + bt_w[ob + j];
            if (bt_cold[cb + 2u * j] != 0xFFFFFFFFu) {
                sum = sum + bitcast<f32>(bt_cold[cb + 2u * j + 1u]);
            }
        }
        if (sum > 0.0) {
            let inv = bt_p.scale / sum;
            for (var j = 0u; j < cnt; j = j + 1u) {
                bt_w[ob + j] = bt_w[ob + j] * inv;
                if (bt_cold[cb + 2u * j] != 0xFFFFFFFFu) {
                    bt_cold[cb + 2u * j + 1u] =
                        bitcast<u32>(bitcast<f32>(bt_cold[cb + 2u * j + 1u]) * inv);
                }
                bt_cold[cb + 2u * k + 2u * j + 1u] =
                    bitcast<u32>(bitcast<f32>(bt_cold[cb + 2u * k + 2u * j + 1u]) * inv);
            }
        }
    }
}

// moe_gate_up_q2tp with grid (row, slot, token).
var<workgroup> btg_pg: array<f32, 64>;
var<workgroup> btg_pu: array<f32, 64>;
@compute @workgroup_size(64)
fn bt_moe_gate_up_q2tp(@builtin(workgroup_id) wid: vec3<u32>,
                       @builtin(local_invocation_index) lid: u32) {
    let row = wid.x;
    let slot = wid.y;
    let t = wid.z;
    let gpr = mg_p.gpr;
    let rows = mg_p.inter;
    let base16 = mg_sel[t * mg_p.slots + slot] * mg_p.mat16;
    let nib16 = base16 + row * gpr * 4u;
    let par16 = base16 + rows * gpr * 4u + row * 2u;
    let cst = (gpr * 5u + 7u) / 8u;
    let cod8 = (base16 + rows * gpr * 4u + rows * 2u) * 2u + row * cst;
    let xb4 = t * gpr * 8u;
    let gl = unpack2x16float(mg_g16(par16) | (mg_g16(par16 + 1u) << 16u));
    let ul = unpack2x16float(mg_u16f(par16) | (mg_u16f(par16 + 1u) << 16u));
    var ag = 0.0;
    var au = 0.0;
    for (var g = lid; g < gpr; g = g + 64u) {
        let bit = g * 5u;
        let cb = bit >> 3u;
        let shf = bit & 7u;
        var cg = mgp_gu8(cod8 + cb);
        var cu = mgp_uu8(cod8 + cb);
        if (shf > 3u) {
            cg = cg | (mgp_gu8(cod8 + cb + 1u) << 8u);
            cu = cu | (mgp_uu8(cod8 + cb + 1u) << 8u);
        }
        let cgv = (cg >> shf) & 31u;
        let cuv = (cu >> shf) & 31u;
        let sg = select(exp2(gl.x + f32(max(cgv, 1u) - 1u) * gl.y), 0.0, cgv == 0u);
        let su = select(exp2(ul.x + f32(max(cuv, 1u) - 1u) * ul.y), 0.0, cuv == 0u);
        let w32 = (nib16 + g * 4u) >> 1u;
        let xq = xb4 + g * 8u;
        let x0 = mg_xv[xq];      let x1 = mg_xv[xq + 1u];
        let x2 = mg_xv[xq + 2u]; let x3 = mg_xv[xq + 3u];
        let x4 = mg_xv[xq + 4u]; let x5 = mg_xv[xq + 5u];
        let x6 = mg_xv[xq + 6u]; let x7 = mg_xv[xq + 7u];
        let dg = mg_dot16v(mg_gw[w32], x0, x1, x2, x3)
               + mg_dot16v(mg_gw[w32 + 1u], x4, x5, x6, x7);
        let du = mg_dot16v(mg_uw[w32], x0, x1, x2, x3)
               + mg_dot16v(mg_uw[w32 + 1u], x4, x5, x6, x7);
        ag = ag + sg * dg;
        au = au + su * du;
    }
    btg_pg[lid] = ag;
    btg_pu[lid] = au;
    workgroupBarrier();
    var stride = 32u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) {
            btg_pg[lid] = btg_pg[lid] + btg_pg[lid + stride];
            btg_pu[lid] = btg_pu[lid] + btg_pu[lid + stride];
        }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        let g = btg_pg[0];
        var gg = g;
        var uu = btg_pu[0];
        if (mg_p.lim > 0.0) {
            uu = clamp(uu, -mg_p.lim, mg_p.lim);
            gg = min(gg, mg_p.lim);
        }
        mg_act[(t * mg_p.slots + slot) * mg_p.inter + row] = (gg / (1.0 + exp(-gg))) * uu;
    }
}

// Four rows of the SwiGLU inputs to one 64-thread workgroup: the x spans
// load once per group and feed all four rows' gate AND up tiles. Each
// row's group order, accumulation and tree are the one-row kernel's, so
// the activations are bit-identical.
var<workgroup> btg4_pg: array<f32, 64>;
var<workgroup> btg4_pu: array<f32, 64>;
@compute @workgroup_size(64)
fn bt_moe_gate_up_q2tp_r4(@builtin(workgroup_id) wid: vec3<u32>,
                          @builtin(local_invocation_index) lid: u32) {
    let row0 = wid.x * 4u;
    let slot = wid.y;
    let t = wid.z;
    let gpr = mg_p.gpr;
    let rows = mg_p.inter;
    let base16 = mg_sel[t * mg_p.slots + slot] * mg_p.mat16;
    let cst = (gpr * 5u + 7u) / 8u;
    let cod0 = (base16 + rows * gpr * 4u + rows * 2u) * 2u;
    let xb4 = t * gpr * 8u;
    var ag0 = 0.0; var au0 = 0.0;
    var ag1 = 0.0; var au1 = 0.0;
    var ag2 = 0.0; var au2 = 0.0;
    var ag3 = 0.0; var au3 = 0.0;
    for (var g = lid; g < gpr; g = g + 64u) {
        let bit = g * 5u;
        let cb0 = bit >> 3u;
        let shf = bit & 7u;
        let xq = xb4 + g * 8u;
        let x0 = mg_xv[xq];      let x1 = mg_xv[xq + 1u];
        let x2 = mg_xv[xq + 2u]; let x3 = mg_xv[xq + 3u];
        let x4 = mg_xv[xq + 4u]; let x5 = mg_xv[xq + 5u];
        let x6 = mg_xv[xq + 6u]; let x7 = mg_xv[xq + 7u];
        for (var r = 0u; r < 4u; r = r + 1u) {
            let row = row0 + r;
            if (row >= rows) { break; }
            let par16 = base16 + rows * gpr * 4u + row * 2u;
            let gl = unpack2x16float(mg_g16(par16) | (mg_g16(par16 + 1u) << 16u));
            let ul = unpack2x16float(mg_u16f(par16) | (mg_u16f(par16 + 1u) << 16u));
            let cod8 = cod0 + row * cst;
            var cg = mgp_gu8(cod8 + cb0);
            var cu = mgp_uu8(cod8 + cb0);
            if (shf > 3u) {
                cg = cg | (mgp_gu8(cod8 + cb0 + 1u) << 8u);
                cu = cu | (mgp_uu8(cod8 + cb0 + 1u) << 8u);
            }
            let cgv = (cg >> shf) & 31u;
            let cuv = (cu >> shf) & 31u;
            let sg = select(exp2(gl.x + f32(max(cgv, 1u) - 1u) * gl.y), 0.0, cgv == 0u);
            let su = select(exp2(ul.x + f32(max(cuv, 1u) - 1u) * ul.y), 0.0, cuv == 0u);
            let nib16 = base16 + row * gpr * 4u;
            let w32 = (nib16 + g * 4u) >> 1u;
            let dg = mg_dot16v(mg_gw[w32], x0, x1, x2, x3)
                   + mg_dot16v(mg_gw[w32 + 1u], x4, x5, x6, x7);
            let du = mg_dot16v(mg_uw[w32], x0, x1, x2, x3)
                   + mg_dot16v(mg_uw[w32 + 1u], x4, x5, x6, x7);
            if (r == 0u) { ag0 = ag0 + sg * dg; au0 = au0 + su * du; }
            if (r == 1u) { ag1 = ag1 + sg * dg; au1 = au1 + su * du; }
            if (r == 2u) { ag2 = ag2 + sg * dg; au2 = au2 + su * du; }
            if (r == 3u) { ag3 = ag3 + sg * dg; au3 = au3 + su * du; }
        }
    }
    for (var r = 0u; r < 4u; r = r + 1u) {
        var ag = ag0; var au = au0;
        if (r == 1u) { ag = ag1; au = au1; }
        if (r == 2u) { ag = ag2; au = au2; }
        if (r == 3u) { ag = ag3; au = au3; }
        btg4_pg[lid] = ag;
        btg4_pu[lid] = au;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) {
                btg4_pg[lid] = btg4_pg[lid] + btg4_pg[lid + stride];
                btg4_pu[lid] = btg4_pu[lid] + btg4_pu[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u && row0 + r < rows) {
            var gg = btg4_pg[0];
            var uu = btg4_pu[0];
            if (mg_p.lim > 0.0) {
                uu = clamp(uu, -mg_p.lim, mg_p.lim);
                gg = min(gg, mg_p.lim);
            }
            mg_act[(t * mg_p.slots + slot) * mg_p.inter + row0 + r] =
                (gg / (1.0 + exp(-gg))) * uu;
        }
        workgroupBarrier();
    }
}

// sparse_attend with grid (head, token): q and out stride by nh*hd, the
// index list by `m` (the uniform carries the per-token STRIDE of the list
// buffer), the attended count comes from a per-token table. The KV cache is
// the layer's and shared — earlier tokens' appends are ordered by the pass.
@group(0) @binding(0) var<storage, read>       bsa_q    : array<f32>;   // b, nh*hd
@group(0) @binding(1) var<storage, read>       bsa_kv   : array<f32>;
@group(0) @binding(2) var<storage, read>       bsa_idx  : array<u32>;   // b, stride
@group(0) @binding(3) var<storage, read>       bsa_sink : array<f32>;   // nh
@group(0) @binding(4) var<storage, read_write> bsa_out  : array<f32>;   // b, nh*hd
@group(0) @binding(5) var<uniform>             bsa_p    : SaP;          // m = idx stride
@group(0) @binding(6) var<storage, read>       bsa_m    : array<u32>;   // b
var<workgroup> bsa_red: array<f32, 256>;
var<workgroup> bsa_w: array<f32, 1024>;
var<workgroup> bsa_den: f32;
@compute @workgroup_size(256)
fn bt_sparse_attend(@builtin(workgroup_id) wid: vec3<u32>,
                    @builtin(local_invocation_index) lid: u32) {
    let h = wid.x;
    if (h >= bsa_p.nh) { return; }
    let t = wid.y;
    let hd = bsa_p.hd;
    let m = bsa_m[t];
    let ib = t * bsa_p.m;
    let qb = t * bsa_p.nh * hd + h * hd;
    var mx = bsa_sink[h];
    var i = lid;
    loop {
        if (i >= m) { break; }
        let p = bsa_idx[ib + i];
        var d = 0.0;
        for (var k = 0u; k < hd; k = k + 1u) {
            d = d + bsa_q[qb + k] * bsa_kv[p * hd + k];
        }
        let sc = d * bsa_p.scale;
        bsa_w[i] = sc;
        mx = max(mx, sc);
        i = i + 256u;
    }
    bsa_red[lid] = mx;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { bsa_red[lid] = max(bsa_red[lid], bsa_red[lid + stride]); }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    let mval = bsa_red[0];
    workgroupBarrier();
    var den = 0.0;
    i = lid;
    loop {
        if (i >= m) { break; }
        let w = exp(bsa_w[i] - mval);
        bsa_w[i] = w;
        den = den + w;
        i = i + 256u;
    }
    bsa_red[lid] = den;
    workgroupBarrier();
    stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { bsa_red[lid] = bsa_red[lid] + bsa_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { bsa_den = bsa_red[0] + exp(bsa_sink[h] - mval); }
    workgroupBarrier();
    let inv = 1.0 / bsa_den;
    var k = lid;
    loop {
        if (k >= hd) { break; }
        var acc = 0.0;
        for (var j = 0u; j < m; j = j + 1u) {
            acc = acc + bsa_w[j] * bsa_kv[bsa_idx[ib + j] * hd + k];
        }
        bsa_out[qb + k] = acc * inv;
        k = k + 256u;
    }
}


// The whole hyper-connection join, fused, with the token on the grid: the
// four dependent dispatches per half become ONE link of the pass's critical
// path per half. The B=1 trade was negative (measured); at B=5 each link
// carries five tokens and the launch bookkeeping amortizes.
@group(0) @binding(0)  var<storage, read>       bhb_x     : array<f32>;   // b, dim
@group(0) @binding(1)  var<storage, read>       bhb_res   : array<f32>;   // b, hc*dim
@group(0) @binding(2)  var<storage, read_write> bhb_post  : array<f32>;   // b, hc
@group(0) @binding(3)  var<storage, read_write> bhb_comb  : array<f32>;   // b, hc*hc
@group(0) @binding(4)  var<storage, read>       bhb_mixw  : array<f32>;
@group(0) @binding(5)  var<storage, read>       bhb_sc    : array<f32>;
@group(0) @binding(6)  var<storage, read>       bhb_base  : array<f32>;
@group(0) @binding(7)  var<storage, read>       bhb_nw    : array<f32>;
@group(0) @binding(8)  var<storage, read_write> bhb_state : array<f32>;   // b, hc*dim
@group(0) @binding(9)  var<storage, read_write> bhb_fold  : array<f32>;   // b, dim
@group(0) @binding(10) var<storage, read_write> bhb_norm  : array<f32>;   // b, dim
@group(0) @binding(11) var<uniform>             bhb_p     : HbP;
var<workgroup> bhb_red: array<f32, 256>;
var<workgroup> bhb_mix: array<f32, 64>;
var<workgroup> bhb_pre: array<f32, 8>;
var<workgroup> bhb_cmb: array<f32, 64>;
var<workgroup> bhb_rsq: f32;
@compute @workgroup_size(256)
fn bt_hc_block(@builtin(workgroup_id) wid: vec3<u32>,
               @builtin(local_invocation_index) lid: u32) {
    let tok = wid.x;
    let xb = tok * bhb_p.dim;
    let rb = tok * bhb_p.hc * bhb_p.dim;
    let sb = rb;
    let fb = xb;
    let pb = tok * bhb_p.hc;
    let cb = tok * bhb_p.hc * bhb_p.hc;

    let hc = bhb_p.hc;
    let dim = bhb_p.dim;
    let n = hc * dim;
    let mh = bhb_p.mix_hc;

    // ── 1. expand ─────────────────────────────────────────────────────────
    var i = lid;
    loop {
        if (i >= n) { break; }
        let j = i / dim;
        let d = i % dim;
        var y = bhb_post[pb + j] * bhb_x[xb + d];
        for (var k = 0u; k < hc; k = k + 1u) {
            y = y + bhb_comb[cb + k * hc + j] * bhb_res[rb + k * dim + d];
        }
        bhb_state[sb + i] = y;
        i = i + 256u;
    }
    workgroupBarrier();

    // ── 2. the mix projection, all outputs at once ────────────────────────
    // Eight threads to an output, striding the shared axis: one barrier for
    // the lot instead of one tree reduction per output.
    let o = lid / HB_LANES;
    let sub = lid % HB_LANES;
    var acc = 0.0;
    if (o < mh) {
        let base = o * n;
        var t = sub;
        loop {
            if (t >= n) { break; }
            acc = acc + bhb_mixw[base + t] * bhb_state[sb + t];
            t = t + HB_LANES;
        }
    }
    bhb_red[lid] = acc;
    workgroupBarrier();
    if (sub == 0u && o < mh) {
        var sm = 0.0;
        for (var t = 0u; t < HB_LANES; t = t + 1u) { sm = sm + bhb_red[o * HB_LANES + t]; }
        bhb_mix[o] = sm;
    }

    // ── 3. rsqrt(mean(state^2) + eps), over ALL copies ────────────────────
    var acc2 = 0.0;
    var i2 = lid;
    loop {
        if (i2 >= n) { break; }
        let v = bhb_state[sb + i2];
        acc2 = acc2 + v * v;
        i2 = i2 + 256u;
    }
    workgroupBarrier();           // bhb_red is being reused
    bhb_red[lid] = acc2;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { bhb_red[lid] = bhb_red[lid] + bhb_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { bhb_rsq = inverseSqrt(bhb_red[0] / f32(n) + bhb_p.eps); }
    workgroupBarrier();
    let rsq = bhb_rsq;

    // ── 4. pre / post / comb, thread 0: hc is 4 and the Sinkhorn is a
    //       sequential fixed point over a 4x4.
    if (lid == 0u) {
        for (var j = 0u; j < hc; j = j + 1u) {
            let m = bhb_mix[j] * rsq * bhb_sc[0] + bhb_base[j];
            bhb_pre[j] = 1.0 / (1.0 + exp(-m)) + bhb_p.eps;
            let m2 = bhb_mix[hc + j] * rsq * bhb_sc[1] + bhb_base[hc + j];
            bhb_post[pb + j] = 2.0 / (1.0 + exp(-m2));
        }
        for (var j = 0u; j < hc; j = j + 1u) {
            var mx = -1e30;
            for (var k = 0u; k < hc; k = k + 1u) {
                let v = bhb_mix[2u * hc + j * hc + k] * rsq * bhb_sc[2]
                      + bhb_base[2u * hc + j * hc + k];
                bhb_cmb[j * hc + k] = v;
                mx = max(mx, v);
            }
            var sum = 0.0;
            for (var k = 0u; k < hc; k = k + 1u) {
                let e = exp(bhb_cmb[j * hc + k] - mx);
                bhb_cmb[j * hc + k] = e;
                sum = sum + e;
            }
            for (var k = 0u; k < hc; k = k + 1u) {
                bhb_cmb[j * hc + k] = bhb_cmb[j * hc + k] / sum + bhb_p.eps;
            }
        }
        for (var k = 0u; k < hc; k = k + 1u) {
            var sum = 0.0;
            for (var j = 0u; j < hc; j = j + 1u) { sum = sum + bhb_cmb[j * hc + k]; }
            for (var j = 0u; j < hc; j = j + 1u) {
                bhb_cmb[j * hc + k] = bhb_cmb[j * hc + k] / (sum + bhb_p.eps);
            }
        }
        for (var it = 1u; it < bhb_p.iters; it = it + 1u) {
            for (var j = 0u; j < hc; j = j + 1u) {
                var sum = 0.0;
                for (var k = 0u; k < hc; k = k + 1u) { sum = sum + bhb_cmb[j * hc + k]; }
                for (var k = 0u; k < hc; k = k + 1u) {
                    bhb_cmb[j * hc + k] = bhb_cmb[j * hc + k] / (sum + bhb_p.eps);
                }
            }
            for (var k = 0u; k < hc; k = k + 1u) {
                var sum = 0.0;
                for (var j = 0u; j < hc; j = j + 1u) { sum = sum + bhb_cmb[j * hc + k]; }
                for (var j = 0u; j < hc; j = j + 1u) {
                    bhb_cmb[j * hc + k] = bhb_cmb[j * hc + k] / (sum + bhb_p.eps);
                }
            }
        }
        for (var j = 0u; j < hc * hc; j = j + 1u) { bhb_comb[cb + j] = bhb_cmb[j]; }
    }
    workgroupBarrier();

    // ── 5. fold, then the norm over it ────────────────────────────────────
    var acc3 = 0.0;
    var d = lid;
    loop {
        if (d >= dim) { break; }
        var y = 0.0;
        for (var j = 0u; j < hc; j = j + 1u) {
            y = y + bhb_pre[j] * bhb_state[sb + j * dim + d];
        }
        bhb_fold[fb + d] = y;
        acc3 = acc3 + y * y;
        d = d + 256u;
    }
    bhb_red[lid] = acc3;
    workgroupBarrier();
    stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { bhb_red[lid] = bhb_red[lid] + bhb_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    let inv = inverseSqrt(bhb_red[0] / f32(dim) + bhb_p.eps);
    var d2 = lid;
    loop {
        if (d2 >= dim) { break; }
        bhb_norm[fb + d2] = bhb_fold[fb + d2] * inv * bhb_nw[d2];
        d2 = d2 + 256u;
    }
}


// The compressor's pending append for a RUN of tokens, one dispatch: every
// token in a fold-free segment writes a distinct slot of the pending
// stream, so the writes commute. The score picks up its in-window position
// bias here, exactly where the per-token step added it.
struct BcaP { width: u32, t0: u32, n: u32, flags: u32 }; // flags: 1 = overlap (add ape), bits 8.. = pos0 % ratio? no — slot base rides in t0's precomputed slots
@group(0) @binding(0) var<storage, read>       bca_ckv : array<f32>;   // b, width
@group(0) @binding(1) var<storage, read>       bca_csc : array<f32>;   // b, width
@group(0) @binding(2) var<storage, read>       bca_ape : array<f32>;
@group(0) @binding(3) var<storage, read_write> bca_pkv : array<f32>;   // ratio, width
@group(0) @binding(4) var<storage, read_write> bca_psc : array<f32>;
@group(0) @binding(5) var<uniform>             bca_p   : BcaP;
@group(0) @binding(6) var<storage, read>       bca_slot: array<u32>;   // b (slot per token)
@compute @workgroup_size(256)
fn bt_comp_append(@builtin(workgroup_id) wid: vec3<u32>,
                  @builtin(local_invocation_id) lid: vec3<u32>) {
    let i = wid.x * 256u + lid.x;
    if (i >= bca_p.width) { return; }
    let t = bca_p.t0 + wid.y;
    if (wid.y >= bca_p.n) { return; }
    let slot = bca_slot[t];
    let src = t * bca_p.width + i;
    bca_pkv[slot * bca_p.width + i] = bca_ckv[src];
    var sc = bca_csc[src];
    if ((bca_p.flags & 1u) != 0u) {
        sc = sc + bca_ape[slot * bca_p.width + i];
    }
    bca_psc[slot * bca_p.width + i] = sc;
}

// The whole fold-token compressor step as ONE dispatch: append the
// closing token's row, softmax-pool the window (kv_pool's expressions,
// element for element), norm through rmsnorm's exact 1024-thread tree,
// rotate the rope tail, land the entry in the cache and shift pending
// into previous. The per-token path spent seven dependent dispatches on
// this; the math here is the same operations in the same order, so the
// bits are too.
struct BcfP {
    width: u32, ratio: u32, rd: u32, flags: u32,
    dst_off: u32, pos_bits: u32, eps_bits: u32, trow: u32,
};
@group(0) @binding(0) var<storage, read>       bcf_ckv : array<f32>;
@group(0) @binding(1) var<storage, read>       bcf_csc : array<f32>;
@group(0) @binding(2) var<storage, read>       bcf_ape : array<f32>;
@group(0) @binding(3) var<storage, read_write> bcf_pkv : array<f32>;
@group(0) @binding(4) var<storage, read_write> bcf_psc : array<f32>;
@group(0) @binding(5) var<storage, read_write> bcf_qkv : array<f32>;
@group(0) @binding(6) var<storage, read_write> bcf_qsc : array<f32>;
@group(0) @binding(7) var<storage, read>       bcf_nw  : array<f32>;
@group(0) @binding(8) var<storage, read>       bcf_fr  : array<f32>;
@group(0) @binding(9) var<storage, read_write> bcf_dst : array<f32>;
@group(0) @binding(10) var<storage, read>      bcf_p   : BcfP;
var<workgroup> bcf_fold: array<f32, 512>;
var<workgroup> bcf_part: array<f32, 1024>;
@compute @workgroup_size(1024)
fn bt_comp_fold(@builtin(local_invocation_index) lid: u32) {
    let w = bcf_p.width;
    let r = bcf_p.ratio;
    let overlap = (bcf_p.flags & 1u) != 0u;
    let have_prev = (bcf_p.flags & 2u) != 0u;
    let ew = select(w, w / 2u, overlap);
    let slot = r - 1u;
    let trow = bcf_p.trow;
    // ── append the closing token ──
    var i = lid;
    loop {
        if (i >= w) { break; }
        bcf_pkv[slot * w + i] = bcf_ckv[trow * w + i];
        var sc = bcf_csc[trow * w + i];
        if (overlap) { sc = sc + bcf_ape[slot * w + i]; }
        bcf_psc[slot * w + i] = sc;
        i = i + 1024u;
    }
    storageBarrier();
    workgroupBarrier();
    // ── pool: kv_pool per element, w there is ew here ──
    let slots = select(r, 2u * r, overlap);
    var d = lid;
    loop {
        if (d >= ew) { break; }
        var mx = KP_NINF;
        for (var t = 0u; t < slots; t = t + 1u) {
            var sc = KP_NINF;
            if (overlap) {
                if (t < r) {
                    if (have_prev) { sc = bcf_qsc[t * 2u * ew + d]; }
                } else {
                    sc = bcf_psc[(t - r) * 2u * ew + ew + d];
                }
            } else {
                sc = bcf_psc[t * ew + d] + bcf_ape[t * ew + d];
            }
            mx = max(mx, sc);
        }
        var outv = 0.0;
        if (mx > KP_NINF) {
            var den = 0.0;
            var acc = 0.0;
            for (var t = 0u; t < slots; t = t + 1u) {
                var sc = KP_NINF;
                var kv = 0.0;
                if (overlap) {
                    if (t < r) {
                        if (have_prev) {
                            sc = bcf_qsc[t * 2u * ew + d];
                            kv = bcf_qkv[t * 2u * ew + d];
                        }
                    } else {
                        sc = bcf_psc[(t - r) * 2u * ew + ew + d];
                        kv = bcf_pkv[(t - r) * 2u * ew + ew + d];
                    }
                } else {
                    sc = bcf_psc[t * ew + d] + bcf_ape[t * ew + d];
                    kv = bcf_pkv[t * ew + d];
                }
                if (sc > KP_NINF) {
                    let e = exp(sc - mx);
                    den = den + e;
                    acc = acc + e * kv;
                }
            }
            if (den > 0.0) { outv = acc / den; }
        }
        bcf_fold[d] = outv;
        d = d + 1024u;
    }
    workgroupBarrier();
    // ── rmsnorm, the 1024-thread kernel's exact tree ──
    var acc2 = 0.0;
    var j = lid;
    loop {
        if (j >= ew) { break; }
        let v = bcf_fold[j];
        acc2 = acc2 + v * v;
        j = j + 1024u;
    }
    bcf_part[lid] = acc2;
    workgroupBarrier();
    var stride = 512u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { bcf_part[lid] = bcf_part[lid] + bcf_part[lid + stride]; }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let inv = inverseSqrt(bcf_part[0] / f32(ew) + bitcast<f32>(bcf_p.eps_bits));
    workgroupBarrier();
    var k2 = lid;
    loop {
        if (k2 >= ew) { break; }
        bcf_fold[k2] = bcf_fold[k2] * inv * bcf_nw[k2];
        k2 = k2 + 1024u;
    }
    workgroupBarrier();
    // ── rope tail, adjacent pairs at the window's FIRST position ──
    let rd = bcf_p.rd;
    let rbase = ew - rd;
    let pos = bitcast<f32>(bcf_p.pos_bits);
    var t2 = lid;
    loop {
        if (t2 >= rd / 2u) { break; }
        let th = pos * bcf_fr[t2];
        let sn = sin(th);
        let cs = cos(th);
        let a = bcf_fold[rbase + 2u * t2];
        let cc = bcf_fold[rbase + 2u * t2 + 1u];
        bcf_fold[rbase + 2u * t2] = a * cs - cc * sn;
        bcf_fold[rbase + 2u * t2 + 1u] = a * sn + cc * cs;
        t2 = t2 + 1024u;
    }
    workgroupBarrier();
    // ── land the entry, then pending becomes previous ──
    var d2 = lid;
    loop {
        if (d2 >= ew) { break; }
        bcf_dst[bcf_p.dst_off + d2] = bcf_fold[d2];
        d2 = d2 + 1024u;
    }
    if (overlap) {
        storageBarrier();
        workgroupBarrier();
        var m = lid;
        loop {
            if (m >= r * w) { break; }
            bcf_qkv[m] = bcf_pkv[m];
            bcf_qsc[m] = bcf_psc[m];
            m = m + 1024u;
        }
    }
}

// index_scores with grid (entry, token): the query heads and the output
// stride by the token; the weights arrive RAW with the fold factor in the
// uniform, multiplied per head exactly where the walk's axpy used to — same
// rounding, one dispatch fewer. Per-token entry limits ride in a table.
struct BixP { nh: u32, hd: u32, n_pos: u32, factor: f32 };
@group(0) @binding(0) var<storage, read>       bix_q    : array<f32>;   // b, nh*hd
@group(0) @binding(1) var<storage, read>       bix_kv   : array<f32>;
@group(0) @binding(2) var<storage, read>       bix_w    : array<f32>;   // b, nh (raw)
@group(0) @binding(3) var<storage, read_write> bix_out  : array<f32>;   // b, 4096
@group(0) @binding(4) var<uniform>             bix_p    : BixP;
@group(0) @binding(5) var<storage, read>       bix_lim  : array<u32>;   // b
var<workgroup> bix_red: array<f32, 256>;
@compute @workgroup_size(256)
fn bt_index_scores(@builtin(workgroup_id) wid: vec3<u32>,
                   @builtin(local_invocation_index) lid: u32) {
    let t = wid.x;
    let tok = wid.y;
    if (t >= bix_p.n_pos) { return; }
    let limit = bix_lim[tok];
    let ob = tok * 4096u;
    if (t >= limit) {
        if (lid == 0u) { bix_out[ob + t] = KP_NINF; }
        return;
    }
    let hd = bix_p.hd;
    let kb = t * hd;
    let qtb = tok * bix_p.nh * hd;
    let wtb = tok * bix_p.nh;
    var acc = 0.0;
    var h = lid;
    loop {
        if (h >= bix_p.nh) { break; }
        var dot = 0.0;
        let qb = qtb + h * hd;
        for (var i = 0u; i < hd; i = i + 1u) {
            dot = dot + bix_q[qb + i] * bix_kv[kb + i];
        }
        let hw = bix_w[wtb + h] * bix_p.factor;
        acc = acc + max(dot, 0.0) * hw;
        h = h + 256u;
    }
    bix_red[lid] = acc;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) { bix_red[lid] = bix_red[lid] + bix_red[lid + stride]; }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) { bix_out[ob + t] = bix_red[0]; }
}

// top_k_index with grid (token): scores, picks and counts stride by the
// token; the entry count comes from the same limit table.
struct BtkP { kmax: u32, _a: u32, _b: u32, _c: u32 };
@group(0) @binding(0) var<storage, read>       btk_s   : array<f32>;   // b, 4096
@group(0) @binding(1) var<storage, read_write> btk_idx : array<u32>;   // b, kmax
@group(0) @binding(2) var<storage, read_write> btk_cnt : array<u32>;   // b
@group(0) @binding(3) var<uniform>             btk_p   : BtkP;
@group(0) @binding(4) var<storage, read>       btk_lim : array<u32>;   // b
var<workgroup> btk_keep: array<u32, 4096>;
@compute @workgroup_size(1024)
fn bt_top_k(@builtin(workgroup_id) wid: vec3<u32>,
            @builtin(local_invocation_index) lid: u32) {
    let tok = wid.x;
    let n = btk_lim[tok];
    let sb = tok * 4096u;
    var i = lid;
    loop {
        if (i >= n) { break; }
        let si = btk_s[sb + i];
        var keep = 0u;
        if (si > KP_NINF) {
            var rank = 0u;
            for (var j = 0u; j < n; j = j + 1u) {
                let sj = btk_s[sb + j];
                if (sj > KP_NINF) {
                    if (sj > si || (sj == si && j < i)) { rank = rank + 1u; }
                }
            }
            if (rank < btk_p.kmax) { keep = 1u; }
        }
        btk_keep[i] = keep;
        i = i + 1024u;
    }
    workgroupBarrier();
    var m = lid;
    loop {
        if (m >= n) { break; }
        if (btk_keep[m] == 1u) {
            var before = 0u;
            for (var j = 0u; j < m; j = j + 1u) { before = before + btk_keep[j]; }
            btk_idx[tok * btk_p.kmax + before] = m;
        }
        m = m + 1024u;
    }
    workgroupBarrier();
    if (lid == 0u) {
        var total = 0u;
        for (var j = 0u; j < n; j = j + 1u) { total = total + btk_keep[j]; }
        btk_cnt[tok] = total;
    }
}

// The staged attended-position list, per token: the visible tail of the
// UNSLID window first (the batch never slides mid-pass; its own rows sit in
// staging at the cache's tail), then this token's staged rows including its
// own, then the compressed positions — the indexer's picks, or all of them
// on the layers without one (seq flag).
struct BibP { window: u32, kmax: u32, seq: u32, srow0: u32 };
@group(0) @binding(0) var<storage, read>       bib_pick : array<u32>;   // b, kmax
@group(0) @binding(1) var<storage, read_write> bib_out  : array<u32>;   // b, 1024
@group(0) @binding(2) var<uniform>             bib_p    : BibP;
@group(0) @binding(3) var<storage, read>       bib_meta : array<vec4<u32>>; // b: [win_start, win_n, staged_n, k]
@compute @workgroup_size(256)
fn bt_idx_build_staged(@builtin(workgroup_id) wid: vec3<u32>,
                       @builtin(local_invocation_id) lid: vec3<u32>) {
    let tok = wid.y;
    let i = wid.x * 256u + lid.x;
    let m = bib_meta[tok];
    let ob = tok * 1024u;
    if (i < m.y) {
        bib_out[ob + i] = m.x + i;
        return;
    }
    let s = i - m.y;
    if (s < m.z) {
        bib_out[ob + i] = bib_p.srow0 + s;
        return;
    }
    let j = i - m.y - m.z;
    if (j < m.w) {
        if (bib_p.seq != 0u) {
            bib_out[ob + i] = bib_p.window + j;
        } else {
            bib_out[ob + i] = bib_p.window + bib_pick[tok * bib_p.kmax + j];
        }
    }
}

// The grouped low-rank output projection with grid (row-quad, token): the
// same four-rows-per-workgroup walk as o_lora_a_m, the activation window and
// the output sliding with the token.
var<workgroup> obt_part: array<f32, 256>;
var<workgroup> obt_lad: array<f32, 128>;
@compute @workgroup_size(256)
fn bt_o_lora_a(@builtin(workgroup_id) wid: vec3<u32>,
               @builtin(num_workgroups) nwg: vec3<u32>,
               @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let lora = q1p._p0;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let t = wid.y;
    let xtok = t * (rows / lora) * gpr * 32u;
    let sub = lid / 64u;
    let lane = lid % 64u;
    var row = wid.x * 4u + sub;
    loop {
        if (row >= rows) { break; }
        if (lane < 32u) {
            let pr = unpack2x16float(q1w[params_w + row]);
            obt_lad[sub * 32u + lane] = exp2(pr.x + f32(lane) * pr.y);
        }
        workgroupBarrier();
        let xoff = xtok + (row / lora) * gpr * 32u;
        var acc = 0.0;
        var g = lane;
        loop {
            if (g >= gpr) { break; }
            let bit = g * 5u;
            let cb = codes_b + row * cstride + (bit >> 3u);
            let sh = bit & 7u;
            var cv = q4tp_byte(cb);
            if (sh > 3u) { cv = cv | (q4tp_byte(cb + 1u) << 8u); }
            let scale = obt_lad[sub * 32u + ((cv >> sh) & 31u)];
            let wv = q1wv[row * gpr + g];
            let xb = xoff + g * 32u;
            var gsum = 0.0;
            gsum = gsum + q4b_dot8(wv.x, xb);
            gsum = gsum + q4b_dot8(wv.y, xb + 8u);
            gsum = gsum + q4b_dot8(wv.z, xb + 16u);
            gsum = gsum + q4b_dot8(wv.w, xb + 24u);
            acc = acc + scale * gsum;
            g = g + 64u;
        }
        obt_part[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lane < stride) { obt_part[lid] = obt_part[lid] + obt_part[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lane == 0u) { q1y[t * rows + row] = obt_part[sub * 64u]; }
        workgroupBarrier();
        row = row + nwg.x * 4u;
    }
}


// The grouped projection, four rows to a 64-thread workgroup: the x span
// loads once per group iteration and feeds all four rows from registers,
// and the scale comes straight from the row params — the shared ladder
// with its barriers is what made the original spend 0.4 ms on 30 MB.
// Each row keeps the original's lane assignment, group order and 64-wide
// tree, so its sum is bit-identical.
var<workgroup> obt_p4: array<f32, 64>;
@compute @workgroup_size(64)
fn bt_o_lora_a4(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let lora = q1p._p0;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let t = wid.y;
    let row0 = wid.x * 4u;
    let xtok = t * (rows / lora) * gpr * 32u;
    let xoff = xtok + (row0 / lora) * gpr * 32u;
    var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
    var g = lid;
    loop {
        if (g >= gpr) { break; }
        let xb = xoff + g * 32u;
        let v0 = q1xv[(xb >> 2u)];      let v1 = q1xv[(xb >> 2u) + 1u];
        let v2 = q1xv[(xb >> 2u) + 2u]; let v3 = q1xv[(xb >> 2u) + 3u];
        let v4 = q1xv[(xb >> 2u) + 4u]; let v5 = q1xv[(xb >> 2u) + 5u];
        let v6 = q1xv[(xb >> 2u) + 6u]; let v7 = q1xv[(xb >> 2u) + 7u];
        let bit = g * 5u;
        let cb0 = bit >> 3u;
        let sh = bit & 7u;
        for (var r = 0u; r < 4u; r = r + 1u) {
            let row = row0 + r;
            if (row >= rows) { break; }
            let pr = unpack2x16float(q1w[params_w + row]);
            let cb = codes_b + row * cstride + cb0;
            var cv = q4tp_byte(cb);
            if (sh > 3u) { cv = cv | (q4tp_byte(cb + 1u) << 8u); }
            let code = (cv >> sh) & 31u;
            let scale = exp2(pr.x + f32(code) * pr.y);
            let wv = q1wv[row * gpr + g];
            var gsum = 0.0;
            gsum = q1_dot8v(wv.x, v0, v1)
                 + q1_dot8v(wv.y, v2, v3)
                 + q1_dot8v(wv.z, v4, v5)
                 + q1_dot8v(wv.w, v6, v7);
            let vsc = scale * gsum;
            if (r == 0u) { a0 = a0 + vsc; }
            if (r == 1u) { a1 = a1 + vsc; }
            if (r == 2u) { a2 = a2 + vsc; }
            if (r == 3u) { a3 = a3 + vsc; }
        }
        g = g + 64u;
    }
    for (var r = 0u; r < 4u; r = r + 1u) {
        var acc = a0;
        if (r == 1u) { acc = a1; }
        if (r == 2u) { acc = a2; }
        if (r == 3u) { acc = a3; }
        obt_p4[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lid < stride) { obt_p4[lid] = obt_p4[lid] + obt_p4[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lid == 0u && row0 + r < rows) { q1y[t * rows + row0 + r] = obt_p4[0]; }
        workgroupBarrier();
    }
}

// The grouped projection with the group's x span staged through shared
// memory: all four sub-rows of a workgroup belong to one o-group and
// re-read the same span — the stage cuts that traffic four-fold. Values
// and per-lane order are the un-staged kernel's, so the sums are
// bit-identical. Created only when the device's workgroup storage fits
// the 4608-float span (18 KB + change); the host also refuses shapes
// that straddle groups.
var<workgroup> obt_xs: array<f32, 4608>;
@compute @workgroup_size(256)
fn bt_o_lora_a2(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(num_workgroups) nwg: vec3<u32>,
                @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let lora = q1p._p0;
    let params_w = rows * gpr * 4u;
    let codes_b = rows * gpr * 16u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let t = wid.y;
    let xtok = t * (rows / lora) * gpr * 32u;
    let sub = lid / 64u;
    let lane = lid % 64u;
    var row = wid.x * 4u + sub;
    loop {
        if (row >= rows) { break; }
        let span = gpr * 32u;
        let xoff = xtok + (row / lora) * span;
        for (var j = lid; j < span; j = j + 256u) { obt_xs[j] = q1x[xoff + j]; }
        if (lane < 32u) {
            let pr = unpack2x16float(q1w[params_w + row]);
            obt_lad[sub * 32u + lane] = exp2(pr.x + f32(lane) * pr.y);
        }
        workgroupBarrier();
        var acc = 0.0;
        var g = lane;
        loop {
            if (g >= gpr) { break; }
            let bit = g * 5u;
            let cb = codes_b + row * cstride + (bit >> 3u);
            let sh = bit & 7u;
            var cv = q4tp_byte(cb);
            if (sh > 3u) { cv = cv | (q4tp_byte(cb + 1u) << 8u); }
            let scale = obt_lad[sub * 32u + ((cv >> sh) & 31u)];
            let base = (row * gpr + g) * 4u;
            let xb = g * 32u;
            var gsum = 0.0;
            for (var k = 0u; k < 4u; k = k + 1u) {
                let w = q1w[base + k];
                let x8 = xb + 8u * k;
                gsum = gsum + (f32(w & 0xFu) - 8.0) * obt_xs[x8]
                     + (f32((w >> 4u) & 0xFu) - 8.0) * obt_xs[x8 + 1u]
                     + (f32((w >> 8u) & 0xFu) - 8.0) * obt_xs[x8 + 2u]
                     + (f32((w >> 12u) & 0xFu) - 8.0) * obt_xs[x8 + 3u]
                     + (f32((w >> 16u) & 0xFu) - 8.0) * obt_xs[x8 + 4u]
                     + (f32((w >> 20u) & 0xFu) - 8.0) * obt_xs[x8 + 5u]
                     + (f32((w >> 24u) & 0xFu) - 8.0) * obt_xs[x8 + 6u]
                     + (f32((w >> 28u) & 0xFu) - 8.0) * obt_xs[x8 + 7u];
            }
            acc = acc + scale * gsum;
            g = g + 64u;
        }
        obt_part[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (lane < stride) { obt_part[lid] = obt_part[lid] + obt_part[lid + stride]; }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (lane == 0u) { q1y[t * rows + row] = obt_part[sub * 64u]; }
        workgroupBarrier();
        row = row + nwg.x * 4u;
    }
}

// BF16 activation materialisation for the V4.1 attention frame. WGSL has no
// portable BF16 storage scalar, so keep the value in f32 and apply the same
// round-to-nearest-even bit operation as the host path at the three tensor
// boundaries that the frame crosses.
struct Bf16P { n: u32, _a: u32, _b: u32, _c: u32 };
@group(0) @binding(0) var<storage, read_write> bf16_x : array<f32>;
@group(0) @binding(1) var<uniform> bf16_p : Bf16P;

fn bf16_round(v: f32) -> f32 {
    let bits = bitcast<u32>(v);
    if ((bits & 0x7F800000u) == 0x7F800000u) { return v; }
    let rounded = bits + 0x7FFFu + ((bits >> 16u) & 1u);
    return bitcast<f32>(rounded & 0xFFFF0000u);
}

@compute @workgroup_size(256)
fn bf16_round_buffer(@builtin(global_invocation_id) gid: vec3<u32>) {
    if (gid.x < bf16_p.n) { bf16_x[gid.x] = bf16_round(bf16_x[gid.x]); }
}

// Qwen Image's two streams have different Q/K/V projections, qk-norm
// weights, and RoPE tables.  This exact join keeps those differences in one
// position-wise pass: the six resident token-major projection panels become
// joint head-major Q/K/V planes in `[text, image]` order.  The attention
// kernels can then consume the planes without the host concat/repack seam.
struct QwenRopeP {
    image_n: u32, text_n: u32, heads: u32, hd: u32,
    hidden: u32, total: u32, pairs: u32, _p: u32,
};
@group(0) @binding(0) var<storage, read>       qwi_q : array<f32>;
@group(0) @binding(1) var<storage, read>       qwi_k : array<f32>;
@group(0) @binding(2) var<storage, read>       qwi_v : array<f32>;
@group(0) @binding(3) var<storage, read>       qwt_q : array<f32>;
@group(0) @binding(4) var<storage, read>       qwt_k : array<f32>;
@group(0) @binding(5) var<storage, read>       qwt_v : array<f32>;
@group(0) @binding(6) var<storage, read>       qwi_bias : array<f32>;
@group(0) @binding(7) var<storage, read>       qwt_bias : array<f32>;
@group(0) @binding(8) var<storage, read>       qwi_qn : array<f32>;
@group(0) @binding(9) var<storage, read>       qwi_kn : array<f32>;
@group(0) @binding(10) var<storage, read>      qwt_qn : array<f32>;
@group(0) @binding(11) var<storage, read>      qwt_kn : array<f32>;
@group(0) @binding(12) var<storage, read>      qwi_cos : array<f32>;
@group(0) @binding(13) var<storage, read>      qwi_sin : array<f32>;
@group(0) @binding(14) var<storage, read>      qwt_cos : array<f32>;
@group(0) @binding(15) var<storage, read>      qwt_sin : array<f32>;
@group(0) @binding(16) var<storage, read_write> qwo_q : array<f32>;
@group(0) @binding(17) var<storage, read_write> qwo_k : array<f32>;
@group(0) @binding(18) var<storage, read_write> qwo_v : array<f32>;
@group(0) @binding(19) var<uniform>             qwr_p : QwenRopeP;
var<workgroup> qwr_q: array<f32, 256>;
var<workgroup> qwr_k: array<f32, 256>;
var<workgroup> qwr_v: array<f32, 256>;
var<workgroup> qwr_qr: array<f32, 256>;
var<workgroup> qwr_kr: array<f32, 256>;

@compute @workgroup_size(256)
fn qwen_rope_pack(@builtin(workgroup_id) wid: vec3<u32>,
                  @builtin(local_invocation_index) lid: u32) {
    let job = wid.y * 65535u + wid.x;
    let jobs = qwr_p.total * qwr_p.heads;
    if (job >= jobs || qwr_p.hd > 256u || qwr_p.hd == 0u) { return; }
    let p = job / qwr_p.heads;
    let h = job - p * qwr_p.heads;
    let text = p < qwr_p.text_n;
    var tok = p;
    if (!text) { tok = p - qwr_p.text_n; }
    let row = tok * qwr_p.hidden + h * qwr_p.hd;
    let base = h * qwr_p.total * qwr_p.hd + p * qwr_p.hd;
    var angle_tok = tok;
    if (text) {
        angle_tok = p;
    }
    var i = lid;
    loop {
        if (i >= qwr_p.hd) { break; }
        var qv = 0.0;
        var kv = 0.0;
        var vv = 0.0;
        var qbias = 0.0;
        var kbias = 0.0;
        var vbias = 0.0;
        if (text) {
            qv = qwt_q[row + i];
            kv = qwt_k[row + i];
            vv = qwt_v[row + i];
            qbias = qwt_bias[h * qwr_p.hd + i];
            kbias = qwt_bias[qwr_p.hidden + h * qwr_p.hd + i];
            vbias = qwt_bias[2u * qwr_p.hidden + h * qwr_p.hd + i];
        } else {
            qv = qwi_q[row + i];
            kv = qwi_k[row + i];
            vv = qwi_v[row + i];
            qbias = qwi_bias[h * qwr_p.hd + i];
            kbias = qwi_bias[qwr_p.hidden + h * qwr_p.hd + i];
            vbias = qwi_bias[2u * qwr_p.hidden + h * qwr_p.hd + i];
        }
        qwr_q[i] = qv + qbias;
        qwr_k[i] = kv + kbias;
        qwr_qr[i] = qwr_q[i] * qwr_q[i];
        qwr_kr[i] = qwr_k[i] * qwr_k[i];
        qwr_v[i] = vv + vbias;
        i = i + 256u;
    }
    if (lid >= qwr_p.hd) {
        qwr_qr[lid] = 0.0;
        qwr_kr[lid] = 0.0;
    }
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) {
            qwr_qr[lid] = qwr_qr[lid] + qwr_qr[lid + stride];
            qwr_kr[lid] = qwr_kr[lid] + qwr_kr[lid + stride];
        }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let qi = inverseSqrt(qwr_qr[0] / f32(qwr_p.hd) + 0.000001);
    let ki = inverseSqrt(qwr_kr[0] / f32(qwr_p.hd) + 0.000001);
    i = lid;
    loop {
        if (i >= qwr_p.hd) { break; }
        var qnw = 0.0;
        var knw = 0.0;
        if (text) {
            qnw = qwt_qn[i];
            knw = qwt_kn[i];
        } else {
            qnw = qwi_qn[i];
            knw = qwi_kn[i];
        }
        qwr_q[i] = qwr_q[i] * qi * qnw;
        qwr_k[i] = qwr_k[i] * ki * knw;
        i = i + 256u;
    }
    // The rotation below reads both members of each pair.  Norm writes are
    // distributed across lanes, so synchronize before a lane reads its
    // neighbour's normalized value.
    workgroupBarrier();
    i = lid;
    loop {
        if (i >= qwr_p.hd) { break; }
        let pair = i / 2u;
        var qv = qwr_q[i];
        var kv = qwr_k[i];
        var cs = 1.0;
        var sn = 0.0;
        if (text) {
            cs = qwt_cos[angle_tok * qwr_p.pairs + pair];
            sn = qwt_sin[angle_tok * qwr_p.pairs + pair];
        } else {
            cs = qwi_cos[angle_tok * qwr_p.pairs + pair];
            sn = qwi_sin[angle_tok * qwr_p.pairs + pair];
        }
        // Qwen's native path rotates adjacent dimensions: [0,1], [2,3],
        // ... .  The earlier half-split form reused pair 0 for output index
        // `pairs`, which only becomes visible once head_dim exceeds 2.
        if (i % 2u == 0u) {
            qv = qwr_q[2u * pair] * cs - qwr_q[2u * pair + 1u] * sn;
            kv = qwr_k[2u * pair] * cs - qwr_k[2u * pair + 1u] * sn;
        } else {
            qv = qwr_q[2u * pair] * sn + qwr_q[2u * pair + 1u] * cs;
            kv = qwr_k[2u * pair] * sn + qwr_k[2u * pair + 1u] * cs;
        }
        qwo_q[base + i] = qv;
        qwo_k[base + i] = kv;
        qwo_v[base + i] = qwr_v[i];
        i = i + 256u;
    }
}
"#;

/// Signed, normalized FWHT used by the Prism activation boundary.  This is
/// deliberately a separate f16-capable module: devices without SHADER_F16
/// keep the conservative per-op/CPU route rather than silently changing the
/// trained boundary.  One workgroup owns one 1024-wide block, with f32
/// butterflies and an explicit f16 round at the output.
const FWHT_SRC: &str = r#"
enable f16;

struct FwhtP {
    width: u32,
    block: u32,
    sign_offset: u32,
    inverse: u32,
    round16: u32,
    // Zero means one row for the original token-graph entry point.  A
    // positive value admits the same transform over a contiguous batch,
    // with the sign table reused for every row (the prefill graph path).
    rows: u32,
    _pad1: u32,
    _pad2: u32,
};

@group(0) @binding(0) var<storage, read> fwht_x: array<f32>;
@group(0) @binding(1) var<storage, read_write> fwht_y: array<f32>;
@group(0) @binding(2) var<storage, read> fwht_signs: array<f32>;
@group(0) @binding(3) var<uniform> fwht_p: FwhtP;

var<workgroup> fwht_s: array<f32, 1024>;

// Implement the declared Prism f16 boundary in bits.  On the target Vulkan
// path, both a source-level f32(f16(v)) round-trip and pack/unpack2x16float
// may be legally folded back to f32 when the intermediate has no observable
// f16 storage.  Keeping the RNE conversion explicit avoids silently running
// the graph at f32 while retaining the same bit contract as the CPU oracle.
fn prism_half_to_f32(h: u32) -> f32 {
    let sign = (h & 0x8000u) << 16u;
    let exp = (h >> 10u) & 0x1Fu;
    let frac = h & 0x03FFu;
    if (exp == 0u) {
        if (frac == 0u) {
            return bitcast<f32>(sign);
        }
        let mag = f32(frac) * 0.000000059604644775390625; // 2^-24
        return select(mag, -mag, sign != 0u);
    }
    if (exp == 31u) {
        return bitcast<f32>(sign | 0x7F800000u | (frac << 13u));
    }
    return bitcast<f32>(sign | ((exp + 112u) << 23u) | (frac << 13u));
}

fn prism_round16(v: f32) -> f32 {
    let bits = bitcast<u32>(v);
    let sign = (bits >> 16u) & 0x8000u;
    let abits = bits & 0x7FFFFFFFu;
    let exp = (abits >> 23u) & 0xFFu;
    let frac = abits & 0x007FFFFFu;
    if (exp == 0xFFu) {
        // The activation path is finite, but preserve IEEE specials for the
        // component boundary rather than turning a diagnostic NaN into zero.
        return v;
    }
    if (exp == 0u) {
        return bitcast<f32>(sign << 16u);
    }
    let unbiased = i32(exp) - 127;
    if (unbiased < -14) {
        let scaled = v * select(16777216.0, -16777216.0, sign != 0u);
        var q = u32(floor(abs(scaled)));
        let rem = abs(scaled) - f32(q);
        if (rem > 0.5 || (rem == 0.5 && (q & 1u) != 0u)) {
            q = q + 1u;
        }
        if (q >= 1024u) {
            return prism_half_to_f32(sign | 0x0400u);
        }
        return prism_half_to_f32(sign | q);
    }
    if (unbiased > 15) {
        return prism_half_to_f32(sign | 0x7C00u);
    }
    let significand = frac | 0x00800000u;
    var half_frac = (significand >> 13u) & 0x03FFu;
    let discarded = significand & 0x1FFFu;
    if (discarded > 0x1000u || (discarded == 0x1000u && (half_frac & 1u) != 0u)) {
        half_frac = half_frac + 1u;
    }
    var half_exp = u32(unbiased + 15);
    if (half_frac == 0x0400u) {
        half_frac = 0u;
        half_exp = half_exp + 1u;
    }
    if (half_exp >= 31u) {
        return prism_half_to_f32(sign | 0x7C00u);
    }
    return prism_half_to_f32(sign | (half_exp << 10u) | half_frac);
}

@compute @workgroup_size(256)
fn fwht(@builtin(workgroup_id) wid: vec3<u32>,
        @builtin(local_invocation_index) lid: u32) {
    let blocks_per_row = fwht_p.width / fwht_p.block;
    let rows = max(fwht_p.rows, 1u);
    let row = wid.x / blocks_per_row;
    let block = wid.x % blocks_per_row;
    if (row >= rows) { return; }
    let base = row * fwht_p.width + block * fwht_p.block;
    let sign_base = fwht_p.sign_offset + block * fwht_p.block;
    var i = lid;
    loop {
        if (i >= fwht_p.block) { break; }
        let si = sign_base + i;
        let sign = fwht_signs[si];
        let raw = fwht_x[base + i];
        // Forward is D·H; inverse is H·D.  H is self-inverse after the
        // normalized 1/sqrt(block) scale, so the sign placement is the
        // only operator distinction.
        fwht_s[i] = select(raw, raw * sign, fwht_p.inverse == 0u);
        i = i + 256u;
    }
    workgroupBarrier();
    var stride = 1u;
    loop {
        if (stride >= fwht_p.block) { break; }
        let span = stride * 2u;
        var j = lid;
        loop {
            if (j >= fwht_p.block / 2u) { break; }
            let group = j / stride;
            let lane = j % stride;
            let a = group * span + lane;
            let b = a + stride;
            let va = fwht_s[a];
            let vb = fwht_s[b];
            fwht_s[a] = va + vb;
            fwht_s[b] = va - vb;
            j = j + 256u;
        }
        workgroupBarrier();
        stride = span;
    }
    var inv = lid;
    loop {
        if (inv >= fwht_p.block) { break; }
        var v = fwht_s[inv] * inverseSqrt(f32(fwht_p.block));
        if (fwht_p.inverse != 0u) {
            v = v * fwht_signs[sign_base + inv];
        }
        // The source runtime's Prism boundary is explicitly f16.  The
        // f32 storage keeps all downstream existing kernels unchanged.
        if (fwht_p.round16 != 0u) {
            v = prism_round16(v);
        }
        fwht_y[base + inv] = v;
        inv = inv + 256u;
    }
}
"#;

/// The bake FFN chain's middle link, in its own module (the main module's
/// binding slots are all taken): act[r][j] = silu(g)·u·scale[j], where g
/// and u are the two halves of the fused gate+up GEMM's row. Runs between
/// the two cooperative GEMMs so the intermediate plane never crosses PCIe.
/// `stride` linearizes a 2-D dispatch — one dim overflows at n=3072.
const BAKE_SILU_SRC: &str = r#"
struct BsP { inter: u32, total: u32, scaled: u32, stride: u32 };
@group(0) @binding(0) var<storage, read> bs_b : array<f32>;
@group(0) @binding(1) var<storage, read> bs_s : array<f32>;
@group(0) @binding(2) var<storage, read_write> bs_a : array<f32>;
@group(0) @binding(3) var<uniform> bs_p : BsP;
@compute @workgroup_size(256)
fn bake_silu_mul(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.y * bs_p.stride + gid.x;
    if (i >= bs_p.total) { return; }
    let r = i / bs_p.inter;
    let j = i % bs_p.inter;
    let g = bs_b[r * 2u * bs_p.inter + j];
    let u = bs_b[r * 2u * bs_p.inter + bs_p.inter + j];
    var a = (g / (1.0 + exp(-g))) * u;
    if (bs_p.scaled != 0u) { a = a * bs_s[j]; }
    bs_a[i] = a;
}
"#;

/// The backward twin of the silu link, in its own module: reads the
/// RESIDENT gate+up plane the forward chain parked on the card and turns
/// dact into the concatenated dgu — silu·mul backward, exactly
/// `ops::silu_bwd`'s σ(x)·(1 + x·(1−σ(x))).
const BAKE_SILU_BWD_SRC: &str = r#"
struct BwP { inter: u32, total: u32, _a: u32, stride: u32 };
@group(0) @binding(0) var<storage, read> bw_plane : array<f32>;
@group(0) @binding(1) var<storage, read> bw_dact : array<f32>;
@group(0) @binding(2) var<storage, read_write> bw_dgu : array<f32>;
@group(0) @binding(3) var<uniform> bw_p : BwP;
@compute @workgroup_size(256)
fn bake_silu_bwd(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.y * bw_p.stride + gid.x;
    if (i >= bw_p.total) { return; }
    let r = i / bw_p.inter;
    let j = i % bw_p.inter;
    let g = bw_plane[r * 2u * bw_p.inter + j];
    let u = bw_plane[r * 2u * bw_p.inter + bw_p.inter + j];
    let da = bw_dact[i];
    let s = 1.0 / (1.0 + exp(-g));
    bw_dgu[r * 2u * bw_p.inter + j] = da * u * (s * (1.0 + g * (1.0 - s)));
    bw_dgu[r * 2u * bw_p.inter + bw_p.inter + j] = da * (g * s);
}
"#;

/// The bake attention chain, stage 1: from the fused qkv plane the coop
/// GEMM left on the card — bias, per-head [q; gate] split, qk-RMSNorm
/// (inv saved for the backward), RoPE from a HOST-precomputed cos/sin
/// table (host trigonometry is f64; recomputing angles in f32 on device
/// is exactly the parity drift the table avoids). One thread per
/// (row, unit) where units run q-heads, then k-heads, then v-heads.
const BAKE_QKR_SRC: &str = r#"
struct QkrP {
    n: u32, t: u32, nh: u32, nkv: u32,
    hd: u32, qrows: u32, half: u32, flags: u32,
    eps: f32, _a: u32, _b: u32, _c: u32,
};
@group(0) @binding(0) var<storage, read> qk_plane : array<f32>;
@group(0) @binding(1) var<storage, read> qk_qnorm : array<f32>;
@group(0) @binding(2) var<storage, read> qk_knorm : array<f32>;
@group(0) @binding(3) var<storage, read> qk_bias : array<f32>;
@group(0) @binding(4) var<storage, read> qk_rope : array<f32>;
@group(0) @binding(5) var<storage, read_write> qk_qrot : array<f32>;
@group(0) @binding(6) var<storage, read_write> qk_krot : array<f32>;
@group(0) @binding(7) var<storage, read_write> qk_vproj : array<f32>;
@group(0) @binding(8) var<storage, read_write> qk_gate : array<f32>;
@group(0) @binding(9) var<storage, read_write> qk_qinv : array<f32>;
@group(0) @binding(10) var<storage, read_write> qk_kinv : array<f32>;
@group(0) @binding(11) var<uniform> qk_p : QkrP;

const FL_GATED: u32 = 1u;
const FL_QNORM: u32 = 2u;
const FL_KNORM: u32 = 4u;
const FL_GEMMA: u32 = 8u;

var<private> head_buf: array<f32, 256>;

fn qk_rms(len: u32, w_is_q: bool) -> f32 {
    var ss = 0.0;
    for (var j = 0u; j < len; j = j + 1u) {
        ss = ss + head_buf[j] * head_buf[j];
    }
    let inv = 1.0 / sqrt(ss / f32(len) + qk_p.eps);
    for (var j = 0u; j < len; j = j + 1u) {
        var w: f32;
        if (w_is_q) { w = qk_qnorm[j]; } else { w = qk_knorm[j]; }
        if ((qk_p.flags & FL_GEMMA) != 0u) { w = 1.0 + w; }
        head_buf[j] = head_buf[j] * inv * w;
    }
    return inv;
}

fn qk_rope_apply(pos: u32) {
    for (var i = 0u; i < qk_p.half; i = i + 1u) {
        let c = qk_rope[(pos * qk_p.half + i) * 2u];
        let s = qk_rope[(pos * qk_p.half + i) * 2u + 1u];
        let x0 = head_buf[i];
        let x1 = head_buf[i + qk_p.half];
        head_buf[i] = x0 * c - x1 * s;
        head_buf[i + qk_p.half] = x0 * s + x1 * c;
    }
}

@compute @workgroup_size(64)
fn bake_qkr(@builtin(global_invocation_id) gid: vec3<u32>) {
    let units_per_row = qk_p.nh + 2u * qk_p.nkv;
    let total = qk_p.n * units_per_row;
    if (gid.x >= total) { return; }
    let r = gid.x / units_per_row;
    let u = gid.x % units_per_row;
    let pos = r % qk_p.t;
    let hd = qk_p.hd;
    let kvdim = qk_p.nkv * hd;
    let fused = qk_p.qrows + 2u * kvdim;
    let base_row = r * fused;
    if (u < qk_p.nh) {
        // q-head u: [q(hd); gate(hd)] when gated, plain q otherwise.
        let h = u;
        var src: u32;
        if ((qk_p.flags & FL_GATED) != 0u) { src = base_row + 2u * h * hd; }
        else { src = base_row + h * hd; }
        var boff: u32;
        if ((qk_p.flags & FL_GATED) != 0u) { boff = 2u * h * hd; } else { boff = h * hd; }
        for (var j = 0u; j < hd; j = j + 1u) {
            head_buf[j] = qk_plane[src + j] + qk_bias[boff + j];
        }
        var inv = 1.0;
        if ((qk_p.flags & FL_QNORM) != 0u) { inv = qk_rms(hd, true); }
        qk_qinv[r * qk_p.nh + h] = inv;
        qk_rope_apply(pos);
        let dst = r * qk_p.nh * hd + h * hd;
        for (var j = 0u; j < hd; j = j + 1u) {
            qk_qrot[dst + j] = head_buf[j];
        }
        if ((qk_p.flags & FL_GATED) != 0u) {
            for (var j = 0u; j < hd; j = j + 1u) {
                qk_gate[dst + j] = qk_plane[src + hd + j] + qk_bias[boff + hd + j];
            }
        }
    } else if (u < qk_p.nh + qk_p.nkv) {
        // k-head
        let g = u - qk_p.nh;
        let src = base_row + qk_p.qrows + g * hd;
        for (var j = 0u; j < hd; j = j + 1u) {
            head_buf[j] = qk_plane[src + j] + qk_bias[qk_p.qrows + g * hd + j];
        }
        var inv = 1.0;
        if ((qk_p.flags & FL_KNORM) != 0u) { inv = qk_rms(hd, false); }
        qk_kinv[r * qk_p.nkv + g] = inv;
        qk_rope_apply(pos);
        let dst = r * kvdim + g * hd;
        for (var j = 0u; j < hd; j = j + 1u) {
            qk_krot[dst + j] = head_buf[j];
        }
    } else {
        // v-head: bias only.
        let g = u - qk_p.nh - qk_p.nkv;
        let src = base_row + qk_p.qrows + kvdim + g * hd;
        let dst = r * kvdim + g * hd;
        for (var j = 0u; j < hd; j = j + 1u) {
            qk_vproj[dst + j] = qk_plane[src + j] + qk_bias[qk_p.qrows + kvdim + g * hd + j];
        }
    }
}
"#;

/// Stage 2: causal max-softmax attention, one workgroup per (query
/// position, sequence·head). Phase one spreads the ≤t score dots over
/// the lanes; the reduction and the exp/denominator run on lane 0 (a
/// few hundred scalar ops — not worth a tree); phase two flips the
/// axis: each lane owns one output component and walks j SEQUENTIALLY,
/// which is the host loop's accumulation order per component.
const BAKE_ATTN_SRC: &str = r#"
struct AtP { t: u32, nh: u32, nkv: u32, hd: u32, b: u32, _a: u32, _b: u32, _c: u32 };
@group(0) @binding(0) var<storage, read> at_q : array<f32>;
@group(0) @binding(1) var<storage, read> at_k : array<f32>;
@group(0) @binding(2) var<storage, read> at_v : array<f32>;
@group(0) @binding(3) var<storage, read_write> at_o : array<f32>;
@group(0) @binding(4) var<uniform> at_p : AtP;

var<workgroup> at_scores: array<f32, 1024>;
var<workgroup> at_stat: array<f32, 2>;

@compute @workgroup_size(128)
fn bake_attn_head(@builtin(workgroup_id) wid: vec3<u32>,
                  @builtin(local_invocation_index) lid: u32) {
    let ti = wid.x;
    let unit = wid.y;
    let bi = unit / at_p.nh;
    let h = unit % at_p.nh;
    let rep = at_p.nh / at_p.nkv;
    let g = h / rep;
    let hd = at_p.hd;
    let qdim = at_p.nh * hd;
    let kvdim = at_p.nkv * hd;
    let scale = 1.0 / sqrt(f32(hd));
    let qbase = (bi * at_p.t + ti) * qdim + h * hd;
    // Phase 1: scores for j ≤ ti.
    for (var j = lid; j <= ti; j = j + 128u) {
        let kbase = (bi * at_p.t + j) * kvdim + g * hd;
        var s = 0.0;
        for (var c = 0u; c < hd; c = c + 1u) {
            s = s + at_q[qbase + c] * at_k[kbase + c];
        }
        at_scores[j] = s * scale;
    }
    workgroupBarrier();
    if (lid == 0u) {
        var mx = at_scores[0];
        for (var j = 1u; j <= ti; j = j + 1u) {
            mx = max(mx, at_scores[j]);
        }
        var den = 0.0;
        for (var j = 0u; j <= ti; j = j + 1u) {
            let e = exp(at_scores[j] - mx);
            at_scores[j] = e;
            den = den + e;
        }
        at_stat[0] = den;
    }
    workgroupBarrier();
    // Phase 2: lane c owns output component c, walks j in order.
    let den = at_stat[0];
    if (lid < hd) {
        var acc = 0.0;
        for (var j = 0u; j <= ti; j = j + 1u) {
            let p = at_scores[j] / den;
            acc = acc + p * at_v[(bi * at_p.t + j) * kvdim + g * hd + lid];
        }
        at_o[(bi * at_p.t + ti) * qdim + h * hd + lid] = acc;
    }
}
"#;

/// Stage 3: the output gate — ao·σ(gate), elementwise, 2-D dispatch.
const BAKE_AGATE_SRC: &str = r#"
struct AgP { total: u32, stride: u32, _a: u32, _b: u32 };
@group(0) @binding(0) var<storage, read> ag_a : array<f32>;
@group(0) @binding(1) var<storage, read> ag_g : array<f32>;
@group(0) @binding(2) var<storage, read_write> ag_o : array<f32>;
@group(0) @binding(3) var<uniform> ag_p : AgP;
@compute @workgroup_size(256)
fn bake_attn_gate(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.y * ag_p.stride + gid.x;
    if (i >= ag_p.total) { return; }
    let g = ag_g[i];
    ag_o[i] = ag_a[i] * (1.0 / (1.0 + exp(-g)));
}
"#;

// Split-K decode attention (its own module: the main module's at_* binding
// slots are taken, and WGSL forbids two resource vars on one binding).
// `gqa_attend_part` runs the flash-decoding loop over ONE ck-position chunk
// per workgroup — grid (nh, nchunks) instead of nh, which left a discrete GPU
// at 16 resident workgroups and latency-bound at depth — and stores each
// chunk's unnormalized accumulator plus its (m, l) softmax frame.
// `gqa_attend_merge` (grid nh) rescales the chunk frames into the global max
// and normalizes. Same math as `gqa_attend` up to one extra merge rounding.
const ATTEND_SPLIT_SRC: &str = r#"
struct ApP { nh: u32, hpk: u32, hd: u32, cap: u32, n: u32, ck: u32, nc: u32, scale: f32 };
@group(0) @binding(0) var<storage, read>       ap_q  : array<vec4<f32>>;
@group(0) @binding(1) var<storage, read>       ap_k  : array<vec4<f32>>;
@group(0) @binding(2) var<storage, read>       ap_v  : array<vec4<f32>>;
@group(0) @binding(3) var<storage, read_write> ap_acc: array<f32>;
@group(0) @binding(4) var<storage, read_write> ap_ml : array<vec2<f32>>;
@group(0) @binding(5) var<uniform>             ap_p  : ApP;
@group(0) @binding(6) var<storage, read_write> ap_o  : array<f32>;
var<workgroup> app_acc: array<f32, 8224>;
var<workgroup> app_m: array<f32, 32>;
var<workgroup> app_l: array<f32, 32>;
@compute @workgroup_size(32)
fn gqa_attend_part(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let ch = wid.y;
    let lane = lid.x;
    if (h >= ap_p.nh) { return; }
    let hd = ap_p.hd;
    let hd4 = hd / 4u;
    let p0 = ch * ap_p.ck;
    let pend = min(ap_p.n, p0 + ap_p.ck);
    let kbase = (h / ap_p.hpk) * ap_p.cap * hd4;
    let qbase = h * hd4;
    let scale = ap_p.scale;
    let base = lane * 257u;
    for (var d = 0u; d < hd; d = d + 1u) { app_acc[base + d] = 0.0; }
    var m = -1e30;
    var l = 0.0;
    var p = p0 + lane;
    loop {
        if (p >= pend) { break; }
        let krow = kbase + p * hd4;
        var dot4 = vec4<f32>(0.0);
        for (var d = 0u; d < hd4; d = d + 1u) { dot4 = dot4 + ap_q[qbase + d] * ap_k[krow + d]; }
        let dot = (dot4.x + dot4.y + dot4.z + dot4.w) * scale;
        let mp = max(m, dot);
        let f = exp(m - mp);
        let w = exp(dot - mp);
        l = l * f + w;
        for (var d = 0u; d < hd4; d = d + 1u) {
            let vv = ap_v[krow + d] * w;
            let a = base + d * 4u;
            app_acc[a]      = app_acc[a]      * f + vv.x;
            app_acc[a + 1u] = app_acc[a + 1u] * f + vv.y;
            app_acc[a + 2u] = app_acc[a + 2u] * f + vv.z;
            app_acc[a + 3u] = app_acc[a + 3u] * f + vv.w;
        }
        m = mp;
        p = p + 32u;
    }
    app_m[lane] = m;
    app_l[lane] = l;
    workgroupBarrier();
    var stride = 16u;
    loop {
        if (stride == 0u) { break; }
        if (lane < stride) {
            let o = lane + stride;
            let m1 = app_m[lane];
            let m2 = app_m[o];
            let mm = max(m1, m2);
            let f1 = exp(m1 - mm);
            let f2 = exp(m2 - mm);
            app_l[lane] = app_l[lane] * f1 + app_l[o] * f2;
            let bo = o * 257u;
            for (var d = 0u; d < hd; d = d + 1u) {
                app_acc[base + d] = app_acc[base + d] * f1 + app_acc[bo + d] * f2;
            }
            app_m[lane] = mm;
        }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let idx = h * ap_p.nc + ch;
    for (var d = lane; d < hd; d = d + 32u) {
        ap_acc[idx * hd + d] = app_acc[d];
    }
    if (lane == 0u) {
        ap_ml[idx] = vec2<f32>(app_m[0], app_l[0]);
    }
}
// hd <= 128 twin of gqa_attend_part at stride 129 (16.5 KB workgroup
// memory — fits the 32 KB mobile/Metal limit; see gqa_attend_s).
var<workgroup> app_acc_s: array<f32, 4128>;
@compute @workgroup_size(32)
fn gqa_attend_part_s(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let ch = wid.y;
    let lane = lid.x;
    if (h >= ap_p.nh) { return; }
    let hd = ap_p.hd;
    let hd4 = hd / 4u;
    let p0 = ch * ap_p.ck;
    let pend = min(ap_p.n, p0 + ap_p.ck);
    let kbase = (h / ap_p.hpk) * ap_p.cap * hd4;
    let qbase = h * hd4;
    let scale = ap_p.scale;
    let base = lane * 129u;
    for (var d = 0u; d < hd; d = d + 1u) { app_acc_s[base + d] = 0.0; }
    var m = -1e30;
    var l = 0.0;
    var p = p0 + lane;
    loop {
        if (p >= pend) { break; }
        let krow = kbase + p * hd4;
        var dot4 = vec4<f32>(0.0);
        for (var d = 0u; d < hd4; d = d + 1u) { dot4 = dot4 + ap_q[qbase + d] * ap_k[krow + d]; }
        let dot = (dot4.x + dot4.y + dot4.z + dot4.w) * scale;
        let mp = max(m, dot);
        let f = exp(m - mp);
        let w = exp(dot - mp);
        l = l * f + w;
        for (var d = 0u; d < hd4; d = d + 1u) {
            let vv = ap_v[krow + d] * w;
            let a = base + d * 4u;
            app_acc_s[a]      = app_acc_s[a]      * f + vv.x;
            app_acc_s[a + 1u] = app_acc_s[a + 1u] * f + vv.y;
            app_acc_s[a + 2u] = app_acc_s[a + 2u] * f + vv.z;
            app_acc_s[a + 3u] = app_acc_s[a + 3u] * f + vv.w;
        }
        m = mp;
        p = p + 32u;
    }
    app_m[lane] = m;
    app_l[lane] = l;
    workgroupBarrier();
    var stride = 16u;
    loop {
        if (stride == 0u) { break; }
        if (lane < stride) {
            let o = lane + stride;
            let m1 = app_m[lane];
            let m2 = app_m[o];
            let mm = max(m1, m2);
            let f1 = exp(m1 - mm);
            let f2 = exp(m2 - mm);
            app_l[lane] = app_l[lane] * f1 + app_l[o] * f2;
            let bo = o * 129u;
            for (var d = 0u; d < hd; d = d + 1u) {
                app_acc_s[base + d] = app_acc_s[base + d] * f1 + app_acc_s[bo + d] * f2;
            }
            app_m[lane] = mm;
        }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let idx = h * ap_p.nc + ch;
    for (var d = lane; d < hd; d = d + 32u) {
        ap_acc[idx * hd + d] = app_acc_s[d];
    }
    if (lane == 0u) {
        ap_ml[idx] = vec2<f32>(app_m[0], app_l[0]);
    }
}
@compute @workgroup_size(32)
fn gqa_attend_merge(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
    let h = wid.x;
    let lane = lid.x;
    if (h >= ap_p.nh) { return; }
    let hd = ap_p.hd;
    let nc = (ap_p.n + ap_p.ck - 1u) / ap_p.ck;
    var mg = -1e30;
    for (var ci = 0u; ci < nc; ci = ci + 1u) { mg = max(mg, ap_ml[h * ap_p.nc + ci].x); }
    var lg = 0.0;
    for (var ci = 0u; ci < nc; ci = ci + 1u) {
        let ml = ap_ml[h * ap_p.nc + ci];
        lg = lg + ml.y * exp(ml.x - mg);
    }
    let invl = select(0.0, 1.0 / lg, lg > 0.0);
    for (var d = lane; d < hd; d = d + 32u) {
        var a = 0.0;
        for (var ci = 0u; ci < nc; ci = ci + 1u) {
            let idx = h * ap_p.nc + ci;
            a = a + ap_acc[idx * hd + d] * exp(ap_ml[idx].x - mg);
        }
        ap_o[h * hd + d] = a * invl;
    }
}

// GQA-SHARED split-K decode attend: ONE workgroup per (kv head, chunk of
// AG_CK positions) serves EVERY query head of that group, so a K/V row
// is read once for all of them instead of once per query head. Qwen3.8
// puts six query heads on each kv head: at 12k context the per-head
// kernel streamed ~600 MB a layer against a 100 MB cache — the token
// went from 20 ms at short context to 37 at 10k, and this is the
// difference. Lanes are POSITIONS for the score pass and DIMENSIONS for
// the value pass (hd <= 256, hpk <= 8); one chunk per workgroup, so the
// online-softmax frame is just the chunk's (m, l). Partials land in the
// [h*nc + ch] layout `gqa_attend_merge` already reads, with ck = AG_CK.
// Same math as the per-head split kernel up to reduction order.
const AG_CK: u32 = 256u;
var<workgroup> ag_q: array<vec4<f32>, 512>;   // [hq*hd4 + d4], hpk<=8, hd<=256
var<workgroup> ag_sc: array<f32, 2048>;       // [hq*256 + p]
var<workgroup> ag_red: array<f32, 2048>;      // [p*8 + hq]
@compute @workgroup_size(256)
fn gqa_attend_gpart(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_index) lid: u32) {
    let g = wid.x;
    let ch = wid.y;
    let hpk = ap_p.hpk;
    let nkv = ap_p.nh / hpk;
    if (g >= nkv) { return; }
    let hd = ap_p.hd;
    let hd4 = hd / 4u;
    let p0 = ch * AG_CK;
    let pend = min(ap_p.n, p0 + AG_CK);
    let cn = select(0u, pend - p0, pend > p0);
    let kbase = g * ap_p.cap * hd4;
    let scale = ap_p.scale;
    // Stage this group's queries (zeros for the unused head slots so the
    // unrolled dots below need no guards).
    for (var i = lid; i < 512u; i = i + 256u) {
        let hq = i / hd4;
        let d4 = i % hd4;
        var qv = vec4<f32>(0.0);
        if (hq < hpk && i < hpk * hd4) { qv = ap_q[(g * hpk + hq) * hd4 + d4]; }
        ag_q[i] = qv;
    }
    workgroupBarrier();
    // Score pass: lane = position, eight dot chains off ONE K row.
    var d0 = vec4<f32>(0.0); var d1 = vec4<f32>(0.0);
    var d2 = vec4<f32>(0.0); var d3 = vec4<f32>(0.0);
    var d4 = vec4<f32>(0.0); var d5 = vec4<f32>(0.0);
    var d6 = vec4<f32>(0.0); var d7 = vec4<f32>(0.0);
    if (lid < cn) {
        let krow = kbase + (p0 + lid) * hd4;
        for (var d = 0u; d < hd4; d = d + 1u) {
            let kv = ap_k[krow + d];
            d0 = d0 + ag_q[d] * kv;
            d1 = d1 + ag_q[hd4 + d] * kv;
            d2 = d2 + ag_q[2u * hd4 + d] * kv;
            d3 = d3 + ag_q[3u * hd4 + d] * kv;
            d4 = d4 + ag_q[4u * hd4 + d] * kv;
            d5 = d5 + ag_q[5u * hd4 + d] * kv;
            d6 = d6 + ag_q[6u * hd4 + d] * kv;
            d7 = d7 + ag_q[7u * hd4 + d] * kv;
        }
    }
    let live = lid < cn;
    var sc: array<f32, 8>;
    sc[0] = select(-1e30, (d0.x + d0.y + d0.z + d0.w) * scale, live);
    sc[1] = select(-1e30, (d1.x + d1.y + d1.z + d1.w) * scale, live);
    sc[2] = select(-1e30, (d2.x + d2.y + d2.z + d2.w) * scale, live);
    sc[3] = select(-1e30, (d3.x + d3.y + d3.z + d3.w) * scale, live);
    sc[4] = select(-1e30, (d4.x + d4.y + d4.z + d4.w) * scale, live);
    sc[5] = select(-1e30, (d5.x + d5.y + d5.z + d5.w) * scale, live);
    sc[6] = select(-1e30, (d6.x + d6.y + d6.z + d6.w) * scale, live);
    sc[7] = select(-1e30, (d7.x + d7.y + d7.z + d7.w) * scale, live);
    for (var h = 0u; h < 8u; h = h + 1u) {
        ag_sc[h * 256u + lid] = sc[h];
        ag_red[lid * 8u + h] = sc[h];
    }
    workgroupBarrier();
    // Chunk max, all eight heads through one tree.
    var st = 128u;
    loop {
        if (st == 0u) { break; }
        if (lid < st) {
            for (var h = 0u; h < 8u; h = h + 1u) {
                ag_red[lid * 8u + h] = max(ag_red[lid * 8u + h], ag_red[(lid + st) * 8u + h]);
            }
        }
        workgroupBarrier();
        st = st >> 1u;
    }
    var cm: array<f32, 8>;
    for (var h = 0u; h < 8u; h = h + 1u) { cm[h] = ag_red[h]; }
    workgroupBarrier();
    // Weights into shared, denominators through the same tree.
    for (var h = 0u; h < 8u; h = h + 1u) {
        let w = select(0.0, exp(sc[h] - cm[h]), live);
        ag_sc[h * 256u + lid] = w;
        ag_red[lid * 8u + h] = w;
    }
    workgroupBarrier();
    st = 128u;
    loop {
        if (st == 0u) { break; }
        if (lid < st) {
            for (var h = 0u; h < 8u; h = h + 1u) {
                ag_red[lid * 8u + h] = ag_red[lid * 8u + h] + ag_red[(lid + st) * 8u + h];
            }
        }
        workgroupBarrier();
        st = st >> 1u;
    }
    // Value pass: lane = output dim, ONE V row read per position for all heads.
    if (lid < hd) {
        var a0 = 0.0; var a1 = 0.0; var a2 = 0.0; var a3 = 0.0;
        var a4 = 0.0; var a5 = 0.0; var a6 = 0.0; var a7 = 0.0;
        let dw = lid >> 2u;
        let dc = lid & 3u;
        for (var p = 0u; p < cn; p = p + 1u) {
            let v = ap_v[kbase + (p0 + p) * hd4 + dw][dc];
            a0 = a0 + ag_sc[p] * v;
            a1 = a1 + ag_sc[256u + p] * v;
            a2 = a2 + ag_sc[512u + p] * v;
            a3 = a3 + ag_sc[768u + p] * v;
            a4 = a4 + ag_sc[1024u + p] * v;
            a5 = a5 + ag_sc[1280u + p] * v;
            a6 = a6 + ag_sc[1536u + p] * v;
            a7 = a7 + ag_sc[1792u + p] * v;
        }
        var acc: array<f32, 8>;
        acc[0] = a0; acc[1] = a1; acc[2] = a2; acc[3] = a3;
        acc[4] = a4; acc[5] = a5; acc[6] = a6; acc[7] = a7;
        for (var h = 0u; h < hpk; h = h + 1u) {
            let idx = (g * hpk + h) * ap_p.nc + ch;
            ap_acc[idx * hd + lid] = acc[h];
        }
    }
    if (lid == 0u) {
        for (var h = 0u; h < hpk; h = h + 1u) {
            let idx = (g * hpk + h) * ap_p.nc + ch;
            ap_ml[idx] = vec2<f32>(cm[h], ag_red[h]);
        }
    }
}
"#;

/// Positions per GQA-shared split-K chunk (`gqa_attend_gpart`): one
/// workgroup of 256 lanes owns exactly one chunk.
const ATTEND_GCK: usize = 256;

/// Positions per split-K attend chunk; the split path engages past
/// `ATTEND_SPLIT_MIN` cached positions (below it the single-workgroup
/// kernel's one dispatch wins).
const ATTEND_CK: usize = 128;
const ATTEND_SPLIT_MIN: usize = 256;

/// Experimental global q2tp ladder-cache component.  The production q2tp
/// shader keeps its row-local ladder in workgroup memory; this opt-in arm
/// interns exact `(f16 lo, f16 step)` pairs once and addresses the resulting
/// 32-slot F32 ladders with a per-row u32 id.  It is deliberately separate
/// from the main module so a validation failure cannot poison the ordinary
/// dtype16/affine path.
const Q2TP_LADDER_CACHE_BUILD_SRC: &str = r#"
struct LadderParams { pairs: u32, _p0: u32, _p1: u32, _p2: u32 };
@group(0) @binding(0) var<storage, read> raw_keys : array<u32>;
@group(0) @binding(1) var<storage, read_write> ladders : array<f32>;
@group(0) @binding(2) var<uniform> lp : LadderParams;

@compute @workgroup_size(256)
fn q2_ladder_build(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= lp.pairs * 32u) { return; }
    let rung = i & 31u;
    if (rung == 0u) {
        ladders[i] = 0.0;
    } else {
        let pr = unpack2x16float(raw_keys[i >> 5u]);
        // Keep this operation/order identical to q2tp_matvec16w's row-local
        // ladder.  The component gate checks the resulting F32 bit pattern.
        ladders[i] = exp2(pr.x + f32(rung - 1u) * pr.y);
    }
}
"#;

const Q2TP_LADDER_CACHE_MV_SRC: &str = r#"
struct Q1Params { np: u32, rows: u32, _p0: u32, _p1: u32 };
@group(0) @binding(0) var<storage, read> q1w : array<u32>;
@group(0) @binding(2) var<storage, read_write> q1y : array<f32>;
@group(0) @binding(3) var<uniform> q1p : Q1Params;
@group(0) @binding(5) var<storage, read> q4v_x : array<vec4<f32>>;
@group(0) @binding(6) var<storage, read> q2_cache : array<f32>;
@group(0) @binding(7) var<storage, read> q2_row_ids : array<u32>;

var<workgroup> partial_q2_cache : array<vec4<f32>, 256>;

fn q2tp_byte(off: u32) -> u32 {
    return (q1w[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
}

fn q2v_c2(w: u32, sh: u32, affine: u32) -> f32 {
    return bitcast<f32>((((w >> sh) & 3u) << 1u) | 0x4B000000u)
        - select(8388611.0, 8388610.0, affine != 0u);
}

fn q2v_d16(w: u32, a: vec4<f32>, b: vec4<f32>, c: vec4<f32>, d: vec4<f32>, affine: u32) -> f32 {
    return (q2v_c2(w, 0u, affine) * a.x
         + q2v_c2(w, 2u, affine) * a.y
         + q2v_c2(w, 4u, affine) * a.z
         + q2v_c2(w, 6u, affine) * a.w
         + q2v_c2(w, 8u, affine) * b.x
         + q2v_c2(w, 10u, affine) * b.y
         + q2v_c2(w, 12u, affine) * b.z
         + q2v_c2(w, 14u, affine) * b.w
         + q2v_c2(w, 16u, affine) * c.x
         + q2v_c2(w, 18u, affine) * c.y
         + q2v_c2(w, 20u, affine) * c.z
         + q2v_c2(w, 22u, affine) * c.w
         + q2v_c2(w, 24u, affine) * d.x
         + q2v_c2(w, 26u, affine) * d.y
         + q2v_c2(w, 28u, affine) * d.z
         + q2v_c2(w, 30u, affine) * d.w) * 0.5;
}

@compute @workgroup_size(256)
fn q2tp_matvec16w_ladder_cache(@builtin(workgroup_id) wid: vec3<u32>,
                               @builtin(num_workgroups) nwg: vec3<u32>,
                               @builtin(local_invocation_index) lid: u32) {
    let gpr = q1p.np;
    let rows = q1p.rows;
    let params_w = rows * gpr * 2u;
    let codes_b = rows * gpr * 8u + rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let sub = lid >> 6u;
    let l = lid & 63u;
    let blocks = (rows + 15u) / 16u;
    var wb = wid.x;
    loop {
        if (wb >= blocks) { break; }
        let base = wb * 16u;
        let r0 = base + sub;
        let r1 = base + sub + 4u;
        let r2 = base + sub + 8u;
        let r3 = base + sub + 12u;
        var acc = vec4<f32>(0.0);
        if (r0 < rows) {
            let c0 = codes_b + r0 * cstride;
            let c1 = codes_b + r1 * cstride;
            let c2 = codes_b + r2 * cstride;
            let c3 = codes_b + r3 * cstride;
            let l1 = r1 < rows;
            let l2 = r2 < rows;
            let l3 = r3 < rows;
            let id0 = q2_row_ids[q1p._p0 + r0] * 32u;
            let id1 = q2_row_ids[q1p._p0 + r1] * 32u;
            let id2 = q2_row_ids[q1p._p0 + r2] * 32u;
            let id3 = q2_row_ids[q1p._p0 + r3] * 32u;
            var g = l;
            loop {
                if (g >= gpr) { break; }
                let bit = g * 5u;
                let cbo = bit >> 3u;
                let sh = bit & 7u;
                let x0 = g * 8u;
                let xa = q4v_x[x0];      let xb = q4v_x[x0 + 1u];
                let xc = q4v_x[x0 + 2u]; let xd = q4v_x[x0 + 3u];
                let xe = q4v_x[x0 + 4u]; let xf = q4v_x[x0 + 5u];
                let xg = q4v_x[x0 + 6u]; let xh = q4v_x[x0 + 7u];
                var cv = q2tp_byte(c0 + cbo);
                if (sh > 3u) { cv = cv | (q2tp_byte(c0 + cbo + 1u) << 8u); }
                var wi = (r0 * gpr + g) * 2u;
                acc.x = acc.x + q2_cache[id0 + ((cv >> sh) & 31u)]
                    * (q2v_d16(q1w[wi], xa, xb, xc, xd, q1p._p1)
                     + q2v_d16(q1w[wi + 1u], xe, xf, xg, xh, q1p._p1));
                if (l1) {
                    cv = q2tp_byte(c1 + cbo);
                    if (sh > 3u) { cv = cv | (q2tp_byte(c1 + cbo + 1u) << 8u); }
                    wi = (r1 * gpr + g) * 2u;
                    acc.y = acc.y + q2_cache[id1 + ((cv >> sh) & 31u)]
                        * (q2v_d16(q1w[wi], xa, xb, xc, xd, q1p._p1)
                         + q2v_d16(q1w[wi + 1u], xe, xf, xg, xh, q1p._p1));
                }
                if (l2) {
                    cv = q2tp_byte(c2 + cbo);
                    if (sh > 3u) { cv = cv | (q2tp_byte(c2 + cbo + 1u) << 8u); }
                    wi = (r2 * gpr + g) * 2u;
                    acc.z = acc.z + q2_cache[id2 + ((cv >> sh) & 31u)]
                        * (q2v_d16(q1w[wi], xa, xb, xc, xd, q1p._p1)
                         + q2v_d16(q1w[wi + 1u], xe, xf, xg, xh, q1p._p1));
                }
                if (l3) {
                    cv = q2tp_byte(c3 + cbo);
                    if (sh > 3u) { cv = cv | (q2tp_byte(c3 + cbo + 1u) << 8u); }
                    wi = (r3 * gpr + g) * 2u;
                    acc.w = acc.w + q2_cache[id3 + ((cv >> sh) & 31u)]
                        * (q2v_d16(q1w[wi], xa, xb, xc, xd, q1p._p1)
                         + q2v_d16(q1w[wi + 1u], xe, xf, xg, xh, q1p._p1));
                }
                g = g + 64u;
            }
        }
        partial_q2_cache[lid] = acc;
        workgroupBarrier();
        var stride = 32u;
        loop {
            if (stride == 0u) { break; }
            if (l < stride) {
                partial_q2_cache[lid] = partial_q2_cache[lid] + partial_q2_cache[lid + stride];
            }
            workgroupBarrier();
            stride = stride >> 1u;
        }
        if (l == 0u) {
            let r = partial_q2_cache[sub << 6u];
            if (r0 < rows) { q1y[r0] = r.x; }
            if (r1 < rows) { q1y[r1] = r.y; }
            if (r2 < rows) { q1y[r2] = r.z; }
            if (r3 < rows) { q1y[r3] = r.w; }
        }
        workgroupBarrier();
        wb = wb + nwg.x;
    }
}
"#;

struct Q2LadderCache {
    model_uid: u64,
    pairs: usize,
    table_bytes: u64,
    row_id_bytes: u64,
    row_id_base: HashMap<usize, u32>,
    row_counts: HashMap<usize, usize>,
    ladders: wgpu::Buffer,
    row_ids: wgpu::Buffer,
}

struct Ctx {
    /// The instance and adapter the device came from, kept for exactly as
    /// long as the device — which Vulkan requires and we were not doing.
    /// They were locals in `init()`, so `VkInstance` was destroyed while a
    /// `VkDevice` made from it lived on in the static below. AddressSanitizer
    /// caught what that costs: a 48-byte block allocated by libEGL is freed
    /// on the way out of `init`, and freed a second time by the NVIDIA driver
    /// at process exit — the `double free or corruption` / `corrupted
    /// double-linked list` abort that has been landing after correct answers.
    _instance: wgpu::Instance,
    _adapter: wgpu::Adapter,
    device: wgpu::Device,
    queue: wgpu::Queue,
    /// Compiled pipelines, kept between runs where the driver supports
    /// it. `None` when it does not, or when the cache is switched off.
    pipeline_cache: Option<wgpu::PipelineCache>,
    /// The adapter's identity, so a later save writes to the same file
    /// the load came from.
    adapter_info: wgpu::AdapterInfo,
    matvec: wgpu::ComputePipeline,
    /// Mobile arm: activations staged in workgroup memory. `CMF_Q8MV=tiled`.
    matvec_tiled: wgpu::ComputePipeline,
    /// Token-graph matvec for q8_2f (both scale planes inside the tensor).
    q8_2f_mv: wgpu::ComputePipeline,
    matmat: wgpu::ComputePipeline,
    mul_mm: wgpu::ComputePipeline,
    q1_mm: wgpu::ComputePipeline,
    silu: wgpu::ComputePipeline,
    axpy: wgpu::ComputePipeline,
    colscale: wgpu::ComputePipeline,
    gate_mul: wgpu::ComputePipeline,
    zero: wgpu::ComputePipeline,
    q1: wgpu::ComputePipeline,
    /// Rows handled by this adapter's specialized q1 workgroup.
    q1_rows: u32,
    q1t: wgpu::ComputePipeline,
    q4b: wgpu::ComputePipeline,
    q4t_mv: wgpu::ComputePipeline,
    /// Hyper-connection fold (with the Sinkhorn) and expand — the join
    /// between blocks in DeepSeek-V4, where an ordinary model has a residual.
    /// Attention over an index list with a learned sink — DeepSeek-V4's,
    /// not the canonical sliding window.
    /// Per-head RMS and the rope tail, forward or inverse.
    rope_heads: wgpu::ComputePipeline,
    /// In-place BF16 materialisation used by the V4.1 attention tail.
    bf16_round: wgpu::ComputePipeline,
    o_lora_a: wgpu::ComputePipeline,
    kv_pool: wgpu::ComputePipeline,
    index_scores: wgpu::ComputePipeline,
    top_k_index: wgpu::ComputePipeline,
    hc_block: wgpu::ComputePipeline,
    f32_matvec_w: wgpu::ComputePipeline,
    o_lora_a_w: wgpu::ComputePipeline,
    f32_matvec_x: wgpu::ComputePipeline,
    f32_mv_split: wgpu::ComputePipeline,
    f32_mv_merge: wgpu::ComputePipeline,
    o_lora_a_m: wgpu::ComputePipeline,
    moe_gu_q2tp_m: wgpu::ComputePipeline,
    moe_dn_q4tp_m: wgpu::ComputePipeline,
    sa_part: wgpu::ComputePipeline,
    sa_merge: wgpu::ComputePipeline,
    blit: wgpu::ComputePipeline,
    idx_build: wgpu::ComputePipeline,
    moe_route: wgpu::ComputePipeline,
    sparse_attend: wgpu::ComputePipeline,
    sa_scores: wgpu::ComputePipeline,
    sa_apply: wgpu::ComputePipeline,
    hc_pre_fold: wgpu::ComputePipeline,
    hc_post_expand: wgpu::ComputePipeline,
    // Token-axis twins for the batched dsv4 frame.
    bt_rope_heads: wgpu::ComputePipeline,
    bt_hc_pre_fold: wgpu::ComputePipeline,
    bt_hc_block: wgpu::ComputePipeline,
    bt_comp_append: wgpu::ComputePipeline,
    bt_comp_fold: wgpu::ComputePipeline,
    bt_hc_post_expand: wgpu::ComputePipeline,
    bt_f32_matvec_w: wgpu::ComputePipeline,
    bt_f32_matvec_x: wgpu::ComputePipeline,
    bt_moe_route: wgpu::ComputePipeline,
    bt_moe_gate_up_q2tp: wgpu::ComputePipeline,
    bt_moe_gate_up_q2tp_r4: wgpu::ComputePipeline,
    bt_sparse_attend: wgpu::ComputePipeline,
    bt_index_scores: wgpu::ComputePipeline,
    bt_top_k: wgpu::ComputePipeline,
    bt_idx_build_staged: wgpu::ComputePipeline,
    bt_o_lora_a: wgpu::ComputePipeline,
    bt_o_lora_a4: wgpu::ComputePipeline,
    bt_o_lora_a2: Option<wgpu::ComputePipeline>,
    q4tp_mv: wgpu::ComputePipeline,
    /// Tall-matrix q4tp matvec (4 rows/workgroup, vec4 nibble loads); the
    /// per-row math is byte-identical to `q4tp_mv`. `CMF_MV4=0` reverts.
    q4tp_mv4: wgpu::ComputePipeline,
    q4tp_mv4_u2: wgpu::ComputePipeline,
    use_mv_u2: bool,
    q4tp_mv4_sg: Option<wgpu::ComputePipeline>,
    q4tp_mv4_nored: wgpu::ComputePipeline,
    use_mv_nored: bool,
    q4tp_mv4_dual: wgpu::ComputePipeline,
    use_mv_dual: bool,
    q4tp_mv4_dsilu: wgpu::ComputePipeline,
    /// Bind groups reused across tokens: (site, layer, kv_id) → (epoch,
    /// group). Valid while no buffer under it was reallocated — the
    /// epoch bumps on every scratch grow and weight eviction. Behind
    /// CMF_GRAPH_BGCACHE=1 until the parity suite blesses it: the host
    /// spends 13.7 of a 23.5 ms token CREATING these objects.
    graph_bgs: Mutex<HashMap<(u32, usize, u64), (u64, wgpu::BindGroup)>>,
    use_bgcache: bool,
    /// The same eight rows, but blocked over a SMALL BATCH in registers:
    /// the weight is read from DRAM once for the whole batch instead of
    /// once per element. What makes a speculative verify of k positions
    /// cost about one position's bandwidth. `CMF_MV_BK=0` reverts to the
    /// (row-block × batch) dispatch above.
    q4tp_mv4_bk: wgpu::ComputePipeline,
    /// The same, with the nibble unpack shared across the batch instead
    /// of repeated per element (`CMF_MV_BK=2` selects it).
    q4tp_mv4_bku: wgpu::ComputePipeline,
    /// 0 = the historical (row-block × batch) dispatch, 1 = batch inside
    /// the workgroup, 2 (default) = that plus a shared unpack.
    use_mv_bk: usize,
    /// `CMF_MV_PROBE=1|2`: the quad-row kernel with the arithmetic taken
    /// out (1) and the activation loads too (2). The ANSWER IS GARBAGE;
    /// the point is the time, which says how much of the matvec is the
    /// memory system and how much is the unpack.
    q4tp_mv16w_probe: Option<wgpu::ComputePipeline>,
    mv_probe: usize,
    use_mv4: bool,
    q4tp_mv16: wgpu::ComputePipeline,
    q4t_mv8: wgpu::ComputePipeline,
    q4b_mv8: wgpu::ComputePipeline,
    q4tp_mm: wgpu::ComputePipeline,
    q2tp_mm: wgpu::ComputePipeline,
    /// Optional Q2TP affine cooperative GEMM. It has its own four-binding
    /// layout because the uniform word is the affine descriptor bit, not the
    /// Q4 activation-scale/sentinel field.
    q2tp_mm_coop: Option<wgpu::ComputePipeline>,
    q4tp_mm_coop: Option<wgpu::ComputePipeline>,
    /// Same kernel, the scale read from a device buffer (binding 4).
    q4tp_mm_coop_s: Option<wgpu::ComputePipeline>,
    /// Dequantize-once path: q4tp plane → packed f16, then a pure f16
    /// coop GEMM. Both or neither.
    ffn_silu_packed: wgpu::ComputePipeline,
    /// max|x| of a device panel → the scale the coop GEMM multiplies by.
    act_absmax: Option<wgpu::ComputePipeline>,
    /// The DiT's attention GEMMs on the matrix units.
    dit_gemm_coop: Option<wgpu::ComputePipeline>,
    /// Interleaved qkv → head-major q/k/v, on the device.
    dit_qkv_split: Option<wgpu::ComputePipeline>,
    vae_im2col: Option<wgpu::ComputePipeline>,
    conv1d_im2col: Option<wgpu::ComputePipeline>,
    music3_glu: Option<wgpu::ComputePipeline>,
    /// v [h][n][hd] → [h][hd][n], so PV is an NT product.
    dit_v_transpose: Option<wgpu::ComputePipeline>,
    /// qk-norm + RoPE + head-major scatter, one of q/k per dispatch.
    dit_qknorm: Option<wgpu::ComputePipeline>,
    /// Two-stage form of the same: partials, then a fold.
    act_amax_part: Option<wgpu::ComputePipeline>,
    act_amax_fold: Option<wgpu::ComputePipeline>,
    q4tp_dq_f16: Option<wgpu::ComputePipeline>,
    /// The int8 twin of `q4tp_dq_f16` — same plane, same GEMM after it.
    q8_dq_f16: Option<wgpu::ComputePipeline>,
    q4tp_mm_coop_f16: Option<wgpu::ComputePipeline>,
    /// The bake's f32-operand forward GEMM on the matrix units; None off
    /// tensor-core devices, and `gemm_nt_f32` falls back to the scalar arm.
    gemm_nt_coop: Option<wgpu::ComputePipeline>,
    /// Its backward twin (reduction over w's leading axis) for `gemm_dx_f32`.
    gemm_nn_coop: Option<wgpu::ComputePipeline>,
    /// silu(g)·u·scale between the chain's two GEMMs (plain f32 module).
    bake_silu: wgpu::ComputePipeline,
    /// Its backward: dact + the resident plane → concatenated dgu.
    bake_silu_bwd: wgpu::ComputePipeline,
    /// The bake attention chain's middle: split+norm+rope, causal
    /// softmax·V per head, output gate (plain f32 modules).
    bake_qkr: wgpu::ComputePipeline,
    bake_attn_head: wgpu::ComputePipeline,
    bake_attn_gate: wgpu::ComputePipeline,
    /// Per-layer gate+up planes the forward chain parks for the backward
    /// one — the resident graph's memory. Grow-only per layer; ~84 MB a
    /// layer at chunk size, discrete cards only.
    bake_planes: Mutex<HashMap<usize, (wgpu::Buffer, u64)>>,
    argmax_part: wgpu::ComputePipeline,
    gdn_step_par: wgpu::ComputePipeline,
    gdn_step_norm: wgpu::ComputePipeline,
    gdn_par: bool,
    ts_query: Option<(wgpu::QuerySet, wgpu::Buffer, wgpu::Buffer)>,
    ts_period: f32,
    argmax_final: wgpu::ComputePipeline,
    embed_gather_q4tp: wgpu::ComputePipeline,
    silu_down: wgpu::ComputePipeline,
    q1t_mm: wgpu::ComputePipeline,
    q4t_mm: wgpu::ComputePipeline,
    dit_qk: wgpu::ComputePipeline,
    dit_pv: wgpu::ComputePipeline,
    dit_softmax: wgpu::ComputePipeline,
    dit_unstack: wgpu::ComputePipeline,
    ffn_silu: wgpu::ComputePipeline,
    /// Qwen Image's exact tanh-GELU+bias epilogue for a resident FFN.
    qwen_gelu_bias: wgpu::ComputePipeline,
    /// Qwen Image's affine-free LayerNorm plus shift/scale modulation.
    qwen_layernorm_mod: wgpu::ComputePipeline,
    /// Qwen Image's bias-then-gated residual epilogue.
    qwen_gated_residual: wgpu::ComputePipeline,
    /// Qwen Image's exact two-stream qk-norm/RoPE/join pass.  It is kept in
    /// the portable main shader module because it uses only f32 storage and
    /// no optional device feature.
    qwen_rope_pack: wgpu::ComputePipeline,
    q1t_ovmm: wgpu::ComputePipeline,
    rmsnorm: wgpu::ComputePipeline,
    add_rmsnorm: wgpu::ComputePipeline,
    rmsnorm_b: wgpu::ComputePipeline,
    add_rmsnorm_b: wgpu::ComputePipeline,
    attn_rope: wgpu::ComputePipeline,
    kv_append: wgpu::ComputePipeline,
    gqa_attend: wgpu::ComputePipeline,
    gqa_attend_s: wgpu::ComputePipeline,
    attend_part: wgpu::ComputePipeline,
    attend_part_s: wgpu::ComputePipeline,
    attend_merge: wgpu::ComputePipeline,
    /// The GQA-shared split-K attend (one workgroup per kv head and
    /// chunk, K/V read once for all its query heads). Needs 24.6 KB of
    /// workgroup storage, so it exists where `big_attend` does.
    /// `CMF_ATTEND_GQA=0` returns to the per-head split kernel.
    attend_gpart: Option<wgpu::ComputePipeline>,
    layout_attend_gpart: Option<wgpu::BindGroupLayout>,
    /// Max head_dim the attend kernels can serve on this device: 256
    /// when 33 KB of workgroup storage fits (desktop), 128 on 32 KB
    /// devices (Adreno/Mali/wgpu-Metal) where only the stride-129
    /// kernels exist.
    hd_cap: usize,
    /// 32 KB+ of workgroup storage: the split-K attend parts need it.
    big_attend: bool,
    gdn_step: wgpu::ComputePipeline,
    gdn_conv: wgpu::ComputePipeline,
    sconv_step: wgpu::ComputePipeline,
    f32_matvec: wgpu::ComputePipeline,
    f32_matvec_b: wgpu::ComputePipeline,
    layout_f32b: wgpu::BindGroupLayout,
    o1_far: wgpu::ComputePipeline,
    o1_push: wgpu::ComputePipeline,
    o1_attend: wgpu::ComputePipeline,
    layout_o1_far: wgpu::BindGroupLayout,
    layout_o1_push: wgpu::BindGroupLayout,
    layout_o1_attend: wgpu::BindGroupLayout,
    /// Device o1 state per (kv_id, layer); re-uploaded when the seal
    /// epoch changes (each generate seals fresh CPU state).
    o1m: Mutex<HashMap<(u64, usize), O1Dev>>,
    moe_select: wgpu::ComputePipeline,
    /// Two independent projections of one input in a single dispatch.
    matvec_pair: wgpu::ComputePipeline,
    layout_mv2: wgpu::BindGroupLayout,
    moe_gate_up: wgpu::ComputePipeline,
    moe_down: wgpu::ComputePipeline,
    /// q4tp twins — same bindings, ladder-plane scale decode.
    moe_gate_up_q4tp: wgpu::ComputePipeline,
    moe_down_q4tp: wgpu::ComputePipeline,
    /// Batched (token-axis) MoE for the batch graph; select_b carries the
    /// f32 router matvec inside itself.
    moe_select_b: wgpu::ComputePipeline,
    gdn_conv_k: wgpu::ComputePipeline,
    q4tp_mv_k: wgpu::ComputePipeline,
    q4tp_mv16w: wgpu::ComputePipeline,
    /// INT8-activation batched matvec (default; `CMF_VERIFY_I8=0`), one pipeline per
    /// batch 2..=8 (index = batch), and its quantizer.
    q4tp_mv4_bk8: Vec<wgpu::ComputePipeline>,
    x_quant_i8: wgpu::ComputePipeline,
    /// Scratch for the int8 activations: (x8 packed, per-group scale/sum).
    i8x: std::sync::Mutex<Option<(wgpu::Buffer, wgpu::Buffer, u64)>>,
    /// Two wide q4tp projections of one input in one dispatch (gate+up,
    /// GDN qkv+z, attention k+v). `CMF_MV_X2=0` splits them again.
    q4tp_mv16w_x2: wgpu::ComputePipeline,
    use_mv_x2: bool,
    /// FFN gate+up+SiLU in one dispatch (the x2 body with the same rows
    /// of both weights per workgroup and the SiLU in the epilogue).
    /// `CMF_MV_GU=0` returns to gate, up, silu as three dispatches.
    q4tp_mv16w_gu: wgpu::ComputePipeline,
    use_mv_gu: bool,
    /// The batched (bku) two-weight kernel for the batch graph.
    q4tp_mv4_bku_x2: wgpu::ComputePipeline,
    /// The dense 2-bit (q2tp profile) decode matvec, kind 9.
    q2tp_mv16w: wgpu::ComputePipeline,
    /// Opt-in affine q2tp NB=1 Q8/DP4A decode.  This stays separate from
    /// q4tp's batched kernel because q2tp has signed ternary symbols and no
    /// q4tp zero-sum correction.
    q2tp_mv1_i8: wgpu::ComputePipeline,
    /// Experimental global q2tp ladder cache, constructed only when
    /// CMF_Q2_LADDER_CACHE=1. The ordinary row-local shader remains the
    /// default and is the fail-closed fallback if this module is rejected.
    q2_ladder_build: Option<wgpu::ComputePipeline>,
    q2_ladder_ref: Option<wgpu::ComputePipeline>,
    q2_ladder_mv: Option<wgpu::ComputePipeline>,
    /// Optional subgroup-reduction q2tp matvec.  It is never selected unless
    /// SUBGROUP is present and CMF_Q2TP_SG=1; the tree kernel above remains
    /// the default/fallback.
    q2tp_mv16w_sg: Option<wgpu::ComputePipeline>,
    /// Descriptor-aware signed FWHT for the resident Prism graph.  Kept
    /// optional so adapters without shader f16 fail closed to the per-op
    /// implementation instead of changing activation-boundary precision.
    fwht: Option<wgpu::ComputePipeline>,
    f32_gemm_dx: wgpu::ComputePipeline,
    vae_conv: wgpu::ComputePipeline,
    dit_ropepack: wgpu::ComputePipeline,
    dit_gres: wgpu::ComputePipeline,
    dit_rmsmod: wgpu::ComputePipeline,
    gdn_step_par_k: wgpu::ComputePipeline,
    gdn_step_norm_k: wgpu::ComputePipeline,
    gdn_step_k: wgpu::ComputePipeline,
    moe_gate_up_q4tp_b: wgpu::ComputePipeline,
    moe_gate_up_q4tp_b_r4: wgpu::ComputePipeline,
    moe_down_q4tp_b: wgpu::ComputePipeline,
    /// Genuine model-wide Q4TP expert cache. These live in a separate
    /// binding-array module and are absent on backends without descriptor
    /// indexing; the old exact per-layer path remains the fallback there.
    dsv4_global_gu: Option<wgpu::ComputePipeline>,
    dsv4_global_gu_q2: Option<wgpu::ComputePipeline>,
    dsv4_global_dn: Option<wgpu::ComputePipeline>,
    /// Optional S16 binding-array pipelines for the V4.1 large-card profile.
    /// Generic models and unsupported adapters never construct these.
    dsv4_global_gu_s16: Option<wgpu::ComputePipeline>,
    dsv4_global_gu_q2_s16: Option<wgpu::ComputePipeline>,
    dsv4_global_dn_s16: Option<wgpu::ComputePipeline>,
    moe_down_q4tp_b2: wgpu::ComputePipeline,
    moe_down_q4tp_part: wgpu::ComputePipeline,
    moe_down_q4tp_b4: wgpu::ComputePipeline,
    moe_down_q2tp_b: wgpu::ComputePipeline,
    moe_down_q4tp_red: wgpu::ComputePipeline,
    layout_moe_sel_b: wgpu::BindGroupLayout,
    layout_moe_gu_b: wgpu::BindGroupLayout,
    layout_moe_dn_b: wgpu::BindGroupLayout,
    layout: wgpu::BindGroupLayout,
    layout_mm: wgpu::BindGroupLayout,
    layout_mmm: wgpu::BindGroupLayout,
    layout_q1mm: wgpu::BindGroupLayout,
    layout_silu: wgpu::BindGroupLayout,
    layout_axpy: wgpu::BindGroupLayout,
    layout_colscale: wgpu::BindGroupLayout,
    layout_gate_mul: wgpu::BindGroupLayout,
    layout_zero: wgpu::BindGroupLayout,
    layout_q1: wgpu::BindGroupLayout,
    layout_rmsnorm: wgpu::BindGroupLayout,
    layout_add_rmsnorm: wgpu::BindGroupLayout,
    layout_rmsnorm_b: wgpu::BindGroupLayout,
    layout_add_rmsnorm_b: wgpu::BindGroupLayout,
    layout_attn_rope: wgpu::BindGroupLayout,
    layout_kv: wgpu::BindGroupLayout,
    layout_attend: wgpu::BindGroupLayout,
    layout_attend_s: wgpu::BindGroupLayout,
    layout_attend_part: wgpu::BindGroupLayout,
    layout_attend_part_s: wgpu::BindGroupLayout,
    layout_attend_merge: wgpu::BindGroupLayout,
    layout_gdn: wgpu::BindGroupLayout,
    layout_gdn_conv: wgpu::BindGroupLayout,
    layout_sconv: wgpu::BindGroupLayout,
    layout_f32: wgpu::BindGroupLayout,
    layout_silu_down: wgpu::BindGroupLayout,
    layout_moe_sel: wgpu::BindGroupLayout,
    layout_moe_gu: wgpu::BindGroupLayout,
    layout_moe_dn: wgpu::BindGroupLayout,
    /// wgpu treats an auto-derived layout as exclusive to the pipeline it
    /// came from, so the q4tp twins need their own even though the
    /// binding lists are identical.
    layout_moe_gu_q4tp: wgpu::BindGroupLayout,
    moe_gate_up_q2tp: wgpu::ComputePipeline,
    layout_moe_gu_q2tp: wgpu::BindGroupLayout,
    moe_gate_up_q2tp_f: wgpu::ComputePipeline,
    moe_down_q4tp_f: wgpu::ComputePipeline,
    gqa_attend_dec: wgpu::ComputePipeline,
    moe_select_sg: Option<wgpu::ComputePipeline>,
    gdn_step_par2: wgpu::ComputePipeline,
    gdn_step_norm2: wgpu::ComputePipeline,
    gdn_inline: bool,
    attend_dec: bool,
    foldsel: bool,
    layout_moe_dn_q4tp: wgpu::BindGroupLayout,
    /// Discrete card (PCIe VRAM) vs UMA — thresholds and budgets differ.
    discrete: bool,
    /// Weight-residency budget in bytes (CMF_GPU_VRAM_MB override). On a
    /// 24 GB card holding a 35 GB model, the first-touched tensors (=
    /// the first layers, decode touches them in order) stay resident and
    /// the rest honestly fall back to CPU — ngl-style offload without an
    /// explicit layer list, and no OOM.
    vram_budget: u64,
    /// Bytes currently resident in `weight_bufs`.
    resident: std::sync::atomic::AtomicU64,
    /// Pooled per-op scratch (grow-only): xs upload, y output, uniform
    /// params, readback staging. Every op used to CREATE all four (plus
    /// a bind group) and map_async-poll a fresh staging buffer — pure
    /// allocator traffic on the hot path. The lock is held across the
    /// whole op (encode → submit → poll): ops already serialize on the
    /// single queue.
    scratch: Mutex<Scratch>,
    /// One GEMM at a time per device. The activation, result and stage
    /// slots are ONE buffer each per context, written under the scratch
    /// lock but read at submit time after it is released — so two
    /// threads could upload into the same slot and one would compute on
    /// the other's operand. The scratch lock cannot cover it: `readback`
    /// takes that lock itself, and std mutexes do not re-enter.
    mm_gate: Mutex<()>,
    /// One run-owned sidecar (the component gate exercises one real matrix).
    /// A mismatched model/tensor rebuilds it rather than reusing stale row ids.
    q2_ladder: Mutex<Option<Q2LadderCache>>,
    /// One unpacked f16 plane PER WEIGHT, kept across calls — the
    /// scratch-slot version re-unpacked every weight before every GEMM.
    planes: Mutex<std::collections::HashMap<(usize, usize), (wgpu::Buffer, u64)>>,
    /// Resident quant weights in VRAM — the WHOLE tensor is loaded once
    /// (key (base_ptr, idx)); ranges/batches address it by offset.
    ///
    /// Residency is DEMAND-DRIVEN and evicting: nothing is preloaded, the
    /// set grows as the model routes, and when the budget is full the
    /// least-valuable tensor makes room. Without eviction the first
    /// tensors to arrive owned the device for the process's life — a
    /// model that switched from prose to code kept the prose experts and
    /// ran the code ones on the CPU forever. The two working sets overlap
    /// at a Jaccard of 0.095, measured, so that is not a corner case.
    weight_bufs: Mutex<HashMap<(usize, usize), Resident>>,
    /// DeepSeek-V4's attention cache, per (kv id, layer), living on the card
    /// between tokens. Re-uploading it each token costs more than the frame
    /// saves: at 4K context it is megabytes per layer.
    dsv4_kv: Mutex<HashMap<(u64, usize), (wgpu::Buffer, usize)>>,
    /// The frame's working buffers, keyed by role and length. Their sizes do
    /// not change from token to token, and allocating ten of them per layer
    /// per token — 430 allocations a token on the release — is most of what a
    /// submission costs. Created once, written thereafter.
    dsv4_scratch: Mutex<HashMap<(u8, usize, usize), wgpu::Buffer>>,
    /// One compressor's streams, alive across tokens ON THE DEVICE:
    /// `[pending_kv, pending_score, prev_kv, prev_score]` keyed by
    /// (kind, kv_id, layer) — kind 0 is the attention compressor, 1 the
    /// indexer's own. Keeping these here is what lets a token advance the
    /// compressor without the host reading anything back: the streams are
    /// the only thing that carried state across the seam.
    dsv4_comp: Mutex<HashMap<(u8, u64, usize), [wgpu::Buffer; 4]>>,
    /// The indexer's compressed cache per layer, grown as the sequence is.
    dsv4_ixkv: Mutex<HashMap<(u64, usize), (wgpu::Buffer, usize)>>,
    /// Stable per-(tag, kv, layer) uniform slots for sequence-varying scalars.
    dsv4_uni: Mutex<HashMap<(u8, u64, usize), wgpu::Buffer>>,
    dsv4_store: Mutex<HashMap<(u8, u64, usize), wgpu::Buffer>>,
    /// CMF_DSV4_SLOT_CHECK: writes per slot key since the last submission.
    /// A slot may be written ONCE per submission — queue writes do not
    /// interleave with passes, so a second write reaches every dispatch of
    /// the first. That is the bug that made a run of layers route with the
    /// last layer's router bias, and a tag collision between two call sites
    /// looks exactly the same from here.
    slot_writes: Mutex<HashMap<(u8, u64, usize), u32>>,
    /// (epoch, bind groups) for the dsv4 frames — see `cached_bind`.
    dsv4_binds: Mutex<(u64, HashMap<(u8, u64, usize), wgpu::BindGroup>)>,
    /// Access clock for the aging above — one tick per weight lookup.
    res_clock: std::sync::atomic::AtomicU64,
    /// row_scale buffer per (model uid, (tensor idx, row0)) — small, cached.
    /// The model component is explicit so image-stage cleanup can remove one
    /// model without touching another model's scales.
    rs_bufs: Mutex<HashMap<(usize, (usize, usize)), wgpu::Buffer>>,
    /// Device K/V cache mirror per (kv_id, layer) for the token graph:
    /// [nkv, cap, hd] each, persists across decode tokens. `synced` counts
    /// the positions already resident (prefill sync + graph appends).
    attn_kv: Mutex<HashMap<(u64, usize), KvMirror>>,
    /// GDN recurrent state per (kv_id, layer): (conv ring, S), persists across
    /// decode tokens (created zeroed on first touch).
    gdn_state: Mutex<HashMap<(u64, usize), (wgpu::Buffer, wgpu::Buffer)>>,
    /// Shape and absolute next-position cursor for each resident GDN state.
    /// The state buffers are recurrent, so a gap or reordered batch cannot be
    /// repaired by choosing a different RoPE position.
    gdn_cursor: Mutex<HashMap<(u64, usize), GdnCursor>>,
    /// Speculative verify: per (kv_id, layer) snapshot buffer holding the
    /// GDN (ring, S) after every batch position, so a partial acceptance
    /// restores the recurrent state to the last position that was real.
    /// Value: (buffer, ring bytes, state bytes, slots).
    gdn_snap: Mutex<HashMap<(u64, usize), (wgpu::Buffer, u64, u64, usize, usize)>>,
    /// Per-layer concatenated MoE expert weights (gate_all, up_all, down_all)
    /// keyed by (file base ptr, first gate idx) — every routed expert plus the
    /// shared one as the trailing block, uploaded once, addressed by expert id
    /// inside the kernels. Counted against `resident` like any weight.
    moe_expw: Mutex<HashMap<(usize, usize), (wgpu::Buffer, wgpu::Buffer, wgpu::Buffer)>>,
    /// One bank cache for every `(layer, expert)` pair of a model. Buffers
    /// are physically segmented only because a Vulkan storage binding tops
    /// out at 4 GiB; slot numbering and eviction are common across segments.
    dsv4_global_moe: Mutex<HashMap<u64, Arc<Dsv4GlobalMoeBufs>>>,
    /// Immutable [rows,cols,…] uniforms cached by content — the ~800 matvec
    /// param buffers per token are token-invariant, so uploading them once
    /// keeps them off the per-token encode critical path.
    uniforms: Mutex<HashMap<[u32; 4], wgpu::Buffer>>,
    uniforms8: Mutex<HashMap<[u32; 8], wgpu::Buffer>>,
    /// Immutable norm/small weight buffers cached by (data ptr, len), each
    /// carrying a content fingerprint — the ~200 per-layer norm uploads per
    /// token are token-invariant, but the ADDRESS is not a stable identity:
    /// a reloaded model's mmap lands where the dropped one's was, and the
    /// bake's streaming replica re-dequantizes layers into recycled Vecs.
    /// A hit whose fingerprint disagrees rewrites the same buffer in place
    /// (never a new object — cached bind groups hold the old handle).
    /// Scoring model B after model A in one process used to inherit A's
    /// norms wherever B's mapping overlapped and report PPL in the
    /// millions for an intact file. Sentinel key (0, n) holds shared zero
    /// buffers (fingerprint 0, content never changes).
    const_bufs: Mutex<HashMap<(usize, usize), (wgpu::Buffer, u64)>>,
    /// `gemm_nt_f32`'s w-side cache, keyed on the CONTENT and not on the
    /// address. Every batched attention in this engine hands that
    /// function a per-head SCRATCH buffer as `w` — one `kh`/`vt` pair
    /// allocated per call and refilled per head — so the address is
    /// stable across heads while the matrix is not, and an
    /// address-keyed entry silently serves head 0's keys to every other
    /// head. (`CmfModel::uid` documents the same footgun on the model
    /// mapping; this is the same mistake one file over.) The
    /// fingerprint is a sequential read of a buffer we would otherwise
    /// push over PCIe, so a genuinely stable weight still uploads once.
    /// `None` buffer = a fingerprint on probation: seen once, not yet
    /// proven to recur. See `bake_weight` — adopting every matrix on
    /// sight let phase B's per-step activations grow this cache past the
    /// card.
    gemm_w_bufs: Mutex<HashMap<(usize, usize), (Option<wgpu::Buffer>, u64)>>,
    /// Per-role scratch for the fused DiT block, kept between calls. The
    /// block used to allocate its sixteen intermediates fresh every time —
    /// three of them 77 MB at 512x512 — which is ~300 MB of driver
    /// allocation per block, 1560 blocks per image. Keyed by role, so two
    /// roles never share a buffer inside one submission, and every call
    /// waits for its own readback before the next reuses them.
    dit_pool: Mutex<HashMap<&'static str, (wgpu::Buffer, u64)>>,
    /// Pooled graph scratch: eliminates per-token buffer allocations in the
    /// whole-token graph path (the dominant decode cost on Vulkan/DX12).
    graph_scratch: Mutex<GraphScratch>,
}

struct Dsv4GlobalMoeBufs {
    gate: Vec<wgpu::Buffer>,
    up: Vec<wgpu::Buffer>,
    down: Vec<wgpu::Buffer>,
    capacity: usize,
    segment_slots: usize,
    gu_len: usize,
    d_len: usize,
    gu_q2: bool,
    segments: usize,
}

struct KvMirror {
    k: wgpu::Buffer,
    v: wgpu::Buffer,
    synced: usize,
    cap: usize,
}

#[derive(Clone, Copy)]
struct GdnCursor {
    /// (nv, nk, dk, dv, kernel, packed qkv width)
    dims: (usize, usize, usize, usize, usize, usize),
    next_pos: usize,
}

/// Device KV mirrors grow with the conversation instead of reserving the
/// model's advertised maximum context at first token. Granite 4.2 advertises
/// 131k: eagerly allocating `[64 layers, K+V, 8 heads, 131k, 128]` asks for
/// about 64 GiB of KV before one token is decoded. Power-of-two growth keeps
/// short prompts small and makes the occasional resize logarithmic.
fn kv_capacity(limit: usize, need: usize) -> usize {
    let floor = std::env::var("CMF_GPU_KV_MIN")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(512)
        .max(1);
    need.max(1)
        .checked_next_power_of_two()
        .unwrap_or(limit)
        .max(floor.min(limit))
        .min(limit)
}

/// Return a mirror with at least `want_cap` positions. Growth copies the
/// already-device-resident history head by head into the new stride, so batch
/// prefill and speculative verification do not need a CPU round trip.
fn kv_mirror_ensure<'a>(
    c: &Ctx,
    mirrors: &'a mut HashMap<(u64, usize), KvMirror>,
    key: (u64, usize),
    nkv: usize,
    hd: usize,
    want_cap: usize,
) -> &'a mut KvMirror {
    let grow = mirrors.get(&key).is_none_or(|m| m.cap < want_cap);
    if grow {
        let old = mirrors.remove(&key);
        let sz = (nkv * want_cap * hd * 4) as u64;
        let mk = || {
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("kv-mirror"),
                size: sz.max(4),
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_DST
                    | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            })
        };
        let fresh = KvMirror {
            k: mk(),
            v: mk(),
            synced: old.as_ref().map_or(0, |m| m.synced.min(want_cap)),
            cap: want_cap,
        };
        if let Some(old) = old.filter(|m| m.synced > 0) {
            let copy_pos = old.synced.min(want_cap);
            let nbytes = (copy_pos * hd * 4) as u64;
            let mut enc = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("kv-grow"),
                });
            for h in 0..nkv {
                let src = (h * old.cap * hd * 4) as u64;
                let dst = (h * want_cap * hd * 4) as u64;
                enc.copy_buffer_to_buffer(&old.k, src, &fresh.k, dst, nbytes);
                enc.copy_buffer_to_buffer(&old.v, src, &fresh.v, dst, nbytes);
            }
            submit(c, finish_enc(enc));
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        }
        mirrors.insert(key, fresh);
    }
    mirrors.get_mut(&key).unwrap()
}

#[derive(Default)]
struct Scratch {
    xs: Option<(wgpu::Buffer, u64)>,
    y: Option<(wgpu::Buffer, u64)>,
    stage: Option<(wgpu::Buffer, u64)>,
    /// The chain's paired readback (folded vector + hyper-connection state).
    /// Its own slot: `stage` is sized for one of them and sharing would make
    /// every token recreate whichever was asked for second.
    stage2: Option<(wgpu::Buffer, u64)>,
    params: Option<wgpu::Buffer>,
    /// Fused-FFN intermediates (gate / up panels).
    g: Option<(wgpu::Buffer, u64)>,
    u: Option<(wgpu::Buffer, u64)>,
    /// q4tp weight plane dequantized to f16 for the tensor-core GEMM
    /// (one grow-only slot: caching planes per tensor would want tens
    /// of GB on a DiT).
    dqw: Option<(wgpu::Buffer, u64)>,
    /// Second plane slot: a fused FFN unpacks BOTH its weights into the
    /// same command buffer, so one slot cannot serve them.
    dqw2: Option<(wgpu::Buffer, u64)>,
    /// Interleaved qkv upload for the device-side split.
    dqkv: Option<(wgpu::Buffer, u64)>,
    /// v transposed per head, for the NT form of PV.
    dvt: Option<(wgpu::Buffer, u64)>,
    /// Partial maxima of the two-stage activation reduction.
    amaxp: Option<(wgpu::Buffer, u64)>,
    /// The folded activation scale (one f32) the batched-prefill coop
    /// GEMMs read; every GEMM in the chunk rewrites it in queue order.
    amax1: Option<(wgpu::Buffer, u64)>,
    /// DiT attention: Q/K/V uploads, scores, panel, output, staging.
    dq: Option<(wgpu::Buffer, u64)>,
    dk: Option<(wgpu::Buffer, u64)>,
    dv: Option<(wgpu::Buffer, u64)>,
    dsc: Option<(wgpu::Buffer, u64)>,
    dpan: Option<(wgpu::Buffer, u64)>,
    m3act: Option<(wgpu::Buffer, u64)>,
    dout: Option<(wgpu::Buffer, u64)>,
    dstage: Option<(wgpu::Buffer, u64)>,
    dpar: Option<wgpu::Buffer>,
    /// The bake GEMMs' per-call activations and readback staging. Fresh
    /// 20-260 MB allocations per call cost ~100 ms each on a discrete
    /// card — more than the tensor-core kernel they served — so these are
    /// grow-only like everything else here.
    bx: Option<(wgpu::Buffer, u64)>,
    /// Left operand of a TRANSIENT gemm: an accumulation whose operand
    /// changes every call has no business in the resident weight ledger
    /// (and its fingerprint is a full pass over 178 MB per call).
    bwt: Option<(wgpu::Buffer, u64)>,
    by: Option<(wgpu::Buffer, u64)>,
    vcx: Option<(wgpu::Buffer, u64)>,
    vcc: Option<(wgpu::Buffer, u64)>,
    vcy: Option<(wgpu::Buffer, u64)>,
    vcs: Option<(wgpu::Buffer, u64)>,
    vcb: Option<(wgpu::Buffer, u64)>,
    bst: Option<(wgpu::Buffer, u64)>,
    /// The FFN chain's device-resident middle (gate+up plane, activation).
    bb: Option<(wgpu::Buffer, u64)>,
    ba: Option<(wgpu::Buffer, u64)>,
    /// The attention chain's per-call planes (qrot/krot/vproj/gate/ao/
    /// ao_eff and the two inv vectors).
    bqr: Option<(wgpu::Buffer, u64)>,
    bkr: Option<(wgpu::Buffer, u64)>,
    bvp: Option<(wgpu::Buffer, u64)>,
    bgp: Option<(wgpu::Buffer, u64)>,
    bao: Option<(wgpu::Buffer, u64)>,
    bae: Option<(wgpu::Buffer, u64)>,
    biq: Option<(wgpu::Buffer, u64)>,
    bik: Option<(wgpu::Buffer, u64)>,
}

impl Scratch {
    /// Grow-only slot: reuse when big enough, else recreate.
    fn ensure(
        dev: &wgpu::Device,
        slot: &mut Option<(wgpu::Buffer, u64)>,
        need: u64,
        usage: wgpu::BufferUsages,
        label: &str,
    ) -> wgpu::Buffer {
        match slot {
            Some((b, cap)) if *cap >= need => b.clone(),
            _ => {
                crate::gpu::probe_note_cold();
                let cap = need.next_power_of_two().max(4096);
                let b = dev.create_buffer(&wgpu::BufferDescriptor {
                    label: Some(label),
                    size: cap,
                    usage,
                    mapped_at_creation: false,
                });
                *slot = Some((b.clone(), cap));
                b
            }
        }
    }
}

/// Pooled scratch for the whole-token graph path. Grow-only: each slot is
/// allocated once (or grown) and reused across tokens — eliminates the ~20
/// Vulkan buffer allocations per token that dominated decode latency.
#[derive(Default)]
struct GraphScratch {
    h: Option<(wgpu::Buffer, u64)>,
    n1: Option<(wgpu::Buffer, u64)>,
    qraw: Option<(wgpu::Buffer, u64)>,
    kb: Option<(wgpu::Buffer, u64)>,
    vb: Option<(wgpu::Buffer, u64)>,
    qout: Option<(wgpu::Buffer, u64)>,
    gout: Option<(wgpu::Buffer, u64)>,
    attn: Option<(wgpu::Buffer, u64)>,
    ob: Option<(wgpu::Buffer, u64)>,
    gbuf: Option<(wgpu::Buffer, u64)>,
    ubuf: Option<(wgpu::Buffer, u64)>,
    abuf: Option<(wgpu::Buffer, u64)>,
    /// Reusable transform output for one Prism activation row.  It is large
    /// enough for the widest graph input (usually the FFN intermediate).
    rot: Option<(wgpu::Buffer, u64)>,
    // GDN intermediates
    qkv_b: Option<(wgpu::Buffer, u64)>,
    // Short-conv intermediates: the fused (B,C,x) projection and the gated y
    sc_bcx: Option<(wgpu::Buffer, u64)>,
    sc_y: Option<(wgpu::Buffer, u64)>,
    cq_b: Option<(wgpu::Buffer, u64)>,
    z_b: Option<(wgpu::Buffer, u64)>,
    a_b: Option<(wgpu::Buffer, u64)>,
    b_b: Option<(wgpu::Buffer, u64)>,
    gdo_b: Option<(wgpu::Buffer, u64)>,
    // Split-K attend partials: [nh·nchunks·hd] accumulators + [nh·nchunks] (m,l)
    apacc: Option<(wgpu::Buffer, u64)>,
    apml: Option<(wgpu::Buffer, u64)>,
    // MoE routing intermediates: router logits, shared-gate logit, selected
    // expert ids + weights, per-slot activations
    m_logit: Option<(wgpu::Buffer, u64)>,
    m_slog: Option<(wgpu::Buffer, u64)>,
    m_sel: Option<(wgpu::Buffer, u64)>,
    m_wt: Option<(wgpu::Buffer, u64)>,
    m_act: Option<(wgpu::Buffer, u64)>,
    // Logits output + readback staging
    logits: Option<(wgpu::Buffer, u64)>,
    stage: Option<(wgpu::Buffer, u64)>,
    // Position-dependent uniforms (fixed size, write_buffer each token)
    kv_u: Option<wgpu::Buffer>,   // 16 bytes: [nkv, hd, cap, position]
    at_u: Option<wgpu::Buffer>,   // 32 bytes: [nh, nh/nkv, hd, cap, pos+1, 0, 0, 0]
    rope_u: Option<wgpu::Buffer>, // 32 bytes: [nh, nkv, hd, rd, pos, flags, eps, 0]
    // Multi-step slots: one uniform PER STEP with a stable identity, so the
    // attention bind groups survive across chunks (write_buffer runs at
    // submit — a single shared uniform would collapse every step to the
    // last position written).
    kv_us: Vec<wgpu::Buffer>,
    at_us: Vec<wgpu::Buffer>,
    rope_us: Vec<wgpu::Buffer>,
    ids: Option<(wgpu::Buffer, u64)>,
    ids_stage: Option<(wgpu::Buffer, u64)>,
    am_pv: Option<(wgpu::Buffer, u64)>,
    am_pi: Option<(wgpu::Buffer, u64)>,
}

impl GraphScratch {
    fn ensure(
        dev: &wgpu::Device,
        slot: &mut Option<(wgpu::Buffer, u64)>,
        need: u64,
        usage: wgpu::BufferUsages,
        label: &str,
    ) -> wgpu::Buffer {
        match slot {
            Some((b, cap)) if *cap >= need => b.clone(),
            _ => {
                let cap = need.next_power_of_two().max(256);
                let b = dev.create_buffer(&wgpu::BufferDescriptor {
                    label: Some(label),
                    size: cap,
                    usage,
                    mapped_at_creation: false,
                });
                *slot = Some((b.clone(), cap));
                b
            }
        }
    }
    /// Pooled uniform buffer of `size` bytes (created once, write_buffer'd each token).
    fn ensure_uniform(
        dev: &wgpu::Device,
        slot: &mut Option<wgpu::Buffer>,
        size: u64,
    ) -> wgpu::Buffer {
        match slot {
            Some(b) => b.clone(),
            None => {
                let b = dev.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("g-unif"),
                    size,
                    usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                });
                *slot = Some(b.clone());
                b
            }
        }
    }
}

/// One context PER GPU, built on first touch of that device. A Ctx owns
/// its weight buffers, KV mirrors and scratch, so keying contexts by
/// device is what keys those caches by device — the alternative (one
/// global context) is why a second card used to be unreachable from the
/// same process.
static CTXS: OnceLock<std::sync::Mutex<std::collections::HashMap<usize, Option<Box<Ctx>>>>> =
    OnceLock::new();

/// Whether the wgpu path is selected (the facade asks before `enabled()`):
/// `CMF_GPU=wgpu` — always; `CMF_GPU=1` (≠0) — only on non-macOS, where
/// there is no native Metal (on macOS `=1` goes to Metal). UNSET selects
/// the path by default on Linux/Windows when the feature is compiled —
/// init failure is a clean CPU fallback, so a box without a driver loses
/// nothing. `CMF_GPU=0` forces the CPU.
pub fn selected() -> bool {
    match std::env::var("CMF_GPU") {
        Ok(v) if v == "wgpu" => true,
        Ok(v) if v != "0" && v != "off" => !cfg!(target_os = "macos"),
        Ok(_) => false,
        Err(_) => {
            crate::pipeline::GLOBAL_USE_GPU.load(std::sync::atomic::Ordering::Relaxed)
                || cfg!(any(target_os = "linux", target_os = "windows"))
        }
    }
}

fn ctx() -> Option<&'static Ctx> {
    ctx_for(crate::gpu::current_device())
}

fn ctx_for(dev: usize) -> Option<&'static Ctx> {
    if !selected() {
        return None;
    }
    let map = CTXS.get_or_init(|| std::sync::Mutex::new(std::collections::HashMap::new()));
    let mut g = map.lock().unwrap();
    if let Some(slot) = g.get(&dev) {
        let ptr = slot.as_deref().map(|c| c as *const Ctx);
        drop(g);
        // The owning Box remains in CTXS until `shutdown`; callers only keep
        // this reference for the process-lifetime runtime phase.  Publishing
        // an owned context lets teardown drop Vulkan objects in order instead
        // of leaving the NVIDIA background thread behind at process exit.
        return ptr.map(|p| unsafe { &*p });
    }
    let built = match init(dev) {
        Ok(c) => Some(Box::new(c)),
        Err(e) => {
            // Tests install no subscriber, so a tracing-only report makes
            // an init failure look exactly like "no GPU here".
            tracing::warn!("wgpu init failed on device {dev} — CPU fallback: {e}");
            if std::env::var("CMF_GPU_DEBUG").is_ok()
                || std::env::var("CMF_DSV4_FRAME_DEBUG").is_ok()
            {
                eprintln!("wgpu init не удался (устройство {dev}) — откат на CPU: {e}");
            }
            None
        }
    };
    let ptr = built.as_deref().map(|c| c as *const Ctx);
    g.insert(dev, built);
    drop(g);
    ptr.map(|p| unsafe { &*p })
}

/// Gracefully release process-wide wgpu contexts before the runtime and
/// tracing dispatcher disappear.  Vulkan drivers commonly keep a background
/// worker for pipeline compilation; leaking `Ctx` avoids use-after-free during
/// inference but leaves that worker observing half-destroyed state at exit.
/// The caller must invoke this only after all model work has stopped.
pub fn shutdown() {
    let Some(map) = CTXS.get() else { return };
    let mut owned = Vec::new();
    {
        let mut all = map.lock().unwrap();
        for (_, slot) in all.drain() {
            if let Some(ctx) = slot {
                owned.push(ctx);
            }
        }
    }
    // Keep a device handle alive across the context drop.  Timestamp query
    // sets/readback buffers are driver-owned asynchronous resources; polling
    // before dropping the context is not enough because their final release
    // can enqueue work after the poll.  Cloning the devices gives the owned
    // shutdown path a live poll target for that last release without leaking
    // the full Ctx (or relying on process-exit driver cleanup).
    let devices: Vec<wgpu::Device> = owned.iter().map(|ctx| ctx.device.clone()).collect();
    for device in &devices {
        let _ = device.poll(wgpu::PollType::wait_indefinitely());
    }
    drop(owned);
    for device in &devices {
        let _ = device.poll(wgpu::PollType::wait_indefinitely());
    }
    drop(devices);
}

/// Look up an already-created context without initializing a device.  Stage
/// cleanup must be inert for CPU-only callers and for an explicitly selected
/// non-wgpu backend.
fn existing_ctx() -> Option<&'static Ctx> {
    let map = CTXS.get()?;
    let ptr = map
        .lock()
        .unwrap()
        .get(&crate::gpu::current_device())
        .and_then(|slot| slot.as_deref().map(|c| c as *const Ctx));
    ptr.map(|p| unsafe { &*p })
}

/// Weight budget of the current device, in bytes (0 when there is no
/// device). What `serve --gpus` asks before deciding between replicas
/// and a layer split.
pub fn device_vram_budget() -> u64 {
    ctx().map(|c| c.vram_budget).unwrap_or(0)
}

/// Logical bytes currently held by the weight-residency arena on the active
/// adapter.  A sequence's KV/O(1) buffers are intentionally excluded; callers
/// report those through the pipeline and `o1_device_stats`.
pub fn resident_bytes() -> u64 {
    ctx()
        .map(|c| c.resident.load(std::sync::atomic::Ordering::Relaxed))
        .unwrap_or(0)
}

/// Physical live-set ceiling used while an over-size graph is assembled.
/// On the measured Vulkan path Q8_2F needs about 1.6× its packed payload in
/// driver allocations (19 Granite-30B layers: 15,384 MiB; 30 layers:
/// 23,832 MiB). The adjusted ceiling still keeps a full graph whenever its
/// packed payload fits below it (3B/8B on the A40, for example); only a Q8_2F
/// stack whose physical expansion would cross the card becomes a prefix.
fn graph_live_weight_budget(c: &Ctx, model: &Arc<CmfModel>) -> u64 {
    let (mut total, mut q8) = (0u64, 0u64);
    for e in &model.tensors {
        total = total.saturating_add(e.nbytes);
        if e.dtype == cortiq_core::TensorDtype::Q8_2f {
            q8 = q8.saturating_add(e.nbytes);
        }
    }
    if q8 > total / 2 {
        // 5/8 matched the packed-to-physical slope but left less than one
        // GiB for the fixed context/pipeline allocations on an A40. 9/16
        // keeps ~3 GiB of measured headroom while retaining the largest safe
        // prefix (and is still substantially better than a blanket 50%).
        c.vram_budget.saturating_mul(9) / 16
    } else {
        c.vram_budget
    }
}

/// Automatic prefix for ordinary per-op layer walks (most importantly the
/// batched prefill). A graph that does not fit already chooses a device
/// prefix, but the fallback prefill used to stream every later CPU-tail layer
/// through the GPU arena. Vulkan allocators retain freed blocks, so a 14 GiB
/// logical arena reached a 25.4 GiB physical peak on Granite 30B Q8_2F.
///
/// Non-layer tensors are reserved first because lm_head/embed operations live
/// outside `set_layer`. `None` means the complete stack fits; `Some(0)` is a
/// valid all-CPU layer stack with the head still independently GPU-eligible.
pub fn automatic_layer_prefix(
    model: &Arc<CmfModel>,
    num_layers: usize,
    physical_layers: usize,
) -> Option<usize> {
    if std::env::var_os("CMF_GPU_LAYERS").is_some() {
        return None; // an explicit split always wins
    }
    let c = ctx()?;
    if c.vram_budget == u64::MAX || num_layers == 0 || physical_layers == 0 {
        return None;
    }
    let live_budget = graph_live_weight_budget(c, model);
    let mut layer_bytes = vec![0u64; physical_layers];
    let mut outside = 0u64;
    for e in &model.tensors {
        let li = layer_of_name(&e.name);
        if li != u16::MAX && (li as usize) < physical_layers {
            layer_bytes[li as usize] = layer_bytes[li as usize].saturating_add(e.nbytes);
        } else {
            outside = outside.saturating_add(e.nbytes);
        }
    }
    let usable = live_budget.saturating_sub(outside.min(live_budget));
    let mut used = 0u64;
    let mut prefix = 0usize;
    for li in 0..num_layers {
        let need = layer_bytes[li % physical_layers];
        if used.saturating_add(need) > usable {
            break;
        }
        used += need;
        prefix += 1;
    }
    (prefix < num_layers).then_some(prefix)
}

/// How many adapters this build can see (0 when the backend is off).
pub fn adapter_count() -> usize {
    if !selected() {
        return 0;
    }
    static N: OnceLock<usize> = OnceLock::new();
    *N.get_or_init(|| {
        let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
            backends: backends_from_env(),
            flags: wgpu::InstanceFlags::default(),
            memory_budget_thresholds: Default::default(),
            backend_options: Default::default(),
            display: None,
        });
        pollster::block_on(instance.enumerate_adapters(backends_from_env()))
            .iter()
            .filter(|a| a.get_info().device_type != wgpu::DeviceType::Cpu)
            .count()
    })
}

/// Backend mask from the standard WGPU_BACKEND env (shared by init and
/// the adapter census, so both count the same list).
fn backends_from_env() -> wgpu::Backends {
    std::env::var("WGPU_BACKEND")
        .ok()
        .map(|v| match v.to_lowercase().as_str() {
            "vulkan" | "vk" => wgpu::Backends::VULKAN,
            "dx12" | "d3d12" => wgpu::Backends::DX12,
            "metal" | "mtl" => wgpu::Backends::METAL,
            "gl" | "gles" => wgpu::Backends::GL,
            _ => wgpu::Backends::all(),
        })
        .unwrap_or(wgpu::Backends::all())
}

/// Total device-local memory of a Vulkan adapter, from the driver's own heap
/// report. None on other backends or when the query is unavailable. Only
/// platforms where wgpu carries the Vulkan backend at all; elsewhere the
/// stub answers None and the conservative default budget stands.
#[cfg(any(target_os = "linux", target_os = "windows", target_os = "android"))]
fn vulkan_vram_total(adapter: &wgpu::Adapter) -> Option<u64> {
    if adapter.get_info().backend != wgpu::Backend::Vulkan {
        return None;
    }
    unsafe {
        let hal = adapter.as_hal::<wgpu::hal::api::Vulkan>()?;
        let phd = hal.raw_physical_device();
        let mem = hal
            .shared_instance()
            .raw_instance()
            .get_physical_device_memory_properties(phd);
        let mut total = 0u64;
        for heap in mem.memory_heaps[..mem.memory_heap_count as usize].iter() {
            // DEVICE_LOCAL is bit 0 of VkMemoryHeapFlags. The LARGEST such
            // heap is the card's memory; summing would double-count the
            // small host-visible BAR window some drivers report separately.
            if heap.flags.as_raw() & 0x1 != 0 {
                total = total.max(heap.size);
            }
        }
        (total > 0).then_some(total)
    }
}

#[cfg(not(any(target_os = "linux", target_os = "windows", target_os = "android")))]
fn vulkan_vram_total(_adapter: &wgpu::Adapter) -> Option<u64> {
    None
}

fn init(dev: usize) -> Result<Ctx, String> {
    // Backend selection is automatic (wgpu picks the platform's best:
    // DX12 on Windows, Vulkan on Linux, Metal on macOS), but the
    // standard WGPU_BACKEND env (vulkan|dx12|metal|gl) forces one.
    let backends = backends_from_env();
    let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
        backends,
        flags: wgpu::InstanceFlags::default(),
        memory_budget_thresholds: Default::default(),
        backend_options: Default::default(),
        display: None,
    });
    // CMF_GPU_ADAPTER pins the card: an index into the `cortiq gpu`
    // listing or a case-insensitive name substring. Without it every
    // process on the host gets the same "best" adapter — the pin is
    // what lets a local worker take the SECOND GPU (`run --gpus N`).
    // Device index first (the multi-GPU registry asks for a specific
    // card), then the CMF_GPU_ADAPTER pin (index or name substring),
    // then wgpu's own "best". A named pin only makes sense for the
    // default device — a request for card 2 by index means card 2.
    let want_env = std::env::var("CMF_GPU_ADAPTER").ok();
    let by_index = (dev != 0 || want_env.is_none()).then_some(dev);
    let adapter = if by_index.is_some() || want_env.is_some() {
        let mut all = pollster::block_on(instance.enumerate_adapters(backends));
        let pick = by_index.filter(|&i| i < all.len()).or_else(|| {
            let want = want_env.as_deref().unwrap_or_default();
            want.parse::<usize>()
                .ok()
                .filter(|&i| i < all.len())
                .or_else(|| {
                    let w = want.to_lowercase();
                    all.iter()
                        .position(|a| a.get_info().name.to_lowercase().contains(&w))
                })
        });
        match pick {
            Some(i) => all.swap_remove(i),
            None => {
                return Err(format!(
                    "device {dev} / CMF_GPU_ADAPTER={:?} matches none of the {} adapters \
                     (see `cortiq gpu`)",
                    want_env,
                    all.len()
                ));
            }
        }
    } else {
        pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::HighPerformance,
            force_fallback_adapter: false,
            compatible_surface: None,
            apply_limit_buckets: false,
        }))
        .map_err(|e| format!("no adapter: {e}"))?
    };

    // A SOFTWARE adapter is not a GPU and taking it is a loss, not a
    // fallback. Mesa ships lavapipe/llvmpipe in most container images —
    // Hugging Face Spaces, CI runners, cloud VMs — and with no real card
    // present wgpu hands it over as the best available: `request_adapter`
    // honours `force_fallback_adapter: false` by not PREFERRING a
    // fallback, not by refusing one. The engine would then report "GPU
    // path: on" and run every shader through an LLVM rasteriser on the
    // same cores its native kernels were already using, which is the same
    // silicon plus an emulation layer.
    //
    // So decline it and let the caller keep the CPU path, exactly as any
    // other init failure does. `CMF_GPU_SOFTWARE=1` takes it anyway —
    // validating a shader against a reference rasteriser is the one job
    // this adapter is genuinely good at.
    let adapter_info = adapter.get_info();
    if adapter_info.device_type == wgpu::DeviceType::Cpu
        && std::env::var("CMF_GPU_SOFTWARE").ok().as_deref() != Some("1")
    {
        return Err(format!(
            "only a software rasteriser is available ({} / {:?}); staying on the CPU path \
             — set CMF_GPU_SOFTWARE=1 to use it anyway",
            adapter_info.name, adapter_info.backend
        ));
    }

    // Take the card's maximum limits — large tensors (lm_head ≈ 254 MB
    // int8) require a raised storage buffer; a discrete card handles GB.
    let limits = adapter.limits();
    // 33 152 B = the stride-257 attend kernels' workgroup footprint.
    // Adreno/Mali/wgpu-Metal report 32 768 — there only the stride-129
    // (hd <= 128) kernels are created.
    let big_attend = limits.max_compute_workgroup_storage_size >= 33_152;
    let wg_storage = limits.max_compute_workgroup_storage_size;
    // GPU timestamps (CMF_GPU_TS=1): ask for the query features when the
    // adapter has them — the frame profiler below is the only consumer.
    // Two tiers, and conflating them cost the dsv4 profiler its clock on
    // Metal: a timestamp PAIR on a pass descriptor needs TIMESTAMP_QUERY and
    // nothing else, while write_timestamp inside an encoder or a pass needs
    // the other two. Demanding all three meant a device that offers the
    // first got no query set at all — which reads exactly like "the profiler
    // printed nothing", not like "this device cannot do the fine one".
    let ts_basic = wgpu::Features::TIMESTAMP_QUERY;
    let ts_fine_features = wgpu::Features::TIMESTAMP_QUERY_INSIDE_ENCODERS
        | wgpu::Features::TIMESTAMP_QUERY_INSIDE_PASSES;
    let have_basic = adapter.features().contains(ts_basic);
    let have_fine = adapter.features().contains(ts_basic | ts_fine_features);
    let ts_features = if have_fine {
        ts_basic | ts_fine_features
    } else {
        ts_basic
    };
    let want_ts = have_basic;
    let want_sg = adapter.features().contains(wgpu::Features::SUBGROUP);
    // History, kept because it explains the shape check below: when the
    // only configuration wgpu exposed was 8x8 f32, this path was opt-in
    // and off — asked for 8x8 f32 the NVIDIA driver fell back to
    // something so slow a render went from 0.68 to 52 seconds a step.
    // That is no longer the default. It is ON wherever the card reports
    // 16x16x16 with f16 operands and an f32 accumulator, and `CMF_COOP=0`
    // is what turns it off. Two things follow that a reader should know:
    // an f32 GEMM on this path accumulates at tf32-class precision
    // (measured on an A100 against an f64 reference: 2.8e-4 relative
    // against 7.7e-7 on the host, both arms of P·V alike, and `CMF_COOP=0`
    // restores the floor), and the shape check is not optional — wgpu
    // raises the feature flag on weaker configurations too, and a shader
    // compiled against a shape the hardware does not have is a silently
    // wrong image, not an error.
    // Tensor cores, and only where the card reports the shape the kernel
    // is written against — f16 in, f32 accumulator, 16x16x16. The feature
    // flag alone is not enough: wgpu raises it on a weaker configuration
    // too, and a shader compiled against a shape the hardware does not
    // have is a silently wrong image, not an error.
    let coop_shape = adapter.cooperative_matrix_properties().iter().any(|c| {
        c.m_size == 16
            && c.n_size == 16
            && c.k_size == 16
            && c.ab_type == wgpu::CooperativeScalarType::F16
            && c.cr_type == wgpu::CooperativeScalarType::F32
    });
    let want_coop = adapter
        .features()
        .contains(wgpu::Features::EXPERIMENTAL_COOPERATIVE_MATRIX)
        && adapter.features().contains(wgpu::Features::SHADER_F16)
        && coop_shape
        && std::env::var("CMF_COOP").map(|v| v != "0").unwrap_or(true);
    // Half precision, where the card has it: the matrix units take f16
    // operands, and without this the shader cannot even name the type.
    let want_f16 = adapter.features().contains(wgpu::Features::SHADER_F16);
    // Keeping compiled pipelines between runs; see `pipeline_cache_path`.
    let want_pcache = adapter.features().contains(wgpu::Features::PIPELINE_CACHE);
    // The unified DSV4 slot pool needs a dynamically indexed array of
    // storage buffers: one logical bank, segmented at Vulkan's 4-GiB binding
    // limit. Request it only where the complete feature set exists; all
    // other adapters retain the exact per-layer cache.
    let bind_array_features = wgpu::Features::BUFFER_BINDING_ARRAY
        | wgpu::Features::STORAGE_RESOURCE_BINDING_ARRAY
        | wgpu::Features::SAMPLED_TEXTURE_AND_STORAGE_BUFFER_ARRAY_NON_UNIFORM_INDEXING;
    let want_bind_arrays = adapter.features().contains(bind_array_features);
    COOP_OK.store(want_coop, std::sync::atomic::Ordering::Relaxed);
    let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
        label: Some("cortiq-wgpu"),
        required_limits: limits.clone(),
        required_features: if want_ts {
            ts_features
        } else {
            wgpu::Features::empty()
        } | if want_sg {
            wgpu::Features::SUBGROUP
        } else {
            wgpu::Features::empty()
        } | if want_coop {
            wgpu::Features::EXPERIMENTAL_COOPERATIVE_MATRIX
        } else {
            wgpu::Features::empty()
        } | if want_f16 {
            wgpu::Features::SHADER_F16
        } else {
            wgpu::Features::empty()
        } | if want_pcache {
            wgpu::Features::PIPELINE_CACHE
        } else {
            wgpu::Features::empty()
        } | if want_bind_arrays {
            bind_array_features
        } else {
            wgpu::Features::empty()
        },
        // Anything wgpu prefixes with EXPERIMENTAL needs this token, and
        // without it `request_device` fails outright rather than dropping
        // the feature — which is how asking for cooperative matrices took
        // the whole device down and every op with it.
        experimental_features: if want_coop {
            unsafe { wgpu::ExperimentalFeatures::enabled() }
        } else {
            wgpu::ExperimentalFeatures::disabled()
        },
        ..Default::default()
    }))
    .map_err(|e| format!("request_device: {e}"))?;
    // Every shader-module and pipeline validation error below must fail
    // init: an invalid pipeline silently turns its dispatches into
    // no-ops and the graph decodes garbage (seen on phones before this
    // scope existed). Err here = clean CPU fallback.
    let vscope = device.push_error_scope(wgpu::ErrorFilter::Validation);

    let info = adapter.get_info();
    let q1_rows = match std::env::var("CMF_Q1_RPG").as_deref() {
        Ok("16") => 16,
        Ok("8") => 8,
        _ if info.name.to_ascii_lowercase().contains("adreno") => 16,
        _ => 8,
    };
    let discrete = info.device_type == wgpu::DeviceType::DiscreteGpu;
    let vram_budget = std::env::var("CMF_GPU_VRAM_MB")
        .ok()
        .and_then(|v| v.parse::<u64>().ok())
        .map(|mb| mb * 1024 * 1024)
        .unwrap_or(if discrete {
            // No knob: ask the driver. The Vulkan device-local heap is what
            // the card actually has; leave room for KV, graph scratch,
            // pipelines and the upload allocator. Backends that cannot answer
            // (DX12/GL) keep the conservative default — CMF_GPU_VRAM_MB
            // overrides either way.
            match vulkan_vram_total(&adapter) {
                Some(total) => {
                    let gib = 1024 * 1024 * 1024u64;
                    // The former 1/32 margin let a 29 GB q8_2f dense graph fill
                    // an A40 to within 1.4 GiB, then OOM while constructing KV
                    // and graph scratch. Measured 30B Q4TP needs more than a
                    // 2 GiB floor on a 16 GiB card, so scale from 2.5 GiB at
                    // 16 GiB through 3 GiB at 24 GiB to the existing 4 GiB
                    // ceiling on large cards.
                    let reserve = (total / 8).clamp(5 * gib / 2, 4 * gib).min(total / 2);
                    total - reserve
                }
                None => 8 * 1024 * 1024 * 1024,
            }
        } else {
            u64::MAX // UMA: the OS pages shared memory
        });
    // The probe cache is keyed by this: a verdict belongs to one piece
    // of silicon on one backend, and must not be adopted by another.
    crate::gpu::probe_set_device(&format!("{}/{:?}", info.name, info.backend));
    tracing::info!(
        "wgpu GPU path: on ({} / {:?}, {}, weight budget {}, q1 {} rows/{} threads)",
        info.name,
        info.backend,
        if discrete { "discrete" } else { "uma" },
        if vram_budget == u64::MAX {
            "unlimited".to_string()
        } else {
            format!("{} MB", vram_budget / 1024 / 1024)
        },
        q1_rows,
        q1_rows * 16,
    );

    let msc = device.push_error_scope(wgpu::ErrorFilter::Validation);
    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("q8"),
        source: wgpu::ShaderSource::Wgsl(wgsl_main_source().into()),
    });
    if let Some(e) = pollster::block_on(msc.pop()) {
        // The module's own diagnostic names the offending line; without
        // it every pipeline built from the module reports the same
        // opaque "Validation Error".
        tracing::warn!("wgpu shader module rejected: {e}");
    }
    // Compiled pipelines from an earlier run, if this driver keeps them.
    let pcache = pipeline_cache_load(&device, &info, want_pcache);
    // Auto layout: the bind group layout is inferred from the shader.
    let pipe = |ep: &str| {
        device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some(ep),
            layout: None, // auto: layout is inferred from the shader
            module: &module,
            entry_point: Some(ep),
            compilation_options: Default::default(),
            cache: pcache.as_ref(),
        })
    };
    let matvec = pipe("q8_matvec");
    let matvec_tiled = pipe("q8_matvec_tiled");
    let q8_2f_mv = pipe("q8_2f_matvec");
    let matmat = pipe("q8_matmat");
    let mul_mm = pipe("q8_mul_mm");
    let q1_mm = pipe("q1_mul_mm");
    let silu = pipe("silu_mul_pre");
    let axpy = pipe("axpy");
    let colscale = pipe("colscale");
    let gate_mul = pipe("gate_mul");
    let zero = pipe("fill_zero");
    let q1_constants = [
        ("Q1_WG", (q1_rows * 16) as f64),
        ("Q1_ROWS", q1_rows as f64),
    ];
    let q1 = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
        label: Some("q1_matvec"),
        layout: None,
        module: &module,
        entry_point: Some("q1_matvec"),
        compilation_options: wgpu::PipelineCompilationOptions {
            constants: &q1_constants,
            ..Default::default()
        },
        cache: pcache.as_ref(),
    });
    let q1t = pipe("q1t_matvec");
    let q4b = pipe("q4b_matvec");
    let q4t_mv = pipe("q4t_matvec");
    let rope_heads = pipe("rope_heads");
    let bf16_round = pipe("bf16_round_buffer");
    let o_lora_a = pipe("o_lora_a");
    let kv_pool = pipe("kv_pool");
    let index_scores = pipe("index_scores");
    let top_k_index = pipe("top_k_index");
    let hc_block = pipe("hc_block");
    let f32_matvec_w = pipe("f32_matvec_w");
    let o_lora_a_w = pipe("o_lora_a_w");
    let f32_matvec_x = pipe("f32_matvec_x");
    let f32_mv_split = pipe("f32_matvec_split");
    let f32_mv_merge = pipe("f32_matvec_merge");
    let o_lora_a_m = pipe("o_lora_a_m");
    let moe_gu_q2tp_m = pipe("moe_gate_up_q2tp_m");
    let moe_dn_q4tp_m = pipe("moe_down_q4tp_m");
    let sa_part = pipe("sparse_attend_part");
    let sa_merge = pipe("sparse_attend_merge");
    let blit = pipe("blit");
    let idx_build = pipe("idx_build");
    let moe_route = pipe("moe_route");
    let sparse_attend = pipe("sparse_attend");
    let sa_scores = pipe("sa_scores");
    let sa_apply = pipe("sa_apply");
    let hc_pre_fold = pipe("hc_pre_fold");
    let hc_post_expand = pipe("hc_post_expand");
    let bt_rope_heads = pipe("bt_rope_heads");
    let bt_hc_pre_fold = pipe("bt_hc_pre_fold");
    let bt_hc_block = pipe("bt_hc_block");
    let bt_comp_append = pipe("bt_comp_append");
    let bt_comp_fold = pipe("bt_comp_fold");
    let bt_hc_post_expand = pipe("bt_hc_post_expand");
    let bt_f32_matvec_w = pipe("bt_f32_matvec_w");
    let bt_f32_matvec_x = pipe("bt_f32_matvec_x");
    let bt_moe_route = pipe("bt_moe_route");
    let bt_moe_gate_up_q2tp = pipe("bt_moe_gate_up_q2tp");
    let bt_moe_gate_up_q2tp_r4 = pipe("bt_moe_gate_up_q2tp_r4");
    let bt_sparse_attend = pipe("bt_sparse_attend");
    let bt_index_scores = pipe("bt_index_scores");
    let bt_top_k = pipe("bt_top_k");
    let bt_idx_build_staged = pipe("bt_idx_build_staged");
    let bt_o_lora_a = pipe("bt_o_lora_a");
    let bt_o_lora_a4 = pipe("bt_o_lora_a4");
    // The staged twin's 4608-float span needs 18 KB and change of
    // workgroup storage; a 16 KB device simply never creates it.
    let bt_o_lora_a2 = (wg_storage >= 19_500).then(|| pipe("bt_o_lora_a2"));
    let q4tp_mv = pipe("q4tp_matvec");
    let q4tp_mv4 = pipe("q4tp_matvec4");
    let q4tp_mv4_u2 = pipe("q4tp_matvec4_u2");
    let q4tp_mv4_nored = pipe("q4tp_matvec4_nored");
    let q4tp_mv4_dual = pipe("q4tp_matvec4_dual");
    let q4tp_mv4_dsilu = pipe("q4tp_matvec4_dsilu");
    let use_mv_dual = std::env::var("CMF_MV_DUAL").as_deref() == Ok("1");
    let use_mv_nored = std::env::var("CMF_MV_NORED").as_deref() == Ok("1");
    // CMF_MV_U2=1: the 2-way unrolled decode matvec — an experiment in
    // in-flight loads, adopted only if the bench says so.
    let use_mv_u2 = std::env::var("CMF_MV_U2").as_deref() == Ok("1");
    let q4tp_mv4_bk = pipe("q4tp_matvec4_bk");
    let q4tp_mv4_bku = pipe("q4tp_matvec4_bku");
    // TIMING ONLY, and only when asked: the arithmetic-free twin of the
    // quad-row kernel. Built lazily so a normal run never compiles it.
    let mv_probe: usize = std::env::var("CMF_MV_PROBE")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(0);
    let q4tp_mv16w_probe = (mv_probe > 0).then(|| pipe("q4tp_matvec16w_probe"));
    // 2 = batch inside the workgroup with the unpack shared. Measured on
    // Qwen3.6-27B / RTX 5090, k=2 speculative verify: the batched FFN is
    // 15.05 / 13.86 / 11.15 ms at arm 0 / 1 / 2 and the verify round
    // 53.3 / 52.1 / 45.9 ms, which is what finally makes a speculative
    // token cheaper than a plain one (51.0 tok/s against 49.3, medians
    // of three; arm 0 was 43.6 — a 11% LOSS).
    let use_mv_bk: usize = std::env::var("CMF_MV_BK")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(2);
    let use_mv4 = std::env::var("CMF_MV4").map(|v| v != "0").unwrap_or(true);
    let q4tp_mv16 = pipe("q4tp_matvec16");
    let q4t_mv8 = pipe("q4t_matvec8");
    let q4b_mv8 = pipe("q4b_matvec8");
    let q4tp_mm = pipe("q4tp_mul_mm");
    let q2tp_mm = pipe("q2tp_mul_mm");
    // The tensor-core GEMM lives in its own module: its `enable` directive
    // does not parse without the feature, so a device that lacks it must
    // never be handed the source.
    let q4tp_mm_coop = want_coop
        .then(|| {
            let sc = device.push_error_scope(wgpu::ErrorFilter::Validation);
            let m = device.create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("q4tp-coop"),
                source: wgpu::ShaderSource::Wgsl(COOP_MM_SRC.into()),
            });
            if let Some(e) = pollster::block_on(sc.pop()) {
                tracing::warn!("cooperative-matrix module rejected: {e}");
                COOP_OK.store(false, std::sync::atomic::Ordering::Relaxed);
                return None;
            }
            let sc = device.push_error_scope(wgpu::ErrorFilter::Validation);
            let p = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("q4tp_mm_coop"),
                layout: None,
                module: &m,
                entry_point: Some("q4tp_mm_coop"),
                compilation_options: Default::default(),
                cache: pcache.as_ref(),
            });
            if let Some(e) = pollster::block_on(sc.pop()) {
                tracing::warn!("cooperative-matrix pipeline rejected: {e}");
                COOP_OK.store(false, std::sync::atomic::Ordering::Relaxed);
                return None;
            }
            Some(p)
        })
        .flatten();
    // Dequantize-once twin: the plane kernel and the f16 GEMM. Guarded
    // like every other coop pipeline — a validation error costs the
    // pipeline, never the device, and the in-kernel path stays.
    let mk_coop = |src: &str, mod_label: &'static str, entry: &'static str| {
        let sc = device.push_error_scope(wgpu::ErrorFilter::Validation);
        let m = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some(mod_label),
            source: wgpu::ShaderSource::Wgsl(src.into()),
        });
        if let Some(e) = pollster::block_on(sc.pop()) {
            tracing::warn!("{mod_label} module rejected: {e}");
            // A rejected optional module can leave a validation event queued
            // on NVIDIA Vulkan. Drain it before dropping the failed module;
            // otherwise a short component-test process may race the driver's
            // update worker during teardown (the ordinary q2tp path remains
            // entirely independent of this experimental shader).
            let _ = device.poll(wgpu::PollType::wait_indefinitely());
            return None;
        }
        let sc = device.push_error_scope(wgpu::ErrorFilter::Validation);
        let p = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some(entry),
            layout: None,
            module: &m,
            entry_point: Some(entry),
            compilation_options: Default::default(),
            cache: pcache.as_ref(),
        });
        if let Some(e) = pollster::block_on(sc.pop()) {
            tracing::warn!("{entry} pipeline rejected: {e}");
            let _ = device.poll(wgpu::PollType::wait_indefinitely());
            return None;
        }
        Some(p)
    };
    // Q2TP affine cooperative GEMM is an explicit A/B arm.  Keep it out of
    // normal initialization until requested: its F16 operand boundary is a
    // separate precision profile, and failure must leave the scalar Q2TP
    // module entirely usable.
    let want_q2_coop = want_coop && std::env::var("CMF_Q2_COOP").as_deref() == Ok("1");
    let q2tp_mm_coop = want_q2_coop
        .then(|| mk_coop(COOP_Q2_MM_SRC, "q2tp-mm-coop", "q2tp_mm_coop"))
        .flatten();
    let ffn_silu_packed = pipe("ffn_silu_mul_packed");
    let act_absmax = mk_coop(COOP_AMAX_SRC, "act-absmax", "act_absmax");
    let dit_qkv_split = mk_coop(DIT_SPLIT_SRC, "dit-split", "dit_qkv_split");
    let vae_im2col = mk_coop(VAE_IM2COL_SRC, "vae-im2col", "vae_im2col");
    let conv1d_im2col = mk_coop(CONV1D_IM2COL_SRC, "conv1d-im2col", "conv1d_im2col");
    let music3_glu = mk_coop(MUSIC3_GLU_SRC, "music3-glu", "music3_glu");
    let dit_v_transpose = mk_coop(DIT_VT_SRC, "dit-vt", "dit_v_transpose");
    let dit_qknorm = mk_coop(DIT_QKNORM_SRC, "dit-qknorm", "dit_qknorm_rope");
    let dit_gemm_coop = (q4tp_mm_coop.is_some())
        .then(|| mk_coop(DIT_COOP_SRC, "dit-coop", "dit_gemm_coop"))
        .flatten();
    let act_amax_part = mk_coop(COOP_AMAX_SRC, "act-absmax-p", "act_absmax_part");
    let act_amax_fold = mk_coop(COOP_AMAX_SRC, "act-absmax-f", "act_absmax_fold");
    let q4tp_dq_f16 = (want_f16 && q4tp_mm_coop.is_some())
        .then(|| mk_coop(COOP_DQ_SRC, "q4tp-dq", "q4tp_dq_f16"))
        .flatten();
    let q8_dq_f16 = q4tp_dq_f16
        .is_some()
        .then(|| mk_coop(COOP_DQ8_SRC, "q8-dq", "q8_dq_f16"))
        .flatten();
    let q4tp_mm_coop_f16 = (q4tp_dq_f16.is_some())
        .then(|| mk_coop(COOP_MM_F16_SRC, "q4tp-mm-f16", "q4tp_mm_coop_f16"))
        .flatten();
    // The device-scale entry point of the in-kernel coop module (the
    // batched prefill's operands never reach the host).
    let q4tp_mm_coop_s = (q4tp_mm_coop.is_some())
        .then(|| mk_coop(COOP_MM_SRC, "q4tp-coop-s", "q4tp_mm_coop_s"))
        .flatten();

    // The bake's f32 forward/backward GEMMs on the same units, guarded the
    // same way: a validation error must cost the pipeline, never the device.
    let bake_coop = |src: &str, mod_label: &'static str, entry: &'static str| {
        let sc = device.push_error_scope(wgpu::ErrorFilter::Validation);
        let m = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some(mod_label),
            source: wgpu::ShaderSource::Wgsl(src.into()),
        });
        if let Some(e) = pollster::block_on(sc.pop()) {
            tracing::warn!("{mod_label} module rejected: {e}");
            return None;
        }
        let sc = device.push_error_scope(wgpu::ErrorFilter::Validation);
        let p = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some(entry),
            layout: None,
            module: &m,
            entry_point: Some(entry),
            compilation_options: Default::default(),
            cache: pcache.as_ref(),
        });
        if let Some(e) = pollster::block_on(sc.pop()) {
            tracing::warn!("{entry} pipeline rejected: {e}");
            return None;
        }
        Some(p)
    };
    let bake_ok = want_coop && q4tp_mm_coop.is_some();
    let gemm_nt_coop = bake_ok
        .then(|| bake_coop(COOP_NT_SRC, "bake-nt-coop", "gemm_nt_coop"))
        .flatten();
    let gemm_nn_coop = bake_ok
        .then(|| bake_coop(COOP_NN_SRC, "bake-nn-coop", "gemm_nn_coop"))
        .flatten();
    let plain = |src: &str, mod_label: &'static str, entry: &'static str| {
        let m = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some(mod_label),
            source: wgpu::ShaderSource::Wgsl(src.into()),
        });
        device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some(entry),
            layout: None,
            module: &m,
            entry_point: Some(entry),
            compilation_options: Default::default(),
            cache: pcache.as_ref(),
        })
    };
    let bake_silu = plain(BAKE_SILU_SRC, "bake-silu", "bake_silu_mul");
    let bake_silu_bwd = plain(BAKE_SILU_BWD_SRC, "bake-silu-bwd", "bake_silu_bwd");
    let bake_qkr = plain(BAKE_QKR_SRC, "bake-qkr", "bake_qkr");
    let bake_attn_head = plain(BAKE_ATTN_SRC, "bake-attn", "bake_attn_head");
    let bake_attn_gate = plain(BAKE_AGATE_SRC, "bake-agate", "bake_attn_gate");
    let argmax_part = pipe("argmax_part");
    let gdn_step_par = pipe("gdn_step_par");
    let gdn_step_par2 = pipe("gdn_step_par2");
    let gdn_step_norm2 = pipe("gdn_step_norm2");
    // Measured -1 tok/s on RTX PRO 6000: every dv-workgroup of a head
    // recomputes the conv reads, 128-fold traffic amplification against
    // one saved hop. Kept for narrow-dv models; CMF_GDN_INLINE=1 enables.
    let gdn_inline = std::env::var("CMF_GDN_INLINE").as_deref() == Ok("1");
    let gdn_step_norm = pipe("gdn_step_norm");
    let gdn_par = std::env::var("CMF_GDN_PAR")
        .map(|v| v != "0")
        .unwrap_or(true);
    // Frame profiler (CMF_GPU_TS=1): 256 timestamp slots + resolve/stage
    // buffers. A bounded batch-kernel profile reserves the query range after
    // the coarse slots in the same set. Keep that range large enough for all
    // Q2/GDN/attention projections in one full batch; the profiler reports
    // any attempted pair that still overflows rather than replaying a prior
    // frame's first-window totals.
    let graph_ts_all = std::env::var("CMF_GRAPH_TS_ALL").as_deref() == Ok("1");
    let ts_query = if want_ts && matches!(std::env::var("CMF_GPU_TS").as_deref(), Ok("1") | Ok("2"))
    {
        let count: u32 = if graph_ts_all
            || std::env::var("CMF_BATCH_KERNEL_TS").as_deref() == Ok("1")
        {
            4096
        } else {
            256
        };
        let qs = device.create_query_set(&wgpu::QuerySetDescriptor {
            label: Some("g-ts"),
            ty: wgpu::QueryType::Timestamp,
            count,
        });
        let resolve = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("g-ts-resolve"),
            size: count as u64 * 8,
            usage: wgpu::BufferUsages::QUERY_RESOLVE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let stage = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("g-ts-stage"),
            size: count as u64 * 8,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        Some((qs, resolve, stage))
    } else {
        None
    };
    let ts_period = queue.get_timestamp_period();
    if std::env::var("CMF_TS_DEBUG").is_ok() {
        eprintln!(
            "[ts] базовый={have_basic} точный={have_fine} набор={} период={ts_period}",
            ts_query.is_some()
        );
    }
    let argmax_final = pipe("argmax_final");
    let embed_gather_q4tp = pipe("embed_gather_q4tp");
    let silu_down = pipe("silu_down_matvec");
    let q1t_mm = pipe("q1t_mul_mm");
    let q4t_mm = pipe("q4t_mul_mm");
    let dit_qk = pipe("dit_qk");
    let dit_pv = pipe("dit_pv");
    let dit_softmax = pipe("dit_softmax");
    let dit_unstack = pipe("dit_unstack");
    let ffn_silu = pipe("ffn_silu_mul");
    let qwen_gelu_bias = pipe("qwen_gelu_bias");
    let qwen_layernorm_mod = pipe("qwen_layernorm_mod");
    let qwen_gated_residual = pipe("qwen_gated_residual");
    let qwen_rope_pack = pipe("qwen_rope_pack");
    let q1t_ovmm = pipe("q1t_overlay_mm");
    let rmsnorm = pipe("rmsnorm");
    let add_rmsnorm = pipe("add_rmsnorm");
    let rmsnorm_b = pipe("rmsnorm_b");
    let add_rmsnorm_b = pipe("add_rmsnorm_b");
    let attn_rope = pipe("attn_rope_qkn");
    let kv_append = pipe("kv_append");
    let gqa_attend_s = pipe("gqa_attend_s");
    // 32 lanes where the workgroup budget allows it, 16 lanes where it does
    // not: both cover head_dim 256, so hd_cap is 256 everywhere now.
    let gqa_attend = if big_attend {
        pipe("gqa_attend")
    } else {
        pipe("gqa_attend_w16")
    };
    let gdn_step = pipe("gdn_step");
    let gdn_conv = pipe("gdn_conv");
    let sconv_step = pipe("sconv_step");
    let layout_sconv = sconv_step.get_bind_group_layout(0);
    let f32_matvec = pipe("f32_matvec");
    let f32_matvec_b = pipe("f32_matvec_b");
    let layout_f32b = f32_matvec_b.get_bind_group_layout(0);
    let o1_far = pipe("o1_far");
    let o1_push = pipe("o1_push");
    let o1_attend = pipe("o1_attend");
    let layout_o1_far = o1_far.get_bind_group_layout(0);
    let layout_o1_push = o1_push.get_bind_group_layout(0);
    let layout_o1_attend = o1_attend.get_bind_group_layout(0);
    let matvec_pair = pipe("matvec_pair");
    let layout_mv2 = matvec_pair.get_bind_group_layout(0);
    let moe_select = pipe("moe_select");
    let moe_gate_up = pipe("moe_gate_up");
    let moe_down = pipe("moe_down");
    let moe_gate_up_q4tp = pipe("moe_gate_up_q4tp");
    let moe_down_q4tp = pipe("moe_down_q4tp");
    let moe_select_b = pipe("moe_select_b");
    let gdn_conv_k = pipe("gdn_conv_k");
    let q4tp_mv_k = pipe("q4tp_matvec4_k");
    let q4tp_mv16w = pipe("q4tp_matvec16w");
    let q4tp_mv4_bk8: Vec<wgpu::ComputePipeline> = (0..=8u32)
        .map(|nb| {
            let cs = [("NB8", nb.clamp(2, 8) as f64)];
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("q4tp_matvec4_bk8"),
                layout: None,
                module: &module,
                entry_point: Some("q4tp_matvec4_bk8"),
                compilation_options: wgpu::PipelineCompilationOptions {
                    constants: &cs,
                    ..Default::default()
                },
                cache: pcache.as_ref(),
            })
        })
        .collect();
    let x_quant_i8 = pipe("x_quant_i8");
    let q4tp_mv16w_x2 = pipe("q4tp_matvec16w_x2");
    let use_mv_x2 = std::env::var("CMF_MV_X2").map(|v| v != "0").unwrap_or(true);
    let q4tp_mv16w_gu = pipe("q4tp_matvec16w_gu");
    let use_mv_gu = std::env::var("CMF_MV_GU").map(|v| v != "0").unwrap_or(true);
    let q4tp_mv4_bku_x2 = pipe("q4tp_matvec4_bku_x2");
    let q2tp_mv16w = pipe("q2tp_matvec16w");
    let q2tp_mv1_i8 = pipe("q2tp_matvec1_i8");
    let want_q2_ladder = std::env::var("CMF_Q2_LADDER_CACHE").as_deref() == Ok("1");
    let q2_ladder_build = want_q2_ladder
        .then(|| mk_coop(Q2TP_LADDER_CACHE_BUILD_SRC, "q2-ladder-build", "q2_ladder_build"))
        .flatten();
    // A second GPU dispatch using the baseline expression is used only for
    // the exact-bit admission check. CPU libm exp2 differs from Vulkan's F32
    // implementation by a one-ulp result on this adapter, so it is retained
    // as a diagnostic, never treated as the GPU reference.
    let q2_ladder_ref = want_q2_ladder
        .then(|| mk_coop(Q2TP_LADDER_CACHE_BUILD_SRC, "q2-ladder-ref", "q2_ladder_build"))
        .flatten();
    let q2_ladder_mv = want_q2_ladder
        .then(|| mk_coop(Q2TP_LADDER_CACHE_MV_SRC, "q2-ladder-mv", "q2tp_matvec16w_ladder_cache"))
        .flatten();
    let q2tp_mv16w_sg = if q2tp_sg_env_admitted() {
        // Keep the subgroup module isolated: a validation error here must
        // return None and preserve the ordinary q2tp pipeline/device. Naga
        // 30 accepts subgroup builtins through the requested capability, but
        // rejects the WGSL `enable subgroups;` directive itself.
        let adapter_sg = want_sg;
        let device_sg = device.features().contains(wgpu::Features::SUBGROUP);
        if !adapter_sg {
            let msg = format!(
                "feature=false adapter_subgroup=false device_subgroup={device_sg}"
            );
            eprintln!("q2tp subgroup admission: {msg}");
            set_q2tp_sg_diag(msg);
            None
        } else if !device_sg {
            let msg = "feature=false adapter_subgroup=true device_subgroup=false".to_string();
            eprintln!("q2tp subgroup admission: {msg}");
            set_q2tp_sg_diag(msg);
            None
        } else {
            let sc = device.push_error_scope(wgpu::ErrorFilter::Validation);
            let m = device.create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("cmf-q2tp-mv-sg"),
                source: wgpu::ShaderSource::Wgsl(Q2TP_SG_SRC.into()),
            });
            if let Some(e) = pollster::block_on(sc.pop()) {
                let msg = format!(
                    "module_error adapter_subgroup={adapter_sg} device_subgroup={device_sg}: {e}"
                );
                eprintln!("q2tp subgroup admission: {msg}");
                set_q2tp_sg_diag(msg);
                let _ = device.poll(wgpu::PollType::wait_indefinitely());
                None
            } else {
                let sc = device.push_error_scope(wgpu::ErrorFilter::Validation);
                let p = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                    label: Some("q2tp_matvec16w_sg"),
                    layout: None,
                    module: &m,
                    entry_point: Some("q2tp_matvec16w_sg"),
                    compilation_options: Default::default(),
                    cache: pcache.as_ref(),
                });
                if let Some(e) = pollster::block_on(sc.pop()) {
                    let msg = format!(
                        "pipeline_error adapter_subgroup={adapter_sg} device_subgroup={device_sg}: {e}"
                    );
                    eprintln!("q2tp subgroup admission: {msg}");
                    set_q2tp_sg_diag(msg);
                    let _ = device.poll(wgpu::PollType::wait_indefinitely());
                    None
                } else {
                    let msg = format!(
                        "admitted=true adapter_subgroup={adapter_sg} device_subgroup={device_sg}"
                    );
                    eprintln!("q2tp subgroup admission: {msg}");
                    set_q2tp_sg_diag(msg);
                    Some(p)
                }
            }
        }
    } else {
        set_q2tp_sg_diag("not_requested");
        None
    };
    let fwht = want_f16
        .then(|| mk_coop(FWHT_SRC, "prism-fwht", "fwht"))
        .flatten();
    let f32_gemm_dx = pipe("f32_gemm_dx");
    // Per-kernel validation while these are new: a scope around each
    // names the shader the driver rejected, where the module-wide scope
    // only says that one of them was.
    let named = |name: &'static str| {
        let sc = device.push_error_scope(wgpu::ErrorFilter::Validation);
        let p = pipe(name);
        if let Some(e) = pollster::block_on(sc.pop()) {
            tracing::warn!("wgpu kernel {name} rejected: {e}");
        }
        p
    };
    let vae_conv = named("vae_conv");
    let dit_ropepack = named("dit_ropepack");
    let dit_gres = named("dit_gated_residual");
    let dit_rmsmod = named("dit_rmsmod");
    let gdn_step_par_k = pipe("gdn_step_par_k");
    let gdn_step_norm_k = pipe("gdn_step_norm_k");
    let gdn_step_k = pipe("gdn_step_k");
    let moe_gate_up_q4tp_b = pipe("moe_gate_up_q4tp_b");
    let moe_gate_up_q4tp_b_r4 = pipe("moe_gate_up_q4tp_b_r4");
    let moe_down_q4tp_b = pipe("moe_down_q4tp_b");
    // Build one global shader/layout family per descriptor-array geometry.  S8
    // remains the default used by every model; S16 is constructed only for an
    // explicitly opted-in V4.1 profile on adapters whose binding-array limits
    // cover the complete six-binding gate/up group and two-bank shader.
    let make_global = |segments: usize| {
        let source = dsv4_global_moe_shader_source(segments)
            .expect("global MoE shader geometry must be one of the supported sizes");
        let suffix = if segments == DSV4_GLOBAL_MOE_SEGMENTS {
            ""
        } else {
            "-s16"
        };
        let gm = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some(if segments == DSV4_GLOBAL_MOE_SEGMENTS {
                "dsv4-global-moe"
            } else {
                "dsv4-global-moe-s16"
            }),
            source: wgpu::ShaderSource::Wgsl(source.into()),
        });
        let storage =
            |binding: u32, read_only: bool, count: Option<u32>| wgpu::BindGroupLayoutEntry {
                binding,
                visibility: wgpu::ShaderStages::COMPUTE,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Storage { read_only },
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: count.and_then(std::num::NonZeroU32::new),
            };
        let gu0 = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some(if suffix.is_empty() {
                "dsv4-global-gu0"
            } else {
                "dsv4-global-gu0-s16"
            }),
            entries: &[
                storage(0, true, Some(segments as u32)),
                storage(1, true, Some(segments as u32)),
                storage(2, true, None),
                storage(3, true, None),
                storage(4, false, None),
                // V4.1 binds route weights here before the BF16 down input;
                // generic callers still provide their mwt buffer, ignored
                // when the BF16 flag is clear.
                storage(5, true, None),
            ],
        });
        let dn0 = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some(if suffix.is_empty() {
                "dsv4-global-dn0"
            } else {
                "dsv4-global-dn0-s16"
            }),
            entries: &[
                storage(0, true, Some(segments as u32)),
                storage(1, true, None),
                storage(2, true, None),
                storage(3, true, None),
                storage(4, false, None),
            ],
        });
        let params = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some(if suffix.is_empty() {
                "dsv4-global-params"
            } else {
                "dsv4-global-params-s16"
            }),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::COMPUTE,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }],
        });
        let gu_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some(if suffix.is_empty() {
                "dsv4-global-gu-layout"
            } else {
                "dsv4-global-gu-layout-s16"
            }),
            bind_group_layouts: &[Some(&gu0), Some(&params)],
            immediate_size: 0,
        });
        let dn_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some(if suffix.is_empty() {
                "dsv4-global-dn-layout"
            } else {
                "dsv4-global-dn-layout-s16"
            }),
            bind_group_layouts: &[Some(&dn0), Some(&params)],
            immediate_size: 0,
        });
        let gp = |entry: &str| {
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some(entry),
                layout: Some(if entry == "dsv4_global_down_q4tp" {
                    &dn_layout
                } else {
                    &gu_layout
                }),
                module: &gm,
                entry_point: Some(entry),
                compilation_options: Default::default(),
                cache: pcache.as_ref(),
            })
        };
        (
            Some(gp("dsv4_global_gate_up_q4tp")),
            Some(gp("dsv4_global_gate_up_q2tp")),
            Some(gp("dsv4_global_down_q4tp")),
        )
    };
    let (dsv4_global_gu, dsv4_global_gu_q2, dsv4_global_dn) = if want_bind_arrays {
        make_global(DSV4_GLOBAL_MOE_SEGMENTS)
    } else {
        (None, None, None)
    };
    let s16_requested = std::env::var("CMF_DSV41_GLOBAL_SEGMENTS").as_deref() == Ok("16");
    let s16_capable = want_bind_arrays
        && limits.max_binding_array_elements_per_shader_stage
            >= (2 * DSV4_GLOBAL_MOE_SEGMENTS_S16) as u32
        && limits.max_bindings_per_bind_group >= 6
        && limits.max_storage_buffers_per_shader_stage >= 6
        && limits.max_storage_buffer_binding_size >= 4
        && limits.max_buffer_size >= 4;
    let (dsv4_global_gu_s16, dsv4_global_gu_q2_s16, dsv4_global_dn_s16) =
        if s16_requested && s16_capable {
            make_global(DSV4_GLOBAL_MOE_SEGMENTS_S16)
        } else {
            (None, None, None)
        };
    if s16_requested && !s16_capable {
        tracing::warn!(
            "CMF_DSV41_GLOBAL_SEGMENTS=16 requested but adapter limits do not support S16; using S8"
        );
    }
    let moe_down_q4tp_b2 = pipe("moe_down_q4tp_b2");
    let moe_down_q4tp_part = pipe("moe_down_q4tp_part");
    let moe_down_q4tp_b4 = pipe("moe_down_q4tp_b4");
    let moe_down_q2tp_b = pipe("moe_down_q2tp_b");
    let moe_down_q4tp_red = pipe("moe_down_q4tp_red");
    let layout_moe_sel_b = moe_select_b.get_bind_group_layout(0);
    let layout_moe_gu_b = moe_gate_up_q4tp_b.get_bind_group_layout(0);
    let layout_moe_dn_b = moe_down_q4tp_b.get_bind_group_layout(0);
    let layout_moe_sel = moe_select.get_bind_group_layout(0);
    let layout_moe_gu = moe_gate_up.get_bind_group_layout(0);
    let layout_moe_dn = moe_down.get_bind_group_layout(0);
    let layout_moe_gu_q4tp = moe_gate_up_q4tp.get_bind_group_layout(0);
    let moe_gate_up_q2tp = pipe("moe_gate_up_q2tp");
    let moe_gate_up_q2tp_f = pipe("moe_gate_up_q2tp_f");
    let moe_down_q4tp_f = pipe("moe_down_q4tp_f");
    let gqa_attend_dec = pipe("gqa_attend_dec");
    let attend_dec = std::env::var("CMF_ATTEND_DEC")
        .map(|v| v != "0")
        .unwrap_or(true);
    // Measured NEGATIVE on RTX PRO 6000 (72.6 vs 79.0 tok/s): the redundant
    // per-workgroup top-k costs more than the retired select hop — in-pass
    // dispatches overlap more than the latency model assumed. Kept for
    // study; CMF_MOE_FOLDSEL=1 enables.
    let foldsel = std::env::var("CMF_MOE_FOLDSEL").as_deref() == Ok("1");
    let layout_moe_gu_q2tp = moe_gate_up_q2tp.get_bind_group_layout(0);
    let layout_moe_dn_q4tp = moe_down_q4tp.get_bind_group_layout(0);
    let layout = matvec.get_bind_group_layout(0);
    let layout_q1 = q1.get_bind_group_layout(0);
    let layout_rmsnorm = rmsnorm.get_bind_group_layout(0);
    let layout_add_rmsnorm = add_rmsnorm.get_bind_group_layout(0);
    let layout_rmsnorm_b = rmsnorm_b.get_bind_group_layout(0);
    let layout_add_rmsnorm_b = add_rmsnorm_b.get_bind_group_layout(0);
    let layout_attn_rope = attn_rope.get_bind_group_layout(0);
    let layout_kv = kv_append.get_bind_group_layout(0);
    let layout_attend = gqa_attend.get_bind_group_layout(0);
    let layout_attend_s = gqa_attend_s.get_bind_group_layout(0);
    let split_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("cmf-attend-split"),
        source: wgpu::ShaderSource::Wgsl(ATTEND_SPLIT_SRC.into()),
    });
    let pipe_split = |ep: &str| {
        device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some(ep),
            layout: None,
            module: &split_module,
            entry_point: Some(ep),
            compilation_options: Default::default(),
            cache: pcache.as_ref(),
        })
    };
    let attend_part_s = pipe_split("gqa_attend_part_s");
    let attend_part = if big_attend {
        pipe_split("gqa_attend_part")
    } else {
        attend_part_s.clone()
    };
    let attend_merge = pipe_split("gqa_attend_merge");
    let attend_gpart = (big_attend && std::env::var("CMF_ATTEND_GQA").as_deref() != Ok("0"))
        .then(|| pipe_split("gqa_attend_gpart"));
    let layout_attend_gpart = attend_gpart.as_ref().map(|p| p.get_bind_group_layout(0));
    // Subgroup select: its own module — `enable subgroups` must never
    // reach a device without the feature.
    // CMF_MV_SG=1: the barrier-light decode matvec — the experiment
    // the bandwidth test funded. Own module, subgroup feature required.
    let q4tp_mv4_sg = if want_sg && std::env::var("CMF_MV_SG").as_deref() == Ok("1") {
        // Through mk_coop for the error scopes: a rejected module logs
        // its reason and degrades to None — the first version created
        // the pipeline bare, and one naga error killed the device.
        mk_coop(MV_SG_SRC, "cmf-mv-sg", "q4tp_matvec4_sg")
    } else {
        None
    };
    let moe_select_sg = if want_sg && std::env::var("CMF_SELECT_SG").as_deref() != Ok("0") {
        let m = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("cmf-select-sg"),
            source: wgpu::ShaderSource::Wgsl(SELECT_SG_SRC.into()),
        });
        Some(
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("moe_select_sg"),
                layout: None,
                module: &m,
                entry_point: Some("moe_select_sg"),
                compilation_options: Default::default(),
                cache: pcache.as_ref(),
            }),
        )
    } else {
        None
    };
    let layout_attend_part = attend_part.get_bind_group_layout(0);
    let layout_attend_part_s = attend_part_s.get_bind_group_layout(0);
    let layout_attend_merge = attend_merge.get_bind_group_layout(0);
    let layout_gdn = gdn_step.get_bind_group_layout(0);
    let layout_gdn_conv = gdn_conv.get_bind_group_layout(0);
    let layout_f32 = f32_matvec.get_bind_group_layout(0);
    let layout_silu_down = silu_down.get_bind_group_layout(0);
    let layout_mm = matmat.get_bind_group_layout(0);
    let layout_mmm = mul_mm.get_bind_group_layout(0);
    let layout_q1mm = q1_mm.get_bind_group_layout(0);
    let layout_silu = silu.get_bind_group_layout(0);
    let layout_axpy = axpy.get_bind_group_layout(0);
    let layout_colscale = colscale.get_bind_group_layout(0);
    let layout_gate_mul = gate_mul.get_bind_group_layout(0);
    let layout_zero = zero.get_bind_group_layout(0);

    if let Some(e) = pollster::block_on(vscope.pop()) {
        // The Display impl stops at "Validation Error"; the cause chain
        // carries the shader and the line, which is the only part that
        // tells you WHICH kernel to fix.
        let mut detail = format!("{e}");
        let mut src: &dyn std::error::Error = &e;
        while let Some(c) = std::error::Error::source(src) {
            detail.push_str(&format!(" | {c}"));
            src = c;
        }
        return Err(format!("wgpu pipeline validation: {detail}"));
    }

    // Pipeline creation and deferred driver work may outlive the synchronous
    // `create_compute_pipeline` calls.  Drain it before publishing the
    // process-lifetime context; otherwise the short GPU parity harness can
    // reach exit with driver work still queued and the NVIDIA Vulkan loader
    // may sporadically fault during its late cleanup.
    let _ = device.poll(wgpu::PollType::wait_indefinitely());

    // Everything the driver compiled during init — written once, read by
    // every process after this one.
    pipeline_cache_store(pcache.as_ref(), &info);

    Ok(Ctx {
        _instance: instance,
        _adapter: adapter,
        device,
        queue,
        pipeline_cache: pcache,
        adapter_info: info,
        matvec,
        matvec_tiled,
        q8_2f_mv,
        matmat,
        mul_mm,
        q1_mm,
        silu,
        axpy,
        colscale,
        gate_mul,
        zero,
        q1,
        q1_rows,
        q1t,
        q4b,
        q4t_mv,
        rope_heads,
        bf16_round,
        o_lora_a,
        kv_pool,
        index_scores,
        top_k_index,
        hc_block,
        f32_matvec_w,
        o_lora_a_w,
        f32_matvec_x,
        f32_mv_split,
        f32_mv_merge,
        o_lora_a_m,
        moe_gu_q2tp_m,
        moe_dn_q4tp_m,
        sa_part,
        sa_merge,
        blit,
        idx_build,
        moe_route,
        sparse_attend,
        sa_scores,
        sa_apply,
        hc_pre_fold,
        hc_post_expand,
        bt_rope_heads,
        bt_hc_pre_fold,
        bt_hc_block,
        bt_comp_append,
        bt_comp_fold,
        bt_hc_post_expand,
        bt_f32_matvec_w,
        bt_f32_matvec_x,
        bt_moe_route,
        bt_moe_gate_up_q2tp,
        bt_moe_gate_up_q2tp_r4,
        bt_sparse_attend,
        bt_index_scores,
        bt_top_k,
        bt_idx_build_staged,
        bt_o_lora_a,
        bt_o_lora_a4,
        bt_o_lora_a2,
        q4tp_mv,
        q4tp_mv4,
        q4tp_mv4_u2,
        use_mv_u2,
        q4tp_mv4_sg,
        q4tp_mv4_nored,
        use_mv_nored,
        q4tp_mv4_dual,
        use_mv_dual,
        q4tp_mv4_dsilu,
        graph_bgs: Mutex::new(HashMap::new()),
        use_bgcache: std::env::var("CMF_GRAPH_BGCACHE").as_deref() == Ok("1"),
        q4tp_mv4_bk,
        q4tp_mv4_bku,
        use_mv_bk,
        q4tp_mv16w_probe,
        mv_probe,
        use_mv4,
        q4tp_mv16,
        q4t_mv8,
        q4b_mv8,
        q4tp_mm,
        q2tp_mm,
        q2tp_mm_coop,
        q4tp_mm_coop,
        q4tp_mm_coop_s,
        ffn_silu_packed,
        act_absmax,
        dit_gemm_coop,
        dit_qkv_split,
        vae_im2col,
        conv1d_im2col,
        music3_glu,
        dit_v_transpose,
        dit_qknorm,
        act_amax_part,
        act_amax_fold,
        q4tp_dq_f16,
        q8_dq_f16,
        q4tp_mm_coop_f16,
        gemm_nt_coop,
        gemm_nn_coop,
        bake_silu,
        bake_silu_bwd,
        bake_qkr,
        bake_attn_head,
        bake_attn_gate,
        bake_planes: Mutex::new(HashMap::new()),
        argmax_part,
        gdn_step_par,
        gdn_step_par2,
        gdn_step_norm2,
        gdn_inline,
        gdn_step_norm,
        gdn_par,
        ts_query,
        ts_period,
        argmax_final,
        embed_gather_q4tp,
        silu_down,
        q1t_mm,
        q4t_mm,
        dit_qk,
        dit_pv,
        dit_softmax,
        dit_unstack,
        ffn_silu,
        qwen_gelu_bias,
        qwen_layernorm_mod,
        qwen_gated_residual,
        qwen_rope_pack,
        q1t_ovmm,
        rmsnorm,
        add_rmsnorm,
        rmsnorm_b,
        add_rmsnorm_b,
        attn_rope,
        kv_append,
        gqa_attend,
        gqa_attend_s,
        attend_part,
        attend_part_s,
        attend_merge,
        attend_gpart,
        layout_attend_gpart,
        hd_cap: 256,
        big_attend,
        gdn_step,
        gdn_conv,
        sconv_step,
        f32_matvec,
        f32_matvec_b,
        layout_f32b,
        o1_far,
        o1_push,
        o1_attend,
        layout_o1_far,
        layout_o1_push,
        layout_o1_attend,
        o1m: Mutex::new(HashMap::new()),
        matvec_pair,
        layout_mv2,
        moe_select,
        moe_gate_up,
        moe_down,
        moe_gate_up_q4tp,
        moe_down_q4tp,
        moe_select_b,
        gdn_conv_k,
        q4tp_mv_k,
        q4tp_mv16w,
        q4tp_mv4_bk8,
        x_quant_i8,
        i8x: std::sync::Mutex::new(None),
        q4tp_mv16w_x2,
        use_mv_x2,
        q4tp_mv16w_gu,
        use_mv_gu,
        q4tp_mv4_bku_x2,
        q2tp_mv16w,
        q2tp_mv1_i8,
        q2_ladder_build,
        q2_ladder_ref,
        q2_ladder_mv,
        q2tp_mv16w_sg,
        fwht,
        f32_gemm_dx,
        vae_conv,
        dit_ropepack,
        dit_gres,
        dit_rmsmod,
        gdn_step_par_k,
        gdn_step_norm_k,
        gdn_step_k,
        moe_gate_up_q4tp_b,
        moe_gate_up_q4tp_b_r4,
        moe_down_q4tp_b,
        dsv4_global_gu,
        dsv4_global_gu_q2,
        dsv4_global_dn,
        dsv4_global_gu_s16,
        dsv4_global_gu_q2_s16,
        dsv4_global_dn_s16,
        moe_down_q4tp_b2,
        moe_down_q4tp_part,
        moe_down_q4tp_b4,
        moe_down_q2tp_b,
        moe_down_q4tp_red,
        layout_moe_sel_b,
        layout_moe_gu_b,
        layout_moe_dn_b,
        layout,
        layout_mm,
        layout_mmm,
        layout_q1mm,
        layout_silu,
        layout_axpy,
        layout_colscale,
        layout_gate_mul,
        layout_zero,
        layout_q1,
        layout_rmsnorm,
        layout_add_rmsnorm,
        layout_rmsnorm_b,
        layout_add_rmsnorm_b,
        layout_attn_rope,
        layout_kv,
        layout_attend,
        layout_attend_s,
        layout_attend_part,
        layout_attend_part_s,
        layout_attend_merge,
        layout_gdn,
        layout_gdn_conv,
        layout_sconv,
        layout_f32,
        layout_silu_down,
        layout_moe_sel,
        layout_moe_gu,
        layout_moe_dn,
        layout_moe_gu_q4tp,
        moe_gate_up_q2tp,
        moe_gate_up_q2tp_f,
        moe_down_q4tp_f,
        gqa_attend_dec,
        moe_select_sg,
        attend_dec,
        foldsel,
        layout_moe_gu_q2tp,
        layout_moe_dn_q4tp,
        discrete,
        vram_budget,
        resident: std::sync::atomic::AtomicU64::new(0),
        scratch: Mutex::new(Scratch::default()),
        mm_gate: Mutex::new(()),
        q2_ladder: Mutex::new(None),
        planes: Mutex::new(std::collections::HashMap::new()),
        weight_bufs: Mutex::new(HashMap::new()),
        dsv4_kv: Mutex::new(HashMap::new()),
        dsv4_scratch: Mutex::new(HashMap::new()),
        dsv4_comp: Mutex::new(HashMap::new()),
        dsv4_ixkv: Mutex::new(HashMap::new()),
        dsv4_uni: Mutex::new(HashMap::new()),
        dsv4_store: Mutex::new(HashMap::new()),
        slot_writes: Mutex::new(HashMap::new()),
        dsv4_binds: Mutex::new((0, HashMap::new())),
        res_clock: std::sync::atomic::AtomicU64::new(0),
        uniforms: Mutex::new(HashMap::new()),
        uniforms8: Mutex::new(HashMap::new()),
        const_bufs: Mutex::new(HashMap::new()),
        gemm_w_bufs: Mutex::new(HashMap::new()),
        dit_pool: Mutex::new(HashMap::new()),
        rs_bufs: Mutex::new(HashMap::new()),
        attn_kv: Mutex::new(HashMap::new()),
        gdn_state: Mutex::new(HashMap::new()),
        gdn_cursor: Mutex::new(HashMap::new()),
        gdn_snap: Mutex::new(HashMap::new()),
        moe_expw: Mutex::new(HashMap::new()),
        dsv4_global_moe: Mutex::new(HashMap::new()),
        graph_scratch: Mutex::new(GraphScratch::default()),
    })
}

/// Is the active adapter a discrete card? (facade: threshold policy)
pub fn is_discrete() -> bool {
    ctx().map(|c| c.discrete).unwrap_or(false)
}

/// Resident quant weights of the WHOLE tensor in VRAM (loaded once per
/// (file, idx)), guarded by the VRAM budget: once the budget is spent,
/// new tensors return None and their ops run on the CPU. Decode touches
/// layers in order, so the resident set is deterministically the first
/// layers — ngl-style offload without configuration.
/// One tensor living on the device, with what the eviction policy needs.
struct Resident {
    buf: wgpu::Buffer,
    bytes: u64,
    /// Which model layer this tensor belongs to (u16::MAX = not a layer
    /// tensor). Eviction balances ACROSS layers with this: a global pool
    /// under a per-token layer-by-layer sweep is the textbook cyclic
    /// pattern that turns plain LRU into 0% — measured on DeepSeek-V4 as
    /// a 4.7% arena hit rate while the same trace replayed through a
    /// PER-LAYER LRU of equal total size hits ~70%.
    layer: u16,
    /// Never evict. Set for the weights of layers the decode loop has
    /// committed to running on the card: their caches live there, so losing
    /// one mid-sequence is not a slower token but a refused one, and at a
    /// budget near the working set that happened every token.
    pinned: bool,
    /// Use count, aged lazily: the score at time `t` is `uses ·
    /// DECAY^(t − last)`. Plain frequency ossifies — an expert that was
    /// popular during the first prompt outranks one being used right now,
    /// forever.
    uses: f32,
    last: u64,
}

/// Per-tick multiplier for the aged use count. 0.999 halves a score over
/// ~700 lookups: long enough that a steady working set is never disturbed,
/// short enough that a change of task migrates within a prompt or two.
const RES_DECAY: f32 = 0.999;

/// A tensor touched within this many lookups is not evicted, whatever its
/// score. Without it a budget slightly smaller than the working set evicts
/// and re-uploads on every token, which is slower than never using the
/// device at all.
const RES_HYSTERESIS: u64 = 512;

fn res_score(e: &Resident, now: u64) -> f32 {
    e.uses * RES_DECAY.powi((now.saturating_sub(e.last)).min(4096) as i32)
}

/// Pin a layer's weights on the card for the rest of the sequence.
///
/// The residency cache evicts by score to make room, which is right while
/// the set is still being chosen and wrong the moment the decode loop has
/// committed to a set: an evicted layer drops off the card, its caches stay
/// there, and the loop refuses the whole fast path rather than read state
/// from two sides. On a budget near the working set that fired every token —
/// 125 times in a 48-token run on an emulated 24 GB card.
pub fn pin_weights(model: &Arc<CmfModel>, idxs: &[usize]) -> usize {
    let Some(c) = ctx() else { return 0 };
    let uid = model.uid() as usize;
    let mut map = c.weight_bufs.lock().unwrap();
    let mut n = 0;
    for &i in idxs {
        if let Some(e) = map.get_mut(&(uid, i)) {
            if !e.pinned {
                e.pinned = true;
                n += 1;
            }
        }
    }
    n
}

/// Write out whatever the driver has compiled so far.
///
/// Deliberately NOT at the end of `init`: on this Adreno the context
/// comes up in 1.5 s while the compiling costs ~200 s, so the driver is
/// clearly building at first USE, and a blob saved before any dispatch
/// is empty. The engine calls this once, after a generation has run.
pub fn pipeline_cache_flush() {
    let Some(c) = ctx() else { return };
    pipeline_cache_store(c.pipeline_cache.as_ref(), &c.adapter_info);
}

/// Where this device's compiled pipelines are remembered between runs.
///
/// Measured on a Snapdragon 778G (Adreno 642L, Vulkan): building the
/// compute pipelines costs about **200 seconds, once per process**, and
/// it lands in whatever the caller thinks is warmup — in the phone app
/// that was the entire first answer, 209.6 s for 25 tokens against 10.5
/// on the CPU path. The cost is the same for a 4-token run and a
/// 40-token one, so it is not the work; it is the driver's compiler, and
/// nothing was keeping what it produced.
///
/// The key includes the driver string and the engine version: a blob
/// compiled by another driver is not merely stale, it is something the
/// driver may refuse or crash on, which is why `create_pipeline_cache`
/// is unsafe in the first place. `CMF_PIPELINE_CACHE=0` opts out.
fn pipeline_cache_path(info: &wgpu::AdapterInfo) -> Option<std::path::PathBuf> {
    match std::env::var("CMF_PIPELINE_CACHE") {
        Ok(v) if v == "0" => return None,
        Ok(v) => return Some(std::path::PathBuf::from(v)),
        Err(_) => {}
    }
    let mut h = std::collections::hash_map::DefaultHasher::new();
    use std::hash::{Hash, Hasher};
    (
        &info.name,
        &info.driver,
        &info.driver_info,
        info.device,
        info.vendor,
        env!("CARGO_PKG_VERSION"),
    )
        .hash(&mut h);
    Some(crate::gpu::cache_dir_pub().join(format!("cortiq-pipelines-{:016x}.bin", h.finish())))
}

/// Load the blob and hand it to the driver. Unsafe by wgpu's contract —
/// the data goes straight to the driver — which the key above bounds:
/// only this build on this driver can produce a matching file name.
fn pipeline_cache_load(
    device: &wgpu::Device,
    info: &wgpu::AdapterInfo,
    supported: bool,
) -> Option<wgpu::PipelineCache> {
    if !supported {
        return None;
    }
    let path = pipeline_cache_path(info)?;
    let data = std::fs::read(&path).ok();
    let had = data.as_ref().map(|d| d.len()).unwrap_or(0);
    let cache = unsafe {
        device.create_pipeline_cache(&wgpu::PipelineCacheDescriptor {
            label: Some("cortiq-pipelines"),
            data: data.as_deref(),
            // A blob the driver rejects must not take the run down with
            // it: fall back to compiling, which is the old behaviour.
            fallback: true,
        })
    };
    tracing::info!(
        "pipeline cache: {} ({})",
        if had > 0 {
            format!("{had} B loaded")
        } else {
            "empty, will compile".to_string()
        },
        path.display()
    );
    Some(cache)
}

/// Persist what the driver produced. Best-effort: a read-only temp dir
/// costs a recompile next time, never a failure now.
fn pipeline_cache_store(cache: Option<&wgpu::PipelineCache>, info: &wgpu::AdapterInfo) {
    let (Some(cache), Some(path)) = (cache, pipeline_cache_path(info)) else {
        return;
    };
    let Some(data) = cache.get_data() else {
        return;
    };
    // Write-and-rename: a half-written blob handed to a driver next run
    // is exactly the crash this whole path is careful about.
    let tmp = path.with_extension("tmp");
    if std::fs::write(&tmp, &data).is_ok() && std::fs::rename(&tmp, &path).is_ok() {
        tracing::debug!("pipeline cache: {} B saved", data.len());
    }
}

/// `CMF_Q8MV=tiled` picks the workgroup-staged matvec. Off by default:
/// it is a mobile-GPU arm and has to earn its place per device, which is
/// what the flag is for.
fn q8mv_tiled() -> bool {
    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *V.get_or_init(|| std::env::var("CMF_Q8MV").is_ok_and(|v| v == "tiled"))
}

/// Residency counters for the expert-arena regime (a MoE that outsizes
/// VRAM routes thousands of small tensors through here per second). Read
/// them with `residency_stats`; `CMF_MOE_RES=1` logs a line every 512
/// misses.
pub static RES_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static RES_MISSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static RES_MISS_BYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static RES_EVICTS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
/// Misses that fetched bytes and then found no room even after eviction.
pub static RES_REFUSED: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
/// Peak logical weight-residency bytes in this process. This excludes Vulkan
/// allocator slack, KV, and graph scratch; it is the honest arena ceiling.
pub static RES_PEAK_BYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

fn note_resident_peak(c: &Ctx) {
    use std::sync::atomic::Ordering;
    let now = c.resident.load(Ordering::Relaxed);
    let mut peak = RES_PEAK_BYTES.load(Ordering::Relaxed);
    while now > peak {
        match RES_PEAK_BYTES.compare_exchange_weak(
            peak,
            now,
            Ordering::Relaxed,
            Ordering::Relaxed,
        ) {
            Ok(_) => break,
            Err(observed) => peak = observed,
        }
    }
}

/// (hits, misses, miss_bytes, evictions) since process start.
pub fn residency_stats() -> (u64, u64, u64, u64) {
    use std::sync::atomic::Ordering::Relaxed;
    (
        RES_HITS.load(Relaxed),
        RES_MISSES.load(Relaxed),
        RES_MISS_BYTES.load(Relaxed),
        RES_EVICTS.load(Relaxed),
    )
}

/// Logical weight-residency profile for a bounded benchmark. The peak does
/// not pretend to be total physical Vulkan allocation; it is paired with the
/// configured arena budget so a run can report that boundary explicitly.
pub fn residency_profile_report(label: &str) {
    let (hits, misses, miss_bytes, evicts) = residency_stats();
    eprintln!(
        "resident profile: label={label} current_mb={:.1} peak_mb={:.1} budget_mb={:.1} hits={hits} misses={misses} miss_mb={:.1} evicts={evicts}",
        resident_bytes() as f64 / 1e6,
        RES_PEAK_BYTES.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6,
        device_vram_budget() as f64 / 1e6,
        miss_bytes as f64 / 1e6,
    );
}

/// The host tier of the expert residency: a pinned-RAM cache of raw
/// weight bytes between the VRAM arena and storage.
///
/// The measurement that makes it load-bearing: on DeepSeek-V4-Flash the
/// VRAM arena alone hit 4.5% and every miss walked to the network volume
/// at 0.34 GB/s — 420 GB streamed for 60 tokens. FreeToken survives the
/// same regime because its misses land in host RAM over PCIe; this tier
/// is that landing pad. A miss now costs the storage read ONCE while the
/// bytes stay resident here, and re-uploads to the card run at memcpy
/// rate. `CMF_RAM_TIER_MB` sizes it (0/unset = off); the ceiling to
/// respect is the CGROUP's, not the host's — `free` lies in a container.
struct HostTier {
    map:
        std::sync::Mutex<std::collections::HashMap<(usize, usize), (std::sync::Arc<Vec<u8>>, u64)>>,
    bytes: std::sync::atomic::AtomicU64,
    clock: std::sync::atomic::AtomicU64,
    budget: u64,
}

pub static TIER_HITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static TIER_FILLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

fn host_tier() -> Option<&'static HostTier> {
    static T: std::sync::OnceLock<Option<HostTier>> = std::sync::OnceLock::new();
    T.get_or_init(|| {
        let mb: u64 = std::env::var("CMF_RAM_TIER_MB").ok()?.parse().ok()?;
        (mb > 0).then(|| HostTier {
            map: std::sync::Mutex::new(std::collections::HashMap::new()),
            bytes: std::sync::atomic::AtomicU64::new(0),
            clock: std::sync::atomic::AtomicU64::new(0),
            budget: mb * 1024 * 1024,
        })
    })
    .as_ref()
}

/// The projection this system was designed around, made real: pin the
/// working set into the RAM tier with ONE SEQUENTIAL sweep of the file,
/// BEFORE decoding — instead of discovering it miss by miss in random
/// order at seek speed. With a task mask the working set is the masked
/// experts (52 GB of DeepSeek-V4-Flash at cover 0.95); without one it is
/// as much of the expert population as the tier holds, hottest file
/// order. Storage misses during decode then start at zero rather than
/// converging toward it.
///
/// Runs on a background thread; decode proceeds meanwhile and simply
/// starts hitting the tier as the sweep passes it.
pub fn prefetch_tier(model: &Arc<CmfModel>, keep: &dyn Fn(&str) -> bool) {
    let Some(t) = host_tier() else { return };
    // (offset, len, idx) of every kept tensor, in FILE ORDER — the whole
    // point is one sequential pass.
    let mut plan: Vec<(usize, usize, usize)> = Vec::new();
    for (idx, e) in model.tensors.iter().enumerate() {
        if !keep(&e.name) {
            continue;
        }
        let Some(abs) = model.entry_abs_offset(e) else {
            continue;
        };
        plan.push((abs, e.nbytes as usize, idx));
    }
    // Admission order. Default: file order — ONE sequential sweep at
    // streaming rate. With `CMF_TIER_HOT=<stats.json>` (the moe-mask
    // statistics file: {"layer": [counts]}), hot experts go FIRST, so a
    // budget-capped tier keeps the experts the router actually reaches —
    // admission policy only, routing and quality untouched. The sweep
    // then seeks, but a tier that holds the right bytes beats one that
    // streamed the wrong ones (measured: 60.7 slot fills a token sourcing
    // at disk speed with file-order admission).
    let hot: Option<std::collections::HashMap<u16, Vec<u64>>> = std::env::var("CMF_TIER_HOT")
        .ok()
        .and_then(|path| std::fs::read_to_string(path).ok())
        .and_then(|text| {
            let v: serde_json::Value = serde_json::from_str(&text).ok()?;
            let mut m = std::collections::HashMap::new();
            for (k, arr) in v.as_object()? {
                let li: u16 = k.parse().ok()?;
                let counts: Vec<u64> = arr
                    .as_array()?
                    .iter()
                    .map(|x| x.as_u64().unwrap_or(0))
                    .collect();
                m.insert(li, counts);
            }
            Some(m)
        });
    match &hot {
        Some(m) => {
            let score = |idx: usize| -> u64 {
                let name = &model.tensors[idx].name;
                let li = layer_of_name(name);
                let e: usize = name
                    .find(".experts.")
                    .and_then(|i| {
                        let r = &name[i + 9..];
                        r[..r.find('.').unwrap_or(r.len())].parse().ok()
                    })
                    .unwrap_or(usize::MAX);
                m.get(&li).and_then(|c| c.get(e)).copied().unwrap_or(0)
            };
            plan.sort_by_key(|p| std::cmp::Reverse(score(p.2)));
        }
        None => plan.sort_unstable_by_key(|p| p.0),
    }
    let total: u64 = plan.iter().map(|p| p.1 as u64).sum();
    let budget = t.budget;
    tracing::info!(
        "tier prefetch: {} tensors, {:.1} GB planned into a {:.1} GB tier",
        plan.len(),
        total as f64 / 1e9,
        budget as f64 / 1e9
    );
    let model = model.clone();
    // A 4-token bench exits while this thread is still sweeping tens of
    // gigabytes; libc exit() racing a thread mid-allocation was a
    // reproducible SIGSEGV that ate the process's buffered stdout. The
    // sweep therefore honors a stop flag, and an atexit hook raises it
    // and joins — at most one 13 MB read of latency at exit.
    static STOP: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
    static HANDLE: Mutex<Option<std::thread::JoinHandle<()>>> = Mutex::new(None);
    extern "C" fn prefetch_atexit() {
        STOP.store(true, std::sync::atomic::Ordering::SeqCst);
        if let Some(h) = HANDLE.lock().unwrap().take() {
            let _ = h.join();
        }
    }
    unsafe extern "C" {
        // The C runtime's own atexit: present on every target the release
        // ships to, without dragging the libc crate into iOS/Windows deps.
        fn atexit(f: extern "C" fn()) -> i32;
    }
    static HOOK: std::sync::Once = std::sync::Once::new();
    HOOK.call_once(|| unsafe {
        atexit(prefetch_atexit);
    });
    let handle = std::thread::spawn(move || {
        let Ok(f) = std::fs::File::open(&model.path) else {
            return;
        };
        let t0 = std::time::Instant::now();
        let mut done = 0u64;
        for (abs, n, idx) in plan {
            if STOP.load(std::sync::atomic::Ordering::Relaxed) {
                return;
            }
            let Some(t) = host_tier() else { return };
            if t.bytes.load(std::sync::atomic::Ordering::Relaxed) + n as u64 > t.budget {
                break; // tier full — the sweep stops, LRU owns the rest
            }
            let key = (model.uid() as usize, idx);
            if host_tier_get(key).is_some() {
                continue;
            }
            let mut v = vec![0u8; n];
            if read_at(&f, &mut v, abs as u64).is_err() {
                return;
            }
            host_tier_put(key, std::sync::Arc::new(v));
            done += n as u64;
        }
        tracing::info!(
            "tier prefetch: {:.1} GB resident in {:.0}s ({:.2} GB/s)",
            done as f64 / 1e9,
            t0.elapsed().as_secs_f64(),
            done as f64 / 1e9 / t0.elapsed().as_secs_f64().max(1e-9)
        );
    });
    *HANDLE.lock().unwrap() = Some(handle);
}

#[inline]
fn host_tier_cache_admissible(len: usize) -> bool {
    host_tier().is_some_and(|t| (len as u64) <= t.budget)
}

fn host_tier_get(key: (usize, usize)) -> Option<std::sync::Arc<Vec<u8>>> {
    let t = host_tier()?;
    let mut m = t.map.lock().unwrap();
    let now = t.clock.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    let e = m.get_mut(&key)?;
    e.1 = now;
    TIER_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    Some(e.0.clone())
}

fn host_tier_put(key: (usize, usize), bytes: std::sync::Arc<Vec<u8>>) {
    let Some(t) = host_tier() else { return };
    use std::sync::atomic::Ordering;
    let len = bytes.len() as u64;
    if len > t.budget {
        return;
    }
    let mut m = t.map.lock().unwrap();
    if m.contains_key(&key) {
        return;
    }
    // Sampled LRU, same discipline as the VRAM arena: stride-sample,
    // evict the oldest of the sample, never scan the whole map per miss.
    let now = t.clock.fetch_add(1, Ordering::Relaxed);
    let mut rounds = 0;
    while t.bytes.load(Ordering::Relaxed) + len > t.budget && rounds < 128 {
        rounds += 1;
        let n = m.len();
        if n == 0 {
            break;
        }
        let stride = (n / 64).max(1);
        let start = (now as usize).wrapping_mul(0x9E37_79B9) % stride;
        let oldest = m
            .iter()
            .skip(start)
            .step_by(stride)
            .min_by_key(|(_, e)| e.1)
            .map(|(k, e)| (*k, e.0.len() as u64));
        let Some((k, b)) = oldest else { break };
        m.remove(&k);
        t.bytes.fetch_sub(b, Ordering::Relaxed);
    }
    if t.bytes.load(Ordering::Relaxed) + len <= t.budget {
        m.insert(key, (bytes, now));
        t.bytes.fetch_add(len, Ordering::Relaxed);
        TIER_FILLS.fetch_add(1, Ordering::Relaxed);
    }
}

/// Remove one model's entries from a UID-prefixed cache and return the exact
/// bytes owned by the removed values.  The caller holds any surrounding cache
/// lock; keeping the predicate here makes every scoped cleanup use the same
/// owner test and prevents an unrelated model from being dropped by a broad
/// cache clear.
fn release_uid_entries<K, V, F>(
    map: &mut std::collections::HashMap<(usize, K), V>,
    uid: usize,
    mut bytes: F,
) -> u64
where
    K: std::cmp::Eq + std::hash::Hash,
    F: FnMut(&V) -> u64,
{
    let mut released = 0u64;
    map.retain(|(owner, _), value| {
        if *owner == uid {
            released = released.saturating_add(bytes(value));
            false
        } else {
            true
        }
    });
    released
}

#[inline]
fn resident_sub(c: &Ctx, bytes: u64) {
    if bytes == 0 {
        return;
    }
    c.resident
        .fetch_update(
            std::sync::atomic::Ordering::AcqRel,
            std::sync::atomic::Ordering::Relaxed,
            |current| Some(current.saturating_sub(bytes)),
        )
        .ok();
}

/// A synchronous image stage owns the model UID it registered.  The context
/// is process-lived, but the model-specific device buffers are disposable.
/// Wait for every submitted command before dropping handles; shared scratch,
/// content caches, and other model UIDs deliberately remain untouched.
pub(crate) struct ImageStageGuard {
    active: bool,
    model_uid: Option<u64>,
}

pub(crate) fn image_stage_scope() -> ImageStageGuard {
    ImageStageGuard {
        // `selected` only reads configuration.  In particular, this guard
        // must not initialize a GPU for a CPU-only stage or release a stale
        // wgpu cache when the caller explicitly selected Metal/CPU.
        active: selected(),
        model_uid: None,
    }
}

impl ImageStageGuard {
    pub(crate) fn track_model(&mut self, uid: u64) {
        if self.active {
            self.model_uid = Some(uid);
        }
    }
}

impl Drop for ImageStageGuard {
    fn drop(&mut self) {
        if self.active {
            if let Some(uid) = self.model_uid {
                release_idle_model_buffers(uid);
            }
        }
    }
}

/// Drop only model-owned wgpu cache entries after all submitted work drains.
/// This is intentionally narrower than a device/cache reset: shared scratch,
/// content-keyed buffers, KV state, and every other model UID survive.
pub(crate) fn release_idle_model_buffers(uid: u64) {
    let Some(c) = existing_ctx() else { return };
    let owner = uid as usize;
    // The image pipeline's GPU calls are synchronous, but a few reusable
    // paths leave a submission pending while returning a device buffer.  A
    // completed poll is the lifetime barrier before cache-owned handles go.
    let _gate = c.mm_gate.lock().unwrap();
    if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        tracing::warn!("wgpu image-stage cache release skipped after device poll failure");
        return;
    }

    let weight_bytes = {
        let mut map = c.weight_bufs.lock().unwrap();
        let bytes = release_uid_entries(&mut map, owner, |entry| entry.bytes);
        resident_sub(c, bytes);
        bytes
    };
    let moe_bytes = {
        let mut map = c.moe_expw.lock().unwrap();
        let bytes = release_uid_entries(&mut map, owner, |(gate, up, down)| {
            gate.size()
                .saturating_add(up.size())
                .saturating_add(down.size())
        });
        resident_sub(c, bytes);
        bytes
    };
    let global_moe_bytes = {
        let mut map = c.dsv4_global_moe.lock().unwrap();
        let mut bytes = 0u64;
        map.retain(|owner, bufs| {
            if *owner != uid {
                return true;
            }
            for buffer in bufs.gate.iter().chain(&bufs.up).chain(&bufs.down) {
                bytes = bytes.saturating_add(buffer.size());
            }
            false
        });
        resident_sub(c, bytes);
        bytes
    };
    let plane_count = {
        let mut map = c.planes.lock().unwrap();
        let before = map.len();
        let _ = release_uid_entries(&mut map, owner, |(_, bytes)| *bytes);
        before - map.len()
    };
    let row_scale_count = {
        let mut map = c.rs_bufs.lock().unwrap();
        let before = map.len();
        let _ = release_uid_entries(&mut map, owner, |_| 0);
        before - map.len()
    };
    {
        let mut registry = layer_registry().lock().unwrap();
        let _ = release_uid_entries(&mut registry, owner, |_| 0);
    }
    let host_bytes = if let Some(tier) = host_tier() {
        let mut map = tier.map.lock().unwrap();
        let bytes = release_uid_entries(&mut map, owner, |(data, _)| data.len() as u64);
        if bytes != 0 {
            tier.bytes
                .fetch_sub(bytes, std::sync::atomic::Ordering::AcqRel);
        }
        bytes
    } else {
        0
    };
    let host_bank_bytes = {
        let mut map = host_banks().lock().unwrap();
        let mut bytes = 0u64;
        map.retain(|(bank_owner, _), banks| {
            if *bank_owner != owner {
                return true;
            }
            bytes = bytes.saturating_add(
                (banks.g.len() as u64)
                    .saturating_add(banks.u.len() as u64)
                    .saturating_add(banks.d.len() as u64),
            );
            false
        });
        bytes
    };
    if weight_bytes != 0
        || moe_bytes != 0
        || global_moe_bytes != 0
        || plane_count != 0
        || row_scale_count != 0
        || host_bytes != 0
        || host_bank_bytes != 0
    {
        tracing::debug!(
            uid,
            weight_bytes,
            moe_bytes,
            global_moe_bytes,
            plane_count,
            row_scale_count,
            host_bytes,
            host_bank_bytes,
            "released wgpu image-stage model caches"
        );
    }
}

#[cfg(test)]
mod image_stage_cache_tests {
    use super::release_uid_entries;
    use std::collections::HashMap;

    #[test]
    fn uid_release_sums_only_owned_entries() {
        let mut cache = HashMap::from([
            ((11usize, 0usize), 17u64),
            ((11usize, 1usize), 23u64),
            ((12usize, 0usize), 41u64),
        ]);
        let released = release_uid_entries(&mut cache, 11, |bytes| *bytes);
        assert_eq!(released, 40);
        assert_eq!(cache.len(), 1);
        assert_eq!(cache.get(&(12, 0)), Some(&41));
    }
}

/// On a network filesystem an mmap MISS is the death of this path: every
/// 4-128 KB page faults through FUSE one round trip at a time, which is
/// the measured "1% CPU, looks hung" failure on MooseFS volumes. With
/// `CMF_WEIGHT_PREAD=1` a residency miss reads its byte range with ONE
/// explicit pread instead — sequential, at the volume's streaming rate.
/// Hits never touch bytes at all, so this only prices the misses.
fn pread_range(model: &Arc<CmfModel>, abs: usize, len: usize) -> Option<Vec<u8>> {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    if !*ON.get_or_init(|| std::env::var("CMF_WEIGHT_PREAD").as_deref() == Ok("1")) {
        return None;
    }
    let f = std::fs::File::open(&model.path).ok()?;
    let mut v = vec![0u8; len];
    read_at(&f, &mut v, abs as u64).ok()?;
    Some(v)
}

/// Is this tensor already resident? A cheap peek so callers can decide
/// how to source the bytes for a MISS without paying for the hit path.
fn weight_resident(c: &Ctx, key: (usize, usize)) -> bool {
    c.weight_bufs.lock().unwrap().contains_key(&key)
}

/// `layers.N.` parsed out of a tensor name; u16::MAX when absent.
fn layer_of_name(name: &str) -> u16 {
    let Some(i) = name.find("layers.") else {
        return u16::MAX;
    };
    let rest = &name[i + 7..];
    let end = rest.find('.').unwrap_or(rest.len());
    rest[..end].parse().unwrap_or(u16::MAX)
}

/// (model, tensor) -> layer, filled by the wrappers that still know the
/// tensor's NAME before it becomes a bare key inside the dispatchers.
/// The eviction policy needs the layer; threading it through every
/// dispatcher signature would touch a dozen call chains for one u16.
fn layer_registry() -> &'static std::sync::Mutex<std::collections::HashMap<(usize, usize), u16>> {
    static R: std::sync::OnceLock<
        std::sync::Mutex<std::collections::HashMap<(usize, usize), u16>>,
    > = std::sync::OnceLock::new();
    R.get_or_init(|| std::sync::Mutex::new(std::collections::HashMap::new()))
}

fn note_layer(key: (usize, usize), name: &str) {
    let l = layer_of_name(name);
    if l != u16::MAX {
        layer_registry().lock().unwrap().entry(key).or_insert(l);
    }
}

fn weight_buffer(c: &Ctx, key: (usize, usize), full_quant: &[u8]) -> Option<wgpu::Buffer> {
    let lay = layer_registry()
        .lock()
        .unwrap()
        .get(&key)
        .copied()
        .unwrap_or(u16::MAX);
    weight_buffer_l(c, key, full_quant, lay)
}

fn weight_buffer_l(
    c: &Ctx,
    key: (usize, usize),
    full_quant: &[u8],
    incoming_layer: u16,
) -> Option<wgpu::Buffer> {
    use std::sync::atomic::Ordering;
    let now = c.res_clock.fetch_add(1, Ordering::Relaxed);
    let len = full_quant.len() as u64;
    let mut map = c.weight_bufs.lock().unwrap();
    if let Some(e) = map.get(&key) {
        if e.bytes >= len {
            let e = map.get_mut(&key).unwrap();
            e.uses = res_score(e, now) + 1.0;
            e.last = now;
            RES_HITS.fetch_add(1, Ordering::Relaxed);
            return Some(e.buf.clone());
        }
        // The same tensor key can first be touched by a body-only q8 GEMM
        // and later by q8_2f's packed graph kernel, which also reads the row
        // and column scale planes. A resident hit is valid only when it is
        // at least as large as the requested payload; upgrade shorter cache
        // entries instead of returning a buffer whose tail does not exist.
        let old = e.bytes;
        map.remove(&key);
        c.resident.fetch_sub(old, Ordering::Relaxed);
        crate::gpu::probe_note_cold();
    }
    if len > c.vram_budget {
        return None; // one tensor larger than the whole budget
    }
    res_who_key(key);
    RES_MISSES.fetch_add(1, Ordering::Relaxed);
    RES_MISS_BYTES.fetch_add(len, Ordering::Relaxed);
    // The RAM tier serves the bytes on a miss when it has them — one
    // memcpy instead of re-faulting mmap pages (or re-reading a disk
    // range). Filled below on BOTH outcomes, insert and refusal: a
    // refused tensor is the one most likely to be missed again next
    // token, and refetching it from storage every time was measured as
    // 65% of all fetch traffic on DeepSeek-V4.
    // A body-only cache fill may also have populated the host tier. It is
    // useful for this request only if it contains the complete payload.
    let tier_bytes = host_tier_get(key).filter(|v| v.len() >= full_quant.len());
    let full_quant: &[u8] = tier_bytes.as_deref().map_or(full_quant, |v| v);
    if std::env::var("CMF_MOE_RES").is_ok() {
        let m = RES_MISSES.load(Ordering::Relaxed);
        if m % 512 == 0 {
            let (h, mm, mb, ev) = residency_stats();
            let th = TIER_HITS.load(Ordering::Relaxed);
            let tf = TIER_FILLS.load(Ordering::Relaxed);
            // Diagnostics for the layer-balance policy: how many resident
            // entries actually KNOW their layer, and how much is pinned.
            tracing::info!(
                "residency: {h} hits / {mm} misses ({:.1}% hit), {:.2} GB fetched, {ev} ev {} ref | tier {th}/{tf} | res {} = {:.1} GB (big {}) reg {}",
                h as f64 / (h + mm).max(1) as f64 * 100.0,
                mb as f64 / 1e9,
                RES_REFUSED.load(Ordering::Relaxed),
                map.len(),
                map.values().map(|e| e.bytes).sum::<u64>() as f64 / 1e9,
                map.values().filter(|e| e.bytes > (64 << 20)).count(),
                layer_registry().lock().unwrap().len()
            );
        }
    }
    // Make room by evicting the least valuable, skipping anything touched
    // recently. Failing to free enough is not an error: the tensor stays on
    // the CPU, which is the pressure valve that keeps a too-small budget
    // from thrashing the bus.
    //
    // Eviction picks from a SAMPLE, not a full scan. The full
    // collect-and-sort was sized for a dense model's dozens of tensors; a
    // MoE arena holds thousands of experts, and a per-miss O(n log n) over
    // them at ~120 misses/token is a decode-scale cost. Stride-sampling 64
    // candidates and evicting the worst of them is the Redis discipline:
    // within a few percent of true-LRU quality at O(1) cost.
    if c.resident.load(Ordering::Relaxed) + len > c.vram_budget {
        let mut freed = 0u64;
        let mut rounds = 0;
        while c.resident.load(Ordering::Relaxed) - freed + len > c.vram_budget && rounds < 64 {
            rounds += 1;
            let n = map.len();
            if n == 0 {
                break;
            }
            let stride = (n / 64).max(1);
            let start = (now as usize).wrapping_mul(0x9E37_79B9) % stride.max(1);
            // Never evict the incoming tensor's own layer while any OTHER
            // layer still holds entries: under the per-token sweep every
            // layer refills in turn, and stealing from your own layer is
            // exactly the cyclic-LRU collapse. With only own-layer entries
            // left (a model with one giant layer), fall back to any.
            let pick = |own: bool| {
                map.iter()
                    .skip(start)
                    .step_by(stride)
                    .filter(|(_, e)| !e.pinned && now.saturating_sub(e.last) > RES_HYSTERESIS)
                    .filter(|(_, e)| own || e.layer != incoming_layer || e.layer == u16::MAX)
                    .min_by(|a, b| {
                        res_score(a.1, now)
                            .partial_cmp(&res_score(b.1, now))
                            .unwrap_or(std::cmp::Ordering::Equal)
                    })
                    .map(|(k, e)| (*k, e.bytes))
            };
            let worst = pick(false).or_else(|| pick(true));
            let Some((k, bytes)) = worst else { break };
            map.remove(&k);
            RES_EVICTS.fetch_add(1, Ordering::Relaxed);
            freed += bytes;
        }
        if freed > 0 {
            c.resident.fetch_sub(freed, Ordering::Relaxed);
            crate::gpu::probe_note_cold();
        }
        if c.resident.load(Ordering::Relaxed) + len > c.vram_budget {
            RES_REFUSED.fetch_add(1, Ordering::Relaxed);
            if tier_bytes.is_none() {
                host_tier_put(key, std::sync::Arc::new(full_quant.to_vec()));
            }
            return None; // still no room — honest CPU
        }
    }
    crate::gpu::probe_note_cold(); // first touch = upload, not a steady sample
                                   // DEVICE-LOCAL residency: create_buffer_init maps at creation → the buffer
                                   // lands in a HOST_VISIBLE heap and every matvec streams its weights over
                                   // PCIe (~25 GB/s) every token. A plain create_buffer + staged write_buffer
                                   // lets the allocator pick DEVICE_LOCAL VRAM (~1 TB/s on a 4090). This is
                                   // THE discrete-GPU decode fix; on UMA it's a wash.
    let buf = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("q1-weights"),
        // Rounded up: write_buffer refuses a size that is not a multiple of
        // four, and a small q4tp payload need not be one. This only ever
        // surfaced once the preparation's weights joined the preflight —
        // until then every tensor through here happened to be aligned.
        size: len.next_multiple_of(4),
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    let t_up = std::time::Instant::now();
    if upload_staged() {
        // Map and copy bounded chunks. `queue.write_buffer` goes through
        // wgpu's belt, which for a 92 GB expert stack measured ~48 MB/s; one
        // staging buffer per WHOLE tensor was fast but left a size-fragmented
        // second model in the Vulkan allocator. A fixed chunk class keeps the
        // speed and bounds the transient peak.
        let chunk = upload_chunk_bytes().next_multiple_of(4);
        for (ci, src) in full_quant.chunks(chunk).enumerate() {
            let n = src.len().next_multiple_of(4) as u64;
            let stg = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("weight-staging"),
                // Every allocation has the SAME size, including the last
                // short chunk of a tensor. Giving Vulkan one differently
                // sized tail per projection made its allocator retain a
                // fragmented second heap: a 14 GiB Granite Q8_2F prefix
                // peaked at 23.8 GiB despite all final weights fitting the
                // arena. Only `n` bytes are copied below, so the unused tail
                // never needs to fit in the destination buffer.
                size: chunk as u64,
                usage: wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: true,
            });
            {
                let Ok(mut view) = stg.slice(..).get_mapped_range_mut() else {
                    return None;
                };
                // Write-only by construction: mapped upload memory may be
                // uncached, so wgpu hands out a slice you may write but not read.
                view.slice(..src.len()).copy_from_slice(src);
                if src.len() < n as usize {
                    view.slice(src.len()..).fill(0);
                }
            }
            stg.unmap();
            let mut enc = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("wup") });
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(&stg, 0, &buf, (ci * chunk) as u64, n);
            submit(c, finish_enc(enc));
            // The staging buffer must outlive the copy; polling before it is
            // dropped also lets the allocator reuse this same chunk class.
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            drop(stg);
        }
    } else if full_quant.len() % 4 == 0 {
        c.queue.write_buffer(&buf, 0, full_quant);
    } else {
        let mut padded = full_quant.to_vec();
        padded.resize(padded.len().next_multiple_of(4), 0);
        c.queue.write_buffer(&buf, 0, &padded);
    }
    UPLOAD_NS.fetch_add(t_up.elapsed().as_nanos() as u64, Ordering::Relaxed);
    UPLOAD_BYTES.fetch_add(len, Ordering::Relaxed);
    c.resident.fetch_add(len, Ordering::Relaxed);
    note_resident_peak(c);
    map.insert(
        key,
        Resident {
            buf: buf.clone(),
            bytes: len,
            uses: 1.0,
            last: now,
            pinned: false,
            layer: incoming_layer,
        },
    );
    if tier_bytes.is_none() {
        host_tier_put(key, std::sync::Arc::new(full_quant.to_vec()));
    }
    Some(buf)
}

/// One MoE layer's expert weights as three concatenated device buffers
/// (gate_all, up_all, down_all), q4t payloads back to back in `experts`
/// order (routed experts then the shared one) — the kernels address
/// expert e at u16 offset e·mat16. Uploaded once per layer (keyed by the
/// first gate idx), budget-guarded like every resident weight; the copy
/// walks the per-tensor directory, so no file-order contiguity is assumed.
/// One-shot reason the whole-token graph declined. Without it the fallback
/// to the per-op path is invisible, which is how a q4tp model looked
/// GPU-accelerated while every layer walked the host.
/// `CMF_GRAPH_SPLIT=N` — how many pieces the token graph's submission is
/// cut into (default 10, the sweep's plateau: 2→110.4, 6→119.9,
/// 10-13→122.3, 40→120.1 on the 35B against 97.9 unsplit; 0/1 = the
/// historical single submit).
/// Workgroup budget for the q4tp matvecs (`CMF_MV_GRID`, 0 = one
/// workgroup per row block, the historical grid). The kernels already
/// loop `wb += num_workgroups`, so a SMALLER grid makes them persistent:
/// each workgroup walks many row blocks instead of streaming 2.5 groups
/// and exiting.
///
/// MEASURED AND NULL — kept so the idea is not re-opened. Qwen3.6-27B on
/// an RTX 5090, GPU timestamps (steady to 0.3%): grid 0/512/680/1024/
/// 2048 give 17.66/17.68/17.67/17.74/17.71 ms a token, and 170 (one
/// workgroup per SM) is WORSE at 20.81. Launch and teardown are not what
/// this kernel is paying. The stripped probe holds 1056 GB/s under every
/// grid, and two decode processes on one card aggregate 52.8 tok/s
/// against a single process's 48.8 — the card is ~92% saturated by one
/// stream, so that 1056 is the bus, not the kernel.
fn mv_grid_cap() -> u32 {
    static N: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
    *N.get_or_init(|| {
        std::env::var("CMF_MV_GRID")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(0)
    })
}

/// Bind group + grid for the two-weight 16w kernel: weight A → ya over
/// rows_a, weight B → yb over rows_b, both against `xs` (cols wide).
#[allow(clippy::too_many_arguments)]
fn mv_x2_bind(
    c: &Ctx,
    wa: &wgpu::Buffer,
    wb: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    ya: &wgpu::Buffer,
    yb: &wgpu::Buffer,
    rows_a: usize,
    rows_b: usize,
    cols: usize,
) -> (wgpu::BindGroup, u32) {
    let gpr = cols / 32;
    let p_buf = uniform_u32x4(c, [gpr as u32, rows_a as u32, 1, rows_b as u32]);
    let layout = c.q4tp_mv16w_x2.get_bind_group_layout(0);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("mv-x2"),
        layout: &layout,
        entries: &[
            bind_buf(0, wa),
            bind_buf(2, ya),
            bind_buf(3, &p_buf),
            bind_buf(4, wa),
            bind_buf(5, xs),
            bind_buf(6, wb),
            bind_buf(7, wb),
            bind_buf(8, yb),
        ],
    });
    let wg = (rows_a as u32).div_ceil(16) + (rows_b as u32).div_ceil(16);
    (bind, mv_grid(wg))
}

#[inline]
fn q2tp_sg_env_admitted() -> bool {
    std::env::var("CMF_Q2TP_SG").as_deref() == Ok("1")
        && std::env::var("CMF_Q2TP_SG_LINEAR").as_deref() == Ok("1")
        && matches!(
            std::env::var("CMF_Q2TP_SG_WIDTH")
                .ok()
                .and_then(|v| v.parse::<u32>().ok()),
            Some(32) | Some(64)
        )
}

// Runtime-owned selection counters make the subgroup A/B observable on the
// resident graph rather than relying on the standalone component probe. A
// lookup is counted where the bind layout/pipeline is selected; graph callers
// use the same selector immediately before their dispatch.
static Q2TP_SG_PIPELINE_LOOKUPS: std::sync::atomic::AtomicU64 =
    std::sync::atomic::AtomicU64::new(0);
static Q2TP_TREE_PIPELINE_LOOKUPS: std::sync::atomic::AtomicU64 =
    std::sync::atomic::AtomicU64::new(0);
// Resident-graph ladder selection is tracked separately from the ordinary
// tree/subgroup counters.  The graph must prove that its own prep/emat route
// actually selected the sidecar; component admission alone is insufficient.
static Q2TP_LADDER_GRAPH_LOOKUPS: std::sync::atomic::AtomicU64 =
    std::sync::atomic::AtomicU64::new(0);
// Number of Prism GDN projection preparations intentionally skipped because
// the transformed output makes the fused whole-chain arm ineligible.  This is
// a run-owned aggregate counter; it does not affect dispatch selection.
static GDN_PRISM_SKIPPED_PREPS: std::sync::atomic::AtomicU64 =
    std::sync::atomic::AtomicU64::new(0);

/// Return resident-graph q2tp pipeline-selection counts for a bounded A/B.
/// This is a run-owned diagnostic hook, not runtime behavior.
pub fn q2tp_selection_snapshot() -> (u64, u64) {
    use std::sync::atomic::Ordering;
    (
        Q2TP_SG_PIPELINE_LOOKUPS.load(Ordering::Relaxed),
        Q2TP_TREE_PIPELINE_LOOKUPS.load(Ordering::Relaxed),
    )
}

/// Emit one concise admission/selection line when explicitly requested by a
/// run-owned benchmark helper. Default runtime output is unchanged.
pub fn q2tp_selection_report(label: &str) {
    let (subgroup, tree) = q2tp_selection_snapshot();
    let diagnostics = q2tp_sg_diag();
    use std::sync::atomic::Ordering;
    let ladder = Q2TP_LADDER_GRAPH_LOOKUPS.load(Ordering::Relaxed);
    let gdn_skipped = GDN_PRISM_SKIPPED_PREPS.load(Ordering::Relaxed);
    eprintln!(
        "q2tp resident selection: label={label} requested={} admitted={} subgroup_lookups={subgroup} tree_lookups={tree} ladder_lookups={ladder} gdn_prism_skipped_preps={gdn_skipped} diagnostics={diagnostics}",
        q2tp_sg_env_admitted(),
        diagnostics.starts_with("admitted="),
    );
}

/// Run-owned ladder-cache gate: `(unique_pairs, table_bytes, row_id_bytes)`
/// after the exact GPU precompute/readback check has admitted the sidecar.
/// `None` means the optional experiment was not requested or failed closed.
pub fn q2tp_ladder_cache_snapshot() -> Option<(usize, u64, u64)> {
    let c = ctx()?;
    let cache = c.q2_ladder.lock().unwrap();
    cache
        .as_ref()
        .map(|v| (v.pairs, v.table_bytes, v.row_id_bytes))
}

#[inline]
fn q2tp_pipeline(c: &Ctx) -> &wgpu::ComputePipeline {
    if q2tp_sg_env_admitted() {
        if let Some(p) = c.q2tp_mv16w_sg.as_ref() {
            use std::sync::atomic::Ordering;
            Q2TP_SG_PIPELINE_LOOKUPS.fetch_add(1, Ordering::Relaxed);
            return p;
        }
    }
    use std::sync::atomic::Ordering;
    Q2TP_TREE_PIPELINE_LOOKUPS.fetch_add(1, Ordering::Relaxed);
    &c.q2tp_mv16w
}

/// Opt into the original-engine-style Q8/DP4A decode only for the explicit
/// affine q2tp experiment.  The default remains the validated scalar kernel;
/// this guard is intentionally independent of CMF_VERIFY_I8, whose q4tp
/// batched path has a different zero-sum contract.
#[inline]
fn q2tp_dp4a_on() -> bool {
    std::env::var("CMF_Q2_DP4A").as_deref() == Ok("1")
}

/// One q2tp (2-bit plane, kind 9) matvec dispatch in its own pass.
fn encode_q2tp_mv16w(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    affine: bool,
) {
    let (bind, wg) = q2tp_mv_bind(c, weight, xs, y, rows, cols, affine);
    let mut pass = begin_pass(enc);
    pass.set_pipeline(q2tp_pipeline(c));
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(wg, 1, 1);
}

/// One affine q2tp vector through the existing Q8 activation quantizer and a
/// dedicated signed-ternary DP4A kernel.  Unlike q4tp's NB=2..8 helper this
/// has a genuine NB=1 specialization and never applies a q4 zero-sum term.
#[allow(clippy::too_many_arguments)]
fn encode_q2tp_mv1_i8(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
) {
    let gpr = cols / 32;
    let need = gpr as u64;
    let (x8, xsb) = {
        let mut g = c.i8x.lock().unwrap();
        let grow = match g.as_ref() {
            Some((_, _, cap)) => *cap < need,
            None => true,
        };
        if grow {
            let cap = need.max(8 * 1024);
            let x8 = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("q2-i8-x8"),
                size: cap * 32,
                usage: wgpu::BufferUsages::STORAGE,
                mapped_at_creation: false,
            });
            let xsb = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("q2-i8-xs"),
                size: cap * 8,
                usage: wgpu::BufferUsages::STORAGE,
                mapped_at_creation: false,
            });
            *g = Some((x8, xsb, cap));
        }
        let (a, b, _) = g.as_ref().unwrap();
        (a.clone(), b.clone())
    };
    // x_quant_i8 consumes the same post-FWHT/f16-boundary activation buffer as
    // the scalar q2 path. It writes one (x8, scale) pair per 32-value group.
    let qp = uniform_u32x4(c, [gpr as u32, 1, 0, 0]);
    let ql = c.x_quant_i8.get_bind_group_layout(0);
    let qbind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q2-i8-quant"),
        layout: &ql,
        entries: &[
            bind_buf(3, &qp),
            bind_buf(5, xs),
            bind_buf(11, &x8),
            bind_buf(12, &xsb),
        ],
    });
    // `_p1=1` is the already-validated affine center selector. This function
    // is only called from an affine descriptor gate; it is not a retagging
    // mechanism for ordinary q2tp tensors.
    let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, 1, 1]);
    let layout = c.q2tp_mv1_i8.get_bind_group_layout(0);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q2-mv1-i8"),
        layout: &layout,
        entries: &[
            bind_buf(0, weight),
            bind_buf(2, y),
            bind_buf(3, &p_buf),
            bind_buf(9, &x8),
            bind_buf(10, &xsb),
        ],
    });
    let mut pass = begin_pass_with(enc, Some("q2-mv1-i8"), None);
    pass.set_pipeline(&c.x_quant_i8);
    pass.set_bind_group(0, &qbind, &[]);
    pass.dispatch_workgroups((gpr as u32).div_ceil(256), 1, 1);
    pass.set_pipeline(&c.q2tp_mv1_i8);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(mv_grid((rows as u32).div_ceil(16)), 1, 1);
}

/// Bind group + grid for the q2tp decode matvec.
fn q2tp_mv_bind(
    c: &Ctx,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    affine: bool,
) -> (wgpu::BindGroup, u32) {
    let gpr = cols / 32;
    // `_p1` is the descriptor-aware center selector: 0 = raw q2tp
    // `(c-1.5)·s`, 1 = q2tp_affine `(c-1)·s`. The payload is unchanged.
    let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, 1, affine as u32]);
    let layout = q2tp_pipeline(c).get_bind_group_layout(0);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("mv-q2"),
        layout: &layout,
        entries: &[
            bind_buf(0, weight),
            bind_buf(2, y),
            bind_buf(3, &p_buf),
            bind_buf(5, xs),
        ],
    });
    (bind, mv_grid((rows as u32).div_ceil(16)))
}

/// Build the run-owned global q2tp ladder sidecar.  The table is intentionally
/// derived from all Q2 rows in the loaded model rather than from the one
/// benchmark tensor: the gate must exercise the exact 197,295-pair inventory
/// and its u32 row-id plane, not a per-matrix duplicate that changes the
/// accounting.  A strict bit check against the same F32 `exp2` expression
/// rejects the cache before it can affect a projection.
fn ensure_q2_ladder_cache(
    c: &Ctx,
    model: &Arc<CmfModel>,
    idx: usize,
    rows: usize,
    cols: usize,
) -> Option<(wgpu::Buffer, wgpu::Buffer, u32)> {
    let _ = cols;
    let build = c.q2_ladder_build.as_ref()?;
    let baseline = c.q2_ladder_ref.as_ref()?;
    c.q2_ladder_mv.as_ref()?;
    {
        let cache = c.q2_ladder.lock().unwrap();
        if let Some(cache) = cache.as_ref().filter(|v| {
            v.model_uid == model.uid()
                && v.row_counts.get(&idx).is_some_and(|&count| count >= rows)
        }) {
            return Some((
                cache.ladders.clone(),
                cache.row_ids.clone(),
                *cache.row_id_base.get(&idx).unwrap(),
            ));
        }
    }

    let t0 = std::time::Instant::now();
    let mut pair_ids: HashMap<u32, u32> = HashMap::with_capacity(197_295);
    let mut raw_keys: Vec<u32> = Vec::with_capacity(197_295);
    let mut row_ids: Vec<u32> = Vec::new();
    let mut target_row_base: Option<u32> = None;
    let mut row_bases: HashMap<usize, u32> = HashMap::new();
    let mut row_counts: HashMap<usize, usize> = HashMap::new();
    for (ti, entry) in model.tensors.iter().enumerate() {
        if entry.dtype != cortiq_core::TensorDtype::Q2TiledP {
            continue;
        }
        let Some(&entry_cols) = entry.shape.get(1) else {
            return None;
        };
        let entry_rows = entry.shape.first().copied().unwrap_or(0);
        if entry_rows == 0 || entry_cols == 0 || entry_cols % cortiq_core::quant::GROUP_SIZE != 0 {
            return None;
        }
        let (params_off, _, _) = cortiq_core::quant::q2tp_sections(entry_rows, entry_cols);
        let bytes = model.entry_bytes(entry);
        let params_end = params_off.checked_add(entry_rows.checked_mul(4)?)?;
        if params_end > bytes.len() {
            return None;
        }
        row_bases.insert(ti, row_ids.len() as u32);
        row_counts.insert(ti, entry_rows);
        for r in 0..entry_rows {
            let o = params_off + r * 4;
            let key = u32::from_le_bytes(bytes[o..o + 4].try_into().ok()?);
            let id = if let Some(&id) = pair_ids.get(&key) {
                id
            } else {
                let id = raw_keys.len() as u32;
                pair_ids.insert(key, id);
                raw_keys.push(key);
                id
            };
            if ti == idx && r == 0 {
                target_row_base = Some(row_ids.len() as u32);
            }
            row_ids.push(id);
        }
    }
    let Some(row_id_base) = target_row_base else {
        eprintln!("q2 ladder cache rejected: target tensor row ids missing");
        return None;
    };
    let target_row_end = row_id_base as usize + rows;
    if target_row_end > row_ids.len() || raw_keys.is_empty() {
        eprintln!(
            "q2 ladder cache rejected: target rows {}/{} pairs {}",
            target_row_end.saturating_sub(row_id_base as usize),
            rows,
            raw_keys.len()
        );
        return None;
    }

    let pairs = raw_keys.len();
    let table_len = pairs.checked_mul(32)?;
    let table_bytes = (table_len * std::mem::size_of::<f32>()) as u64;
    let row_id_bytes = (row_ids.len() * std::mem::size_of::<u32>()) as u64;
    let expected: Vec<f32> = raw_keys
        .iter()
        .flat_map(|&key| {
            let lo = cortiq_core::quant::f16_to_f32((key & 0xffff) as u16);
            let step = cortiq_core::quant::f16_to_f32((key >> 16) as u16);
            (0..32)
                .map(move |r| if r == 0 { 0.0 } else { (lo + (r - 1) as f32 * step).exp2() })
        })
        .collect();
    let raw_buf = c.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some("q2-ladder-raw-keys"),
        contents: bytemuck::cast_slice(&raw_keys),
        usage: wgpu::BufferUsages::STORAGE,
    });
    let ladders = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("q2-ladder-global"),
        size: table_bytes.max(4),
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        mapped_at_creation: false,
    });
    let row_ids_buf = c.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some("q2-ladder-row-ids"),
        contents: bytemuck::cast_slice(&row_ids),
        usage: wgpu::BufferUsages::STORAGE,
    });
    let p_buf = c.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some("q2-ladder-params"),
        contents: bytemuck::cast_slice(&[pairs as u32, 0u32, 0u32, 0u32]),
        usage: wgpu::BufferUsages::UNIFORM,
    });
    let build_bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q2-ladder-build-bg"),
        layout: &build.get_bind_group_layout(0),
        entries: &[bind_buf(0, &raw_buf), bind_buf(1, &ladders), bind_buf(2, &p_buf)],
    });
    let baseline_buf = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("q2-ladder-baseline"),
        size: table_bytes.max(4),
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        mapped_at_creation: false,
    });
    let baseline_bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q2-ladder-baseline-bg"),
        layout: &baseline.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, &raw_buf),
            bind_buf(1, &baseline_buf),
            bind_buf(2, &p_buf),
        ],
    });
    let mut enc = c.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
        label: Some("q2-ladder-build"),
    });
    {
        let mut pass = begin_pass_with(&mut enc, Some("q2-ladder-build"), None);
        pass.set_pipeline(build);
        pass.set_bind_group(0, &build_bind, &[]);
        pass.dispatch_workgroups((table_len as u32).div_ceil(256).min(MAX_WG), 1, 1);
    }
    {
        let mut pass = begin_pass_with(&mut enc, Some("q2-ladder-baseline"), None);
        pass.set_pipeline(baseline);
        pass.set_bind_group(0, &baseline_bind, &[]);
        pass.dispatch_workgroups((table_len as u32).div_ceil(256).min(MAX_WG), 1, 1);
    }
    let mut got = vec![0f32; table_len];
    let mut gpu_baseline = vec![0f32; table_len];
    if !readback2(
        c,
        enc,
        (&ladders, &mut got),
        (&baseline_buf, &mut gpu_baseline),
    ) {
        eprintln!("q2 ladder cache rejected: GPU table readback failed");
        return None;
    }
    let cpu_mismatches = got
        .iter()
        .zip(expected.iter())
        .filter(|(a, b)| a.to_bits() != b.to_bits())
        .count();
    let gpu_mismatches = got
        .iter()
        .zip(gpu_baseline.iter())
        .filter(|(a, b)| a.to_bits() != b.to_bits())
        .count();
    if gpu_mismatches != 0 {
        let first = got
            .iter()
            .zip(gpu_baseline.iter())
            .position(|(a, b)| a.to_bits() != b.to_bits())
            .unwrap_or(0);
        eprintln!(
            "q2 ladder cache rejected: gpu_baseline_scale_mismatches={} first={} cache=0x{:08x} baseline=0x{:08x}",
            gpu_mismatches,
            first,
            got[first].to_bits(),
            gpu_baseline[first].to_bits()
        );
        return None;
    }
    let build_ms = t0.elapsed().as_secs_f64() * 1e3;
    eprintln!(
        "q2 ladder cache admitted: pairs={} table_bytes={} row_id_bytes={} total_bytes={} build_ms={:.3} gpu_baseline_scale_mismatches=0 cpu_exp2_diagnostic_mismatches={}",
        pairs,
        table_bytes,
        row_id_bytes,
        table_bytes + row_id_bytes,
        build_ms,
        cpu_mismatches
    );
    let fresh = Q2LadderCache {
        model_uid: model.uid(),
        pairs,
        table_bytes,
        row_id_bytes,
        ladders,
        row_ids: row_ids_buf,
        row_id_base: row_bases,
        row_counts,
    };
    let mut cache = c.q2_ladder.lock().unwrap();
    if let Some(old) = cache.as_ref().filter(|v| {
        v.model_uid == model.uid()
            && v.row_counts.get(&idx).is_some_and(|&count| count >= rows)
    }) {
        return Some((
            old.ladders.clone(),
            old.row_ids.clone(),
            *old.row_id_base.get(&idx).unwrap(),
        ));
    }
    let result = (
        fresh.ladders.clone(),
        fresh.row_ids.clone(),
        *fresh.row_id_base.get(&idx).unwrap(),
    );
    *cache = Some(fresh);
    Some(result)
}

/// Encode one q2tp matvec using the global ladder table and target-row IDs.
/// The symbol plane and affine center selector are unchanged from the
/// production shader; only the scale source is replaced.
fn encode_q2tp_ladder_cache(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    ladders: &wgpu::Buffer,
    row_ids: &wgpu::Buffer,
    row_id_base: u32,
    rows: usize,
    cols: usize,
    affine: bool,
) {
    let Some(pipe) = c.q2_ladder_mv.as_ref() else {
        return;
    };
    let gpr = cols / 32;
    let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, row_id_base, affine as u32]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q2-ladder-mv-bg"),
        layout: &pipe.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, weight),
            bind_buf(2, y),
            bind_buf(3, &p_buf),
            bind_buf(5, xs),
            bind_buf(6, ladders),
            bind_buf(7, row_ids),
        ],
    });
    let mut pass = begin_pass_with(enc, Some("q2-ladder-mv"), None);
    pass.set_pipeline(pipe);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(mv_grid((rows as u32).div_ceil(16)), 1, 1);
}

// An OCT-row variant (8 rows a lane, 32 a workgroup, bit-identical rows)
// was written and MEASURED on the RTX 5090 (Qwen3.8-27B): gate/up 995
// against 1056 GB/s, down 846 against 1107, decode 44.0 against 48.7
// tok/s; the 2-bit plane the same (42.5 against 46.8). Half the
// workgroups is the wrong direction on a 170-SM card. Removed; the
// x-LSU theory it tested is dead for the one-vector kernel.

/// Bind group + grid for the fused gate+up+SiLU kernel: gate and up
/// (both [inter x cols] q4tp) against `xs`, activations into `act`.
fn mv_gu_bind(
    c: &Ctx,
    gate: &wgpu::Buffer,
    up: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    act: &wgpu::Buffer,
    inter: usize,
    cols: usize,
) -> (wgpu::BindGroup, u32) {
    let gpr = cols / 32;
    // `_p0` carries the swiglu limit as f32 bits; the token graph's dense
    // FFN has none (0.0), which is what silu_mul_pre gets there too.
    let p_buf = uniform_u32x4(c, [gpr as u32, inter as u32, 0, inter as u32]);
    let layout = c.q4tp_mv16w_gu.get_bind_group_layout(0);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("mv-gu"),
        layout: &layout,
        entries: &[
            bind_buf(0, gate),
            bind_buf(2, act),
            bind_buf(3, &p_buf),
            bind_buf(4, gate),
            bind_buf(5, xs),
            bind_buf(6, up),
            bind_buf(7, up),
        ],
    });
    (bind, mv_grid((inter as u32).div_ceil(8)))
}

/// The main module's source, with the nibble unpack swapped back to
/// integer-to-float conversions under `CMF_MAGIC_UNPACK=0` — the A/B for
/// the magic-mantissa unpack. Same values to the bit either way.
fn wgsl_main_source() -> String {
    if std::env::var("CMF_MAGIC_UNPACK").as_deref() != Ok("0") {
        return WGSL.to_string();
    }
    WGSL.replace(
        "return bitcast<f32>(((w >> sh) & 0xFu) | 0x4B000000u) - 8388616.0;",
        "return f32((w >> sh) & 0xFu) - 8.0;",
    )
    .replace(
        "return bitcast<f32>((((w >> sh) & 3u) << 1u) | 0x4B000000u) - select(8388611.0, 8388610.0, affine != 0u);",
        "return f32(((w >> sh) & 3u) << 1u) - select(3.0, 2.0, affine != 0u);",
    )
}

/// The batched (verify / small-batch prefill) q4tp matvec on int8
/// activations and dp4a (`q4tp_matvec4_bk8`). ON by default: measured
/// on the RTX 5090 with Qwen3.8-27B q4tp it decodes 76 against the f32
/// verify's 66 tok/s (core) and 56.5 against 51.8 on a code prompt, with
/// the same acceptance; the activations carry a per-32 int8 rounding
/// (the Q8_1 grid), so a near-tie can resolve differently from the plain
/// path — `CMF_VERIFY_I8=0` restores the f32 verify and the bit-exact
/// stream.
pub(crate) fn verify_i8_on() -> bool {
    static N: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *N.get_or_init(|| std::env::var("CMF_VERIFY_I8").as_deref() != Ok("0"))
}

/// Quantize `batch` activation vectors (cols wide, f32) to the packed
/// int8 layout, then the dp4a batched matvec: `y[e*rows + r]`.
#[allow(clippy::too_many_arguments)]
fn encode_q4tp_mv4_b_i8(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    batch: usize,
) {
    let gpr = cols / 32;
    let need = (batch * gpr) as u64;
    let (x8, xsb) = {
        let mut g = c.i8x.lock().unwrap();
        let grow = match g.as_ref() {
            Some((_, _, cap)) => *cap < need,
            None => true,
        };
        if grow {
            let cap = need.max(8 * 1024);
            let x8 = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("i8-x8"),
                size: cap * 32,
                usage: wgpu::BufferUsages::STORAGE,
                mapped_at_creation: false,
            });
            let xsb = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("i8-xs"),
                size: cap * 8,
                usage: wgpu::BufferUsages::STORAGE,
                mapped_at_creation: false,
            });
            *g = Some((x8, xsb, cap));
        }
        let (a, b, _) = g.as_ref().unwrap();
        (a.clone(), b.clone())
    };
    // quantizer: one thread per (element, group)
    let qp = uniform_u32x4(c, [gpr as u32, batch as u32, 0, 0]);
    let ql = c.x_quant_i8.get_bind_group_layout(0);
    let qbind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("i8-quant"),
        layout: &ql,
        entries: &[
            bind_buf(3, &qp),
            bind_buf(5, xs),
            bind_buf(11, &x8),
            bind_buf(12, &xsb),
        ],
    });
    let p_buf = q4tp_mv_params(c, gpr, rows, batch);
    let pipe8 = &c.q4tp_mv4_bk8[batch.clamp(2, 8)];
    let layout = pipe8.get_bind_group_layout(0);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("mv-bk8"),
        layout: &layout,
        entries: &[
            bind_buf(0, weight),
            bind_buf(2, y),
            bind_buf(3, &p_buf),
            bind_buf(4, weight),
            bind_buf(9, &x8),
            bind_buf(10, &xsb),
        ],
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(&c.x_quant_i8);
    pass.set_bind_group(0, &qbind, &[]);
    pass.dispatch_workgroups(((batch * gpr) as u32).div_ceil(256), 1, 1);
    pass.set_pipeline(pipe8);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(mv_grid((rows as u32).div_ceil(16)), 1, 1);
}

/// Whether the wgpu device runs on Metal. The batched verify graph is
/// verified on Vulkan (bit-exact against the plain token, the aquarium
/// identical); on wgpu-over-Metal it produced 0 accepted drafts and
/// garbage text on Qwen3.5-0.8B (measured 16.08) while the plain graph
/// was fine — so speculation does not default on there until that is
/// found. `CMF_GRAPH_SPEC=1` still forces it (for the investigation).
pub(crate) fn wgpu_backend_is_metal() -> bool {
    ctx().is_some_and(|c| c.adapter_info.backend == wgpu::Backend::Metal)
}

/// `CMF_MV16W=0`: the wide-row q4tp decode matvec takes the 8-row pair
/// kernel (`q4tp_matvec4`) instead of the 16-row quad (`q4tp_matvec16w`).
fn mv16w_on() -> bool {
    static N: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *N.get_or_init(|| std::env::var("CMF_MV16W").as_deref() != Ok("0"))
}

/// `CMF_BATCH_COOP=1`: the batched prefill's wide GEMMs (k > 16) on the
/// cooperative-matrix kernel with a device-computed activation scale.
/// Opt-in until measured on more than one card.
fn batch_coop_on() -> bool {
    static N: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *N.get_or_init(|| std::env::var("CMF_BATCH_COOP").as_deref() == Ok("1"))
}

fn mv_grid(blocks: u32) -> u32 {
    let cap = mv_grid_cap();
    let g = if cap == 0 { blocks } else { blocks.min(cap) };
    g.max(1).min(MAX_WG)
}

fn graph_split_n() -> usize {
    static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
    *N.get_or_init(|| {
        std::env::var("CMF_GRAPH_SPLIT")
            .ok()
            .and_then(|v| v.parse().ok())
            // 16 pieces, which the `.max(4)` floor turns into four-layer
            // chunks on a 64-layer stack. Measured on Qwen3.6-27B / RTX
            // 5090, medians of three: 49.5 tok/s against 48.4 at 10.
            // Anything above 16 is the SAME configuration (the floor
            // binds) — five such settings measured 44.8/44.7/44.7/44.7/
            // 49.4, which is also this stand's honest ±10% spread on a
            // single end-to-end run, and the reason these numbers are
            // medians. GPU timestamps are steady to 0.3%; tok/s is not.
            .unwrap_or(16)
    })
}

/// Loud decline for the honest bench: tracing is invisible on many
/// stacks (the CLI's Run filter, containers without a collector), and
/// a silent `return false` cost a day of KAT diagnosis. CMF_GPU_DEBUG=1
/// prints every decline reason to stderr directly.
fn graph_decline(why: &str) {
    if std::env::var("CMF_GPU_DEBUG").is_ok() {
        eprintln!("wgpu graph decline: {why}");
    }
    tracing::warn!("wgpu token graph declined: {why}");
}

/// Active weight bytes dispatched — the honest floor's numerator, the
/// instrument the Metal campaign ended on: two days of kernel work on
/// two backends argued with floors computed from file size, and the
/// "2x amplification" turned out to be a looped transformer doing its
/// job. Qwen3.8 is NOT looped, so on the 5090 this counter either
/// finds its own amplification (KV? activations?) or the 52% mystery
/// is real. Either way: measure the floor before arguing with it.
pub static WEIGHT_BYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

fn graph_refused(why: &'static str) {
    use std::sync::atomic::{AtomicBool, Ordering};
    static SAID: AtomicBool = AtomicBool::new(false);
    if !SAID.swap(true, Ordering::Relaxed) {
        tracing::warn!("wgpu token graph declined: {why}");
    }
}

// The optional subgroup pipeline is deliberately fail-closed, but a plain
// `None` is not enough for the component gate: it cannot distinguish a card
// without the feature from a WGSL/module/pipeline rejection.  Keep the full
// scoped diagnostic for the test-visible admission line without making the
// ordinary decode path depend on a tracing subscriber.
fn q2tp_sg_diag_slot() -> &'static std::sync::Mutex<Option<String>> {
    static DIAG: std::sync::OnceLock<std::sync::Mutex<Option<String>>> =
        std::sync::OnceLock::new();
    DIAG.get_or_init(|| std::sync::Mutex::new(None))
}

fn set_q2tp_sg_diag(message: impl Into<String>) {
    *q2tp_sg_diag_slot().lock().unwrap() = Some(message.into());
}

fn q2tp_sg_diag() -> String {
    q2tp_sg_diag_slot()
        .lock()
        .unwrap()
        .clone()
        .unwrap_or_else(|| "not_requested".to_string())
}

/// Device bytes one layer's expert stack wants (gate + up + down, all
/// experts). The builder asks BEFORE uploading to decide where the device
/// prefix ends; `moe_expert_bufs` uses the same arithmetic for its budget
/// refusal, so the two never disagree.
fn moe_pack_bytes(
    n_experts: usize,
    inter: usize,
    hidden: usize,
    q4tp: bool,
    gu_q2: bool,
    dn_q2: bool,
) -> Option<u64> {
    let plen = |rows: usize, cols: usize| -> Option<usize> {
        if q4tp {
            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
        } else {
            Some(rows * (cols / 32) * 18)
        }
    };
    let gu_len = if gu_q2 {
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q2TiledP, &[inter, hidden])?
    } else {
        plen(inter, hidden)?
    };
    let d_len = if dn_q2 {
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q2TiledP, &[hidden, inter])?
    } else {
        plen(hidden, inter)?
    };
    Some((n_experts * (2 * gu_len + d_len)) as u64)
}

fn moe_expert_bufs(
    c: &Ctx,
    model: &Arc<CmfModel>,
    experts: &[(usize, usize, usize)],
    inter: usize,
    hidden: usize,
    q4tp: bool,
    gu_q2: bool,
    dn_q2: bool,
) -> Option<(wgpu::Buffer, wgpu::Buffer, wgpu::Buffer)> {
    use std::sync::atomic::Ordering;
    if hidden % 32 != 0 || inter % 32 != 0 {
        graph_refused("moe_expert_bufs: hidden/inter not 32-aligned");
        return None;
    }
    let bytes = model.primary_bytes();
    let key = (model.uid() as usize, experts.first()?.0);
    if let Some(t) = c.moe_expw.lock().unwrap().get(&key) {
        return Some(t.clone());
    }
    // q4t is 18 B a group flat; q4tp is 16 B of nibbles plus the row
    // params and 5-bit code planes, which the format's own accessor sizes.
    let plen = |rows: usize, cols: usize| -> Option<usize> {
        if q4tp {
            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
        } else {
            Some(rows * (cols / 32) * 18)
        }
    };
    let gu_len = if gu_q2 {
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q2TiledP, &[inter, hidden])?
    } else {
        plen(inter, hidden)?
    };
    let d_len = if dn_q2 {
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q2TiledP, &[hidden, inter])?
    } else {
        plen(hidden, inter)?
    };
    let total = (experts.len() * (2 * gu_len + d_len)) as u64;
    if c.resident.load(Ordering::Relaxed) + total > c.vram_budget {
        // Over budget — the whole graph falls to CPU. Say so ONCE with the
        // numbers: the default budget is a conservative 8 GB on discrete
        // cards, so a 32 GB card running a big MoE lands here and every
        // expert quietly walks the host. That refusal used to be silent.
        use std::sync::atomic::AtomicBool;
        static SAID: AtomicBool = AtomicBool::new(false);
        if !SAID.swap(true, Ordering::Relaxed) {
            let mb = |b: u64| b / (1024 * 1024);
            tracing::warn!(
                "MoE experts need {} MB for this layer on top of {} MB resident, \
                 over the {} MB weight budget — the whole-token graph falls back to \
                 the CPU. Raise it with CMF_GPU_VRAM_MB (e.g. {} on this card).",
                mb(total),
                mb(c.resident.load(Ordering::Relaxed)),
                mb(c.vram_budget),
                mb(c.vram_budget) * 3,
            );
        }
        return None;
    }
    crate::gpu::probe_note_cold();
    // Every byte range this layer ships to the card, for the post-upload
    // evict below.
    let uploaded = std::cell::RefCell::new(Vec::<(usize, usize)>::new());
    let mk = |role: &dyn Fn(&(usize, usize, usize)) -> usize,
              rows: usize,
              cols: usize,
              plen: usize|
     -> Option<wgpu::Buffer> {
        // Check every expert BEFORE allocating: a shape mismatch found
        // halfway used to leave a gigabyte-scale buffer behind.
        let mut offs = Vec::with_capacity(experts.len());
        for t in experts {
            let e = model.tensors.get(role(t))?;
            if *e.shape.first()? != rows || *e.shape.get(1)? != cols || e.nbytes as usize != plen {
                return None;
            }
            let abs = model.entry_abs_offset(e)?;
            bytes.get(abs..abs + plen)?; // in range
            offs.push(abs);
        }
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("moe-experts"),
            size: (experts.len() * plen) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        // Straight from the mapping to the queue, expert by expert. Gathering
        // them into one Vec first meant a 2.2 GB allocation and a full extra
        // memcpy PER LAYER — 94 GB of pointless copying across the release,
        // on top of the 94 GB that has to move anyway.
        // Counted, like every other upload. The rate the profile printed —
        // 7850 MB/s — was measured over the SKELETON only, because the
        // expert stack is ninety per cent of the bytes and goes through
        // this loop rather than through `weight_buffer`.
        let t_up = std::time::Instant::now();
        for (i, &abs) in offs.iter().enumerate() {
            c.queue
                .write_buffer(&b, (i * plen) as u64, &bytes[abs..abs + plen]);
        }
        // Bound transient memory by ONE projection, not a whole expert
        // pack. Near the VRAM limit, keeping gate + up + down staging alive
        // together can OOM even though their final device buffers fit.
        submit_empty(c);
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        UPLOAD_NS.fetch_add(t_up.elapsed().as_nanos() as u64, Ordering::Relaxed);
        UPLOAD_BYTES.fetch_add((offs.len() * plen) as u64, Ordering::Relaxed);
        uploaded
            .borrow_mut()
            .extend(offs.iter().map(|&a| (a, plen)));
        Some(b)
    };
    let (g, u, d) = match (
        mk(&|t| t.0, inter, hidden, gu_len),
        mk(&|t| t.1, inter, hidden, gu_len),
        mk(&|t| t.2, hidden, inter, d_len),
    ) {
        (Some(g), Some(u), Some(d)) => (g, u, d),
        _ => {
            graph_refused("moe_expert_bufs: expert tensor shape/nbytes mismatch");
            return None;
        }
    };
    // Flush the write_buffer staging belt NOW: with 40 MoE layers the
    // pending uploads (~17 GB) would otherwise coexist with their device
    // copies until the graph's first submit — twice the expert weights in
    // memory = device OOM on discrete cards. One submit+wait per layer
    // bounds transient staging to this layer's three buffers.
    submit_empty(c);
    let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
    // The card holds these bytes now, so the host copy is dead weight: the
    // page cache otherwise keeps every resident expert a second time, and on
    // a 112 GB model that second copy IS the machine's RAM (measured: 172 of
    // 176 GB cached).
    //
    // Alternated off/on/off/on from a warmed file, 256 tokens each:
    //
    //   off  94 GB resident  1.0 tok/s      off  91 GB  0.5 tok/s
    //   on    5 GB resident  6.9 tok/s      on    5 GB  6.1 tok/s
    //
    // It is not a trade — holding a second copy of the weights is what was
    // costing the speed. At 94 GB resident the machine has nothing left for
    // the pages it does need, and reclaim churns; at 5 GB it does not.
    //
    // Pairs with `open()` skipping its whole-file `WillNeed` when this is on:
    // reading 104 GB ahead only to drop it behind the uploader had the kernel
    // fetching the same bytes twice. `CMF_UPLOAD_EVICT=0` opts out; discrete
    // only, as on UMA the mapping IS the working copy.
    if c.discrete
        && std::env::var("CMF_UPLOAD_EVICT")
            .map(|v| v != "0")
            .unwrap_or(true)
    {
        model.evict_ranges(&uploaded.borrow());
    }
    c.resident.fetch_add(total, Ordering::Relaxed);
    note_resident_peak(c);
    c.moe_expw
        .lock()
        .unwrap()
        .insert(key, (g.clone(), u.clone(), d.clone()));
    Some((g, u, d))
}

/// The draft pack's expert upload with gate/up requantized q4tp → q2tp on
/// the way through — via the encoder the binary registered. The down
/// projection stays q4tp (the q2tp down kernel does not exist). Cached
/// under the same first-gate key as every pack; only the draft reaches
/// these tensors, so the variant is unambiguous per process.
pub fn moe_expert_bufs_requant_gu(
    model: &Arc<CmfModel>,
    experts: &[(usize, usize, usize)],
    inter: usize,
    hidden: usize,
) -> Option<(wgpu::Buffer, wgpu::Buffer, wgpu::Buffer)> {
    use std::sync::atomic::Ordering;
    let c = ctx()?;
    let enc2 = *crate::dsv4::DSPARK_Q2TP_ENCODE.get()?;
    if hidden % 32 != 0 || inter % 32 != 0 {
        return None;
    }
    let bytes = model.primary_bytes();
    let key = (model.uid() as usize, experts.first()?.0);
    if let Some(t) = c.moe_expw.lock().unwrap().get(&key) {
        return Some(t.clone());
    }
    let gu_len =
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q2TiledP, &[inter, hidden])?;
    let d_len =
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[hidden, inter])?;
    let total = (experts.len() * (2 * gu_len + d_len)) as u64;
    if c.resident.load(Ordering::Relaxed) + total > c.vram_budget {
        return None;
    }
    let t0 = std::time::Instant::now();
    let mut vals = vec![0.0f32; inter * hidden];
    let mk_gu = |role: &dyn Fn(&(usize, usize, usize)) -> usize,
                 vals: &mut Vec<f32>|
     -> Option<wgpu::Buffer> {
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("dspark-experts-q2"),
            size: (experts.len() * gu_len) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        for (i, t) in experts.iter().enumerate() {
            let e = model.tensors.get(role(t))?;
            if e.dtype != cortiq_core::TensorDtype::Q4TiledP || e.shape != [inter, hidden] {
                return None;
            }
            let abs = model.entry_abs_offset(e)?;
            let src = bytes.get(abs..abs + e.nbytes as usize)?;
            cortiq_core::quant::dequant_q4tp(src, inter, hidden, vals);
            let q2 = enc2(vals, inter, hidden);
            if q2.len() != gu_len {
                return None;
            }
            c.queue.write_buffer(&b, (i * gu_len) as u64, &q2);
        }
        submit_empty(c);
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        Some(b)
    };
    let mk_d = || -> Option<wgpu::Buffer> {
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("dspark-experts-dn"),
            size: (experts.len() * d_len) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        for (i, t) in experts.iter().enumerate() {
            let e = model.tensors.get(t.2)?;
            if e.dtype != cortiq_core::TensorDtype::Q4TiledP || e.nbytes as usize != d_len {
                return None;
            }
            let abs = model.entry_abs_offset(e)?;
            c.queue
                .write_buffer(&b, (i * d_len) as u64, bytes.get(abs..abs + d_len)?);
        }
        submit_empty(c);
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        Some(b)
    };
    let (g, u, d) = match (
        mk_gu(&|t| t.0, &mut vals),
        mk_gu(&|t| t.1, &mut vals),
        mk_d(),
    ) {
        (Some(g), Some(u), Some(d)) => (g, u, d),
        _ => return None,
    };
    tracing::info!(
        "DSpark: реквант q4tp→q2tp {} экспертов за {:.1} с ({} МБ на карте)",
        experts.len(),
        t0.elapsed().as_secs_f64(),
        total / (1024 * 1024),
    );
    c.resident.fetch_add(total, Ordering::Relaxed);
    note_resident_peak(c);
    c.moe_expw
        .lock()
        .unwrap()
        .insert(key, (g.clone(), u.clone(), d.clone()));
    Some((g, u, d))
}

/// GPU enabled and initialized?
pub fn enabled() -> bool {
    ctx().is_some()
}

/// Probe helper: true — tensor `idx`'s weights are already resident;
/// false — not yet (with `may_upload`, the upload happens NOW within the
/// budget, without a dispatch, so the next touch is warm) or the tensor
/// can't be resolved.
pub fn q8_resident_or_upload(model: &Arc<CmfModel>, idx: usize, may_upload: bool) -> bool {
    let Some(c) = ctx() else { return false };
    let entry = &model.tensors[idx];
    let rows_total = entry.shape.first().copied().unwrap_or(0);
    let cols = entry.shape.get(1).copied().unwrap_or(0);
    if rows_total == 0 || cols == 0 {
        return false;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return false;
    };
    let bytes = model.primary_bytes();
    // q8_2f's graph kernel reads the row and column scale planes from the
    // same device buffer.  The per-op q8 kernel only needs the int8 body,
    // but caching that shorter slice under the shared (model,tensor) key
    // poisoned a later whole-token graph lookup: it found the resident body,
    // read both scale planes out of bounds, and returned all-zero logits for
    // large vocabulary heads.  Keep the complete packed payload resident so
    // both consumers can safely share the cache entry.
    let payload_len = if entry.dtype == cortiq_core::TensorDtype::Q8_2f {
        entry.nbytes as usize
    } else {
        rows_total * cols
    };
    if abs + payload_len > bytes.len() {
        return false;
    }
    let key = (model.uid() as usize, idx);
    if c.weight_bufs
        .lock()
        .unwrap()
        .get(&key)
        .is_some_and(|e| e.bytes >= payload_len as u64)
    {
        return true;
    }
    if may_upload {
        let _ = weight_buffer_l(
            c,
            key,
            &bytes[abs..abs + payload_len],
            layer_of_name(&model.tensors[idx].name),
        );
    }
    false
}

/// q8_row/q8_2f matvec on the GPU, rows [row0, row0+rows). `xs` are already
/// prescaled activations. false = could not (the caller falls back to CPU).
#[allow(clippy::too_many_arguments)]
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 {
    let Some(c) = ctx() else { return false };
    if cols % 4 != 0 || rows == 0 {
        return false;
    }
    let entry = &model.tensors[idx];
    let rows_total = entry.shape.first().copied().unwrap_or(0);
    if rows_total < row0 + rows {
        return false;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return false; // neighboring shard — different mapping; CPU
    };
    let bytes = model.primary_bytes();
    let payload_len = if entry.dtype == cortiq_core::TensorDtype::Q8_2f {
        entry.nbytes as usize
    } else {
        rows_total * cols
    };
    if abs + payload_len > bytes.len() {
        return false;
    }
    let full_quant = &bytes[abs..abs + payload_len];
    let key = (model.uid() as usize, idx);
    dispatch_matvec(
        c,
        Some(key),
        full_quant,
        row0,
        row_scale,
        xs,
        rows,
        cols,
        out,
    )
}

/// matvec kernel: resident weights of the WHOLE tensor + row0 offset, rs, xs,
/// dispatch, readback. `weight_key = None` — no cache (test).
#[allow(clippy::too_many_arguments)]
fn dispatch_matvec(
    c: &Ctx,
    weight_key: Option<(usize, usize)>,
    full_quant: &[u8],
    row0: usize,
    row_scale: &[f32],
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    if row_scale.len() < rows || xs.len() < cols || full_quant.len() < (row0 + rows) * cols {
        return false;
    }
    let q_buf = match weight_key {
        Some(k) => match weight_buffer(c, k, full_quant) {
            Some(b) => b,
            None => return false, // over VRAM budget — honest CPU path
        },
        None => c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("q8-weights"),
                contents: full_quant,
                usage: wgpu::BufferUsages::STORAGE,
            }),
    };
    let make_rs = || {
        c.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("q8-rs"),
                contents: bytemuck::cast_slice(&row_scale[..rows]),
                usage: wgpu::BufferUsages::STORAGE,
            })
    };
    let rs_buf = match weight_key {
        Some((base, idx)) => c
            .rs_bufs
            .lock()
            .unwrap()
            .entry((base, (idx, row0)))
            .or_insert_with(|| {
                crate::gpu::probe_note_cold();
                make_rs()
            })
            .clone(),
        None => make_rs(),
    };

    // Pooled scratch for the whole op (encode → submit → poll).
    let mut sc = c.scratch.lock().unwrap();
    let xs_buf = Scratch::ensure(
        &c.device,
        &mut sc.xs,
        (cols * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "q8-xs",
    );
    c.queue
        .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&xs[..cols]));
    let y_size = (rows * 4) as u64;
    let y_buf = Scratch::ensure(
        &c.device,
        &mut sc.y,
        y_size,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "q8-y",
    );
    let params = [
        (cols / 4) as u32,
        rows as u32,
        (row0 * cols / 4) as u32,
        0u32,
    ];
    let p_buf = match &sc.params {
        Some(b) => b.clone(),
        None => {
            let b = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("q8-params"),
                size: 16,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            sc.params = Some(b.clone());
            b
        }
    };
    c.queue
        .write_buffer(&p_buf, 0, bytemuck::cast_slice(&params));
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        y_size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "q8-stage",
    );

    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q8-bg"),
        layout: &c.layout,
        entries: &[
            bind_buf(0, &q_buf),
            bind_buf(1, &xs_buf),
            bind_buf(2, &rs_buf),
            bind_buf(3, &y_buf),
            bind_buf(4, &p_buf),
        ],
    });

    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("q8") });
    {
        let mut pass = begin_pass_with(&mut enc, Some("q8"), None);
        // The staged-activation arm reads xs from workgroup memory, so
        // it only fits while xs does: 768 vec4 = 3072 columns.
        let tiled = q8mv_tiled() && cols <= 3072;
        pass.set_pipeline(if tiled { &c.matvec_tiled } else { &c.matvec });
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1); // grid-stride over rows
    }
    let ok = readback(c, enc, &y_buf, &stage_buf, y_size, &mut out[..rows]);
    drop(sc);
    ok
}

/// q1t (base+overlay) / q4_block matvec on wgpu — raw f32 x, scales embedded.
/// The kernel decodes bytes out of the u32 weight buffer; params carry
/// (gpr, rows, cols). Weights resident under the shared VRAM budget.
pub fn q1t_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    q1t_like(model, idx, xs, rows, cols, out, false)
}

/// q4_block matvec on wgpu (nibbles + trailing scales, no overlay).
pub fn q4b_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    q1t_like(model, idx, xs, rows, cols, out, true)
}

fn q1t_like(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
    q4: bool,
) -> bool {
    let Some(c) = ctx() else { return false };
    let gpr = cols / 32;
    if rows == 0 || cols % 32 != 0 || xs.len() < cols || out.len() < rows {
        return false;
    }
    let entry = &model.tensors[idx];
    if entry.shape.first().copied().unwrap_or(0) < rows {
        return false;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return false;
    };
    let bytes = model.primary_bytes();
    let plen = entry.nbytes as usize;
    WEIGHT_BYTES.fetch_add(plen as u64, std::sync::atomic::Ordering::Relaxed);
    // sanity: the base must at least fit (q1t base 9 B/group, q4b 18 B/group).
    let min_base = if q4 { rows * gpr * 18 } else { rows * gpr * 9 };
    if plen < min_base || abs + plen > bytes.len() {
        return false;
    }
    let pipeline = if q4 { &c.q4b } else { &c.q1t };
    dispatch_q1t(
        c,
        pipeline,
        {
            note_layer((model.uid() as usize, idx), &model.tensors[idx].name);
            Some((model.uid() as usize, idx))
        },
        &bytes[abs..abs + plen],
        xs,
        rows,
        cols,
        out,
    )
}

#[allow(clippy::too_many_arguments)]
fn dispatch_q1t(
    c: &Ctx,
    pipeline: &wgpu::ComputePipeline,
    weight_key: Option<(usize, usize)>,
    payload: &[u8],
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    let gpr = cols / 32;
    let q_buf = match weight_key {
        Some(k) => match weight_buffer(c, k, payload) {
            Some(b) => b,
            None => return false,
        },
        None => c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("q1t-weights"),
                contents: payload,
                usage: wgpu::BufferUsages::STORAGE,
            }),
    };
    let mut sc = c.scratch.lock().unwrap();
    let xs_buf = Scratch::ensure(
        &c.device,
        &mut sc.xs,
        (cols * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "q1t-xs",
    );
    c.queue
        .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&xs[..cols]));
    let y_size = (rows * 4) as u64;
    let y_buf = Scratch::ensure(
        &c.device,
        &mut sc.y,
        y_size,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "q1t-y",
    );
    let params = [gpr as u32, rows as u32, cols as u32, 0u32];
    let p_buf = match &sc.params {
        Some(b) => b.clone(),
        None => {
            let b = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("q1t-params"),
                size: 16,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            sc.params = Some(b.clone());
            b
        }
    };
    c.queue
        .write_buffer(&p_buf, 0, bytemuck::cast_slice(&params));
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        y_size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "q1t-stage",
    );
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q1t-bg"),
        // Must be THIS pipeline's layout (wgpu treats each pipeline's layout as
        // distinct even when structurally identical to q1's).
        layout: &pipeline.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, &q_buf),
            bind_buf(1, &xs_buf),
            bind_buf(2, &y_buf),
            bind_buf(3, &p_buf),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("q1t") });
    {
        let mut pass = begin_pass_with(&mut enc, Some("q1t"), None);
        pass.set_pipeline(pipeline);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
    }
    let ok = readback(c, enc, &y_buf, &stage_buf, y_size, &mut out[..rows]);
    drop(sc);
    ok
}

/// q1 matvec: raw f32 activations, tile-embedded scales (no rs buffer).
/// Weights resident under the same VRAM budget as q8; false = CPU path.
pub fn q1_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    let gpr = cols / 32;
    if rows == 0 || cols % 32 != 0 || gpr % 2 != 0 || xs.len() < cols || out.len() < rows {
        return false;
    }
    let entry = &model.tensors[idx];
    if entry.shape.first().copied().unwrap_or(0) < rows {
        return false;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return false;
    };
    let bytes = model.primary_bytes();
    let plen = rows * gpr * 6;
    if abs + plen > bytes.len() {
        return false;
    }
    dispatch_q1(
        c,
        {
            note_layer((model.uid() as usize, idx), &model.tensors[idx].name);
            Some((model.uid() as usize, idx))
        },
        &bytes[abs..abs + plen],
        xs,
        rows,
        cols,
        out,
    )
}

/// GPU RMSNorm of one row — the token-graph building block that keeps the
/// hidden state resident across the norm→matvec boundary. One workgroup,
/// direct buffers (no residency cache). Returns false without a GPU context.
pub fn rmsnorm_row(x: &[f32], w: &[f32], out: &mut [f32], gemma: bool, eps: f32) -> bool {
    let Some(c) = ctx() else { return false };
    let n = x.len();
    if n == 0 || w.len() < n || out.len() < n {
        return false;
    }
    let x_b = storage_bytes(c, bytemuck::cast_slice(x));
    let w_b = storage_bytes(c, bytemuck::cast_slice(&w[..n]));
    let o_b = rw_f32(c, n, true);
    let p_buf = uniform_u32x4(c, [n as u32, gemma as u32, eps.to_bits(), 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("rms-bg"),
        layout: &c.layout_rmsnorm,
        entries: &[
            bind_buf(0, &x_b),
            bind_buf(1, &w_b),
            bind_buf(2, &o_b),
            bind_buf(3, &p_buf),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("rms") });
    {
        let mut pass = begin_pass_with(&mut enc, Some("rms"), None);
        pass.set_pipeline(&c.rmsnorm);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(1, 1, 1);
    }
    let size = (n * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "rms-stage",
    );
    let ok = readback(c, enc, &o_b, &stage, size, &mut out[..n]);
    drop(sc);
    ok
}

/// GPU RoPE + qk-norm + gate-split building block (bring-up / parity). One
/// workgroup per head; writes qout[nh·hd], k in place[nkv·hd], gout[nh·hd].
/// qnw/knw must be hd-long (dummy ok if the norm flag is off), invf rd/2-long.
#[allow(clippy::too_many_arguments)]
pub fn attn_rope_qkn_gpu(
    qraw: &[f32],
    k_in: &[f32],
    qnw: &[f32],
    knw: &[f32],
    invf: &[f32],
    nh: usize,
    nkv: usize,
    hd: usize,
    rd: usize,
    pos: usize,
    flags: u32,
    eps: f32,
    qout: &mut [f32],
    k_out: &mut [f32],
    gout: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    let qraw_b = storage_bytes(c, bytemuck::cast_slice(qraw));
    let k_b = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("rq-k"),
            contents: bytemuck::cast_slice(&k_in[..nkv * hd]),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        });
    let qout_b = rw_f32(c, nh * hd, true);
    let gout_b = rw_f32(c, nh * hd, true);
    let qnw_b = storage_bytes(c, bytemuck::cast_slice(qnw));
    let knw_b = storage_bytes(c, bytemuck::cast_slice(knw));
    let invf_b = storage_bytes(c, bytemuck::cast_slice(invf));
    let p_data = [
        nh as u32,
        nkv as u32,
        hd as u32,
        rd as u32,
        pos as u32,
        flags,
        eps.to_bits(),
        0u32,
    ];
    let p_buf = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("rq-p"),
            contents: bytemuck::cast_slice(&p_data),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("rq-bg"),
        layout: &c.layout_attn_rope,
        entries: &[
            bind_buf(0, &qraw_b),
            bind_buf(1, &k_b),
            bind_buf(2, &qout_b),
            bind_buf(3, &gout_b),
            bind_buf(4, &qnw_b),
            bind_buf(5, &knw_b),
            bind_buf(6, &invf_b),
            bind_buf(7, &p_buf),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("rq") });
    {
        let mut pass = begin_pass_with(&mut enc, Some("rq"), None);
        pass.set_pipeline(&c.attn_rope);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((nh + nkv) as u32, 1, 1);
    }
    let mk_stage = |n: usize| {
        c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("rq-stage"),
            size: (n * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        })
    };
    let sq = mk_stage(nh * hd);
    let sk = mk_stage(nkv * hd);
    let sgt = mk_stage(nh * hd);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&qout_b, 0, &sq, 0, (nh * hd * 4) as u64);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&k_b, 0, &sk, 0, (nkv * hd * 4) as u64);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&gout_b, 0, &sgt, 0, (nh * hd * 4) as u64);
    submit(c, finish_enc(enc));
    for s in [&sq, &sk, &sgt] {
        s.slice(..).map_async(wgpu::MapMode::Read, |_| {});
    }
    if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        return false;
    }
    let (Ok(dq), Ok(dk), Ok(dg)) = (
        sq.slice(..).get_mapped_range(),
        sk.slice(..).get_mapped_range(),
        sgt.slice(..).get_mapped_range(),
    ) else {
        return false;
    };
    qout[..nh * hd].copy_from_slice(bytemuck::cast_slice(&dq[..nh * hd * 4]));
    k_out[..nkv * hd].copy_from_slice(bytemuck::cast_slice(&dk[..nkv * hd * 4]));
    gout[..nh * hd].copy_from_slice(bytemuck::cast_slice(&dg[..nh * hd * 4]));
    true
}

/// GPU grouped decode attention (bring-up / parity). K/V caches are laid out
/// [nkv, cap, hd]; attends q[nh·hd] over the first `n` rows, writes out[nh·hd].
#[allow(clippy::too_many_arguments)]
pub fn gqa_attend_gpu(
    q: &[f32],
    kcache: &[f32],
    vcache: &[f32],
    nh: usize,
    hpk: usize,
    hd: usize,
    cap: usize,
    n: usize,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if hd % 4 != 0 || hd > c.hd_cap {
        return false; // vec4 K/V reads; hd_cap = workgroup-storage limit
    }
    let q_b = storage_bytes(c, bytemuck::cast_slice(q));
    let k_b = storage_bytes(c, bytemuck::cast_slice(kcache));
    let v_b = storage_bytes(c, bytemuck::cast_slice(vcache));
    let o_b = rw_f32(c, nh * hd, true);
    let p_buf = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("at-p"),
            contents: bytemuck::cast_slice(&[
                nh as u32,
                hpk as u32,
                hd as u32,
                cap as u32,
                n as u32,
                (1.0 / (hd as f32).sqrt()).to_bits(),
                0u32,
                0u32,
            ]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("at-bg"),
        layout: attend_pipes(c, hd).1,
        entries: &[
            bind_buf(0, &q_b),
            bind_buf(1, &k_b),
            bind_buf(2, &v_b),
            bind_buf(3, &o_b),
            bind_buf(4, &p_buf),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("at") });
    {
        let mut pass = begin_pass_with(&mut enc, Some("at"), None);
        pass.set_pipeline(attend_pipes(c, hd).0);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(nh as u32, 1, 1);
    }
    let size = (nh * hd * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "at-stage",
    );
    let ok = readback(c, enc, &o_b, &stage, size, &mut out[..nh * hd]);
    drop(sc);
    ok
}

/// The split-K decode attend (part + merge) as one call, on either the
/// per-head kernel or the GQA-shared one — the parity harness for the
/// two against the CPU reference and each other.
pub fn gqa_attend_split_gpu(
    q: &[f32],
    kcache: &[f32],
    vcache: &[f32],
    nh: usize,
    hpk: usize,
    hd: usize,
    cap: usize,
    n: usize,
    gqa: bool,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if hd % 4 != 0 || hd > c.hd_cap || nh % hpk != 0 {
        return false;
    }
    if gqa && (c.attend_gpart.is_none() || hpk > 8 || hd > 256) {
        return false;
    }
    let ck = if gqa { ATTEND_GCK } else { ATTEND_CK };
    let nc = cap.div_ceil(ck);
    let nc_used = n.div_ceil(ck);
    let q_b = storage_bytes(c, bytemuck::cast_slice(q));
    let k_b = storage_bytes(c, bytemuck::cast_slice(kcache));
    let v_b = storage_bytes(c, bytemuck::cast_slice(vcache));
    let o_b = rw_f32(c, nh * hd, true);
    let pacc = rw_f32(c, nh * nc * hd, false);
    let pml = rw_f32(c, nh * nc * 2, false);
    let p_buf = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("ap-p"),
            contents: bytemuck::cast_slice(&[
                nh as u32,
                hpk as u32,
                hd as u32,
                cap as u32,
                n as u32,
                ck as u32,
                nc as u32,
                (1.0 / (hd as f32).sqrt()).to_bits(),
            ]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let (pp, pl) = if gqa {
        (
            c.attend_gpart.as_ref().unwrap(),
            c.layout_attend_gpart.as_ref().unwrap(),
        )
    } else {
        attend_part_pipes(c, hd)
    };
    let bg_part = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("ap-bg"),
        layout: pl,
        entries: &[
            bind_buf(0, &q_b),
            bind_buf(1, &k_b),
            bind_buf(2, &v_b),
            bind_buf(3, &pacc),
            bind_buf(4, &pml),
            bind_buf(5, &p_buf),
        ],
    });
    let bg_merge = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("am-bg"),
        layout: &c.layout_attend_merge,
        entries: &[
            bind_buf(3, &pacc),
            bind_buf(4, &pml),
            bind_buf(5, &p_buf),
            bind_buf(6, &o_b),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("ap") });
    {
        let mut pass = begin_pass_with(&mut enc, Some("ap"), None);
        pass.set_pipeline(pp);
        pass.set_bind_group(0, &bg_part, &[]);
        let gx = if gqa { (nh / hpk) as u32 } else { nh as u32 };
        pass.dispatch_workgroups(gx, nc_used as u32, 1);
        pass.set_pipeline(&c.attend_merge);
        pass.set_bind_group(0, &bg_merge, &[]);
        pass.dispatch_workgroups(nh as u32, 1, 1);
    }
    let size = (nh * hd * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "ap-stage",
    );
    let ok = readback(c, enc, &o_b, &stage, size, &mut out[..nh * hd]);
    drop(sc);
    ok
}

/// Resident q1 weight for a model tensor (cached in VRAM by (ptr, idx)).
/// Returns (buffer, rows, cols). None on budget/shape refusal.
/// Is wgpu initialized on a DISCRETE adapter? Gates the whole-token
/// graph default (see gpu::wgpu_graph_default).
pub fn discrete_active() -> bool {
    ctx().map(|c| c.discrete).unwrap_or(false)
}

/// A wgpu context of any kind is up — the graph-default check for
/// desktop-class UMA (Apple silicon), where discreteness is the wrong
/// question.
pub(crate) fn adapter_up() -> bool {
    ctx().is_some()
}

fn q1_weight(c: &Ctx, model: &Arc<CmfModel>, idx: usize) -> Option<(wgpu::Buffer, usize, usize)> {
    let entry = model.tensors.get(idx)?;
    let rows = *entry.shape.first()?;
    let cols = *entry.shape.get(1)?;
    if cols % 32 != 0 {
        return None;
    }
    let abs = model.entry_abs_offset(entry)?;
    let bytes = model.primary_bytes();
    let plen = rows * (cols / 32) * 6;
    if abs + plen > bytes.len() {
        return None;
    }
    let buf = weight_buffer_l(
        c,
        (model.uid() as usize, idx),
        &bytes[abs..abs + plen],
        layer_of_name(&model.tensors[idx].name),
    )?;
    Some((buf, rows, cols))
}

/// A q4_tiled / q4tp weight as one device buffer — the whole tensor, since
/// both layouts keep their scales inside (q4t) or in trailing planes (q4tp)
/// and the kernels index them from the same base.
/// `CMF_RES_WHO=1`: every 256th arena miss prints the tensor NAME — the
/// counters say how much is fetched, this says by whom.
/// The idx half of the arena key names the tensor via the model directory
/// at exit; here it is enough to know WHICH indices carry the traffic.
fn res_who_key(key: (usize, usize)) {
    use std::sync::atomic::{AtomicU64, Ordering};
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    static N: AtomicU64 = AtomicU64::new(0);
    if !*ON.get_or_init(|| std::env::var("CMF_RES_WHO").is_ok()) {
        return;
    }
    if N.fetch_add(1, Ordering::Relaxed) % 256 == 0 {
        tracing::info!("res miss idx {}", key.1);
    }
}

fn res_who(name: &str) {
    use std::sync::atomic::{AtomicU64, Ordering};
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    static N: AtomicU64 = AtomicU64::new(0);
    if !*ON.get_or_init(|| std::env::var("CMF_RES_WHO").is_ok()) {
        return;
    }
    if N.fetch_add(1, Ordering::Relaxed) % 256 == 0 {
        tracing::info!("res miss: {name}");
    }
}

fn tile_weight(c: &Ctx, model: &Arc<CmfModel>, idx: usize) -> Option<(wgpu::Buffer, usize, usize)> {
    let entry = model.tensors.get(idx)?;
    let rows = *entry.shape.first()?;
    let cols = *entry.shape.get(1)?;
    if cols % 32 != 0 {
        return None;
    }
    let abs = model.entry_abs_offset(entry)?;
    let bytes = model.primary_bytes();
    let plen = entry.nbytes as usize;
    if abs + plen > bytes.len() {
        return None;
    }
    if !weight_resident(c, (model.uid() as usize, idx)) {
        res_who(&model.tensors[idx].name);
    }
    let buf = weight_buffer_l(
        c,
        (model.uid() as usize, idx),
        &bytes[abs..abs + plen],
        layer_of_name(&model.tensors[idx].name),
    )?;
    Some((buf, rows, cols))
}

/// Production drop-in for the attention sub-block on the token graph: takes
/// the already-normed hidden and returns the O-projection output (pre-
/// residual) — exactly where `qwen_attention` slots in. QKV/O weights are
/// resident (VRAM cache), the K/V cache is a persistent device mirror keyed
/// by (kv_id, layer) that is synced once from the CPU cache (prefill) then
/// appended to each token. Everything runs in ONE command encoder; only the
/// attention output reads back. false = refusal (caller keeps the CPU path).
#[allow(clippy::too_many_arguments)]
pub fn attn_dropin_gpu(
    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>],
    attn_out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if pos >= cap || hd % 4 != 0 || hd > c.hd_cap {
        return false; // vec4 K/V reads; hd_cap = workgroup-storage limit
    }
    let cap = kv_capacity(cap, pos + 1);
    let (wq, rq, cq) = q1_weight(c, model, wq_idx).unwrap_or((
        c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: 4,
            usage: wgpu::BufferUsages::STORAGE,
            mapped_at_creation: false,
        }),
        0,
        0,
    ));
    if rq != nh * hd || cq != hidden {
        return false; // gated arch (e.g. output_gate doubles rows) → CPU path
    }
    let Some((wk, _, _)) = q1_weight(c, model, wk_idx) else {
        return false;
    };
    let Some((wv, _, _)) = q1_weight(c, model, wv_idx) else {
        return false;
    };
    let Some((wo, ro, co)) = q1_weight(c, model, wo_idx) else {
        return false;
    };
    if ro != hidden || co != nh * hd {
        return false;
    }
    // Device K/V mirror (persist across tokens).
    let mut kvm = c.attn_kv.lock().unwrap();
    let entry = kv_mirror_ensure(c, &mut kvm, (kv_id, layer), nkv, hd, cap);
    // Sync prefill history 0..pos from the CPU cache (once).
    if entry.synced < pos {
        for h in 0..nkv {
            let src_k = &cpu_k[h];
            let src_v = &cpu_v[h];
            let from = entry.synced;
            let take = pos.min(src_k.len() / hd);
            if take > from {
                let off = ((h * cap + from) * hd * 4) as u64;
                c.queue.write_buffer(
                    &entry.k,
                    off,
                    bytemuck::cast_slice(&src_k[from * hd..take * hd]),
                );
                c.queue.write_buffer(
                    &entry.v,
                    off,
                    bytemuck::cast_slice(&src_v[from * hd..take * hd]),
                );
            }
        }
        entry.synced = pos;
    }
    let kbuf = entry.k.clone();
    let vbuf = entry.v.clone();
    drop(kvm);

    let stor = |data: &[u8]| {
        c.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: data,
                usage: wgpu::BufferUsages::STORAGE,
            })
    };
    let dummy = vec![0f32; hd];
    let qnw_b = stor(bytemuck::cast_slice(q_norm.unwrap_or(&dummy)));
    let knw_b = stor(bytemuck::cast_slice(k_norm.unwrap_or(&dummy)));
    let invf_b = stor(bytemuck::cast_slice(invf));
    let normed_b = stor(bytemuck::cast_slice(&normed[..hidden]));
    let qraw_b = rw_f32(c, nh * hd, false);
    let k_b = rw_f32(c, nkv * hd, false);
    let v_b = rw_f32(c, nkv * hd, false);
    let qout_b = rw_f32(c, nh * hd, false);
    let gout_b = rw_f32(c, nh * hd, false);
    let attn_b = rw_f32(c, nh * hd, false);
    let o_b = rw_f32(c, hidden, true);
    let flags = if q_norm.is_some() { 2u32 } else { 0 }
        | if k_norm.is_some() { 4 } else { 0 }
        | if gemma { 8 } else { 0 }
        | if late_qk_norm { 32 } else { 0 };
    let unif = |d: &[u32]| {
        c.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(d),
                usage: wgpu::BufferUsages::UNIFORM,
            })
    };
    let bg = |layout: &wgpu::BindGroupLayout, bufs: &[&wgpu::Buffer]| {
        let e: Vec<_> = bufs
            .iter()
            .enumerate()
            .map(|(i, b)| bind_buf(i as u32, b))
            .collect();
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout,
            entries: &e,
        })
    };
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("attn-dropin"),
        });
    let go =
        |enc: &mut wgpu::CommandEncoder, p: &wgpu::ComputePipeline, b: &wgpu::BindGroup, g: u32| {
            let mut pass = begin_pass(enc);
            pass.set_pipeline(p);
            pass.set_bind_group(0, b, &[]);
            pass.dispatch_workgroups(g, 1, 1);
        };
    encode_matvec_q1(c, &mut enc, &wq, &normed_b, &qraw_b, nh * hd, hidden);
    encode_matvec_q1(c, &mut enc, &wk, &normed_b, &k_b, nkv * hd, hidden);
    encode_matvec_q1(c, &mut enc, &wv, &normed_b, &v_b, nkv * hd, hidden);
    let rq_p = unif(&[
        nh as u32,
        nkv as u32,
        hd as u32,
        rd as u32,
        pos as u32,
        flags,
        eps.to_bits(),
        0,
    ]);
    go(
        &mut enc,
        &c.attn_rope,
        &bg(
            &c.layout_attn_rope,
            &[
                &qraw_b, &k_b, &qout_b, &gout_b, &qnw_b, &knw_b, &invf_b, &rq_p,
            ],
        ),
        (nh + nkv) as u32,
    );
    let kv_p = unif(&[nkv as u32, hd as u32, cap as u32, pos as u32]);
    go(
        &mut enc,
        &c.kv_append,
        &bg(&c.layout_kv, &[&k_b, &v_b, &kbuf, &vbuf, &kv_p]),
        ((nkv * hd) as u32).div_ceil(256),
    );
    let at_p = unif(&[
        nh as u32,
        (nh / nkv) as u32,
        hd as u32,
        cap as u32,
        (pos + 1) as u32,
        (1.0 / (hd as f32).sqrt()).to_bits(),
        0,
        0,
    ]);
    {
        let (ap, al) = attend_pipes(c, hd);
        go(
            &mut enc,
            ap,
            &bg(al, &[&qout_b, &kbuf, &vbuf, &attn_b, &at_p]),
            nh as u32,
        );
    }
    encode_matvec_q1(c, &mut enc, &wo, &attn_b, &o_b, hidden, nh * hd);
    let size = (hidden * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dropin-stage",
    );
    let ok = readback(c, enc, &o_b, &stage, size, &mut attn_out[..hidden]);
    drop(sc);
    if ok {
        c.attn_kv
            .lock()
            .unwrap()
            .get_mut(&(kv_id, layer))
            .map(|m| m.synced = pos + 1);
    }
    ok
}

/// Device bytes held alive by one layer of a whole-token graph.
///
/// This is deliberately independent of the residency counter. While a graph
/// is being assembled every `GMat` keeps its `wgpu::Buffer` alive; removing an
/// entry from the LRU therefore does not release the allocation until the
/// graph finishes. A dense checkpoint larger than the arena used to appear to
/// obey `CMF_GPU_VRAM_MB` while physically retaining the complete model and
/// OOMing the driver. The graph builders use this real live-set size to choose
/// a device prefix (decode) or decline the all-layers batch graph (prefill).
fn graph_layer_payload_bytes(
    model: &Arc<CmfModel>,
    layer: &crate::gpu::GraphLayer<'_>,
) -> Option<u64> {
    let mut tensor_seen = std::collections::HashSet::<usize>::new();
    let mut f32_seen = std::collections::HashSet::<usize>::new();
    let mut total = 0u64;
    let mut add = |w: &crate::gpu::GraphW<'_>| -> Option<()> {
        let bytes = if w.kind == 4 {
            if !f32_seen.insert(w.data.as_ptr() as usize) {
                return Some(());
            }
            w.data.len().checked_mul(4)? as u64
        } else {
            if !tensor_seen.insert(w.idx) {
                return Some(());
            }
            let e = model.tensors.get(w.idx)?;
            // q8_row keeps its decoded f32 row-scale side buffer beside the
            // byte body. Counting it in addition to `nbytes` is slightly
            // conservative (the file has f16 scales), which is the safe side
            // of a hard VRAM boundary.
            (e.nbytes as u64).checked_add(if w.kind == 0 {
                (w.row_scale.len() as u64).checked_mul(4)?
            } else {
                0
            })?
        };
        total = total.checked_add(bytes)?;
        Some(())
    };
    match &layer.attn {
        crate::gpu::GraphAttn::Full { wq, wk, wv, wo, .. } => {
            add(wq)?;
            add(wk)?;
            add(wv)?;
            add(wo)?;
        }
        crate::gpu::GraphAttn::Gdn {
            qkv, z, a, b, out, ..
        } => {
            add(qkv)?;
            add(z)?;
            add(a)?;
            add(b)?;
            add(out)?;
        }
        crate::gpu::GraphAttn::ShortConv { inp, out, .. } => {
            add(inp)?;
            add(out)?;
        }
    }
    match &layer.ffn {
        crate::gpu::GraphFfn::Dense { gate, up, down } => {
            add(gate)?;
            add(up)?;
            add(down)?;
        }
        crate::gpu::GraphFfn::Moe {
            router,
            shared_gate,
            experts,
            has_shared,
            ..
        } => {
            add(router)?;
            if *has_shared {
                add(shared_gate)?;
            }
            for &(g, u, d) in experts {
                for idx in [g, u, d] {
                    if tensor_seen.insert(idx) {
                        total = total.checked_add(model.tensors.get(idx)?.nbytes as u64)?;
                    }
                }
            }
        }
    }
    Some(total)
}

fn graph_stack_payload_bytes(
    model: &Arc<CmfModel>,
    layers: &[crate::gpu::GraphLayer<'_>],
) -> Option<u64> {
    layers.iter().try_fold(0u64, |n, l| {
        n.checked_add(graph_layer_payload_bytes(model, l)?)
    })
}

/// WHOLE-TOKEN decode graph: the entire layer stack (rmsnorm → attention →
/// residual → rmsnorm → SiLU-FFN → residual, every layer) encoded into ONE
/// command buffer with the hidden RESIDENT on the GPU — only the final hidden
/// reads back (one submit/token instead of ~2 per layer). This is what lifts
/// the submit-latency wall. A failed run after sealed O(1) admission is
/// distinct from a preflight refusal because the CPU state is stale.
#[allow(clippy::too_many_arguments)]
pub fn forward_token_graph(
    model: &Arc<CmfModel>,
    kv_id: u64,
    layers: &[crate::gpu::GraphLayer],
    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,
    // Optional final-norm + lm_head fold: (weight, rows). When Some and the
    // weight resolves, the graph rides the final RMSNorm and lm_head in the
    // same submit and reads back `logits` (rows) instead of the hidden — one
    // fewer op + sync per token, and the lm_head stays on-device.
    lm_head: Option<(&crate::gpu::GraphW, usize)>,
    final_norm: &[f32],
    logits: &mut Vec<f32>,
    loop_norm_at: &[usize],
    // Multi-step greedy: encode `steps` whole frames in THIS submit, argmax
    // and re-embed on the device, and return the k winner ids instead of
    // logits. Needs `embed` = (q4tp embedding weight, vocab rows, multiplier).
    steps: usize,
    embed: Option<(&crate::gpu::GraphW, usize, f32)>,
    ids_out: Option<&mut Vec<u32>>,
    // Reports how many leading layers the graph executed. Equal to
    // `layers.len()` on a full run; smaller when the expert budget ended
    // the device prefix early — then `h` holds the boundary hidden, no
    // logits were produced, and the caller owns the remaining layers.
    layers_run: Option<&mut usize>,
    // Absolute index of layers[0] in the model: KV/GDN mirror keys use
    // `layer_base + li` so a span run (network split) and a full run
    // address the SAME per-layer mirrors.
    layer_base: usize,
    // With a fused lm_head, ALSO read the final hidden back into `h` (the
    // MTP draft chain feeds the block's hidden to its next step and wants
    // the head's logits from the same submit).
    hidden_too: bool,
) -> crate::gpu::TokenGraphOutcome {
    let _hp_t0 = std::time::Instant::now(); // CMF_GRAPH_HOSTPROF
    let span_submit_start = SUBMITS.load(std::sync::atomic::Ordering::Relaxed);
    let mut o1_started = false;
    // Persistent GDN/short-conv state makes a submitted graph mutation-sensitive
    // even when O(1) is not enabled.  A failed readback must therefore abort
    // the sequence instead of declining into a stale CPU fallback.
    let mut state_started = false;
    let Some(c) = ctx() else {
        graph_refused("no ctx");
        return token_graph_outcome(o1_started || state_started, false);
    };
    let graph_live_budget = graph_live_weight_budget(c, model);
    if position >= cap || hd % 4 != 0 || hd > c.hd_cap {
        {
            use std::sync::atomic::{AtomicBool, Ordering};
            static SAID: AtomicBool = AtomicBool::new(false);
            if !SAID.swap(true, Ordering::Relaxed) {
                tracing::warn!(
                    "wgpu token graph declined: position {position} >= cap {cap}, or head_dim \
                     {hd} (must be %4 and <= hd_cap {})",
                    c.hd_cap
                );
            }
        }
        return token_graph_outcome(o1_started || state_started, false); // vec4 K/V reads; hd_cap = workgroup-storage limit
    }
    let cap = kv_capacity(cap, position.saturating_add(steps.max(1)));
    let t_start = std::time::Instant::now();
    // A resolved matvec weight: the device-local buffer, (q8 only) its row
    // scales, and the codec kind (0=q8_row 1=q1 2=q4_block 3=q1t 4=f32 5=q4_tiled).
    #[derive(Clone)]
    struct GMat {
        buf: wgpu::Buffer,
        rs: Option<wgpu::Buffer>,
        // Model tensor identity is part of the graph descriptor: the global
        // ladder sidecar indexes row IDs by tensor, never by a guessed shape.
        idx: usize,
        kind: u8,
        prism: crate::gpu::GraphPrismOp,
        affine: bool,
    }
    enum LAttn {
        Full {
            wq: GMat,
            wk: GMat,
            wv: GMat,
            wo: GMat,
        },
        Conv {
            inp: GMat,
            out: GMat,
            kernel: usize,
        },
        Gdn {
            qkv: GMat,
            z: GMat,
            a: GMat,
            b: GMat,
            out: GMat,
            nv: usize,
            nk: usize,
            dk: usize,
            dv: usize,
            kk: usize,
            cdim: usize,
        },
    }
    enum LFfn {
        Dense {
            gate: GMat,
            up: GMat,
            down: GMat,
            /// This LAYER's intermediate width. Not the model's: a pruned
            /// model narrows the FFN per layer (bonsai-1.7b runs 6130 …
            /// 6140 across its 28 layers), and taking the header's number
            /// here asked for weights that do not exist — the graph then
            /// refused for the whole model, on every token, silently.
            width: usize,
        },
        Moe {
            router: GMat,
            sgate: GMat,
            gate_all: wgpu::Buffer,
            up_all: wgpu::Buffer,
            down_all: wgpu::Buffer,
            n_exp: usize,
            top_k: usize,
            inter: usize,
            norm_topk: bool,
            q4tp: bool,
            gu_q2: bool,
            sigmoid: bool,
            bias: Option<wgpu::Buffer>,
            has_shared: bool,
            shared_gated: bool,
            route_scale: f32,
        },
    }
    struct LW {
        attn: LAttn,
        ffn: LFfn,
    }
    // Resolve + cache every layer's weights (q8_row or q1) up front; bail (CPU)
    // on any refusal (budget/shape/dtype).
    let resolve = |gw: &crate::gpu::GraphW, rows: usize, cols: usize| -> Option<GMat> {
        if std::env::var("CMF_GRAPH_RESOLVE").is_ok() {
            eprintln!(
                "graph resolve: idx={} kind={} want={}x{} prism={:?} affine={}",
                gw.idx, gw.kind, rows, cols, gw.prism, gw.affine
            );
        }
        // The resident graph currently has a descriptor-aware affine kernel
        // only for the dense q2tp plane.  Refuse any other declared affine
        // target rather than silently applying the raw center.  Embedding
        // inverse operators belong to the separate embedding path, never a
        // layer projection.
        if gw.prism == crate::gpu::GraphPrismOp::InverseEmbedding
            || (gw.affine && gw.kind != 9)
        {
            return None;
        }
        match gw.kind {
            0 => {
                // q8_row: weight bytes = rows*cols, plus per-row scales.
                if gw.row_scale.len() < rows {
                    return None;
                }
                let b = tensor_weight(c, model, gw.idx, rows, cols)?; // device-local
                                                                      // Row scales are token-invariant — cache by (ptr,rows),
                                                                      // fingerprint-checked (the ptr is not a stable identity).
                let key = (gw.row_scale.as_ptr() as usize, rows);
                let fp = fp_bytes(bytemuck::cast_slice(&gw.row_scale[..rows]));
                let mut cb = c.const_bufs.lock().unwrap();
                let rsb = if let Some((x, f)) = cb.get_mut(&key) {
                    if *f != fp {
                        c.queue
                            .write_buffer(x, 0, bytemuck::cast_slice(&gw.row_scale[..rows]));
                        *f = fp;
                    }
                    x.clone()
                } else {
                    let x = c.device.create_buffer(&wgpu::BufferDescriptor {
                        label: Some("g-rs"),
                        size: (rows * 4) as u64,
                        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                        mapped_at_creation: false,
                    });
                    c.queue
                        .write_buffer(&x, 0, bytemuck::cast_slice(&gw.row_scale[..rows]));
                    cb.insert(key, (x.clone(), fp));
                    x
                };
                Some(GMat {
                    buf: b,
                    rs: Some(rsb),
                    idx: gw.idx,
                    kind: 0,
                    prism: gw.prism,
                    affine: gw.affine,
                })
            }
            1 => {
                let (b, r, cc) = q1_weight(c, model, gw.idx)?;
                if r != rows || cc != cols {
                    return None;
                }
                Some(GMat {
                    buf: b,
                    rs: None,
                    idx: gw.idx,
                    kind: 1,
                    prism: gw.prism,
                    affine: gw.affine,
                })
            }
            2 | 3 | 5 | 6 | 7 | 9 => {
                // The int8 body must end on a word, because the per-ROW f16
                // plane right after it is still word-indexed (`rs0`). Rows
                // are byte-addressed, so an odd `cols` is fine — an odd
                // `rows*cols` is not. The per-COLUMN plane needs no such
                // gate: it starts `rows` half-words later, which is
                // word-aligned only for even `rows`, and `f16x4` reads it
                // from an arbitrary half-word. So this check plus that
                // helper cover every layout the packer can emit.
                if gw.kind == 7 && (rows * cols) % 4 != 0 {
                    return None;
                }
                // q4_block / q1t / q4_tiled / q4tp / q8_2f: the tensor
                // carries its own byte length (tiles, q1t's sparse
                // overlay, q8_2f's two scale planes) — fetch whole,
                // device-local.
                let entry = model.tensors.get(gw.idx)?;
                if *entry.shape.first()? != rows || *entry.shape.get(1)? != cols {
                    // A silent refusal here is what made this model look
                    // like the graph "does not work on mobile" for a day:
                    // say which gate closed.
                    if std::env::var("CMF_GRAPH_DEBUG").is_ok() {
                        eprintln!(
                            "graph refuse: kind {} idx {} shape {:?} but graph wants {}x{}",
                            gw.kind, gw.idx, entry.shape, rows, cols
                        );
                    }
                    return None;
                }
                let abs = model.entry_abs_offset(entry)?;
                let plen = entry.nbytes as usize;
                let bytes = model.primary_bytes();
                if abs + plen > bytes.len() {
                    return None;
                }
                WEIGHT_BYTES.fetch_add(plen as u64, std::sync::atomic::Ordering::Relaxed);
                let b = weight_buffer_l(
                    c,
                    (model.uid() as usize, gw.idx),
                    &bytes[abs..abs + plen],
                    layer_of_name(&model.tensors[gw.idx].name),
                )?;
                Some(GMat {
                    buf: b,
                    rs: None,
                    idx: gw.idx,
                    kind: gw.kind,
                    prism: gw.prism,
                    affine: gw.affine,
                })
            }
            4 => {
                // f32 weight (small unquantized projection, e.g. GDN a/b) —
                // token-invariant: cache device-local by (ptr, rows*cols)
                // instead of re-uploading it every token.
                if gw.data.len() < rows * cols {
                    return None;
                }
                let key = (gw.data.as_ptr() as usize, rows * cols);
                let fp = fp_bytes(bytemuck::cast_slice(&gw.data[..rows * cols]));
                let mut cb = c.const_bufs.lock().unwrap();
                let b = if let Some((x, f)) = cb.get_mut(&key) {
                    if *f != fp {
                        c.queue
                            .write_buffer(x, 0, bytemuck::cast_slice(&gw.data[..rows * cols]));
                        *f = fp;
                    }
                    x.clone()
                } else {
                    let x = c.device.create_buffer(&wgpu::BufferDescriptor {
                        label: Some("g-f32w"),
                        size: (rows * cols * 4) as u64,
                        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                        mapped_at_creation: false,
                    });
                    c.queue
                        .write_buffer(&x, 0, bytemuck::cast_slice(&gw.data[..rows * cols]));
                    cb.insert(key, (x.clone(), fp));
                    x
                };
                Some(GMat {
                    buf: b,
                    rs: None,
                    idx: gw.idx,
                    kind: 4,
                    prism: gw.prism,
                    affine: gw.affine,
                })
            }
            _ => None,
        }
    };
    let mut lws = Vec::with_capacity(layers.len());
    let mut gdn_dims: Option<(usize, usize, usize, usize, usize, usize)> = None; // nv,nk,dk,dv,kk,cdim
                                                                                 // Budget-driven device prefix: the first layer whose live graph buffers
                                                                                 // no longer fit ends the prefix instead of retaining an over-budget dense
                                                                                 // stack behind evicted LRU handles. The caller finishes the remaining
                                                                                 // layers on the host from the boundary hidden: one sync per token at the
                                                                                 // boundary, not per layer.
    let mut prefix = false;
    let mut graph_bytes = 0u64;
    for l in layers {
        let Some(layer_bytes) = graph_layer_payload_bytes(model, l) else {
            graph_decline("cannot size layer payload");
            return token_graph_outcome(o1_started || state_started, false);
        };
        if graph_bytes.saturating_add(layer_bytes) > graph_live_budget {
            if lws.is_empty() {
                graph_decline("one layer exceeds the weight budget");
                return token_graph_outcome(o1_started || state_started, false);
            }
            prefix = true;
            break;
        }
        graph_bytes += layer_bytes;
        let attn = match &l.attn {
            crate::gpu::GraphAttn::Full {
                wq,
                wk,
                wv,
                wo,
                output_gate,
                ..
            } => {
                // Gated attention: wq packs q||gate per head → 2·nh·hd rows.
                let qrows = nh * hd * (1 + *output_gate as usize);
                let (Some(wq), Some(wk), Some(wv), Some(wo)) = (
                    resolve(wq, qrows, hidden),
                    resolve(wk, nkv * hd, hidden),
                    resolve(wv, nkv * hd, hidden),
                    resolve(wo, hidden, nh * hd),
                ) else {
                    graph_decline("attn q/k/v/o resolve");
                    return token_graph_outcome(o1_started || state_started, false);
                };
                LAttn::Full { wq, wk, wv, wo }
            }
            crate::gpu::GraphAttn::Gdn {
                qkv,
                z,
                a,
                b,
                out,
                nv,
                nk,
                dk,
                dv,
                kk,
                ..
            } => {
                let cdim = 2 * nk * dk + nv * dv;
                let dims = (*nv, *nk, *dk, *dv, *kk, cdim);
                if gdn_dims.is_some_and(|prev| prev != dims) {
                    graph_decline("heterogeneous GDN geometry across token-graph layers");
                    return token_graph_outcome(o1_started || state_started, false);
                }
                gdn_dims = Some(dims);
                let (Some(qkv), Some(z), Some(a), Some(b), Some(out)) = (
                    resolve(qkv, cdim, hidden),
                    resolve(z, nv * dv, hidden),
                    resolve(a, *nv, hidden),
                    resolve(b, *nv, hidden),
                    resolve(out, hidden, nv * dv),
                ) else {
                    graph_decline("weight resolve (dtype/shape outside graph contract)");
                    return token_graph_outcome(o1_started || state_started, false);
                };
                LAttn::Gdn {
                    qkv,
                    z,
                    a,
                    b,
                    out,
                    nv: *nv,
                    nk: *nk,
                    dk: *dk,
                    dv: *dv,
                    kk: *kk,
                    cdim,
                }
            }
            crate::gpu::GraphAttn::ShortConv {
                inp, out, kernel, ..
            } => {
                let (Some(inp), Some(out)) = (
                    resolve(inp, 3 * hidden, hidden),
                    resolve(out, hidden, hidden),
                ) else {
                    graph_decline("short-conv resolve (dtype outside graph contract)");
                    return token_graph_outcome(o1_started || state_started, false);
                };
                LAttn::Conv {
                    inp,
                    out,
                    kernel: *kernel,
                }
            }
        };
        let ffn = match &l.ffn {
            crate::gpu::GraphFfn::Dense { gate, up, down } => {
                // The tensor knows its own width; the config only knows
                // the widest. Ask the weight.
                let ffn_w = model
                    .tensors
                    .get(gate.idx)
                    .and_then(|e| e.shape.first().copied())
                    .map(|r| r)
                    .filter(|w| *w > 0 && *w <= inter)
                    .unwrap_or(inter);
                let (Some(gate), Some(up), Some(down)) = (
                    resolve(gate, ffn_w, hidden),
                    resolve(up, ffn_w, hidden),
                    resolve(down, hidden, ffn_w),
                ) else {
                    graph_decline("weight resolve (dtype/shape outside graph contract)");
                    return token_graph_outcome(o1_started || state_started, false);
                };
                LFfn::Dense {
                    gate,
                    up,
                    down,
                    width: ffn_w,
                }
            }
            crate::gpu::GraphFfn::Moe {
                router,
                shared_gate,
                experts,
                n_exp,
                top_k,
                inter: mi,
                norm_topk,
                q4tp,
                gu_q2,
                sigmoid,
                bias,
                has_shared,
                shared_gated,
                route_scale,
            } => {
                // Select kernel: logits live in a 256-slot workgroup array;
                // with a shared expert it rides as the last block.
                let want = n_exp + usize::from(*has_shared);
                if *top_k >= 16 || *n_exp > 256 || experts.len() != want {
                    graph_decline(&format!(
                        "moe shape: top_k {top_k} n_exp {n_exp} experts {} (want {want})",
                        experts.len()
                    ));
                    return token_graph_outcome(o1_started || state_started, false);
                }
                let Some(router) = resolve(router, *n_exp, hidden) else {
                    graph_decline("moe router resolve");
                    return token_graph_outcome(o1_started || state_started, false);
                };
                // Without a shared expert there is no gate to resolve; the
                // kernel never reads it (flags bit 3 off), so the router
                // stands in to keep the binding set total.
                let sgate = if *has_shared {
                    let Some(sg) = resolve(shared_gate, 1, hidden) else {
                        graph_decline("moe shared_gate resolve");
                        return token_graph_outcome(o1_started || state_started, false);
                    };
                    sg
                } else {
                    router.clone()
                };
                // The layer-level live-set check above owns the prefix
                // boundary. Keep this arena check for other resident data
                // that may already occupy the explicit weight budget.
                let cached = experts.first().is_some_and(|e| {
                    c.moe_expw
                        .lock()
                        .unwrap()
                        .contains_key(&(model.uid() as usize, e.0))
                });
                let need = moe_pack_bytes(experts.len(), *mi, hidden, *q4tp, *gu_q2, false)
                    .unwrap_or(u64::MAX);
                if !cached
                    && !lws.is_empty()
                    && c.resident.load(std::sync::atomic::Ordering::Relaxed) + need > c.vram_budget
                {
                    prefix = true;
                    break;
                }
                if experts.iter().flat_map(|&(g, u, d)| [g, u, d]).any(|idx| {
                    model
                        .tensors
                        .get(idx)
                        .is_some_and(|e| crate::prism::is_forward_weight(model, &e.name))
                }) {
                    graph_decline("Prism MoE experts require an unimplemented resident transform");
                    return token_graph_outcome(o1_started || state_started, false);
                }
                let Some((gate_all, up_all, down_all)) =
                    moe_expert_bufs(c, model, experts, *mi, hidden, *q4tp, *gu_q2, false)
                else {
                    graph_decline("moe_expert_bufs (pack/budget/dtype)");
                    return token_graph_outcome(o1_started || state_started, false);
                };
                LFfn::Moe {
                    router,
                    sgate,
                    gate_all,
                    up_all,
                    down_all,
                    n_exp: *n_exp,
                    top_k: *top_k,
                    inter: *mi,
                    norm_topk: *norm_topk,
                    q4tp: *q4tp,
                    gu_q2: *gu_q2,
                    sigmoid: *sigmoid,
                    bias: bias.map(|b| {
                        c.device
                            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                                label: Some("moe-sel-bias"),
                                contents: bytemuck::cast_slice(b),
                                usage: wgpu::BufferUsages::STORAGE,
                            })
                    }),
                    has_shared: *has_shared,
                    shared_gated: *shared_gated,
                    route_scale: *route_scale,
                }
            }
        };
        lws.push(LW { attn, ffn });
    }
    // The tail weight is resolved only after the layer command stream has
    // been encoded. Reserve it now: otherwise a stack that fits by itself
    // can evict one of its already-referenced buffers to admit lm_head and
    // exceed the physical limit despite a correct LRU counter.
    if !prefix {
        let mut tail_bytes = 0u64;
        if let Some((w, _)) = lm_head {
            tail_bytes = tail_bytes.saturating_add(
                model
                    .tensors
                    .get(w.idx)
                    .map(|e| e.nbytes as u64)
                    .unwrap_or(u64::MAX),
            );
        }
        if steps > 1 {
            if let Some((w, _, _)) = embed {
                if lm_head.is_none_or(|(lm, _)| lm.idx != w.idx) {
                    tail_bytes = tail_bytes.saturating_add(
                        model
                            .tensors
                            .get(w.idx)
                            .map(|e| e.nbytes as u64)
                            .unwrap_or(u64::MAX),
                    );
                }
            }
        }
        if graph_bytes.saturating_add(tail_bytes) > graph_live_budget {
            if lws.pop().is_none() {
                graph_decline("layers plus tail exceed the weight budget");
                return token_graph_outcome(o1_started || state_started, false);
            }
            prefix = true;
        }
    }
    if prefix {
        // The multi-step tail (on-device argmax + re-embed) needs the head,
        // which a prefix does not reach — those callers fall back whole.
        if steps > 1 || ids_out.is_some() {
            graph_refused("device prefix cannot serve the multi-step tail");
            return token_graph_outcome(o1_started || state_started, false);
        }
        use std::sync::atomic::{AtomicBool, Ordering};
        static SAID: AtomicBool = AtomicBool::new(false);
        if !SAID.swap(true, Ordering::Relaxed) {
            tracing::info!(
                "wgpu token graph: device prefix {} of {} layers (VRAM budget), \
                 the tail runs on the host from the boundary hidden",
                lws.len(),
                layers.len()
            );
        }
    }
    if let Some(n) = layers_run {
        *n = lws.len();
    }
    let layers = &layers[..lws.len()];
    // DEVICE-LOCAL + content-cached: create_buffer + write_buffer keeps norm
    // weights in VRAM (not the HOST_VISIBLE heap create_buffer_init forces);
    // caching by (ptr,len) uploads each token-invariant norm buffer once,
    // and the fingerprint refreshes it in place when another model's mmap
    // lands on the same address.
    let stor = |data: &[u8]| {
        let key = (data.as_ptr() as usize, data.len());
        let fp = fp_bytes(data);
        let mut cb = c.const_bufs.lock().unwrap();
        if let Some((b, f)) = cb.get_mut(&key) {
            if *f != fp {
                c.queue.write_buffer(b, 0, data);
                *f = fp;
            }
            return b.clone();
        }
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: data.len() as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        c.queue.write_buffer(&b, 0, data);
        cb.insert(key, (b.clone(), fp));
        b
    };
    // Shared zero buffer of `n` f32 (sentinel key (0,n)) — for absent q/k-norms
    // and the silu bias slot, so no per-token zero Vec is allocated/uploaded.
    let zeros = |n: usize| -> wgpu::Buffer {
        let key = (0usize, n * 4);
        let mut cb = c.const_bufs.lock().unwrap();
        if let Some((b, _)) = cb.get(&key) {
            return b.clone();
        }
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("g-zero"),
            size: (n * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        c.queue.write_buffer(&b, 0, &vec![0u8; n * 4]);
        cb.insert(key, (b.clone(), 0));
        b
    };
    let unif = |d: &[u32]| {
        c.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(d),
                usage: wgpu::BufferUsages::UNIFORM,
            })
    };
    let mk_bg = |layout: &wgpu::BindGroupLayout, bufs: &[&wgpu::Buffer]| {
        let e: Vec<_> = bufs
            .iter()
            .enumerate()
            .map(|(i, b)| bind_buf(i as u32, b))
            .collect();
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout,
            entries: &e,
        })
    };
    // The host spends 13.7 of a 23.5 ms token creating these objects
    // (CMF_GRAPH_HOSTPROF). Cache per (site, layer, kv): buffers under a
    // cached group are stable until ANY cold event — a scratch grow or a
    // weight upload — which invalidates everything via the global cold
    // epoch. Positions ride in content-cached uniforms, so a给 (site,
    // layer) the uniform BUFFER changes when position changes — those
    // sites pass site=0, which bypasses the cache.
    let bgc = |site: u32, li: usize, layout: &wgpu::BindGroupLayout, bufs: &[&wgpu::Buffer]| {
        if site == 0 || !c.use_bgcache {
            return mk_bg(layout, bufs);
        }
        let ep = crate::gpu::cold_epoch();
        let key = (site, li, kv_id);
        let mut m = c.graph_bgs.lock().unwrap();
        if let Some((e, g)) = m.get(&key) {
            if *e == ep {
                return g.clone();
            }
        }
        let g = mk_bg(layout, bufs);
        m.insert(key, (ep, g.clone()));
        g
    };
    let bg = |layout: &wgpu::BindGroupLayout, bufs: &[&wgpu::Buffer]| mk_bg(layout, bufs);
    let _ = &bg;
    // ── Pooled scratch: all intermediate buffers are reused across tokens ──
    let mut gs = c.graph_scratch.lock().unwrap();
    let st = wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC; // COPY_SRC: debug taps (CMF_O1_TRACE)
    let h_buf = GraphScratch::ensure(
        &c.device,
        &mut gs.h,
        (hidden * 4) as u64,
        st | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
        "g-h",
    );
    c.queue
        .write_buffer(&h_buf, 0, bytemuck::cast_slice(&h[..hidden]));
    let token_tap_layer = std::env::var("CMF_GRAPH_TAP_LAYER")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .filter(|&li| li < layers.len());
    let token_tap_stage = token_tap_layer.map(|_| c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("token-graph-tap"),
        size: (hidden * 4) as u64,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    }));
    let n1 = GraphScratch::ensure(
        &c.device,
        &mut gs.n1,
        (hidden * 4) as u64,
        st | wgpu::BufferUsages::COPY_SRC,
        "g-n1",
    );
    // Gated attention (Qwen3.5) makes wq emit 2·nh·hd (q||gate per head), so the
    // raw-QKV scratch must hold the widened q output for any gated layer.
    let any_gate = layers.iter().any(|l| {
        matches!(
            &l.attn,
            crate::gpu::GraphAttn::Full {
                output_gate: true,
                ..
            }
        )
    });
    let qraw = GraphScratch::ensure(
        &c.device,
        &mut gs.qraw,
        (nh * hd * (1 + any_gate as usize) * 4) as u64,
        st,
        "g-qraw",
    );
    let kb = GraphScratch::ensure(&c.device, &mut gs.kb, (nkv * hd * 4) as u64, st, "g-kb");
    let vb = GraphScratch::ensure(&c.device, &mut gs.vb, (nkv * hd * 4) as u64, st, "g-vb");
    let qout = GraphScratch::ensure(&c.device, &mut gs.qout, (nh * hd * 4) as u64, st, "g-qout");
    let gout = GraphScratch::ensure(&c.device, &mut gs.gout, (nh * hd * 4) as u64, st, "g-gout");
    let attn = GraphScratch::ensure(&c.device, &mut gs.attn, (nh * hd * 4) as u64, st, "g-attn");
    let ob = GraphScratch::ensure(&c.device, &mut gs.ob, (hidden * 4) as u64, st, "g-ob");
    let gbuf = GraphScratch::ensure(&c.device, &mut gs.gbuf, (inter * 4) as u64, st, "g-gbuf");
    let ubuf = GraphScratch::ensure(&c.device, &mut gs.ubuf, (inter * 4) as u64, st, "g-ubuf");
    let abuf = GraphScratch::ensure(&c.device, &mut gs.abuf, (inter * 4) as u64, st, "g-abuf");
    // MoE routing scratch, sized to the largest MoE layer (absent → skipped).
    let moe_geom = lws
        .iter()
        .filter_map(|lw| match &lw.ffn {
            LFfn::Moe {
                n_exp,
                top_k,
                inter,
                ..
            } => Some((*n_exp, *top_k + 1, *inter)),
            _ => None,
        })
        .reduce(|a, b| (a.0.max(b.0), a.1.max(b.1), a.2.max(b.2)));
    let moe_bufs = moe_geom.map(|(mn, ms, mi)| {
        (
            GraphScratch::ensure(&c.device, &mut gs.m_logit, (mn * 4) as u64, st, "g-mlogit"),
            GraphScratch::ensure(&c.device, &mut gs.m_slog, 4, st, "g-mslog"),
            GraphScratch::ensure(&c.device, &mut gs.m_sel, (ms * 4) as u64, st, "g-msel"),
            GraphScratch::ensure(&c.device, &mut gs.m_wt, (ms * 4) as u64, st, "g-mwt"),
            GraphScratch::ensure(&c.device, &mut gs.m_act, (ms * mi * 4) as u64, st, "g-mact"),
        )
    });
    let invf_b = stor(bytemuck::cast_slice(invf));
    let dummy_hd = zeros(hd);
    // GDN intermediates (sized to the model's GDN geometry; 1 if no GDN layer).
    let (gnv, _gnk, gdk, gdv, _gkk, gcdim) = gdn_dims.unwrap_or((1, 1, 1, 1, 1, 1));
    let qkv_b = GraphScratch::ensure(&c.device, &mut gs.qkv_b, (gcdim * 4) as u64, st, "g-qkv");
    let cq_b = GraphScratch::ensure(&c.device, &mut gs.cq_b, (gcdim * 4) as u64, st, "g-cq");
    let z_b = GraphScratch::ensure(&c.device, &mut gs.z_b, (gnv * gdv * 4) as u64, st, "g-z");
    let a_b = GraphScratch::ensure(&c.device, &mut gs.a_b, (gnv * 4) as u64, st, "g-a");
    let b_b = GraphScratch::ensure(&c.device, &mut gs.b_b, (gnv * 4) as u64, st, "g-b");
    let gdo_b = GraphScratch::ensure(
        &c.device,
        &mut gs.gdo_b,
        (gnv * gdv * 4) as u64,
        st,
        "g-gdo",
    );
    let sc_bcx = GraphScratch::ensure(
        &c.device,
        &mut gs.sc_bcx,
        (3 * hidden * 4) as u64,
        st,
        "g-scbcx",
    );
    let sc_y = GraphScratch::ensure(&c.device, &mut gs.sc_y, (hidden * 4) as u64, st, "g-scy");
    // A single reusable transform destination is enough because the graph
    // emits Prism FWHTs in dependency order.  It is sized for both the
    // hidden input projections and the widest dense-FFN down input.
    let rot_width = hidden.max(inter);
    let rot = GraphScratch::ensure(
        &c.device,
        &mut gs.rot,
        (rot_width * 4) as u64,
        st | wgpu::BufferUsages::COPY_DST,
        "g-prism-rot",
    );
    // The validated Prism header stores one concatenated sign vector per
    // supported width.  Keep that table resident and address the selected
    // width by an explicit offset in the FWHT uniform; no name/dtype guess
    // is permitted in the graph.
    let prism_signs = model
        .header
        .arch
        .prism_hadamard
        .as_ref()
        .map(|cfg| stor(bytemuck::cast_slice(&cfg.signs)));
    let prism_sign_offset = |width: usize| -> Option<usize> {
        let cfg = model.header.arch.prism_hadamard.as_ref()?;
        let mut off = 0usize;
        for &w in &cfg.widths {
            if w == width {
                return Some(off);
            }
            off = off.checked_add(w)?;
        }
        None
    };
    let prism_round16 = model
        .header
        .arch
        .prism_hadamard
        .as_ref()
        .is_some_and(|cfg| cfg.activation_f16);
    // Sync each Full layer's device K/V mirror from the CPU cache (once);
    // GDN layers carry a persistent (ring, S) recurrent state instead.
    let mut kvbufs: Vec<Option<(wgpu::Buffer, wgpu::Buffer)>> = Vec::with_capacity(layers.len());
    let mut gdnbufs: Vec<Option<(wgpu::Buffer, wgpu::Buffer)>> = Vec::with_capacity(layers.len());
    {
        let mut kvm = c.attn_kv.lock().unwrap();
        let mut gsm = c.gdn_state.lock().unwrap();
        let mut gcm = c.gdn_cursor.lock().unwrap();
        for (li, l) in layers.iter().enumerate() {
            match &l.attn {
                crate::gpu::GraphAttn::Full { cpu_k, cpu_v, .. } => {
                    if o1.get(li).is_some_and(|v| v.is_some()) {
                        // o1 replaces this layer's KV attention outright —
                        // no mirror, and no prefill K/V upload (16K of it
                        // at long context) that nothing would read.
                        kvbufs.push(None);
                        gdnbufs.push(None);
                        continue;
                    }
                    let e = kv_mirror_ensure(c, &mut kvm, (kv_id, layer_base + li), nkv, hd, cap);
                    if e.synced > position {
                        graph_refused("KV mirror is ahead of token-graph position");
                        // A resident mirror proves that device state already
                        // belongs to this request.  Falling through to the
                        // host path would consume a stale CPU KV copy.
                        return token_graph_outcome(true, false);
                    }
                    if e.synced < position {
                        for hh in 0..nkv {
                            let take = position.min(cpu_k[hh].len() / hd);
                            if take > e.synced {
                                let off = ((hh * cap + e.synced) * hd * 4) as u64;
                                c.queue.write_buffer(
                                    &e.k,
                                    off,
                                    bytemuck::cast_slice(&cpu_k[hh][e.synced * hd..take * hd]),
                                );
                                c.queue.write_buffer(
                                    &e.v,
                                    off,
                                    bytemuck::cast_slice(&cpu_v[hh][e.synced * hd..take * hd]),
                                );
                            }
                        }
                        e.synced = position;
                    }
                    kvbufs.push(Some((e.k.clone(), e.v.clone())));
                    gdnbufs.push(None);
                }
                crate::gpu::GraphAttn::Gdn {
                    cpu_state,
                    nv,
                    nk,
                    dk,
                    dv,
                    kk,
                    ..
                } => {
                    let key = (kv_id, layer_base + li);
                    let dims = (*nv, *nk, *dk, *dv, *kk, 2 * nk * dk + nv * dv);
                    match (gsm.get(&key), gcm.get(&key)) {
                        (Some(_), Some(cur)) if cur.dims == dims && cur.next_pos == position => {}
                        (Some(_), Some(_)) => {
                            graph_refused("GDN device state position/geometry mismatch");
                            return token_graph_outcome(true, false);
                        }
                        (Some(_), None) => {
                            graph_refused("GDN device state cursor missing");
                            return token_graph_outcome(true, false);
                        }
                        (None, Some(_)) => {
                            graph_refused("GDN cursor exists without device state");
                            return token_graph_outcome(true, false);
                        }
                        (None, None) => {}
                    }
                    let ring_words = gcdim * _gkk.max(1).saturating_sub(1);
                    let state_words = gnv * gdk * gdv;
                    if !gsm.contains_key(&key)
                        && !gcm.contains_key(&key)
                        && position > 0
                        && cpu_state.len() != ring_words.saturating_add(state_words)
                    {
                        graph_refused("GDN CPU seed missing for nonzero token position");
                        return token_graph_outcome(o1_started || state_started, false);
                    }
                    let e = gsm.entry(key).or_insert_with(|| {
                        let ring_sz = ((gcdim * (_gkk.max(1).saturating_sub(1))) * 4) as u64;
                        let s_sz = (gnv * gdk * gdv * 4) as u64;
                        let mk = |sz: u64| {
                            let bf = c.device.create_buffer(&wgpu::BufferDescriptor {
                                label: Some("gdn-state"),
                                size: sz.max(4),
                                usage: wgpu::BufferUsages::STORAGE
                                    | wgpu::BufferUsages::COPY_DST
                                    | wgpu::BufferUsages::COPY_SRC,
                                mapped_at_creation: false,
                            });
                            c.queue.write_buffer(&bf, 0, &vec![0u8; sz.max(4) as usize]);
                            bf
                        };
                        let (ring, sbuf) = (mk(ring_sz), mk(s_sz));
                        // Seed from the CPU recurrence when the host ran the
                        // prefill (o1 collection, CPU fallback). A fresh entry
                        // with an EMPTY cpu_state is the graph-prefill flow —
                        // the graph builds the state itself from position 0.
                        // Zero-initialized device state at decode is the
                        // "coherent but contextless" failure this closes.
                        let want = (ring_sz + s_sz) as usize / 4;
                        if cpu_state.len() == want && want > 0 {
                            let ring_n = ring_sz as usize / 4;
                            c.queue.write_buffer(
                                &ring,
                                0,
                                bytemuck::cast_slice(&cpu_state[..ring_n]),
                            );
                            c.queue.write_buffer(
                                &sbuf,
                                0,
                                bytemuck::cast_slice(&cpu_state[ring_n..]),
                            );
                        }
                        (ring, sbuf)
                    });
                    gcm.entry(key).or_insert(GdnCursor {
                        dims,
                        next_pos: position,
                    });
                    gdnbufs.push(Some((e.0.clone(), e.1.clone())));
                    kvbufs.push(None);
                }
                crate::gpu::GraphAttn::ShortConv {
                    kernel, cpu_state, ..
                } => {
                    let key = (kv_id, layer_base + li);
                    let dims = (hidden, 0, 0, 0, *kernel, hidden);
                    match (gsm.get(&key), gcm.get(&key)) {
                        (Some(_), Some(cur)) if cur.dims == dims && cur.next_pos == position => {}
                        (Some(_), Some(_)) => {
                            graph_refused("short-conv device state position/geometry mismatch");
                            return token_graph_outcome(true, false);
                        }
                        (Some(_), None) => {
                            graph_refused("short-conv device state cursor missing");
                            return token_graph_outcome(true, false);
                        }
                        (None, Some(_)) => {
                            graph_refused("short-conv cursor exists without device state");
                            return token_graph_outcome(true, false);
                        }
                        (None, None) => {}
                    }
                    // The conv ring rides the same state map as GDN; the
                    // second buffer of the pair is a 4-byte placeholder.
                    let ring_sz = ((kernel.saturating_sub(1)) * hidden * 4) as u64;
                    if !gsm.contains_key(&key)
                        && !gcm.contains_key(&key)
                        && position > 0
                        && cpu_state.len() * 4 != ring_sz as usize
                    {
                        graph_refused("short-conv CPU seed missing for nonzero token position");
                        return token_graph_outcome(o1_started || state_started, false);
                    }
                    let e = gsm.entry(key).or_insert_with(|| {
                        let mk = |sz: u64| {
                            let bf = c.device.create_buffer(&wgpu::BufferDescriptor {
                                label: Some("sconv-ring"),
                                size: sz.max(4),
                                usage: wgpu::BufferUsages::STORAGE
                                    | wgpu::BufferUsages::COPY_DST
                                    | wgpu::BufferUsages::COPY_SRC,
                                mapped_at_creation: false,
                            });
                            c.queue.write_buffer(&bf, 0, &vec![0u8; sz.max(4) as usize]);
                            bf
                        };
                        let ring = mk(ring_sz);
                        // Prefill always runs on the host for this mixer
                        // (the batch graph declines it), so the CPU ring is
                        // the truth at the first decode token. Layout is
                        // the host's: [channel][kernel-1], slot 0 newest.
                        if cpu_state.len() * 4 == ring_sz as usize && !cpu_state.is_empty() {
                            c.queue
                                .write_buffer(&ring, 0, bytemuck::cast_slice(*cpu_state));
                        }
                        (ring, mk(4))
                    });
                    gcm.entry(key).or_insert(GdnCursor {
                        dims,
                        next_pos: position,
                    });
                    gdnbufs.push(Some((e.0.clone(), e.1.clone())));
                    kvbufs.push(None);
                }
            }
        }
    }
    let prof = std::env::var("CMF_GRAPH_PROF").is_ok();
    // Group mutually-independent projections (that all read the same normed
    // hidden) into ONE compute pass — the GPU can overlap them, cutting the
    // per-pass barrier bubbles that dominate single-token decode. Default on
    // (measured +5-8% token-identical across q1/q8/GDN); CMF_GPU_GROUP=0 off.
    let group = std::env::var("CMF_GPU_GROUP")
        .map(|v| v != "0")
        .unwrap_or(true);
    // Hand strictly-serial single-dispatch stages to the NEXT pass instead of
    // opening a pass for each. This is opt-in: it defers selected residual/norm
    // producers, and is separate from PassMergeGuard's pass-merging switch.
    // CMF_PASSFUSE=1 is an explicit experimental override; the narrower
    // GRAPH_FUSE_* switches below remain off unless requested. Prism inputs
    // force a deferred FFN norm back before their FWHT consumer.
    let passfuse = std::env::var("CMF_PASSFUSE")
        .map(|v| v != "0")
        .unwrap_or(false);
    // Keep the two residual/norm joins independently switchable. The safe
    // default leaves both producers in their own encoded stage; any explicit
    // fusion must preserve producer-before-transform ordering below.
    let fuse_pre = passfuse
        && std::env::var("CMF_GRAPH_FUSE_PRE")
            .map(|v| v != "0")
            .unwrap_or(false);
    let fuse_tail = passfuse
        && std::env::var("CMF_GRAPH_FUSE_TAIL")
            .map(|v| v != "0")
            .unwrap_or(false);
    // CMF_SKIP_PROBE=moe|gdn — TIMING ONLY, the answer is garbage. Drops a
    // whole stage's dispatches while leaving every buffer, pass and shape
    // in place, so the delta is that stage's real share of the frame. The
    // arithmetic-only probe (CMF_TOPK_PROBE) says MoE math is ~2 ms of 17;
    // this one says where the rest actually goes, which neither dispatch
    // counting nor pass counting predicted correctly.
    let skip = std::env::var("CMF_SKIP_PROBE").unwrap_or_default();
    let (skip_moe, skip_gdn) = (skip.contains("moe"), skip.contains("gdn"));
    // Skeleton pieces, so the 7.5 ms that is neither MoE nor GDN can be
    // attributed instead of guessed at: the four GDN input projections,
    // the GDN output projection, the fused residual+norm, the router, and
    // the whole full-attention chain.
    let skip_proj = skip.contains("proj");
    let skip_outp = skip.contains("outp");
    let skip_norm = skip.contains("norm");
    let skip_router = skip.contains("router");
    let skip_attn = skip.contains("attn");
    // CMF_LAYERS_PROBE=N — TIMING ONLY, the answer is garbage. Encodes just
    // the first N layers, leaving the final norm + lm_head + readback in
    // place. Decode time against N is a straight line whose slope is the
    // per-layer cost and whose intercept is everything that happens once a
    // token: the submit, the ~1 MB logits readback and the lm_head. Neither
    // dispatch counting nor pass counting predicted the frame correctly, so
    // this splits it by measurement instead.
    let layer_cap = std::env::var("CMF_LAYERS_PROBE")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .unwrap_or(usize::MAX);
    let t_enc0 = std::time::Instant::now();
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("token-graph"),
        });
    // Pass merging for this encoder (and its chunk successors — same
    // slot, same key): every begin_pass below hands back one open pass;
    // copies, timestamps, swaps and finishes flush it first.
    let _merge_guard = PassMergeGuard::new(&enc);
    batch_kernel_ts_begin(c);
    let go =
        |enc: &mut wgpu::CommandEncoder, p: &wgpu::ComputePipeline, b: &wgpu::BindGroup, g: u32| {
            let mut pass = begin_pass(enc);
            pass.set_pipeline(p);
            pass.set_bind_group(0, b, &[]);
            pass.dispatch_workgroups(g, 1, 1);
        };
    // Full-frame token microscope.  The normal 256-slot profile intentionally
    // samples only a few early layers; this opt-in route reserves at most the
    // 4096 slots allocated above and stamps every selected dispatch class.
    // It is diagnostics only and never changes the graph's math or grouping.
    // Interior mutability keeps the long-lived prism closure from holding a
    // mutable borrow over the rest of graph encoding.
    let ts_n = std::cell::Cell::new(0u32);
    let ts_lbl = std::cell::RefCell::new(Vec::<(u8, u8)>::new());
    // Bounded architecture discriminator: when requested, only two endpoint
    // timestamps are emitted and all work remains in the ordinary graph
    // passes.  This is deliberately separate from the detailed microscope:
    // its span must not inherit per-dispatch timestamp flushes or labels.
    let device_span_probe = std::env::var("CMF_GRAPH_DEVICE_SPAN").as_deref() == Ok("1")
        && steps == 1;
    // CMF_GRAPH_PREENCODE=1 disables only the historical mid-stack split for
    // this probe.  It does not fuse passes or alter dispatch/math ordering;
    // the caller still submits the finished encoder through the same readback.
    let preencode_frame = std::env::var("CMF_GRAPH_PREENCODE").as_deref() == Ok("1")
        && steps == 1;
    let ts_fine = !device_span_probe && std::env::var("CMF_GPU_TS").as_deref() == Ok("2");
    let ts_full = !device_span_probe && std::env::var("CMF_GRAPH_TS_ALL").as_deref() == Ok("1");
    let ts_cap = if device_span_probe { 2 } else if ts_full { 4096 } else { 255 };
    macro_rules! ts_point {
        ($enc:expr, $stage:expr) => {
            if !device_span_probe && steps == 1 && ts_full {
                if let Some((qs, _, _)) = &c.ts_query {
                    if ts_n.get() < ts_cap {
                        flush_pass(&$enc);
                        $enc.write_timestamp(qs, ts_n.get());
                        ts_lbl.borrow_mut().push(($stage, 9));
                        ts_n.set(ts_n.get() + 1);
                    }
                }
            }
        };
    }
    macro_rules! ts {
        ($enc:expr, $stage:expr, $kind:expr) => {
            if !device_span_probe && steps == 1 {
                if let Some((qs, _, _)) = &c.ts_query {
                    if ts_n.get() < ts_cap {
                        flush_pass(&$enc);
                        $enc.write_timestamp(qs, ts_n.get());
                        ts_lbl.borrow_mut().push(($stage, $kind));
                        ts_n.set(ts_n.get() + 1);
                    }
                }
            }
        };
    }
    macro_rules! tsp {
        ($pass:expr, $on:expr, $stage:expr) => {
            if ts_fine && $on && steps == 1 {
                if let Some((qs, _, _)) = &c.ts_query {
                    if ts_n.get() < ts_cap {
                        $pass.write_timestamp(qs, ts_n.get());
                        ts_lbl.borrow_mut().push(($stage, 9));
                        ts_n.set(ts_n.get() + 1);
                    }
                }
            }
        };
    }
    let write_span_endpoint = |enc: &mut wgpu::CommandEncoder| {
        if device_span_probe {
            if let Some((qs, _, _)) = &c.ts_query {
                if ts_n.get() < 2 {
                    flush_pass(enc);
                    let slot = ts_n.get();
                    enc.write_timestamp(qs, slot);
                    ts_lbl.borrow_mut().push((60 + slot as u8, 0));
                    ts_n.set(slot + 1);
                }
            }
        }
    };
    write_span_endpoint(&mut enc);
    // Resolve the timestamp query on every readback route.  Previously this
    // lived only in the logits/head arm, so hidden-only frames mapped stale
    // query data and reported it as if it belonged to the current graph.
    let resolve_timestamps = |enc: &mut wgpu::CommandEncoder| {
        if let Some((qs, resolve, tstage)) = &c.ts_query {
            if steps == 1 && ts_n.get() > 0 {
                flush_pass(enc);
                enc.resolve_query_set(qs, 0..ts_n.get(), resolve, 0);
                flush_pass(enc);
                enc.copy_buffer_to_buffer(resolve, 0, tstage, 0, ts_n.get() as u64 * 8);
            }
        }
    };
    // Encode a descriptor-aware signed FWHT without a host round trip.  The
    // output is the pooled `rot` buffer and is consumed by the immediately
    // following projection dispatch(es) in queue order.  Refusal is explicit:
    // callers must not fall back to an untransformed Prism matvec.
    let prism_transform =
        |enc: &mut wgpu::CommandEncoder,
         src: &wgpu::Buffer,
         width: usize,
         op: crate::gpu::GraphPrismOp|
         -> Option<wgpu::Buffer> {
            if op == crate::gpu::GraphPrismOp::None {
                return Some(src.clone());
            }
            let pipe = c.fwht.as_ref()?;
            let cfg = model.header.arch.prism_hadamard.as_ref()?;
            let block = cfg.block_size;
            let sign_offset = prism_sign_offset(width)?;
            if block != 1024 || width == 0 || width % block != 0 || width > rot_width {
                return None;
            }
            let signs = prism_signs.as_ref()?;
            let p = uniform_u32x8(
                c,
                [
                    width as u32,
                    block as u32,
                    sign_offset as u32,
                    u32::from(op == crate::gpu::GraphPrismOp::InverseEmbedding),
                    u32::from(prism_round16),
                    0,
                    0,
                    0,
                ],
            );
            let layout = pipe.get_bind_group_layout(0);
            let bg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("prism-fwht"),
                layout: &layout,
                entries: &[
                    bind_buf(0, src),
                    bind_buf(1, &rot),
                    bind_buf(2, signs),
                    bind_buf(3, &p),
                ],
            });
            ts_point!(enc, 50);
            go(enc, pipe, &bg, (width / block) as u32);
            ts_point!(enc, 51);
            Some(rot.clone())
        };
    let prism_input =
        |enc: &mut wgpu::CommandEncoder,
         mats: &[&GMat],
         src: &wgpu::Buffer,
         width: usize|
         -> Option<wgpu::Buffer> {
            let mut op = crate::gpu::GraphPrismOp::None;
            for m in mats {
                if m.prism != crate::gpu::GraphPrismOp::None {
                    if op != crate::gpu::GraphPrismOp::None && op != m.prism {
                        return None;
                    }
                    op = m.prism;
                }
            }
            prism_transform(enc, src, width, op)
        };
    let flags = |qn: bool, kn: bool, late: bool| {
        (if qn { 2u32 } else { 0 })
            | (if kn { 4 } else { 0 })
            | (if gemma { 8 } else { 0 })
            | (if late { 32 } else { 0 })
    };
    // Constant uniforms for the whole token (position is fixed for this call).
    // Token-invariant ones use the content-keyed cache; position-dependent ones
    // use pooled buffers updated via write_buffer (no allocation after first token).
    let g = if gemma { 1u32 } else { 0 };
    let rms_u = uniform_u32x4(c, [hidden as u32, g, eps.to_bits(), 0]);
    let ax_u = uniform_u32x4(c, [1.0f32.to_bits(), hidden as u32, 0, 0]);
    let silu_u = uniform_u32x4(c, [inter as u32, 0, 0, 0]);
    let steps = steps.max(1);
    // One uniform PER STEP, with stable identities: write_buffer lands at
    // submit, so a single shared buffer would collapse every step to the
    // last position written. Slot 0 is the plain single-step path.
    let mku = |v: &mut Vec<wgpu::Buffer>, size: u64, count: usize| {
        while v.len() < count {
            v.push(c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("g-step-u"),
                size,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }));
        }
    };
    mku(&mut gs.kv_us, 16, steps);
    mku(&mut gs.at_us, 32, steps);
    mku(&mut gs.rope_us, 32, steps * layers.len());
    for st in 0..steps {
        let p = position + st;
        c.queue.write_buffer(
            &gs.kv_us[st],
            0,
            bytemuck::cast_slice(&[nkv as u32, hd as u32, cap as u32, p as u32]),
        );
        c.queue.write_buffer(
            &gs.at_us[st],
            0,
            bytemuck::cast_slice(&[
                nh as u32,
                (nh / nkv) as u32,
                hd as u32,
                cap as u32,
                (p + 1) as u32,
                attn_scale.to_bits(),
                0,
                0,
            ]),
        );
    }
    let kv_us = std::mem::take(&mut gs.kv_us);
    let at_us = std::mem::take(&mut gs.at_us);
    let rope_us = std::mem::take(&mut gs.rope_us);
    // The ladder sidecar is selected from the same resolved tensor identity
    // used by the graph's prep/emat descriptors.  Keeping this in one helper
    // prevents the graph path from accidentally using a component-only or
    // shape-only cache binding, and lets every fallback retain the ordinary
    // row-local q2tp kernel.
    let ladder = |m: &GMat,
                  xs: &wgpu::Buffer,
                  y: &wgpu::Buffer,
                  rows: usize,
                  cols: usize|
     -> Option<(&wgpu::ComputePipeline, wgpu::BindGroup, u32)> {
        if m.kind != 9
            || !m.affine
            || std::env::var("CMF_Q2_LADDER_CACHE").as_deref() != Ok("1")
        {
            return None;
        }
        let (ladders, row_ids, row_id_base) = ensure_q2_ladder_cache(c, model, m.idx, rows, cols)?;
        let pipe = c.q2_ladder_mv.as_ref()?;
        let gpr = cols / 32;
        let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, row_id_base, 1]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("graph-q2-ladder-mv-bg"),
            layout: &pipe.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &m.buf),
                bind_buf(2, y),
                bind_buf(3, &p_buf),
                bind_buf(5, xs),
                bind_buf(6, &ladders),
                bind_buf(7, &row_ids),
            ],
        });
        use std::sync::atomic::Ordering;
        Q2TP_LADDER_GRAPH_LOOKUPS.fetch_add(1, Ordering::Relaxed);
        Some((pipe, bind, mv_grid((rows as u32).div_ceil(16))))
    };
    // Encode one matvec, dtype-dispatched: q8_row (encode_matvec + row
    // scales) or q1 (encode_matvec_q1). Each is its own pass.
    //
    // MEASURED on Adreno 642L: the steady graph is one queue submission,
    // 87 compute passes and zero weight uploads per token. After the Adreno
    // q1 workgroup retile, a timestamped 0.992 s frame spent ~604 ms in
    // dense gate/up/down matvecs, ~209 ms in attention and ~178 ms in the
    // final norm/lm_head. Pass grouping is still useful for ordering and
    // encoder overhead, but kernel work is the wall; “token time = pass
    // count × 4.4 ms” was a false model.
    //
    // A/B experiments rejected q1 wave64 (-5.7%) and removing q1's
    // stride-33 padding (~-2%). The q8_2f row sweep was invalidated when the
    // measured artifact proved to contain zero q8_2f tensors. One change did
    // win: 16 rows/256 threads reused each staged q1 activation tile twice
    // as far and raised steady decode 0.589 -> 0.856 tok/s (+45%) on Adreno.
    //
    // Read that +45% with the caveat it deserves: it is WALL-CLOCK, and the
    // 8/128 baseline wandered 0.589/0.602/0.758 across sessions on this
    // phone — the same instability that made the q8_2f sweep look real.
    // Only the sign is safe. The kernel-level confirmation (a CMF_GPU_TS=2
    // frame at CMF_Q1_RPG=8 vs =16; q1 time should be ~1310 ms vs 782 ms)
    // has not been taken yet.
    let emat = |enc: &mut wgpu::CommandEncoder,
                m: &GMat,
                xs: &wgpu::Buffer,
                y: &wgpu::Buffer,
                rows: usize,
                cols: usize| {
        match m.kind {
            0 => encode_matvec(c, enc, &m.buf, xs, m.rs.as_ref().unwrap(), y, rows, cols),
            1 => encode_matvec_q1(c, enc, &m.buf, xs, y, rows, cols),
            2 => encode_q1t_like(c, enc, &c.q4b, &m.buf, xs, y, rows, cols),
            3 => encode_q1t_like(c, enc, &c.q1t, &m.buf, xs, y, rows, cols),
            5 => {
                if c.use_mv4 {
                    let gpr = cols / 32;
                    let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, cols as u32, 0]);
                    let layout = c.q4t_mv8.get_bind_group_layout(0);
                    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &layout,
                        entries: &[
                            bind_buf(0, &m.buf),
                            bind_buf(2, y),
                            bind_buf(3, &p_buf),
                            bind_buf(5, xs),
                        ],
                    });
                    let mut pass = begin_pass(enc);
                    pass.set_pipeline(&c.q4t_mv8);
                    pass.set_bind_group(0, &bind, &[]);
                    pass.dispatch_workgroups((rows as u32).div_ceil(8).min(MAX_WG), 1, 1);
                } else {
                    encode_q1t_like(c, enc, &c.q4t_mv, &m.buf, xs, y, rows, cols)
                }
            }
            6 => {
                if c.use_mv4 {
                    encode_q4tp_mv4(c, enc, &m.buf, xs, y, rows, cols)
                } else {
                    encode_q1t_like(c, enc, &c.q4tp_mv, &m.buf, xs, y, rows, cols)
                }
            }
            7 => {
                // Both scale planes live inside the buffer, so the kernel
                // needs the true `cols` alongside the word count.
                let p_buf = uniform_u32x4(c, [(cols / 4) as u32, rows as u32, cols as u32, 0]);
                let layout = c.q8_2f_mv.get_bind_group_layout(0);
                let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &layout,
                    entries: &[
                        bind_buf(0, &m.buf),
                        bind_buf(1, xs),
                        bind_buf(2, y),
                        bind_buf(3, &p_buf),
                    ],
                });
                let mut pass = begin_pass(enc);
                pass.set_pipeline(&c.q8_2f_mv);
                pass.set_bind_group(0, &bind, &[]);
                pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
            }
            9 => {
                if let Some((pipe, bind, groups)) = ladder(m, xs, y, rows, cols) {
                    let mut pass = begin_pass(enc);
                    pass.set_pipeline(pipe);
                    pass.set_bind_group(0, &bind, &[]);
                    pass.dispatch_workgroups(groups, 1, 1);
                    return;
                }
                if q2tp_dp4a_on() && m.affine && cols % 32 == 0 {
                    encode_q2tp_mv1_i8(c, enc, &m.buf, xs, y, rows, cols);
                    return;
                }
                let gpr = cols / 32;
                let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, 1, m.affine as u32]);
                let layout = q2tp_pipeline(c).get_bind_group_layout(0);
                let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &layout,
                    entries: &[
                        bind_buf(0, &m.buf),
                        bind_buf(2, y),
                        bind_buf(3, &p_buf),
                        bind_buf(5, xs),
                    ],
                });
                let mut pass = begin_pass(enc);
                pass.set_pipeline(q2tp_pipeline(c));
                pass.set_bind_group(0, &bind, &[]);
                pass.dispatch_workgroups(mv_grid((rows as u32).div_ceil(16)), 1, 1);
            }
            _ => encode_f32matvec(c, enc, &m.buf, xs, y, rows, cols),
        }
    };
    // Prep a matvec (pipeline, bind group, workgroups) WITHOUT opening a pass —
    // so several independent ones can share a pass. None = a dtype we don't
    // group (q4t/q1t) → caller falls back to per-op emat. The bind group keeps
    // its uniform buffer alive, so returning it alone is enough.
    let prep = |m: &GMat,
                xs: &wgpu::Buffer,
                y: &wgpu::Buffer,
                rows: usize,
                cols: usize|
     -> Option<(&wgpu::ComputePipeline, wgpu::BindGroup, u32)> {
        match m.kind {
            0 => {
                let p_buf = uniform_u32x4(c, [(cols / 4) as u32, rows as u32, 0, 0]);
                let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.layout,
                    entries: &[
                        bind_buf(0, &m.buf),
                        bind_buf(1, xs),
                        bind_buf(2, m.rs.as_ref().unwrap()),
                        bind_buf(3, y),
                        bind_buf(4, &p_buf),
                    ],
                });
                Some((&c.matvec, bind, (rows as u32).min(MAX_WG)))
            }
            7 => {
                // q8_2f, so the FFN's `down` can ride in the same serialized
                // pass as gate/up/silu. This removes encoder/pass overhead;
                // timestamps show that the matvec arithmetic remains the wall.
                let p_buf = uniform_u32x4(c, [(cols / 4) as u32, rows as u32, cols as u32, 0]);
                let layout = c.q8_2f_mv.get_bind_group_layout(0);
                let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &layout,
                    entries: &[
                        bind_buf(0, &m.buf),
                        bind_buf(1, xs),
                        bind_buf(2, y),
                        bind_buf(3, &p_buf),
                    ],
                });
                Some((&c.q8_2f_mv, bind, (rows as u32).min(MAX_WG)))
            }
            1 => {
                let gpr = cols / 32;
                let p_buf = uniform_u32x4(c, [(gpr / 2) as u32, rows as u32, 0, 0]);
                let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.layout_q1,
                    entries: &[
                        bind_buf(0, &m.buf),
                        bind_buf(1, xs),
                        bind_buf(2, y),
                        bind_buf(3, &p_buf),
                    ],
                });
                Some((&c.q1, bind, (rows as u32).div_ceil(c.q1_rows).min(MAX_WG)))
            }
            4 => {
                let p_buf = uniform_u32x4(c, [cols as u32, rows as u32, 0, 0]);
                let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.layout_f32,
                    entries: &[
                        bind_buf(0, &m.buf),
                        bind_buf(1, xs),
                        bind_buf(2, y),
                        bind_buf(3, &p_buf),
                    ],
                });
                Some((&c.f32_matvec, bind, (rows as u32).min(MAX_WG)))
            }
            // These arms must exist: without them `prep` returned None for
            // every q4t/q4tp projection, `group_mats` fell back to one
            // compute pass PER matvec and the MoE layer took its per-op
            // branch — a pass costs ~60 us on this Vulkan stack against
            // ~2 ms of arithmetic for the whole MoE block.
            2 | 5 | 6 => {
                let gpr = cols / 32;
                let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, cols as u32, 0]);
                if m.kind == 2 && c.use_mv4 {
                    let layout = c.q4b_mv8.get_bind_group_layout(0);
                    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &layout,
                        entries: &[
                            bind_buf(0, &m.buf),
                            bind_buf(2, y),
                            bind_buf(3, &p_buf),
                            bind_buf(4, &m.buf),
                            bind_buf(5, xs),
                        ],
                    });
                    return Some((&c.q4b_mv8, bind, (rows as u32).div_ceil(8).min(MAX_WG)));
                }
                if m.kind == 5 && c.use_mv4 {
                    // q4t's twin takes the same five bindings as q4tp's.
                    let layout = c.q4t_mv8.get_bind_group_layout(0);
                    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &layout,
                        // NO slot 4: q4t assembles weights from u16 halves
                        // (18 B tiles are 2-aligned), so the entry point
                        // never reads the vec4 weight view and its auto
                        // layout does not carry that binding.
                        entries: &[
                            bind_buf(0, &m.buf),
                            bind_buf(2, y),
                            bind_buf(3, &p_buf),
                            bind_buf(5, xs),
                        ],
                    });
                    return Some((&c.q4t_mv8, bind, (rows as u32).div_ceil(8).min(MAX_WG)));
                }
                if m.kind == 6 && c.use_mv4 {
                    // Wide rows: the quad-row kernel, four dot chains per
                    // x fetch (the pair kernel was x-LSU-bound at the
                    // dense FFN widths).
                    // `CMF_MV16W=0` sends the wide rows to the 8-row pair
                    // kernel instead — the SAME shape q4t's decode kernel
                    // runs, for a like-for-like A/B: on the RTX 5090 pod
                    // q4t (8-row pairs) decoded 7% FASTER than q4tp (quad
                    // 16-row) on Qwen3.8-27B despite 7.5% fewer bytes.
                    let (pipe6, per_wg) = if gpr <= 64 {
                        (&c.q4tp_mv16, 16u32)
                    } else if let Some(p) = c.q4tp_mv16w_probe.as_ref() {
                        (p, 16u32) // CMF_MV_PROBE: garbage answers, real timing
                    } else if !mv16w_on() {
                        (&c.q4tp_mv4, 8u32)
                    } else {
                        (&c.q4tp_mv16w, 16u32)
                    };
                    // Word 2 is the batch count for THIS pair and `cols` for
                    // everything else bound to the same struct — the shared
                    // uniform above cannot serve both. Word 3 carries the
                    // probe level when the probe is the one running.
                    let p6 = q4tp_mv_params_w(c, gpr, rows, 1, c.mv_probe);
                    let layout = pipe6.get_bind_group_layout(0);
                    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &layout,
                        entries: &[
                            bind_buf(0, &m.buf),
                            bind_buf(2, y),
                            bind_buf(3, &p6),
                            bind_buf(4, &m.buf),
                            bind_buf(5, xs),
                        ],
                    });
                    return Some((pipe6, bind, mv_grid((rows as u32).div_ceil(per_wg))));
                }
                let pl = match m.kind {
                    5 => &c.q4t_mv,
                    6 => &c.q4tp_mv,
                    _ => &c.q4b,
                };
                let layout = pl.get_bind_group_layout(0);
                let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &layout,
                    entries: &[
                        bind_buf(0, &m.buf),
                        bind_buf(1, xs),
                        bind_buf(2, y),
                        bind_buf(3, &p_buf),
                    ],
                });
                Some((pl, bind, (rows as u32).min(MAX_WG)))
            }
            3 => {
                let gpr = cols / 32;
                let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, cols as u32, 0]);
                let layout = c.q1t.get_bind_group_layout(0);
                let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &layout,
                    entries: &[
                        bind_buf(0, &m.buf),
                        bind_buf(1, xs),
                        bind_buf(2, y),
                        bind_buf(3, &p_buf),
                    ],
                });
                Some((&c.q1t, bind, (rows as u32).min(MAX_WG)))
            }
            9 => {
                // The 2-bit plane of the q2tp profile: its own quad-row kernel.
                if cols % 32 != 0 {
                    return None;
                }
                if let Some((pipe, bind, groups)) = ladder(m, xs, y, rows, cols) {
                    return Some((pipe, bind, groups));
                }
                // The NB=1 Q8/DP4A arm needs a preceding activation-quantize
                // dispatch, so it intentionally declines grouping; `group_mats`
                // then routes this operation through `emat`, which emits the
                // quantizer and the dedicated affine kernel in one pass.
                if q2tp_dp4a_on() && m.affine {
                    return None;
                }
                let gpr = cols / 32;
                let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, 1, m.affine as u32]);
                let layout = q2tp_pipeline(c).get_bind_group_layout(0);
                let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &layout,
                    entries: &[
                        bind_buf(0, &m.buf),
                        bind_buf(2, y),
                        bind_buf(3, &p_buf),
                        bind_buf(5, xs),
                    ],
                });
                Some((q2tp_pipeline(c), bind, mv_grid((rows as u32).div_ceil(16))))
            }
            _ => None,
        }
    };
    // Two projections over the same input in ONE dispatch (the x2
    // kernel): both q4tp, both wide (gpr > 64, the 16w regime), same
    // cols. None when the pair does not qualify — the caller falls back
    // to two `prep`s. Bit-identical per row to the single kernel.
    let prep2 = |a: &GMat,
                 b: &GMat,
                 xs: &wgpu::Buffer,
                 ya: &wgpu::Buffer,
                 yb: &wgpu::Buffer,
                 rows_a: usize,
                 rows_b: usize,
                 cols: usize|
     -> Option<(&wgpu::ComputePipeline, wgpu::BindGroup, u32)> {
        if !c.use_mv_x2 || !c.use_mv4 || a.kind != 6 || b.kind != 6 || cols % 32 != 0 {
            return None;
        }
        let gpr = cols / 32;
        if gpr <= 64 || c.q4tp_mv16w_probe.is_some() {
            return None;
        }
        let (bind, wg) = mv_x2_bind(c, &a.buf, &b.buf, xs, ya, yb, rows_a, rows_b, cols);
        Some((&c.q4tp_mv16w_x2, bind, wg))
    };
    // Emit a set of mutually-INDEPENDENT matvecs. When grouping is on and every
    // one preps, they share a single compute pass (no barrier between them);
    // otherwise each goes through emat as its own pass. Correctness rests on the
    // caller passing only matvecs with no read-after-write among them.
    let group_mats =
        |enc: &mut wgpu::CommandEncoder,
         mats: &[(&GMat, &wgpu::Buffer, &wgpu::Buffer, usize, usize)]| {
            if group {
                let prepped: Vec<_> = mats
                    .iter()
                    .filter_map(|(m, xs, y, r, cc)| prep(m, xs, y, *r, *cc))
                    .collect();
                if prepped.len() == mats.len() {
                    let mut pass = begin_pass(enc);
                    for (p, b, g) in &prepped {
                        pass.set_pipeline(p);
                        pass.set_bind_group(0, b, &[]);
                        pass.dispatch_workgroups(*g, 1, 1);
                    }
                    return;
                }
            }
            for (m, xs, y, r, cc) in mats {
                emat(enc, m, xs, y, *r, *cc);
            }
        };
    // Two projections of the same input under ONE dispatch. `false` = a
    // kind the paired kernel does not cover, caller keeps `group_mats`.
    let _pair_mats = |enc: &mut wgpu::CommandEncoder,
                      a: (&GMat, &wgpu::Buffer, usize, usize),
                      b: (&GMat, &wgpu::Buffer, usize, usize),
                      xs: &wgpu::Buffer|
     -> bool {
        let ok = |k: u8| k == 4 || k == 6;
        if !ok(a.0.kind) || !ok(b.0.kind) {
            return false;
        };
        let p = unif(&[
            a.2 as u32,
            a.3 as u32,
            a.0.kind as u32,
            0,
            b.2 as u32,
            b.3 as u32,
            b.0.kind as u32,
            0,
        ]);
        let bind = bgc(0, 0, &c.layout_mv2, &[&a.0.buf, &b.0.buf, xs, a.1, b.1, &p]);
        go(enc, &c.matvec_pair, &bind, ((a.2 + b.2) as u32).min(MAX_WG));
        true
    };
    let mut o1_dbg: Vec<(usize, wgpu::Buffer)> = Vec::new();
    // ── Multi-step prerequisites: the lm_head fold and a q4tp embedding,
    // both resolved up front. Anything missing refuses the WHOLE call so
    // the pipeline can fall back to single-step.
    // "multi" really means: the DEVICE picks the token(s) and the CPU
    // reads ids, not logits. k=1 rides it too — a 4-byte readback against
    // a megabyte of logits.
    let multi = ids_out.is_some();
    let lm_pre = lm_head.and_then(|(gw, rows)| resolve(gw, rows, hidden).map(|m| (m, rows)));
    let emb_pre =
        embed.and_then(|(gw, rows, mult)| resolve(gw, rows, hidden).map(|m| (m, rows, mult)));
    if multi {
        // The initial resident slice deliberately keeps the inverse embedding
        // on the CPU.  The existing gather kernel has no descriptor/sign
        // binding, so declining here is safer than returning a plausible but
        // untransformed multi-step stream.
        let embed_ok = matches!(&emb_pre, Some((m, _, _)) if m.kind == 6
            && m.prism == crate::gpu::GraphPrismOp::None);
        if lm_pre.is_none() || !embed_ok {
            graph_refused("multi-step needs the lm_head fold and a q4tp embedding");
            return token_graph_outcome(o1_started || state_started, false);
        }
    }
    // O(1) admission is completed before the first command is encoded.  A
    // later readback/submit failure must therefore be reported as Failed,
    // while a malformed or unsealed view remains an ordinary decline.  The
    // caller can clear the sequence on Failed instead of walking stale CPU
    // accumulators beside a partially advanced device state.
    // A budget prefix may deliberately truncate `layers`; the caller's O(1)
    // vector is for the full stack, so only a short vector is malformed here.
    if !o1.is_empty() && o1.len() < layers.len() {
        graph_refused("o1 layer/view count mismatch");
        return token_graph_outcome(o1_started || state_started, false);
    }
    let has_o1 = o1.iter().take(layers.len()).any(Option::is_some);
    if has_o1 {
        for (li, views) in o1.iter().take(layers.len()).enumerate() {
            match (&layers[li].attn, views) {
                (crate::gpu::GraphAttn::Full { .. }, Some(v)) => {
                    if !o1_views_valid(v, nh, nkv, hd) {
                        graph_refused("o1 view failed token admission");
                        return token_graph_outcome(o1_started || state_started, false);
                    }
                }
                (crate::gpu::GraphAttn::Full { .. }, None) => {}
                (crate::gpu::GraphAttn::Gdn { .. }, None) => {}
                (crate::gpu::GraphAttn::Gdn { .. }, Some(_))
                | (crate::gpu::GraphAttn::ShortConv { .. }, Some(_)) => {
                    graph_refused("o1 view attached to non-full attention");
                    return token_graph_outcome(o1_started || state_started, false);
                }
                (crate::gpu::GraphAttn::ShortConv { .. }, None) => {}
            }
        }
        {
            let om = c.o1m.lock().unwrap();
            for (li, views) in o1.iter().take(layers.len()).enumerate() {
                if views.is_some()
                    && om
                        .get(&(kv_id, layer_base + li))
                        .filter(|d| d.epoch == o1_epoch)
                        .and_then(|d| d.next_pos)
                        .is_some_and(|next| next != position)
                {
                    graph_refused("o1 device state position mismatch");
                    return token_graph_outcome(true, false);
                }
            }
        }
        // Set this before the first upload: an upload can allocate earlier
        // mirrors before a later layer rejects the same epoch.
        o1_started = true;
        for (li, views) in o1.iter().take(layers.len()).enumerate() {
            if let Some(views) = views {
                if o1_ensure(c, kv_id, layer_base + li, views, o1_epoch).is_none() {
                    graph_refused("o1 state not portable");
                    return token_graph_outcome(o1_started || state_started, false);
                }
            }
        }
    }
    const AM_PARTS: u32 = 512;
    let lm_rows_pre = lm_pre.as_ref().map(|(_, r)| *r).unwrap_or(0);
    let (lbuf_pre, am_pv, am_pi, ids_buf, ids_stage) = if multi {
        let lsize = (lm_rows_pre * 4) as u64;
        (
            Some(GraphScratch::ensure(
                &c.device,
                &mut gs.logits,
                lsize,
                wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                "g-logits",
            )),
            Some(GraphScratch::ensure(
                &c.device,
                &mut gs.am_pv,
                (AM_PARTS * 4) as u64,
                wgpu::BufferUsages::STORAGE,
                "g-am-pv",
            )),
            Some(GraphScratch::ensure(
                &c.device,
                &mut gs.am_pi,
                (AM_PARTS * 4) as u64,
                wgpu::BufferUsages::STORAGE,
                "g-am-pi",
            )),
            Some(GraphScratch::ensure(
                &c.device,
                &mut gs.ids,
                (steps * 4) as u64,
                wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                "g-ids",
            )),
            Some(GraphScratch::ensure(
                &c.device,
                &mut gs.ids_stage,
                (steps * 4) as u64,
                wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                "g-ids-stage",
            )),
        )
    } else {
        (None, None, None, None, None)
    };
    for stp in 0..steps {
        let kv_u = kv_us[stp].clone();
        let at_u = at_us[stp].clone();
        let position = position + stp;
        // Bootstrap the first layer's attention norm; thereafter each residual is
        // fused with the following norm (add_rmsnorm), saving two dispatches/layer.
        let inw0 = stor(bytemuck::cast_slice(layers[0].input_norm));
        ts!(enc, 0, 0);
        go(
            &mut enc,
            &c.rmsnorm,
            &bgc(0, 0, &c.layout_rmsnorm, &[&h_buf, &inw0, &n1, &rms_u]),
            1,
        );
        // Split the submission mid-stack (single-step decode only): the card
        // starts the first layers while the host still encodes the rest —
        // the DSV4 chain-split trick. Measured +12.8% on the 35B (97.9 ->
        // 110.4): it hides the encode AND the driver's submit latency.
        // Same queue, same order; nothing about the computation changes.
        // `CMF_GRAPH_SPLIT=N` pieces; 0/1 = historical single submit.
        let split_n = if preencode_frame { 1 } else { graph_split_n() };
        let chunk = if steps == 1 && split_n > 1 {
            // Never below four layers a piece: a short device prefix cut
            // into per-layer submits pays more in submissions than it
            // hides in encode.
            lws.len().div_ceil(split_n).max(4)
        } else {
            usize::MAX
        };
        for (li, l) in layers.iter().enumerate() {
            if li >= layer_cap {
                break;
            }
            if li > 0 && chunk != usize::MAX && li % chunk == 0 {
                flush_pass(&enc);
                let full = std::mem::replace(
                    &mut enc,
                    c.device
                        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                            label: Some("token-graph"),
                        }),
                );
                submit(c, finish_enc(full));
            }
            let lw = &lws[li];
            let rope_u = rope_us[stp * layers.len() + li].clone();
            let lkind: u8 = if matches!(lw.attn, LAttn::Full { .. }) {
                1
            } else {
                0
            };
            let pnw = stor(bytemuck::cast_slice(l.post_norm));
            // ── token mixing (attention or GDN) → ob ──
            match (&lw.attn, &l.attn) {
                (
                    LAttn::Full { wq, wk, wv, wo },
                    crate::gpu::GraphAttn::Full {
                        q_norm,
                        k_norm,
                        late_qk_norm,
                        bias,
                        output_gate,
                        ..
                    },
                ) => {
                    let o1_here = o1.get(li).and_then(|v| v.as_ref());
                    // true = the fused short-context arm already ran the output
                    // gate and the O projection inside its pass.
                    let mut attn_done = false;
                    let qnw = q_norm
                        .map(|q| stor(bytemuck::cast_slice(q)))
                        .unwrap_or_else(|| zeros(hd));
                    let knw = k_norm
                        .map(|k| stor(bytemuck::cast_slice(k)))
                        .unwrap_or_else(|| zeros(hd));
                    let gate_flag = if *output_gate { 1u32 } else { 0 };
                    c.queue.write_buffer(
                        &rope_u,
                        0,
                        bytemuck::cast_slice(&[
                            nh as u32,
                            nkv as u32,
                            hd as u32,
                            rd as u32,
                            position as u32,
                            flags(q_norm.is_some(), k_norm.is_some(), *late_qk_norm) | gate_flag,
                            eps.to_bits(),
                            0,
                        ]),
                    );
                    let qkv_in = match prism_input(
                        &mut enc,
                        &[wq, wk, wv],
                        &n1,
                        hidden,
                    ) {
                        Some(x) => x,
                        None => {
                            graph_decline("Prism QKV transform unavailable");
                            return token_graph_outcome(o1_started || state_started, false);
                        }
                    };
                    // Gated wq emits 2·nh·hd (q||gate interleaved per head); the rope
                    // kernel splits it, roping q and passing gate through to `gout`.
                    let qrows = nh * hd * (1 + *output_gate as usize);
                    // k+v in ONE dispatch (the x2 kernel) next to q, when
                    // both are wide q4tp; otherwise the grouped pass.
                    let pkv = if group {
                        prep2(wk, wv, &qkv_in, &kb, &vb, nkv * hd, nkv * hd, hidden)
                    } else {
                        None
                    };
                    match (pkv, prep(wq, &qkv_in, &qraw, qrows, hidden)) {
                        (Some((p2, b2, w2)), Some((pq, bq, wgq))) => {
                            let mut pass = begin_pass(&mut enc);
                            pass.set_pipeline(pq);
                            pass.set_bind_group(0, &bq, &[]);
                            pass.dispatch_workgroups(wgq, 1, 1);
                            pass.set_pipeline(p2);
                            pass.set_bind_group(0, &b2, &[]);
                            pass.dispatch_workgroups(w2, 1, 1);
                        }
                        _ => group_mats(
                            &mut enc,
                            &[
                                (wq, &qkv_in, &qraw, qrows, hidden),
                                (wk, &qkv_in, &kb, nkv * hd, hidden),
                                (wv, &qkv_in, &vb, nkv * hd, hidden),
                            ],
                        ),
                    }
                    if let Some((bq, bk, bv)) = bias {
                        let (bqb, bkb, bvb) = (
                            stor(bytemuck::cast_slice(bq)),
                            stor(bytemuck::cast_slice(bk)),
                            stor(bytemuck::cast_slice(bv)),
                        );
                        let axq = uniform_u32x4(c, [1.0f32.to_bits(), (nh * hd) as u32, 0, 0]);
                        let axkv = uniform_u32x4(c, [1.0f32.to_bits(), (nkv * hd) as u32, 0, 0]);
                        go(
                            &mut enc,
                            &c.axpy,
                            &bgc(13, li, &c.layout_axpy, &[&bqb, &qraw, &axq]),
                            ((nh * hd) as u32).div_ceil(256),
                        );
                        go(
                            &mut enc,
                            &c.axpy,
                            &bgc(14, li, &c.layout_axpy, &[&bkb, &kb, &axkv]),
                            ((nkv * hd) as u32).div_ceil(256),
                        );
                        go(
                            &mut enc,
                            &c.axpy,
                            &bgc(15, li, &c.layout_axpy, &[&bvb, &vb, &axkv]),
                            ((nkv * hd) as u32).div_ceil(256),
                        );
                    }
                    if let Some(views) = o1_here {
                        // O(1) attention: rope as usual, then the three o1
                        // kernels replace kv_append + attend. State mirrors on
                        // the device once per seal epoch; kv mirrors are not
                        // touched for this layer at all.
                        // A network/in-process span uses an absolute layer
                        // base. Keep the O(1) mirror key aligned with the
                        // ordinary KV/GDN mirrors so spans cannot collide
                        // with layer zero or unwrap a missing state.
                        if o1_ensure(c, kv_id, layer_base + li, views, o1_epoch).is_none() {
                            graph_refused("o1 state not portable");
                            return token_graph_outcome(o1_started || state_started, false);
                        }
                        let (
                            dmeta,
                            drk,
                            drv,
                            dsk,
                            dsv,
                            dkt,
                            dqt,
                            dmu,
                            dmz,
                            dth,
                            gg,
                            hh_,
                            mm,
                            ww,
                            nns,
                            sc,
                        ) = {
                            let map = c.o1m.lock().unwrap();
                            let d = map.get(&(kv_id, layer_base + li)).unwrap();
                            (
                                d.meta.clone(),
                                d.ring_k.clone(),
                                d.ring_v.clone(),
                                d.sink_k.clone(),
                                d.sink_v.clone(),
                                d.k_tilde.clone(),
                                d.qt.clone(),
                                d.mu.clone(),
                                d.mz.clone(),
                                d.that.clone(),
                                d.g,
                                d.h,
                                d.m,
                                d.w,
                                d.ns,
                                d.scale,
                            )
                        };
                        let rect_fm = views
                            .first()
                            .and_then(|v| v.heads.first())
                            .is_some_and(|h| h.rect_fm);
                        let o1_u = uniform_u32x8(
                            c,
                            [
                                hh_ as u32,
                                mm as u32,
                                ww as u32,
                                (nns as u32) | (u32::from(rect_fm) << 8),
                                hd as u32,
                                hd as u32,
                                sc.to_bits(),
                                0,
                            ],
                        );
                        let bg_rope = bg(
                            &c.layout_attn_rope,
                            &[&qraw, &kb, &qout, &gout, &qnw, &knw, &invf_b, &rope_u],
                        );
                        let bg_far = bg(
                            &c.layout_o1_far,
                            &[&dmeta, &drk, &drv, &dqt, &dmz, &dth, &o1_u],
                        );
                        let bg_push = bgc(
                            16,
                            li,
                            &c.layout_o1_push,
                            &[&dmeta, &kb, &vb, &drk, &drv, &o1_u],
                        );
                        let bg_att = bg(
                            &c.layout_o1_attend,
                            &[
                                &dmeta, &qout, &drk, &drv, &dsk, &dsv, &dkt, &dmu, &dmz, &dth,
                                &attn, &o1_u,
                            ],
                        );
                        let mut pass = begin_pass(&mut enc);
                        pass.set_pipeline(&c.attn_rope);
                        pass.set_bind_group(0, &bg_rope, &[]);
                        pass.dispatch_workgroups((nh + nkv) as u32, 1, 1);
                        pass.set_pipeline(&c.o1_far);
                        pass.set_bind_group(0, &bg_far, &[]);
                        pass.dispatch_workgroups((gg * hh_ * mm) as u32, 1, 1);
                        pass.set_pipeline(&c.o1_push);
                        pass.set_bind_group(0, &bg_push, &[]);
                        pass.dispatch_workgroups(gg as u32, 1, 1);
                        pass.set_pipeline(&c.o1_attend);
                        pass.set_bind_group(0, &bg_att, &[]);
                        pass.dispatch_workgroups((gg * hh_) as u32, 1, 1);
                        drop(pass);
                        if std::env::var("CMF_O1_TRACE").is_ok() {
                            // Debug-only: stage this layer's o1 attention output
                            // for a post-submit dump (the attn buffer itself is
                            // reused by every later layer).
                            let dbgb = c.device.create_buffer(&wgpu::BufferDescriptor {
                                label: Some("o1-dbg"),
                                size: (nh * hd * 4) as u64,
                                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                                mapped_at_creation: false,
                            });
                            flush_pass(&enc);
                            enc.copy_buffer_to_buffer(&attn, 0, &dbgb, 0, (nh * hd * 4) as u64);
                            o1_dbg.push((li, dbgb));
                            let dbgq = c.device.create_buffer(&wgpu::BufferDescriptor {
                                label: Some("o1-dbg-q"),
                                size: 64,
                                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                                mapped_at_creation: false,
                            });
                            flush_pass(&enc);
                            enc.copy_buffer_to_buffer(&qout, 0, &dbgq, 0, 16);
                            flush_pass(&enc);
                            enc.copy_buffer_to_buffer(&kb, 0, &dbgq, 16, 16);
                            flush_pass(&enc);
                            enc.copy_buffer_to_buffer(&vb, 0, &dbgq, 32, 16);
                            o1_dbg.push((li + 10_000, dbgq));
                        }
                    } else {
                        let (kbuf, vbuf) = kvbufs[li].as_ref().unwrap();
                        // rope + kv_append are independent (both read kb, neither
                        // writes it) — share ONE compute pass to avoid the inter-pass
                        // pipeline flush (~78 μs on NVIDIA Vulkan).
                        let n_ctx = position + 1;
                        // Short context (the decode regime): attend + output gate +
                        // O-projection ride the SAME pass as rope/kv when they prep —
                        // five passes become one, and dispatch order is unchanged.
                        let short_ctx = !(n_ctx > ATTEND_SPLIT_MIN && (hd <= 128 || c.big_attend));
                        if passfuse && short_ctx && !skip_attn {
                            attn_done = true;
                            let (mut ap, al) = attend_pipes(c, hd);
                            let dec_l;
                            let bg_att = if c.attend_dec && hd <= 256 {
                                ap = &c.gqa_attend_dec;
                                // Auto layouts are pipeline-exclusive — the twin's
                                // binding SET matches, its layout object does not.
                                dec_l = c.gqa_attend_dec.get_bind_group_layout(0);
                                bgc(17, li, &dec_l, &[&qout, kbuf, vbuf, &attn, &at_u])
                            } else {
                                bg(al, &[&qout, kbuf, vbuf, &attn, &at_u])
                            };
                            let gm = if *output_gate {
                                let gm_u = uniform_u32x4(c, [(nh * hd) as u32, 0, 0, 0]);
                                Some(bgc(18, li, &c.layout_gate_mul, &[&gout, &attn, &gm_u]))
                            } else {
                                None
                            };
                            // A Prism O projection needs its FWHT after the
                            // attention result exists, so it cannot be
                            // pre-bound into this earlier fused pass.
                            let wo_prep = if wo.prism == crate::gpu::GraphPrismOp::None {
                                prep(wo, &attn, &ob, hidden, nh * hd)
                            } else {
                                None
                            };
                            {
                                let bg_rope = bg(
                                    &c.layout_attn_rope,
                                    &[&qraw, &kb, &qout, &gout, &qnw, &knw, &invf_b, &rope_u],
                                );
                                let bg_kv =
                                    bgc(19, li, &c.layout_kv, &[&kb, &vb, kbuf, vbuf, &kv_u]);
                                let mut pass = begin_pass(&mut enc);
                                let fine = ts_full || li < 4;
                                tsp!(pass, fine, 20); // pass start (after qkv projections)
                                pass.set_pipeline(&c.attn_rope);
                                pass.set_bind_group(0, &bg_rope, &[]);
                                pass.dispatch_workgroups((nh + nkv) as u32, 1, 1);
                                tsp!(pass, fine, 21); // rope
                                                      // Exact Full attention has now admitted a
                                                      // persistent K/V mutation.  If a later
                                                      // dispatch or readback fails, the CPU cache
                                                      // cannot safely resume this sequence.
                                state_started = true;
                                pass.set_pipeline(&c.kv_append);
                                pass.set_bind_group(0, &bg_kv, &[]);
                                pass.dispatch_workgroups(((nkv * hd) as u32).div_ceil(256), 1, 1);
                                tsp!(pass, fine, 22); // kv append
                                pass.set_pipeline(ap);
                                pass.set_bind_group(0, &bg_att, &[]);
                                pass.dispatch_workgroups(nh as u32, 1, 1);
                                tsp!(pass, fine, 23); // attend
                                if let Some(bg_gm) = &gm {
                                    pass.set_pipeline(&c.gate_mul);
                                    pass.set_bind_group(0, bg_gm, &[]);
                                    pass.dispatch_workgroups(
                                        ((nh * hd) as u32).div_ceil(256),
                                        1,
                                        1,
                                    );
                                }
                                tsp!(pass, fine, 24); // gate
                                if let Some((wp, wb, ww)) = &wo_prep {
                                    pass.set_pipeline(wp);
                                    pass.set_bind_group(0, wb, &[]);
                                    pass.dispatch_workgroups(*ww, 1, 1);
                                }
                                tsp!(pass, fine, 25); // o-proj
                            }
                            if wo_prep.is_none() {
                                let Some(wo_in) =
                                    prism_input(&mut enc, &[wo], &attn, nh * hd)
                                else {
                                    graph_decline("Prism O transform unavailable");
                                    return token_graph_outcome(o1_started || state_started, false);
                                };
                                emat(&mut enc, wo, &wo_in, &ob, hidden, nh * hd);
                            }
                        } else {
                            {
                                let bg_rope = bg(
                                    &c.layout_attn_rope,
                                    &[&qraw, &kb, &qout, &gout, &qnw, &knw, &invf_b, &rope_u],
                                );
                                let bg_kv =
                                    bgc(20, li, &c.layout_kv, &[&kb, &vb, kbuf, vbuf, &kv_u]);
                                let mut pass = begin_pass(&mut enc);
                                pass.set_pipeline(&c.attn_rope);
                                pass.set_bind_group(0, &bg_rope, &[]);
                                pass.dispatch_workgroups((nh + nkv) as u32, 1, 1);
                                // See the fused short-context arm above:
                                // kv_append is the Full-attention admission
                                // boundary for the terminal outcome contract.
                                state_started = true;
                                pass.set_pipeline(&c.kv_append);
                                pass.set_bind_group(0, &bg_kv, &[]);
                                pass.dispatch_workgroups(((nkv * hd) as u32).div_ceil(256), 1, 1);
                            }
                            if n_ctx > ATTEND_SPLIT_MIN && (hd <= 128 || c.big_attend) {
                                // Split-K attend: (nh × chunks) part workgroups + a
                                // per-head merge, both in ONE pass (WebGPU orders
                                // dispatches within a pass, so the merge sees the
                                // partials without an inter-pass flush).
                                // GQA-shared kernel: one workgroup per (kv head,
                                // 256-position chunk) serves all its query heads.
                                let hpk = nh / nkv;
                                let gqa = c.attend_gpart.is_some()
                                    && hpk <= 8
                                    && hd <= 256
                                    && hd % 4 == 0
                                    && nh % nkv == 0;
                                let ck = if gqa { ATTEND_GCK } else { ATTEND_CK };
                                let nc = cap.div_ceil(ck);
                                let nc_used = n_ctx.div_ceil(ck);
                                let pacc = GraphScratch::ensure(
                                    &c.device,
                                    &mut gs.apacc,
                                    (nh * nc * hd * 4) as u64,
                                    st,
                                    "g-apacc",
                                );
                                let pml = GraphScratch::ensure(
                                    &c.device,
                                    &mut gs.apml,
                                    (nh * nc * 8) as u64,
                                    st,
                                    "g-apml",
                                );
                                let ap_u = unif(&[
                                    nh as u32,
                                    (nh / nkv) as u32,
                                    hd as u32,
                                    cap as u32,
                                    n_ctx as u32,
                                    ck as u32,
                                    nc as u32,
                                    attn_scale.to_bits(),
                                ]);
                                let (pp, pl) = if gqa {
                                    (
                                        c.attend_gpart.as_ref().unwrap(),
                                        c.layout_attend_gpart.as_ref().unwrap(),
                                    )
                                } else {
                                    attend_part_pipes(c, hd)
                                };
                                let part_x = if gqa { nkv as u32 } else { nh as u32 };
                                let bg_part = bg(pl, &[&qout, kbuf, vbuf, &pacc, &pml, &ap_u]);
                                let bg_merge =
                                    c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                        label: None,
                                        layout: &c.layout_attend_merge,
                                        entries: &[
                                            bind_buf(3, &pacc),
                                            bind_buf(4, &pml),
                                            bind_buf(5, &ap_u),
                                            bind_buf(6, &attn),
                                        ],
                                    });
                                let mut pass = begin_pass(&mut enc);
                                pass.set_pipeline(pp);
                                pass.set_bind_group(0, &bg_part, &[]);
                                pass.dispatch_workgroups(part_x, nc_used as u32, 1);
                                pass.set_pipeline(&c.attend_merge);
                                pass.set_bind_group(0, &bg_merge, &[]);
                                pass.dispatch_workgroups(nh as u32, 1, 1);
                            } else if !skip_attn {
                                let (ap, al) = attend_pipes(c, hd);
                                go(
                                    &mut enc,
                                    ap,
                                    &bg(al, &[&qout, kbuf, vbuf, &attn, &at_u]),
                                    nh as u32,
                                );
                            }
                        } // fused-vs-split attend arms
                          // attn_out *= sigmoid(gate) before the O projection.
                    }
                    if *output_gate && !attn_done {
                        let gm_u = uniform_u32x4(c, [(nh * hd) as u32, 0, 0, 0]);
                        go(
                            &mut enc,
                            &c.gate_mul,
                            &bgc(21, li, &c.layout_gate_mul, &[&gout, &attn, &gm_u]),
                            ((nh * hd) as u32).div_ceil(256),
                        );
                    }
                    if !attn_done {
                        let Some(wo_in) = prism_input(&mut enc, &[wo], &attn, nh * hd) else {
                            graph_decline("Prism O transform unavailable");
                            return token_graph_outcome(o1_started || state_started, false);
                        };
                        emat(&mut enc, wo, &wo_in, &ob, hidden, nh * hd);
                    }
                }
                (
                    LAttn::Gdn {
                        qkv,
                        z,
                        a,
                        b,
                        out,
                        nv,
                        nk,
                        dk,
                        dv,
                        kk,
                        cdim,
                    },
                    crate::gpu::GraphAttn::Gdn {
                        conv1d,
                        a_log,
                        dt_bias,
                        norm,
                        ..
                    },
                ) => {
                    state_started = true;
                    let (ring, s) = gdnbufs[li].as_ref().unwrap();
                    let taps = stor(bytemuck::cast_slice(conv1d));
                    let alog = stor(bytemuck::cast_slice(a_log));
                    let dtb = stor(bytemuck::cast_slice(dt_bias));
                    let gnorm = stor(bytemuck::cast_slice(norm));
                    let gc_p = uniform_u32x4(c, [*cdim as u32, *kk as u32, 0, 0]);
                    let gd_p = unif(&[
                        *nv as u32,
                        *dk as u32,
                        *dv as u32,
                        (nk * dk) as u32,
                        (nv / nk) as u32,
                        *cdim as u32,
                        eps.to_bits(),
                        0,
                    ]);
                    let bg_conv = bgc(
                        22,
                        li,
                        &c.layout_gdn_conv,
                        &[&qkv_b, &taps, ring, &cq_b, &gc_p],
                    );
                    let bg_step = bg(
                        &c.layout_gdn,
                        &[
                            &cq_b, &z_b, &a_b, &b_b, &alog, &dtb, &gnorm, s, &gdo_b, &gd_p,
                        ],
                    );
                    // The parallel step/norm entries use SUBSETS of the gdn
                    // binding set, and an auto layout lists only what its entry
                    // reads — each gets its own bind group (lesson of the day).
                    let bg_par = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &c.gdn_step_par.get_bind_group_layout(0),
                        entries: &[
                            bind_buf(0, &cq_b),
                            bind_buf(2, &a_b),
                            bind_buf(3, &b_b),
                            bind_buf(4, &alog),
                            bind_buf(5, &dtb),
                            bind_buf(7, s),
                            bind_buf(8, &gdo_b),
                            bind_buf(9, &gd_p),
                        ],
                    });
                    let bg_snorm = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &c.gdn_step_norm.get_bind_group_layout(0),
                        entries: &[
                            bind_buf(1, &z_b),
                            bind_buf(6, &gnorm),
                            bind_buf(8, &gdo_b),
                            bind_buf(9, &gd_p),
                        ],
                    });
                    let gi_u = uniform_u32x4(c, [*kk as u32, 0, 0, 0]);
                    let (bg_par2, bg_snorm2) = if c.gdn_inline {
                        (
                            Some(c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                label: None,
                                layout: &c.gdn_step_par2.get_bind_group_layout(0),
                                entries: &[
                                    bind_buf(2, &a_b),
                                    bind_buf(3, &b_b),
                                    bind_buf(4, &alog),
                                    bind_buf(5, &dtb),
                                    bind_buf(7, s),
                                    bind_buf(8, &gdo_b),
                                    bind_buf(9, &gd_p),
                                    bind_buf(10, &qkv_b),
                                    bind_buf(11, ring),
                                    bind_buf(12, &taps),
                                    bind_buf(13, &gi_u),
                                ],
                            })),
                            Some(c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                label: None,
                                layout: &c.gdn_step_norm2.get_bind_group_layout(0),
                                entries: &[
                                    bind_buf(1, &z_b),
                                    bind_buf(6, &gnorm),
                                    bind_buf(8, &gdo_b),
                                    bind_buf(9, &gd_p),
                                    bind_buf(10, &qkv_b),
                                    bind_buf(11, ring),
                                    bind_buf(13, &gi_u),
                                ],
                            })),
                        )
                    } else {
                        (None, None)
                    };
                    // The whole GDN chain in ONE compute pass: projections →
                    // conv → step → out_proj. Each stage reads the previous
                    // stage's output, which is exactly what a pass guarantees
                    // (dispatches inside it are ordered, with memory visible
                    // between them — the same rule the fused SiLU FFN relies on).
                    //
                    // A pass, not a dispatch, is the unit that costs here:
                    // teaching `prep` the q4tp kind collapsed this layer's four
                    // projection passes into one and bought 1.46 ms a token
                    // across 30 layers — ~16 us per pass. Four passes become one.
                    // qkv+z ride ONE dispatch when both are wide q4tp
                    // (the x2 kernel); a and b (tiny f32 rows) ride ONE
                    // dispatch of the pair kernel, whose f32 arm is
                    // `f32_matvec` lane for lane (64-stride, same tree) —
                    // four projection launches become two.
                    let qz_in = match prism_input(&mut enc, &[qkv, z], &n1, hidden) {
                        Some(x) => x,
                        None => {
                            graph_decline("Prism GDN QKV transform unavailable");
                            return token_graph_outcome(o1_started || state_started, false);
                        }
                    };
                    // The GDN output transform depends on the recurrent step,
                    // so a Prism output cannot be pre-bound into the same
                    // projection pass.  The fallback arm below emits it after
                    // the stateful step has written gdo_b. Decide that
                    // eligibility before preparing any of the input
                    // projections: Prism's transformed output necessarily
                    // takes the fallback below, so its first four prep calls
                    // would be discarded and repeated by group_mats.
                    let prism_output = out.prism != crate::gpu::GraphPrismOp::None;
                    let outp = if !prism_output {
                        prep(out, &gdo_b, &ob, hidden, nv * dv)
                    } else {
                        None
                    };
                    let (n_proj, projs) = if prism_output {
                        use std::sync::atomic::Ordering;
                        GDN_PRISM_SKIPPED_PREPS.fetch_add(4, Ordering::Relaxed);
                        (4, std::array::from_fn(|_| None))
                    } else {
                        let pqz = prep2(qkv, z, &qz_in, &qkv_b, &z_b, *cdim, nv * dv, hidden);
                        let pab = if pqz.is_some() && a.kind == 4 && b.kind == 4 && c.use_mv_x2 {
                            let pu = unif(&[
                                *nv as u32,
                                hidden as u32,
                                4,
                                0,
                                *nv as u32,
                                hidden as u32,
                                4,
                                0,
                            ]);
                            let bind = bg(&c.layout_mv2, &[&a.buf, &b.buf, &n1, &a_b, &b_b, &pu]);
                            Some((&c.matvec_pair, bind, ((2 * *nv) as u32).min(MAX_WG)))
                        } else {
                            None
                        };
                        let n_proj = match (pqz.is_some(), pab.is_some()) {
                            (true, true) => 2,
                            (true, false) => 3,
                            _ => 4,
                        };
                        let projs = match (pqz, pab) {
                            (Some(p2), Some(pp)) => [Some(p2), Some(pp), None, None],
                            (Some(p2), None) => [
                                Some(p2),
                                prep(a, &n1, &a_b, *nv, hidden),
                                prep(b, &n1, &b_b, *nv, hidden),
                                None,
                            ],
                            (None, _) => [
                                prep(qkv, &qz_in, &qkv_b, *cdim, hidden),
                                prep(z, &qz_in, &z_b, nv * dv, hidden),
                                prep(a, &n1, &a_b, *nv, hidden),
                                prep(b, &n1, &b_b, *nv, hidden),
                            ],
                        };
                        (n_proj, projs)
                    };
                    let projs_ok = projs.iter().take(n_proj).all(|p| p.is_some());
                    if projs_ok && outp.is_some() {
                        let _ = (skip_proj, skip_outp);
                        let mut pass = begin_pass(&mut enc);
                        if !skip_proj {
                            for p in projs.iter().flatten() {
                                pass.set_pipeline(p.0);
                                pass.set_bind_group(0, &p.1, &[]);
                                pass.dispatch_workgroups(p.2, 1, 1);
                            }
                        }
                        let fine = ts_full || li == 0;
                        tsp!(pass, fine, 10); // after projections
                        if !skip_gdn {
                            if c.gdn_par && c.gdn_inline {
                                pass.set_pipeline(&c.gdn_step_par2);
                                pass.set_bind_group(0, bg_par2.as_ref().unwrap(), &[]);
                                pass.dispatch_workgroups(*nv as u32, *dv as u32, 1);
                                tsp!(pass, fine, 12); // step_par (conv inline)
                                pass.set_pipeline(&c.gdn_step_norm2);
                                pass.set_bind_group(0, bg_snorm2.as_ref().unwrap(), &[]);
                                pass.dispatch_workgroups(*nv as u32, 1, 1);
                                tsp!(pass, fine, 13); // step_norm (+ring shift)
                            } else {
                                pass.set_pipeline(&c.gdn_conv);
                                pass.set_bind_group(0, &bg_conv, &[]);
                                pass.dispatch_workgroups((*cdim as u32).div_ceil(256), 1, 1);
                                tsp!(pass, fine, 11); // conv
                                if c.gdn_par {
                                    pass.set_pipeline(&c.gdn_step_par);
                                    pass.set_bind_group(0, &bg_par, &[]);
                                    pass.dispatch_workgroups(
                                        *nv as u32,
                                        (*dv as u32).div_ceil(4),
                                        1,
                                    );
                                    tsp!(pass, fine, 12); // step_par
                                    pass.set_pipeline(&c.gdn_step_norm);
                                    pass.set_bind_group(0, &bg_snorm, &[]);
                                    pass.dispatch_workgroups(*nv as u32, 1, 1);
                                    tsp!(pass, fine, 13); // step_norm
                                } else {
                                    pass.set_pipeline(&c.gdn_step);
                                    pass.set_bind_group(0, &bg_step, &[]);
                                    pass.dispatch_workgroups(*nv as u32, 1, 1);
                                    tsp!(pass, fine, 12);
                                }
                            } // gdn_inline arms
                        }
                        if !skip_outp {
                            let o = outp.as_ref().unwrap();
                            pass.set_pipeline(o.0);
                            pass.set_bind_group(0, &o.1, &[]);
                            pass.dispatch_workgroups(o.2, 1, 1);
                            tsp!(pass, fine, 14); // out-proj
                        }
                    } else {
                        group_mats(
                            &mut enc,
                            &[
                                (qkv, &qz_in, &qkv_b, *cdim, hidden),
                                (z, &qz_in, &z_b, nv * dv, hidden),
                                (a, &n1, &a_b, *nv, hidden),
                                (b, &n1, &b_b, *nv, hidden),
                            ],
                        );
                        go(
                            &mut enc,
                            &c.gdn_conv,
                            &bg_conv,
                            (*cdim as u32).div_ceil(256),
                        );
                        if c.gdn_par {
                            {
                                let mut pass = begin_pass(&mut enc);
                                pass.set_pipeline(&c.gdn_step_par);
                                pass.set_bind_group(0, &bg_par, &[]);
                                pass.dispatch_workgroups(*nv as u32, (*dv as u32).div_ceil(4), 1);
                                pass.set_pipeline(&c.gdn_step_norm);
                                pass.set_bind_group(0, &bg_snorm, &[]);
                                pass.dispatch_workgroups(*nv as u32, 1, 1);
                            }
                        } else {
                            go(&mut enc, &c.gdn_step, &bg_step, *nv as u32);
                        }
                        let Some(out_in) = prism_input(&mut enc, &[out], &gdo_b, nv * dv) else {
                            graph_decline("Prism GDN output transform unavailable");
                            return token_graph_outcome(o1_started || state_started, false);
                        };
                        emat(&mut enc, out, &out_in, &ob, hidden, nv * dv);
                    }
                }
                (
                    LAttn::Conv { inp, out, kernel },
                    crate::gpu::GraphAttn::ShortConv { taps, .. },
                ) => {
                    state_started = true;
                    let (ring, _) = gdnbufs[li].as_ref().unwrap();
                    let taps_b = stor(bytemuck::cast_slice(*taps));
                    // GcP reused verbatim: cdim = hidden, kk = kernel.
                    let sc_u = uniform_u32x4(c, [hidden as u32, *kernel as u32, 0, 0]);
                    let Some(inp_in) = prism_input(&mut enc, &[inp], &n1, hidden) else {
                        graph_decline("Prism short-conv input transform unavailable");
                        return token_graph_outcome(o1_started || state_started, false);
                    };
                    emat(&mut enc, inp, &inp_in, &sc_bcx, 3 * hidden, hidden);
                    let bg_sc = bgc(
                        33,
                        li,
                        &c.layout_sconv,
                        &[&sc_bcx, &taps_b, ring, &sc_y, &sc_u],
                    );
                    {
                        let mut pass = begin_pass(&mut enc);
                        pass.set_pipeline(&c.sconv_step);
                        pass.set_bind_group(0, &bg_sc, &[]);
                        pass.dispatch_workgroups((hidden as u32).div_ceil(256), 1, 1);
                    }
                    let Some(out_in) = prism_input(&mut enc, &[out], &sc_y, hidden) else {
                        graph_decline("Prism short-conv output transform unavailable");
                        return token_graph_outcome(o1_started || state_started, false);
                    };
                    emat(&mut enc, out, &out_in, &ob, hidden, hidden);
                }
                _ => return token_graph_outcome(o1_started || state_started, false),
            }
            ts!(enc, 1, lkind);
            // token-mix residual + FFN-norm fused: h += ob, n1 = rms(h, post_norm).
            // It used to open its own compute pass. On this Vulkan stack a PASS
            // BOUNDARY is the expensive part (the MoE block is built entirely
            // around that fact), and this one sits between two passes that are
            // strictly serial anyway — so hand it to the FFN pass as a prologue
            // and let within-pass serialization do the same job for free.
            // CMF_PASSFUSE=0 puts it back in its own pass.
            let mut ffn_pre: Option<(&wgpu::ComputePipeline, wgpu::BindGroup, u32)> = None;
            if !skip_norm {
                let nbg = bgc(
                    23,
                    li,
                    &c.layout_add_rmsnorm,
                    &[&h_buf, &ob, &pnw, &n1, &rms_u],
                );
                if fuse_pre {
                    ffn_pre = Some((&c.add_rmsnorm, nbg, 1));
                } else {
                    go(&mut enc, &c.add_rmsnorm, &nbg, 1);
                }
            }
            // …and the layer's TAIL (FFN residual + the next layer's input norm)
            // rides out on the same pass. It reads `ob`, which that pass's last
            // dispatch writes — the same within-pass ordering the block above
            // relies on. With both ends folded in, a layer is TWO passes
            // (token-mix, then FFN) instead of four.
            let simple_tail = fuse_tail && !loop_norm_at.contains(&li);
            let mut ffn_post: Option<(&wgpu::ComputePipeline, wgpu::BindGroup, u32)> = None;
            let mut tail_done = false;
            if simple_tail {
                ffn_post = Some(if li + 1 < layers.len() {
                    let inw_next = stor(bytemuck::cast_slice(layers[li + 1].input_norm));
                    (
                        &c.add_rmsnorm,
                        bg(
                            &c.layout_add_rmsnorm,
                            &[&h_buf, &ob, &inw_next, &n1, &rms_u],
                        ),
                        1,
                    )
                } else {
                    (
                        &c.axpy,
                        bgc(24, li, &c.layout_axpy, &[&ob, &h_buf, &ax_u]),
                        (hidden as u32).div_ceil(256),
                    )
                });
            }
            // SiLU FFN: gate+up matvecs + silu fused in ONE compute pass
            // (dispatches within a pass are serialized — silu safely reads gate/up output).
            match &lw.ffn {
                LFfn::Dense {
                    gate,
                    up,
                    down,
                    width,
                } => {
                    let inter = *width; // this layer's, not the model's
                    let mut continue_ffn = false;
                    // DUAL (CMF_MV_DUAL=1): gate and up in ONE dispatch —
                    // the elimination table's verdict was the wave drain
                    // between serialized dispatches, and this deletes one
                    // of the four per layer. Falls through to the split
                    // path whenever either side is not plain q4tp.
                    if c.use_mv_dual
                        && gate.kind == 2
                        && up.kind == 2
                        && gate.prism == crate::gpu::GraphPrismOp::None
                        && up.prism == crate::gpu::GraphPrismOp::None
                        && hidden % 32 == 0
                    {
                        let gpr = hidden / 32;
                        let p_buf = uniform_u32x4(
                            c,
                            [gpr as u32, inter as u32, hidden as u32, inter as u32],
                        );
                        let layout = c.q4tp_mv4_dual.get_bind_group_layout(0);
                        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                            label: Some("ffn-dual"),
                            layout: &layout,
                            entries: &[
                                bind_buf(0, &gate.buf),
                                bind_buf(2, &gbuf),
                                bind_buf(3, &p_buf),
                                bind_buf(4, &gate.buf),
                                bind_buf(5, &n1),
                                bind_buf(6, &up.buf),
                                bind_buf(7, &up.buf),
                                bind_buf(8, &ubuf),
                            ],
                        });
                        let blocks = ((inter as u32).div_ceil(8) * 2).min(MAX_WG);
                        let bg_silu = bgc(
                            25,
                            li,
                            &c.layout_silu,
                            &[&gbuf, &ubuf, &dummy_hd, &abuf, &silu_u],
                        );
                        let mut pass = begin_pass(&mut enc);
                        if let Some((p, b, w)) = &ffn_pre {
                            pass.set_pipeline(p);
                            pass.set_bind_group(0, b, &[]);
                            pass.dispatch_workgroups(*w, 1, 1);
                        }
                        pass.set_pipeline(&c.q4tp_mv4_dual);
                        pass.set_bind_group(0, &bind, &[]);
                        pass.dispatch_workgroups(blocks, 1, 1);
                        if down.kind == 2 && inter % 32 == 0 {
                            // Fused down: SiLU inline — the whole FFN is
                            // two waves (dual, dsilu-down) instead of five.
                            let gpr_d = inter / 32;
                            let pd_buf =
                                uniform_u32x4(c, [gpr_d as u32, hidden as u32, inter as u32, 0]);
                            let layout = c.q4tp_mv4_dsilu.get_bind_group_layout(0);
                            let bind_d = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                label: Some("ffn-dsilu"),
                                layout: &layout,
                                entries: &[
                                    bind_buf(0, &down.buf),
                                    bind_buf(2, &ob),
                                    bind_buf(3, &pd_buf),
                                    bind_buf(4, &down.buf),
                                    bind_buf(5, &gbuf),
                                    bind_buf(9, &ubuf),
                                ],
                            });
                            pass.set_pipeline(&c.q4tp_mv4_dsilu);
                            pass.set_bind_group(0, &bind_d, &[]);
                            pass.dispatch_workgroups((hidden as u32).div_ceil(8).min(MAX_WG), 1, 1);
                            drop(pass);
                        } else {
                            pass.set_pipeline(&c.silu);
                            pass.set_bind_group(0, &bg_silu, &[]);
                            pass.dispatch_workgroups_flat((inter as u32).div_ceil(256));
                            drop(pass);
                            if let Some((pdp, bg_d, wd)) = prep(down, &abuf, &ob, hidden, inter) {
                                let mut pass = begin_pass(&mut enc);
                                pass.set_pipeline(pdp);
                                pass.set_bind_group(0, &bg_d, &[]);
                                pass.dispatch_workgroups(wd, 1, 1);
                            }
                        }
                        continue_ffn = true;
                    }
                    if !continue_ffn {
                        // gate+up+SiLU in one dispatch when both are wide q4tp
                        // (`prep_gu`); else gate+up in one and SiLU separate
                        // (`prep2`); else the three-dispatch path.
                        // A Prism FFN input consumes `n1` through a device FWHT.
                        // Do not defer the residual+norm producer into the same
                        // pass in that case: the FWHT would be encoded before
                        // its producer and observe stale `n1` regardless of
                        // storage visibility barriers.  Keep the fusion for
                        // ordinary Q2TP, but make the Prism ordering explicit.
                        if (gate.prism != crate::gpu::GraphPrismOp::None
                            || up.prism != crate::gpu::GraphPrismOp::None)
                            && let Some((p, b, w)) = ffn_pre.take()
                        {
                            go(&mut enc, p, &b, w);
                        }
                        let ffn_in = match prism_input(&mut enc, &[gate, up], &n1, hidden) {
                            Some(x) => x,
                            None => {
                                graph_decline("Prism FFN input transform unavailable");
                                return token_graph_outcome(o1_started || state_started, false);
                            }
                        };
                        let gu_ok = c.use_mv_gu
                            && c.use_mv4
                            && gate.kind == 6
                            && up.kind == 6
                            && gate.prism == crate::gpu::GraphPrismOp::None
                            && up.prism == crate::gpu::GraphPrismOp::None
                            && hidden % 32 == 0
                            && hidden / 32 > 64
                            && c.q4tp_mv16w_probe.is_none();
                        let pgu_fused = if gu_ok {
                            let (b, w) =
                                mv_gu_bind(c, &gate.buf, &up.buf, &n1, &abuf, inter, hidden);
                            Some((&c.q4tp_mv16w_gu, b, w))
                        } else {
                            None
                        };
                        let pgu = if pgu_fused.is_some() {
                            None
                        } else {
                                prep2(gate, up, &ffn_in, &gbuf, &ubuf, inter, inter, hidden)
                            };
                        let (pg, pu) = if pgu.is_some() || pgu_fused.is_some() {
                            (None, None)
                        } else {
                            (
                                prep(gate, &ffn_in, &gbuf, inter, hidden),
                                prep(up, &ffn_in, &ubuf, inter, hidden),
                            )
                        };
                        let pd = if down.prism == crate::gpu::GraphPrismOp::None {
                            prep(down, &abuf, &ob, hidden, inter)
                        } else {
                            None
                        };
                        let mut down_rode_along = false;
                        if let Some((pf, bgf, wf)) = pgu_fused {
                            let mut pass = begin_pass(&mut enc);
                            let fine = ts_full || li < 32;
                            if let Some((p, b, w)) = &ffn_pre {
                                pass.set_pipeline(p);
                                pass.set_bind_group(0, b, &[]);
                                pass.dispatch_workgroups(*w, 1, 1);
                            }
                            tsp!(pass, fine, 40);
                            pass.set_pipeline(pf);
                            pass.set_bind_group(0, &bgf, &[]);
                            pass.dispatch_workgroups(wf, 1, 1);
                            tsp!(pass, fine, 43); // gate+up+silu in one dispatch
                            if let Some((pdp, bg_d, wd)) = &pd {
                                pass.set_pipeline(pdp);
                                pass.set_bind_group(0, bg_d, &[]);
                                pass.dispatch_workgroups(*wd, 1, 1);
                                tsp!(pass, fine, 44);
                                if let Some((p, b, w)) = &ffn_post {
                                    pass.set_pipeline(p);
                                    pass.set_bind_group(0, b, &[]);
                                    pass.dispatch_workgroups(*w, 1, 1);
                                    tail_done = true;
                                }
                                tsp!(pass, fine, 45);
                                down_rode_along = true;
                            }
                        } else if let Some((p2, bg2, w2)) = pgu {
                            let bg_silu = bgc(
                                26,
                                li,
                                &c.layout_silu,
                                &[&gbuf, &ubuf, &dummy_hd, &abuf, &silu_u],
                            );
                            let mut pass = begin_pass(&mut enc);
                            let fine = ts_full || li < 32;
                            if let Some((p, b, w)) = &ffn_pre {
                                pass.set_pipeline(p);
                                pass.set_bind_group(0, b, &[]);
                                pass.dispatch_workgroups(*w, 1, 1);
                            }
                            tsp!(pass, fine, 40);
                            pass.set_pipeline(p2);
                            pass.set_bind_group(0, &bg2, &[]);
                            pass.dispatch_workgroups(w2, 1, 1);
                            tsp!(pass, fine, 42); // gate+up in one dispatch
                            pass.set_pipeline(&c.silu);
                            pass.set_bind_group(0, &bg_silu, &[]);
                            pass.dispatch_workgroups_flat((inter as u32).div_ceil(256));
                            tsp!(pass, fine, 43);
                            if let Some((pdp, bg_d, wd)) = &pd {
                                pass.set_pipeline(pdp);
                                pass.set_bind_group(0, bg_d, &[]);
                                pass.dispatch_workgroups(*wd, 1, 1);
                                tsp!(pass, fine, 44);
                                if let Some((p, b, w)) = &ffn_post {
                                    pass.set_pipeline(p);
                                    pass.set_bind_group(0, b, &[]);
                                    pass.dispatch_workgroups(*w, 1, 1);
                                    tail_done = true;
                                }
                                tsp!(pass, fine, 45);
                                down_rode_along = true;
                            }
                        } else if let (Some((pgp, bg_g, wg)), Some((pup, bg_u, wu))) = (pg, pu) {
                            let bg_silu = bgc(
                                26,
                                li,
                                &c.layout_silu,
                                &[&gbuf, &ubuf, &dummy_hd, &abuf, &silu_u],
                            );
                            let mut pass = begin_pass(&mut enc);
                            let fine = ts_full || li < 32;
                            if let Some((p, b, w)) = &ffn_pre {
                                pass.set_pipeline(p);
                                pass.set_bind_group(0, b, &[]);
                                pass.dispatch_workgroups(*w, 1, 1);
                            }
                            tsp!(pass, fine, 40); // residual + FFN norm
                            pass.set_pipeline(pgp);
                            pass.set_bind_group(0, &bg_g, &[]);
                            pass.dispatch_workgroups(wg, 1, 1);
                            tsp!(pass, fine, 41); // gate projection
                            pass.set_pipeline(pup);
                            pass.set_bind_group(0, &bg_u, &[]);
                            pass.dispatch_workgroups(wu, 1, 1);
                            tsp!(pass, fine, 42); // up projection
                            pass.set_pipeline(&c.silu);
                            pass.set_bind_group(0, &bg_silu, &[]);
                            pass.dispatch_workgroups_flat((inter as u32).div_ceil(256));
                            tsp!(pass, fine, 43); // SiLU × up
                                                  // `down` rides here too when its dtype can be
                                                  // prepped: dispatches inside one pass serialize
                                                  // with memory visibility — the same guarantee the
                                                  // MoE arm leans on — so it reads the `abuf` silu
                                                  // just wrote. This saves one pass per layer without
                                                  // pretending that pass count predicts kernel time.
                            if let Some((pdp, bg_d, wd)) = &pd {
                                pass.set_pipeline(pdp);
                                pass.set_bind_group(0, bg_d, &[]);
                                pass.dispatch_workgroups(*wd, 1, 1);
                                tsp!(pass, fine, 44); // down projection
                                if let Some((p, b, w)) = &ffn_post {
                                    pass.set_pipeline(p);
                                    pass.set_bind_group(0, b, &[]);
                                    pass.dispatch_workgroups(*w, 1, 1);
                                    tail_done = true;
                                }
                                tsp!(pass, fine, 45); // residual + next norm
                                down_rode_along = true;
                            }
                        } else {
                            group_mats(
                                &mut enc,
                                &[
                                    (gate, &ffn_in, &gbuf, inter, hidden),
                                    (up, &ffn_in, &ubuf, inter, hidden),
                                ],
                            );
                            go(
                                &mut enc,
                                &c.silu,
                                &bgc(
                                    27,
                                    li,
                                    &c.layout_silu,
                                    &[&gbuf, &ubuf, &dummy_hd, &abuf, &silu_u],
                                ),
                                (inter as u32).div_ceil(256),
                            );
                        }
                        if !down_rode_along {
                            let Some(down_in) =
                                prism_input(&mut enc, &[down], &abuf, inter)
                            else {
                                graph_decline("Prism FFN down transform unavailable");
                                return token_graph_outcome(o1_started || state_started, false);
                            };
                            emat(&mut enc, down, &down_in, &ob, hidden, inter);
                        }
                    } // !continue_ffn
                }
                LFfn::Moe {
                    router,
                    sgate,
                    gate_all,
                    up_all,
                    down_all,
                    n_exp,
                    top_k,
                    inter: mi,
                    norm_topk,
                    q4tp,
                    gu_q2,
                    sigmoid,
                    bias,
                    has_shared,
                    shared_gated,
                    route_scale,
                } => {
                    // The WHOLE MoE FFN — router + shared-gate matvecs, top-k
                    // select, fused gate+up+SiLU over the selected experts, and
                    // the weighted down accumulation into ob — rides in ONE
                    // compute pass: dispatches within a pass serialize with
                    // memory visibility (same guarantee the dense fused FFN
                    // uses), and the inter-pass pipeline flush (~78 µs on
                    // NVIDIA Vulkan) is what dominates a 40-layer decode.
                    let (mlogit, mslog, msel, mwt, mact) = moe_bufs.as_ref().unwrap();
                    let slots = *top_k + usize::from(*has_shared);
                    // sg_kind = 4 tells the select kernel to compute the shared
                    // gate itself; then the sgate matvec below is not encoded.
                    let sg_fold = sgate.kind == 4;
                    // Cached uniform: `unif` mints a fresh buffer per call, and one
                    // per MoE layer per token exhausted the device. hidden and the
                    // fold flag share the spare word.
                    let sel_u = uniform_u32x8(
                        c,
                        [
                            *n_exp as u32,
                            *top_k as u32,
                            // One flags word: bit0 renorm, bit1 sigmoid
                            // scores, bit2 selection bias, bit3 shared
                            // expert present, bit4 shared expert UNGATED
                            // (weight 1). Softmax models pass 0/1 exactly
                            // as before.
                            u32::from(*norm_topk)
                                | (u32::from(*sigmoid) << 1)
                                | (u32::from(bias.is_some()) << 2)
                                | (u32::from(*has_shared) << 3)
                                | (u32::from(*has_shared && !*shared_gated) << 4),
                            ((hidden as u32) << 8) | (u32::from(sg_fold) * 4),
                            // routed_scaling_factor on the routed mix (f32 bits).
                            route_scale.to_bits(),
                            0,
                            0,
                            0,
                        ],
                    );
                    // Per-expert stride in u16 units: q4t is 9 per group flat,
                    // q4tp adds the row params and code planes on top of 8.
                    let mat16 = |rows: usize, cols: usize| -> u32 {
                        let n = if *q4tp {
                            cortiq_core::quant::expected_nbytes(
                                cortiq_core::TensorDtype::Q4TiledP,
                                &[rows, cols],
                            )
                            .unwrap_or(0)
                        } else {
                            rows * (cols / 32) * 18
                        };
                        (n / 2) as u32
                    };
                    // gate/up may be a HALF-WIDTH plane (q2tp experts against a
                    // q4tp down), so its per-expert stride is its own.
                    let gu_mat16 = if *gu_q2 {
                        (cortiq_core::quant::expected_nbytes(
                            cortiq_core::TensorDtype::Q2TiledP,
                            &[*mi, hidden],
                        )
                        .unwrap_or(0)
                            / 2) as u32
                    } else {
                        mat16(*mi, hidden)
                    };
                    // Eight words now: the fifth is the swiglu limit, zero
                    // for every architecture but DeepSeek-V4.
                    let gu_u = uniform_u32x8(
                        c,
                        [
                            (hidden / 32) as u32,
                            *mi as u32,
                            slots as u32,
                            gu_mat16,
                            0,
                            0,
                            0,
                            0,
                        ],
                    );
                    let dn_u = uniform_u32x4(
                        c,
                        [
                            (*mi / 32) as u32,
                            hidden as u32,
                            slots as u32,
                            mat16(hidden, *mi),
                        ],
                    );
                    let (p_gu, p_dn, l_gu, l_dn) = if *gu_q2 {
                        (
                            &c.moe_gate_up_q2tp,
                            &c.moe_down_q4tp,
                            &c.layout_moe_gu_q2tp,
                            &c.layout_moe_dn_q4tp,
                        )
                    } else if *q4tp {
                        (
                            &c.moe_gate_up_q4tp,
                            &c.moe_down_q4tp,
                            &c.layout_moe_gu_q4tp,
                            &c.layout_moe_dn_q4tp,
                        )
                    } else {
                        (
                            &c.moe_gate_up,
                            &c.moe_down,
                            &c.layout_moe_gu,
                            &c.layout_moe_dn,
                        )
                    };
                    let bias_buf = bias.clone().unwrap_or_else(|| {
                        c.device
                            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                                label: Some("moe-sel-bias0"),
                                contents: &[0u8; 4],
                                usage: wgpu::BufferUsages::STORAGE,
                            })
                    });
                    let bg_sel = bg(
                        &c.layout_moe_sel,
                        &[mlogit, mslog, msel, mwt, &sel_u, &sgate.buf, &n1, &bias_buf],
                    );
                    let bg_sel_sg = c.moe_select_sg.as_ref().map(|p| {
                        let l = p.get_bind_group_layout(0);
                        bgc(
                            28,
                            li,
                            &l,
                            &[mlogit, mslog, msel, mwt, &sel_u, &sgate.buf, &n1],
                        )
                    });
                    let bg_gu = bg(l_gu, &[gate_all, up_all, &n1, msel, mact, &gu_u]);
                    let bg_dn = bg(l_dn, &[down_all, mact, msel, mwt, &ob, &dn_u]);
                    let pr = prep(router, &n1, mlogit, *n_exp, hidden);
                    let ps = prep(sgate, &n1, mslog, 1, hidden);
                    let mut continue_moe_std = true;
                    // Fold-select (q2tp + folded shared gate): router feeds the
                    // gu/down twins DIRECTLY — the select hop and the sgate
                    // matvec disappear from the layer's dependency chain.
                    let fold = c.foldsel && *gu_q2 && sg_fold && !skip_router;
                    if fold {
                        if let Some((prp, bgr, wr)) = prep(router, &n1, mlogit, *n_exp, hidden) {
                            let mgf_u = uniform_u32x4(c, [*n_exp as u32, 0, 0, 0]);
                            let l_guf = c.moe_gate_up_q2tp_f.get_bind_group_layout(0);
                            let bg_guf = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                label: None,
                                layout: &l_guf,
                                entries: &[
                                    bind_buf(0, gate_all),
                                    bind_buf(1, up_all),
                                    bind_buf(2, &n1),
                                    bind_buf(3, mlogit),
                                    bind_buf(4, mact),
                                    bind_buf(5, &gu_u),
                                    bind_buf(7, &mgf_u),
                                ],
                            });
                            let mdf_u = uniform_u32x4(
                                c,
                                [
                                    *n_exp as u32,
                                    *top_k as u32,
                                    *norm_topk as u32,
                                    (hidden as u32) << 8 | 4,
                                ],
                            );
                            let l_dnf = c.moe_down_q4tp_f.get_bind_group_layout(0);
                            let bg_dnf = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                label: None,
                                layout: &l_dnf,
                                entries: &[
                                    bind_buf(0, down_all),
                                    bind_buf(1, mact),
                                    bind_buf(2, mlogit),
                                    bind_buf(3, &sgate.buf),
                                    bind_buf(4, &ob),
                                    bind_buf(5, &dn_u),
                                    bind_buf(6, &n1),
                                    bind_buf(7, &mdf_u),
                                ],
                            });
                            let mut pass = begin_pass(&mut enc);
                            if let Some((p, b, w)) = &ffn_pre {
                                pass.set_pipeline(p);
                                pass.set_bind_group(0, b, &[]);
                                pass.dispatch_workgroups(*w, 1, 1);
                            }
                            pass.set_pipeline(prp);
                            pass.set_bind_group(0, &bgr, &[]);
                            pass.dispatch_workgroups(wr, 1, 1);
                            if !skip_moe {
                                pass.set_pipeline(&c.moe_gate_up_q2tp_f);
                                pass.set_bind_group(0, &bg_guf, &[]);
                                pass.dispatch_workgroups(*mi as u32, slots as u32, 1);
                                pass.set_pipeline(&c.moe_down_q4tp_f);
                                pass.set_bind_group(0, &bg_dnf, &[]);
                                pass.dispatch_workgroups(hidden as u32, 1, 1);
                                if let Some((p, b, w)) = &ffn_post {
                                    pass.set_pipeline(p);
                                    pass.set_bind_group(0, b, &[]);
                                    pass.dispatch_workgroups(*w, 1, 1);
                                    tail_done = true;
                                }
                            }
                            drop(pass);
                            flush_pass(&enc);
                            enc.copy_buffer_to_buffer(&n1, 0, &h_buf, 0, 0);
                            // (zero-length copy: keeps the borrow checker shape
                            // identical to the non-fold arm; no-op on device)
                            continue_moe_std = false;
                        }
                    }
                    if continue_moe_std {
                        if let (Some((prp, bgr, wr)), Some((psp, bgs, ws))) = (pr, ps) {
                            let mut pass = begin_pass(&mut enc);
                            if let Some((p, b, w)) = &ffn_pre {
                                pass.set_pipeline(p);
                                pass.set_bind_group(0, b, &[]);
                                pass.dispatch_workgroups(*w, 1, 1);
                            }
                        let fine = ts_full || li == 0;
                            tsp!(pass, fine, 30); // pass start (after prologue norm)
                            if !skip_router {
                                pass.set_pipeline(prp);
                                pass.set_bind_group(0, &bgr, &[]);
                                pass.dispatch_workgroups(wr, 1, 1);
                            }
                            tsp!(pass, fine, 31); // router
                            if !sg_fold {
                                pass.set_pipeline(psp);
                                pass.set_bind_group(0, &bgs, &[]);
                                pass.dispatch_workgroups(ws, 1, 1);
                            }
                            // The subgroup select hard-codes the GATED
                            // shared expert; hy_v3's ungated one stays on
                            // the tree kernel.
                            let plain =
                                !*sigmoid && bias.is_none() && *has_shared && *shared_gated;
                            if let (Some(sgp), true) = (&c.moe_select_sg, plain) {
                                // Same binding ORDER as the tree kernel's bg_sel —
                                // but its OWN layout (auto layouts are exclusive).
                                pass.set_pipeline(sgp);
                                pass.set_bind_group(0, bg_sel_sg.as_ref().unwrap(), &[]);
                                pass.dispatch_workgroups(1, 1, 1);
                            } else {
                                pass.set_pipeline(&c.moe_select);
                                pass.set_bind_group(0, &bg_sel, &[]);
                                pass.dispatch_workgroups(1, 1, 1);
                            }
                            tsp!(pass, fine, 32); // select
                            if !skip_moe {
                                pass.set_pipeline(p_gu);
                                pass.set_bind_group(0, &bg_gu, &[]);
                                pass.dispatch_workgroups(*mi as u32, slots as u32, 1);
                                tsp!(pass, fine, 33); // gate/up experts
                                pass.set_pipeline(p_dn);
                                pass.set_bind_group(0, &bg_dn, &[]);
                                pass.dispatch_workgroups(hidden as u32, 1, 1);
                                tsp!(pass, fine, 34); // down experts
                                if let Some((p, b, w)) = &ffn_post {
                                    pass.set_pipeline(p);
                                    pass.set_bind_group(0, b, &[]);
                                    pass.dispatch_workgroups(*w, 1, 1);
                                    tail_done = true;
                                }
                            }
                        } else {
                            // Un-preppable router dtype: per-op passes (correct, rare).
                            group_mats(
                                &mut enc,
                                &[
                                    (router, &n1, mlogit, *n_exp, hidden),
                                    (sgate, &n1, mslog, 1, hidden),
                                ],
                            );
                            go(&mut enc, &c.moe_select, &bg_sel, 1);
                            {
                                let mut pass = begin_pass(&mut enc);
                                pass.set_pipeline(p_gu);
                                pass.set_bind_group(0, &bg_gu, &[]);
                                pass.dispatch_workgroups(*mi as u32, slots as u32, 1);
                            }
                            go(&mut enc, p_dn, &bg_dn, hidden as u32);
                        }
                    } // continue_moe_std
                }
            }
            // FFN-residual + next layer's attn-norm fused (plain residual on the last).
            // At loop boundaries (Looped Transformer), insert final_norm between the
            // residual and the next iteration's input norm.
            ts!(enc, 2, lkind);
            if tail_done {
                // already emitted at the end of the FFN pass
            } else if li + 1 < layers.len() {
                if loop_norm_at.contains(&li) {
                    // h += ob; n1 = rms(h, final_norm); copy n1→h; n1 = rms(h, next_input_norm)
                    let fnw = stor(bytemuck::cast_slice(final_norm));
                    let inw_next = stor(bytemuck::cast_slice(layers[li + 1].input_norm));
                    go(
                        &mut enc,
                        &c.add_rmsnorm,
                        &bgc(
                            29,
                            li,
                            &c.layout_add_rmsnorm,
                            &[&h_buf, &ob, &fnw, &n1, &rms_u],
                        ),
                        1,
                    );
                    flush_pass(&enc);
                    enc.copy_buffer_to_buffer(&n1, 0, &h_buf, 0, (hidden * 4) as u64);
                    go(
                        &mut enc,
                        &c.rmsnorm,
                        &bgc(30, li, &c.layout_rmsnorm, &[&h_buf, &inw_next, &n1, &rms_u]),
                        1,
                    );
                } else {
                    let inw_next = stor(bytemuck::cast_slice(layers[li + 1].input_norm));
                    go(
                        &mut enc,
                        &c.add_rmsnorm,
                        &bg(
                            &c.layout_add_rmsnorm,
                            &[&h_buf, &ob, &inw_next, &n1, &rms_u],
                        ),
                        1,
                    );
                }
            } else if loop_norm_at.contains(&li) {
                // The span ENDS on a loop boundary: the boundary norm
                // belongs to this side — the wire (and the CPU span it
                // must match) hands over the NORMED hidden. Dropping it
                // here fed the next loop a raw residual: nanbeige split
                // at the default half spoke template noise ("М user").
                let fnw = stor(bytemuck::cast_slice(final_norm));
                go(
                    &mut enc,
                    &c.add_rmsnorm,
                    &bgc(
                        31,
                        li,
                        &c.layout_add_rmsnorm,
                        &[&h_buf, &ob, &fnw, &n1, &rms_u],
                    ),
                    1,
                );
                flush_pass(&enc);
                enc.copy_buffer_to_buffer(&n1, 0, &h_buf, 0, (hidden * 4) as u64);
            } else {
                go(
                    &mut enc,
                    &c.axpy,
                    &bgc(32, li, &c.layout_axpy, &[&ob, &h_buf, &ax_u]),
                    (hidden as u32).div_ceil(256),
                );
            }
            if token_tap_layer == Some(li) {
                if let Some(tap) = token_tap_stage.as_ref() {
                    flush_pass(&enc);
                    enc.copy_buffer_to_buffer(&h_buf, 0, tap, 0, (hidden * 4) as u64);
                }
            }
        }
        // ── Multi-step tail: final norm + lm_head + on-device argmax; the
        // winner's embedding becomes the next step's h. All inside the SAME
        // encoder — one submit carries every step.
        if multi {
            let (lm, lrows) = lm_pre.as_ref().unwrap();
            let lrows = *lrows;
            let lbuf = lbuf_pre.as_ref().unwrap();
            let fnw = stor(bytemuck::cast_slice(final_norm));
            go(
                &mut enc,
                &c.rmsnorm,
                &bgc(0, 0, &c.layout_rmsnorm, &[&h_buf, &fnw, &n1, &rms_u]),
                1,
            );
            let Some(lm_in) = prism_input(&mut enc, &[&lm], &n1, hidden) else {
                graph_decline("Prism lm_head transform unavailable");
                return token_graph_outcome(o1_started || state_started, false);
            };
            emat(&mut enc, lm, &lm_in, lbuf, lrows, hidden);
            let am_u = uniform_u32x4(c, [lrows as u32, AM_PARTS, stp as u32, 0]);
            let l_ap = c.argmax_part.get_bind_group_layout(0);
            let bg_ap = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &l_ap,
                entries: &[
                    bind_buf(0, lbuf),
                    bind_buf(1, am_pv.as_ref().unwrap()),
                    bind_buf(2, am_pi.as_ref().unwrap()),
                    bind_buf(3, &am_u),
                ],
            });
            go(&mut enc, &c.argmax_part, &bg_ap, AM_PARTS);
            let l_af = c.argmax_final.get_bind_group_layout(0);
            let bg_af = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &l_af,
                entries: &[
                    bind_buf(0, am_pv.as_ref().unwrap()),
                    bind_buf(1, am_pi.as_ref().unwrap()),
                    bind_buf(2, ids_buf.as_ref().unwrap()),
                    bind_buf(3, &am_u),
                ],
            });
            go(&mut enc, &c.argmax_final, &bg_af, 1);
            if stp + 1 < steps {
                let (em, e_rows, mult) = emb_pre.as_ref().unwrap();
                let eg_u = uniform_u32x8(
                    c,
                    [
                        hidden as u32,
                        (hidden / 32) as u32,
                        *e_rows as u32,
                        stp as u32,
                        mult.to_bits(),
                        0,
                        0,
                        0,
                    ],
                );
                let l_eg = c.embed_gather_q4tp.get_bind_group_layout(0);
                let bg_eg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &l_eg,
                    entries: &[
                        bind_buf(0, &em.buf),
                        bind_buf(1, ids_buf.as_ref().unwrap()),
                        bind_buf(2, &h_buf),
                        bind_buf(3, &eg_u),
                    ],
                });
                go(
                    &mut enc,
                    &c.embed_gather_q4tp,
                    &bg_eg,
                    (hidden as u32).div_ceil(256),
                );
            }
        }
    } // for stp (multi-step frames)
    let t_enc = t_enc0.elapsed().as_secs_f64() * 1000.0;
    let t_sub0 = std::time::Instant::now();
    // Return the step-slot uniforms to the scratch pool.
    gs.kv_us = kv_us;
    gs.at_us = at_us;
    gs.rope_us = rope_us;
    // ── Multi-step exit: one submit, one k×u32 readback, no logits. ──
    if multi {
        let ids_b = ids_buf.as_ref().unwrap();
        let stage = ids_stage.as_ref().unwrap();
        let sz = (steps * 4) as u64;
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(ids_b, 0, stage, 0, sz);
        submit(c, finish_enc(enc));
        let (tx, rx) = std::sync::mpsc::channel();
        stage.map_async(wgpu::MapMode::Read, ..sz, move |r| {
            let _ = tx.send(r);
        });
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        let ok = rx.recv().map(|r| r.is_ok()).unwrap_or(false);
        if ok {
            let raw = stage.get_mapped_range(..sz).unwrap();
            let ids: &[u32] = bytemuck::cast_slice(&raw);
            if let Some(out) = ids_out {
                out.clear();
                out.extend_from_slice(ids);
            }
            drop(raw);
        }
        stage.unmap();
        if ok {
            let mut kvm = c.attn_kv.lock().unwrap();
            for li in 0..layers.len() {
                if let Some(m) = kvm.get_mut(&(kv_id, layer_base + li)) {
                    m.synced = position + steps;
                }
            }
            // The recurrent buffers advanced on-device.  Keep host-side
            // cursors in lockstep so the next graph call cannot silently
            // reuse a state at a gap or an older absolute position.
            let next = position + steps;
            let mut gcm = c.gdn_cursor.lock().unwrap();
            for li in 0..layers.len() {
                if let Some(cur) = gcm.get_mut(&(kv_id, layer_base + li)) {
                    cur.next_pos = next;
                }
            }
            drop(gcm);
            let mut om = c.o1m.lock().unwrap();
            for (li, views) in o1.iter().take(layers.len()).enumerate() {
                if views.is_some() {
                    if let Some(d) = om.get_mut(&(kv_id, layer_base + li)) {
                        d.next_pos = Some(next);
                    }
                }
            }
        }
        drop(gs);
        if prof {
            let setup = t_enc0.duration_since(t_start).as_secs_f64() * 1000.0;
            eprintln!(
                "token-graph[x{steps}]: setup {setup:.2} ms | encode {t_enc:.2} ms | submit+ids {:.2} ms",
                t_sub0.elapsed().as_secs_f64() * 1000.0
            );
        }
        return token_graph_outcome(o1_started || state_started, ok);
    }
    // h_buf now holds the final hidden. Either ride final-norm + lm_head and
    // read back logits, or (no lm / unresolved weight / device prefix) read
    // back the hidden — a prefix boundary is mid-stack, so no final norm.
    let lm_resolved = if prefix {
        None
    } else {
        lm_head.and_then(|(gw, rows)| resolve(gw, rows, hidden).map(|m| (m, rows)))
    };
    let ok = if let Some((lm, lrows)) = lm_resolved {
        let fnw = stor(bytemuck::cast_slice(final_norm));
        go(
            &mut enc,
            &c.rmsnorm,
            &bgc(0, 0, &c.layout_rmsnorm, &[&h_buf, &fnw, &n1, &rms_u]),
            1,
        );
        let lsize = (lrows * 4) as u64;
        let lbuf = GraphScratch::ensure(
            &c.device,
            &mut gs.logits,
            lsize,
            wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            "g-logits",
        );
        let Some(lm_in) = prism_input(&mut enc, &[&lm], &n1, hidden) else {
            graph_decline("Prism lm_head transform unavailable");
            return token_graph_outcome(o1_started || state_started, false);
        };
        emat(&mut enc, &lm, &lm_in, &lbuf, lrows, hidden);
        if !device_span_probe {
            ts!(enc, 3, 0);
        }
        write_span_endpoint(&mut enc);
        resolve_timestamps(&mut enc);
        logits.resize(lrows, 0.0);
        let hsize = if hidden_too { (hidden * 4) as u64 } else { 0 };
        let stage = GraphScratch::ensure(
            &c.device,
            &mut gs.stage,
            lsize + hsize,
            wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            "g-stage",
        );
        crate::gpu::hostprof_encode_done(_hp_t0);
        let r = if hidden_too {
            readback_two(
                c,
                enc,
                &lbuf,
                lsize,
                &h_buf,
                hsize,
                &stage,
                &mut logits[..lrows],
                &mut h[..hidden],
            )
        } else {
            readback(c, enc, &lbuf, &stage, lsize, &mut logits[..lrows])
        };
        crate::gpu::hostprof_total(_hp_t0);
        drop(gs);
        r
    } else {
        let size = (hidden * 4) as u64;
        let stage = GraphScratch::ensure(
            &c.device,
            &mut gs.stage,
            size,
            wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            "g-stage",
        );
        write_span_endpoint(&mut enc);
        resolve_timestamps(&mut enc);
        let r = readback(c, enc, &h_buf, &stage, size, &mut h[..hidden]);
        drop(gs);
        r
    };
    if ok && ts_n.get() > 1 {
        if let Some((_, _, tstage)) = &c.ts_query {
            let (tx, rx) = std::sync::mpsc::channel();
            tstage.map_async(wgpu::MapMode::Read, ..(ts_n.get() as u64 * 8), move |r| {
                let _ = tx.send(r);
            });
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            if rx.recv().map(|r| r.is_ok()).unwrap_or(false) {
                let raw = tstage.get_mapped_range(..(ts_n.get() as u64 * 8)).unwrap();
                let t: Vec<u64> = bytemuck::cast_slice::<u8, u64>(&raw).to_vec();
                drop(raw);
                let labels = ts_lbl.borrow();
                if device_span_probe {
                    let us = t[1].saturating_sub(t[0]) as f64 * c.ts_period as f64 / 1000.0;
                    let submits = SUBMITS
                        .load(std::sync::atomic::Ordering::Relaxed)
                        .saturating_sub(span_submit_start);
                    eprintln!(
                        "gpu-device-span: preencoded={} span_ms={:.3} slots={} submits={} frames=one",
                        preencode_frame,
                        us / 1000.0,
                        ts_n.get(),
                        submits,
                    );
                } else {
                    // Attribute each delta to the LATER stamp's (stage, kind).
                    let mut agg = std::collections::BTreeMap::<(u8, u8), (f64, u32)>::new();
                    for i in 1..ts_n.get() as usize {
                        let dt = t[i].saturating_sub(t[i - 1]) as f64 * c.ts_period as f64 / 1000.0;
                        let e = agg.entry(labels[i]).or_insert((0.0, 0));
                        e.0 += dt;
                        e.1 += 1;
                    }
                    let name = |k: (u8, u8)| match k {
                    (1, 0) => "gdn-after-fwht-residual",
                    (1, 1) => "attention-after-fwht-residual",
                    (2, 0) => "ffn@gdn",
                    (2, 1) => "ffn@attn",
                    (3, _) => "tail(norm+lm)",
                    (10, _) => "|gdn:proj",
                    (11, _) => "|gdn:conv",
                    (12, _) => "|gdn:step",
                    (13, _) => "|gdn:snorm",
                    (14, _) => "|gdn:outp",
                    (20, _) => "|attn:qkv",
                    (21, _) => "|attn:rope",
                    (22, _) => "|attn:kv",
                    (23, _) => "|attn:attend",
                    (24, _) => "|attn:gate",
                    (25, _) => "|attn:wo",
                    (30, _) => "|moe:pre",
                    (31, _) => "|moe:router",
                    (32, _) => "|moe:select",
                    (33, _) => "|moe:gu",
                    (34, _) => "|moe:dn",
                    (40, _) => "|ffn:pre",
                    (41, _) => "|ffn:gate",
                    (42, _) => "|ffn:up",
                    (43, _) => "|ffn:silu",
                    (44, _) => "|ffn:down",
                    (45, _) => "|ffn:tail",
                    (50, _) => "|pre-fwht-span",
                    (51, _) => "|fwht-dispatch-span",
                    _ => "start",
                };
                    let mut line = String::from("gpu-ts: timeline spans (later-stamp attribution):");
                    let total: f64 = agg.values().map(|v| v.0).sum();
                    for (k, (us, n)) in &agg {
                        line.push_str(&format!(" {}={:.0}us/{}", name(*k), us, n));
                    }
                    line.push_str(&format!(" | timeline_span {:.2} ms | slots {}", total / 1000.0, ts_n.get()));
                    eprintln!("{line}");
                }
            }
            tstage.unmap();
        }
    }
    if ok {
        if let Some(tap) = token_tap_stage.as_ref() {
            let bytes = (hidden * 4) as u64;
            let (tx, rx) = std::sync::mpsc::channel();
            tap.slice(..bytes).map_async(wgpu::MapMode::Read, move |r| {
                let _ = tx.send(r);
            });
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            if rx.recv().map(|r| r.is_ok()).unwrap_or(false) {
                if let Ok(raw) = tap.get_mapped_range(..bytes) {
                    let vals: &[f32] = bytemuck::cast_slice(&raw);
                    let norm = vals.iter().map(|&v| (v as f64) * (v as f64)).sum::<f64>().sqrt();
                    eprintln!(
                        "token-tap layer={} row0={:?} norm={:.9e}",
                        token_tap_layer.unwrap_or(usize::MAX),
                        &vals[..hidden.min(4)],
                        norm
                    );
                    drop(raw);
                }
            }
            tap.unmap();
        }
    }
    if ok {
        for (li, b) in &o1_dbg {
            let (tx, rx) = std::sync::mpsc::channel();
            b.map_async(wgpu::MapMode::Read, .., move |r| {
                let _ = tx.send(r);
            });
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            if rx.recv().map(|r| r.is_ok()).unwrap_or(false) {
                let raw = b.get_mapped_range(..).unwrap();
                let all: &[f32] = bytemuck::cast_slice(&raw);
                let v: Vec<f32> = all[..all.len().min(16)].to_vec();
                drop(raw);
                if *li >= 10_000 {
                    eprintln!(
                        "o1-trace L{} gpu q[..4]={:?} k[..4]={:?} v[..4]={:?}",
                        li - 10_000,
                        &v[..4],
                        &v[4..8],
                        &v[8..12]
                    );
                } else {
                    eprintln!("o1-trace L{li} gpu attn[..8] = {v:?}");
                }
            }
        }
        // The append at `position` is now durable — advance each mirror.
        let mut kvm = c.attn_kv.lock().unwrap();
        for li in 0..layers.len() {
            if let Some(m) = kvm.get_mut(&(kv_id, layer_base + li)) {
                m.synced = position + 1;
            }
        }
        let next = position + 1;
        let mut gcm = c.gdn_cursor.lock().unwrap();
        for li in 0..layers.len() {
            if let Some(cur) = gcm.get_mut(&(kv_id, layer_base + li)) {
                cur.next_pos = next;
            }
        }
        drop(gcm);
        let mut om = c.o1m.lock().unwrap();
        for (li, views) in o1.iter().take(layers.len()).enumerate() {
            if views.is_some() {
                if let Some(d) = om.get_mut(&(kv_id, layer_base + li)) {
                    d.next_pos = Some(next);
                }
            }
        }
    }
    if prof {
        let setup = t_enc0.duration_since(t_start).as_secs_f64() * 1000.0;
        eprintln!(
            "token-graph: setup {setup:.2} ms | encode {t_enc:.2} ms | tail+submit+readback {:.2} ms",
            t_sub0.elapsed().as_secs_f64() * 1000.0
        );
    }
    token_graph_outcome(o1_started || state_started, ok)
}

/// Batched prefill: K prompt positions through the whole layer stack in ONE
/// submit. Projections & FFN run as resident GEMMs (each weight read once per K
/// columns instead of once per position); attention and GDN loop the existing
/// per-position kernels over scratch slices (KV mirror / recurrent S persist).
/// Cuts graph prefill from N whole-graph submits to N/K. Returns `Declined`
/// before device mutation when an unsupported case (bias, q4t/q1t
/// projections) requires the per-position graph; returns `Failed` after a
/// submitted batch mutates persistent state, so callers cannot use stale CPU
/// state. `Completed` provides valid output. `positions[i]` is the absolute
/// sequence position of batch row i (contiguous causal run starting at
/// `positions[0]`); `h` is [k·hidden] in/out.
#[allow(clippy::too_many_arguments)]
pub fn forward_batch_graph(
    model: &Arc<CmfModel>,
    kv_id: u64,
    layers: &[crate::gpu::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,
    // Speculative verify: fold final-norm + lm_head over every position and
    // read the k logit rows back beside the hiddens, snapshotting the GDN
    // state after each position so a partial acceptance can restore it
    // (`gdn_spec_restore`). None = the plain batched prefill.
    o1: &[Option<Vec<crate::nystrom::O1DeviceView<'_>>>],
    o1_epoch: u64,
    mut spec: Option<crate::gpu::SpecTail<'_>>,
) -> crate::gpu::BatchGraphOutcome {
    let t_bfn = std::time::Instant::now();
    // Once sealed O(1) views have been admitted and uploaded, a failed batch
    // cannot safely fall through to the CPU path: the device state may have
    // advanced while the CPU copy is intentionally stale. All refusals before
    // that point remain ordinary declines.
    let mut o1_started = false;
    // A submitted ordinary GDN batch mutates persistent device state just
    // like the O(1) route; failures after that point must be terminal rather
    // than a stale CPU fallback.
    let mut state_started = false;
    let Some(c) = ctx() else {
        bgraph_refused("no ctx");
        return batch_outcome(o1_started || state_started, false);
    };
    // This is the true entry for the resident batch graph.  The token graph
    // has its own reset, but a batch call normally follows a completed
    // readback without passing through that path.  Reset before resolving
    // weights or encoding any timestamped pass so kernel totals belong only
    // to this batch (and never replay the previous first-use window).
    batch_kernel_ts_begin(c);
    let graph_live_budget = graph_live_weight_budget(c, model);
    if k == 0 || positions.len() != k {
        bgraph_refused("k/positions mismatch");
        return batch_outcome(o1_started || state_started, false);
    }
    let pos0 = positions[0];
    let Some(pos_end) = pos0.checked_add(k) else {
        bgraph_refused("position range overflow");
        return batch_outcome(o1_started || state_started, false);
    };
    if positions
        .windows(2)
        .any(|pair| pair[0].checked_add(1) != Some(pair[1]))
    {
        bgraph_refused("batch positions are not contiguous");
        return batch_outcome(o1_started || state_started, false);
    }
    if pos_end > cap || hd % 4 != 0 || hd > c.hd_cap {
        bgraph_refused("pos+k past cap, or head_dim not %4 / over hd_cap");
        return batch_outcome(o1_started || state_started, false); // vec4 K/V reads; hd_cap = workgroup-storage limit
    }
    // Unlike the decode graph this function has no host-tail handoff: every
    // layer must coexist in one command buffer. Refuse before the first
    // upload when that live set cannot fit. Relying on the LRU here is
    // incorrect because the `GMat`s below retain evicted buffers until the
    // batch command finishes, producing a full-model transient allocation.
    let Some(mut live_bytes) = graph_stack_payload_bytes(model, layers) else {
        bgraph_refused("cannot size layer payload");
        return batch_outcome(o1_started || state_started, false);
    };
    if let Some(sp) = spec.as_ref() {
        live_bytes = live_bytes.saturating_add(
            model
                .tensors
                .get(sp.lm.idx)
                .map(|e| e.nbytes as u64)
                .unwrap_or(u64::MAX),
        );
    }
    if live_bytes > graph_live_budget {
        bgraph_refused("all-layer live set exceeds the weight budget");
        return batch_outcome(o1_started || state_started, false);
    }
    let cap = kv_capacity(cap, pos_end);
    struct GMat {
        buf: wgpu::Buffer,
        rs: Option<wgpu::Buffer>,
        kind: u8,
        prism: crate::gpu::GraphPrismOp,
        affine: bool,
    }
    enum LAttn {
        Full {
            wq: GMat,
            wk: GMat,
            wv: GMat,
            wo: GMat,
        },
        Gdn {
            qkv: GMat,
            z: GMat,
            a: GMat,
            b: GMat,
            out: GMat,
            nv: usize,
            nk: usize,
            dk: usize,
            dv: usize,
            kk: usize,
            cdim: usize,
        },
    }
    /// Батчевый FFN слоя. MoE маршрутизируется ПО ТОКЕНАМ, поэтому его
    /// эксперты кодируются в цикле внутри того же submit'а, тогда как
    /// attention и проекции остаются батчевыми GEMM'ами. Раньше здесь
    /// допускался только Dense, и любая MoE-модель уходила на путь
    /// «одна позиция за submit»: префилл 33 tok/s против 54 на декоде,
    /// то есть промпт обрабатывался медленнее, чем генерация.
    enum BFfn {
        Dense {
            gate: GMat,
            up: GMat,
            down: GMat,
            /// This LAYER's intermediate width. Not the model's: a pruned
            /// model narrows the FFN per layer (bonsai-1.7b runs 6130 …
            /// 6140 across its 28 layers), and taking the header's number
            /// here asked for weights that do not exist — the graph then
            /// refused for the whole model, on every token, silently.
            width: usize,
        },
        Moe {
            router: GMat,
            sgate: GMat,
            gate_all: wgpu::Buffer,
            up_all: wgpu::Buffer,
            down_all: wgpu::Buffer,
            n_exp: usize,
            top_k: usize,
            inter: usize,
            norm_topk: bool,
            q4tp: bool,
            /// Mixed 2-bit profile: q2tp gate/up over a q4tp down. The
            /// whole-chunk q4tp fast lane must NOT take these bytes.
            gu_q2: bool,
            sigmoid: bool,
            bias: Option<wgpu::Buffer>,
            shared_gated: bool,
            route_scale: f32,
        },
    }
    struct LW {
        attn: LAttn,
        ffn: BFfn,
    }
    let resolve = |gw: &crate::gpu::GraphW, rows: usize, cols: usize| -> Option<GMat> {
        match gw.kind {
            0 => {
                if gw.row_scale.len() < rows {
                    return None;
                }
                let b = tensor_weight(c, model, gw.idx, rows, cols)?;
                let rsb = c.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("bg-rs"),
                    size: (rows * 4) as u64,
                    usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                });
                c.queue
                    .write_buffer(&rsb, 0, bytemuck::cast_slice(&gw.row_scale[..rows]));
                Some(GMat {
                    buf: b,
                    rs: Some(rsb),
                    kind: 0,
                    prism: gw.prism,
                    affine: gw.affine,
                })
            }
            1 => {
                let (b, r, cc) = q1_weight(c, model, gw.idx)?;
                if r != rows || cc != cols {
                    return None;
                }
                Some(GMat {
                    buf: b,
                    rs: None,
                    kind: 1,
                    prism: gw.prism,
                    affine: gw.affine,
                })
            }
            4 => {
                if gw.data.len() < rows * cols {
                    return None;
                }
                let b = c.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("bg-f32w"),
                    size: (rows * cols * 4) as u64,
                    usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                });
                c.queue
                    .write_buffer(&b, 0, bytemuck::cast_slice(&gw.data[..rows * cols]));
                Some(GMat {
                    buf: b,
                    rs: None,
                    kind: 4,
                    prism: gw.prism,
                    affine: gw.affine,
                })
            }
            // q4_tiled and q4tp: same buffer shape, the kernel differs.
            // Leaving these out is what kept every q4t/q4tp model off the
            // batched path — including its GDN projections, which is where
            // the refusal actually landed.
            k @ (5 | 6 | 9) => {
                let (b, r, cc) = tile_weight(c, model, gw.idx)?;
                if r != rows || cc != cols {
                    return None;
                }
                Some(GMat {
                    buf: b,
                    rs: None,
                    kind: k,
                    prism: gw.prism,
                    affine: gw.affine,
                })
            }
            _ => None, // q1t not batched here → CPU/per-position path
        }
    };
    // GEMM-able projection? (q8_row/q1). f32 (a/b) is per-position; anything else bails.
    // kinds 5/6 (q4_tiled, q4tp) have tile GEMMs too — admitting only 0/1
    // is what kept every q4tp model off the batched path.
    // 9 (the 2-bit plane) has the tile GEMM `q2tp_mm`; without it a q2tp
    // file's speculative verify declined every round and spun.
    let gemmable = |m: &GMat| matches!(m.kind, 0 | 1 | 5 | 6 | 9);
    let mut lws = Vec::with_capacity(layers.len());
    let mut gdn_dims: Option<(usize, usize, usize, usize, usize, usize)> = None;
    for l in layers {
        let attn = match &l.attn {
            crate::gpu::GraphAttn::Full {
                wq,
                wk,
                wv,
                wo,
                output_gate,
                bias,
                ..
            } => {
                if bias.is_some() {
                    bgraph_refused("site:5889");
                    return batch_outcome(o1_started || state_started, false);
                } // batched bias axpy not wired
                let qrows = nh * hd * (1 + *output_gate as usize);
                let (Some(wq), Some(wk), Some(wv), Some(wo)) = (
                    resolve(wq, qrows, hidden),
                    resolve(wk, nkv * hd, hidden),
                    resolve(wv, nkv * hd, hidden),
                    resolve(wo, hidden, nh * hd),
                ) else {
                    bgraph_refused("site:5898");
                    return batch_outcome(o1_started || state_started, false);
                };
                if !(gemmable(&wq) && gemmable(&wk) && gemmable(&wv) && gemmable(&wo)) {
                    bgraph_refused("attention weights not gemmable");
                    return batch_outcome(o1_started || state_started, false);
                }
                LAttn::Full { wq, wk, wv, wo }
            }
            crate::gpu::GraphAttn::Gdn {
                qkv,
                z,
                a,
                b,
                out,
                nv,
                nk,
                dk,
                dv,
                kk,
                ..
            } => {
                let cdim = 2 * nk * dk + nv * dv;
                let dims = (*nv, *nk, *dk, *dv, *kk, cdim);
                if gdn_dims.is_some_and(|prev| prev != dims) {
                    bgraph_refused("heterogeneous GDN geometry across batch layers");
                    return batch_outcome(o1_started || state_started, false);
                }
                gdn_dims = Some(dims);
                let (Some(qkv), Some(z), Some(a), Some(b), Some(out)) = (
                    resolve(qkv, cdim, hidden),
                    resolve(z, nv * dv, hidden),
                    resolve(a, *nv, hidden),
                    resolve(b, *nv, hidden),
                    resolve(out, hidden, nv * dv),
                ) else {
                    bgraph_refused("site:5928");
                    return batch_outcome(o1_started || state_started, false);
                };
                if !(gemmable(&qkv) && gemmable(&z) && gemmable(&out) && a.kind == 4 && b.kind == 4)
                {
                    bgraph_refused("site:5932");
                    return batch_outcome(o1_started || state_started, false);
                }
                LAttn::Gdn {
                    qkv,
                    z,
                    a,
                    b,
                    out,
                    nv: *nv,
                    nk: *nk,
                    dk: *dk,
                    dv: *dv,
                    kk: *kk,
                    cdim,
                }
            }
            crate::gpu::GraphAttn::ShortConv { .. } => {
                // The batch (prefill) graph has no conv-ring kernel that
                // walks positions in order yet; prefill stays per-op.
                return batch_outcome(o1_started || state_started, false);
            }
        };
        let bffn = match &l.ffn {
            crate::gpu::GraphFfn::Dense {
                gate: lg,
                up: lu,
                down: ld,
            } => {
                // Per-layer FFN width, from the weight rather than the
                // header — a pruned model narrows it layer by layer.
                let ffn_w = model
                    .tensors
                    .get(lg.idx)
                    .and_then(|e| e.shape.first().copied())
                    .map(|r| r)
                    .filter(|w| *w > 0 && *w <= inter)
                    .unwrap_or(inter);
                let (Some(gate), Some(up), Some(down)) = (
                    resolve(lg, ffn_w, hidden),
                    resolve(lu, ffn_w, hidden),
                    resolve(ld, hidden, ffn_w),
                ) else {
                    bgraph_refused("site:5960");
                    return batch_outcome(o1_started || state_started, false);
                };
                if !(gemmable(&gate) && gemmable(&up) && gemmable(&down)) {
                    bgraph_refused("dense FFN not gemmable");
                    return batch_outcome(o1_started || state_started, false);
                }
                BFfn::Dense {
                    gate,
                    up,
                    down,
                    width: ffn_w,
                }
            }
            crate::gpu::GraphFfn::Moe {
                router,
                shared_gate,
                experts,
                n_exp,
                top_k,
                inter: mi,
                norm_topk,
                q4tp,
                gu_q2,
                sigmoid,
                bias,
                has_shared,
                shared_gated,
                route_scale,
            } => {
                if *top_k >= 16 || *n_exp > 256 || experts.len() != n_exp + 1 {
                    bgraph_refused("site:5979");
                    return batch_outcome(o1_started || state_started, false);
                }
                let (Some(router), Some(sgate)) = (
                    resolve(router, *n_exp, hidden),
                    resolve(shared_gate, 1, hidden),
                ) else {
                    bgraph_refused("site:5985");
                    return batch_outcome(o1_started || state_started, false);
                };
                let Some((gate_all, up_all, down_all)) =
                    moe_expert_bufs(c, model, experts, *mi, hidden, *q4tp, *gu_q2, false)
                else {
                    bgraph_refused("site:5990");
                    return batch_outcome(o1_started || state_started, false);
                };
                BFfn::Moe {
                    router,
                    sgate,
                    gate_all,
                    up_all,
                    down_all,
                    n_exp: *n_exp,
                    top_k: *top_k,
                    inter: *mi,
                    norm_topk: *norm_topk,
                    q4tp: *q4tp,
                    gu_q2: *gu_q2,
                    sigmoid: *sigmoid,
                    bias: bias.map(|b| {
                        c.device
                            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                                label: Some("bmoe-sel-bias"),
                                contents: bytemuck::cast_slice(b),
                                usage: wgpu::BufferUsages::STORAGE,
                            })
                    }),
                    shared_gated: *shared_gated,
                    route_scale: *route_scale,
                }
            }
        };
        lws.push(LW { attn, ffn: bffn });
    }
    // Recurrent state and exact-KV mirrors must enter this batch at the same
    // absolute position. Validate all cursors before allocating the batch
    // scratch or uploading a sealed O(1) mirror; otherwise a later refusal
    // could leave one layer advanced while the caller still believes the
    // batch declined. A missing cursor is only valid when its state is also
    // absent and the CPU seed can initialize a nonzero position.
    {
        let kvm = c.attn_kv.lock().unwrap();
        let gsm = c.gdn_state.lock().unwrap();
        let gcm = c.gdn_cursor.lock().unwrap();
        let om = c.o1m.lock().unwrap();
        for (li, l) in layers.iter().enumerate() {
            let key = (kv_id, li);
            let o1_here = o1.get(li).is_some_and(|v| v.is_some());
            if o1_here
                && om
                    .get(&key)
                    .filter(|d| d.epoch == o1_epoch)
                    .and_then(|d| d.next_pos)
                    .is_some_and(|next| next != pos0)
            {
                bgraph_refused("o1 device state position mismatch");
                // A resident O(1) state is already authoritative on the
                // device; a host fallback would pair the sealed request with
                // stale accumulators.
                return batch_outcome(true, false);
            }
            match &l.attn {
                crate::gpu::GraphAttn::Full { .. } if !o1_here => {
                    if kvm.get(&key).is_some_and(|m| m.synced > pos0) {
                        bgraph_refused("KV mirror is ahead of batch position");
                        return batch_outcome(true, false);
                    }
                }
                crate::gpu::GraphAttn::Full { .. } => {}
                crate::gpu::GraphAttn::Gdn {
                    cpu_state,
                    nv,
                    nk,
                    dk,
                    dv,
                    kk,
                    ..
                } => {
                    let dims = (*nv, *nk, *dk, *dv, *kk, 2 * nk * dk + nv * dv);
                    match (gsm.get(&key), gcm.get(&key)) {
                        (Some(_), Some(cur)) if cur.dims == dims && cur.next_pos == pos0 => {}
                        (Some(_), Some(_)) => {
                            bgraph_refused("GDN device state position/geometry mismatch");
                            return batch_outcome(true, false);
                        }
                        (Some(_), None) => {
                            bgraph_refused("GDN device state cursor missing");
                            return batch_outcome(true, false);
                        }
                        (None, Some(_)) => {
                            bgraph_refused("GDN cursor exists without device state");
                            return batch_outcome(true, false);
                        }
                        (None, None) => {
                            let cdim = 2 * nk * dk + nv * dv;
                            let want = cdim
                                .saturating_mul(kk.saturating_sub(1))
                                .saturating_add((*nv).saturating_mul(*dk).saturating_mul(*dv));
                            if pos0 > 0 && cpu_state.len() != want {
                                bgraph_refused("GDN CPU seed missing for nonzero batch position");
                                return batch_outcome(o1_started || state_started, false);
                            }
                        }
                    }
                }
                crate::gpu::GraphAttn::ShortConv {
                    kernel, cpu_state, ..
                } => {
                    let dims = (hidden, 0, 0, 0, *kernel, hidden);
                    match (gsm.get(&key), gcm.get(&key)) {
                        (Some(_), Some(cur)) if cur.dims == dims && cur.next_pos == pos0 => {}
                        (Some(_), Some(_)) => {
                            bgraph_refused("short-conv device state position/geometry mismatch");
                            return batch_outcome(true, false);
                        }
                        (Some(_), None) => {
                            bgraph_refused("short-conv device state cursor missing");
                            return batch_outcome(true, false);
                        }
                        (None, Some(_)) => {
                            bgraph_refused("short-conv cursor exists without device state");
                            return batch_outcome(true, false);
                        }
                        (None, None) => {
                            let want = kernel.saturating_sub(1).saturating_mul(hidden);
                            if pos0 > 0 && cpu_state.len() != want {
                                bgraph_refused(
                                    "short-conv CPU seed missing for nonzero batch position",
                                );
                                return batch_outcome(o1_started || state_started, false);
                            }
                        }
                    }
                }
            }
        }
    }
    // O(1) admission is all-or-nothing. Validate every sealed view and every
    // layer pairing before creating the persistent batch scratch or touching a
    // device mirror. A partial set would otherwise make a later CPU fallback
    // consume a stale exact-KV copy beside an already-mutated O(1) state.
    let has_o1 = o1.iter().any(Option::is_some);
    if !o1.is_empty() && o1.len() != layers.len() {
        bgraph_refused("o1 layer/view count mismatch");
        return batch_outcome(o1_started || state_started, false);
    }
    if has_o1 {
        if spec.is_some() {
            bgraph_refused("o1 batch does not support speculative tail");
            return batch_outcome(o1_started || state_started, false);
        }
        for (li, views) in o1.iter().enumerate() {
            match (&layers[li].attn, views) {
                (crate::gpu::GraphAttn::Full { .. }, Some(v)) => {
                    if !o1_views_valid(v, nh, nkv, hd) {
                        bgraph_refused("o1 view failed batch admission");
                        return batch_outcome(o1_started || state_started, false);
                    }
                }
                (crate::gpu::GraphAttn::Full { .. }, None) => {}
                (crate::gpu::GraphAttn::Gdn { .. }, None) => {}
                (crate::gpu::GraphAttn::Gdn { .. }, Some(_))
                | (crate::gpu::GraphAttn::ShortConv { .. }, Some(_)) => {
                    bgraph_refused("o1 view attached to non-full attention");
                    return batch_outcome(o1_started || state_started, false);
                }
                (crate::gpu::GraphAttn::ShortConv { .. }, None) => {}
            }
        }
        // Upload every layer's sealed state before the first command dispatch.
        // `o1_started` is set before the first call because a later upload can
        // still fail after earlier layers have allocated their mirrors.
        o1_started = true;
        for (li, views) in o1.iter().enumerate() {
            if let Some(views) = views {
                if o1_ensure(c, kv_id, li, views, o1_epoch).is_none() {
                    graph_refused("o1 state not portable");
                    return batch_outcome(o1_started || state_started, false);
                }
            }
        }
    }
    // Content-cached, exactly as the token graph's: norm weights are
    // token-invariant, and this minted a fresh device buffer for every one
    // of them on every call — a speculative verify walks 64 layers and
    // asks for several apiece. Same `const_bufs` map, same (ptr,len) key,
    // same fingerprint so another model's mmap landing on the address
    // refreshes rather than aliases.
    let stor = |data: &[u8]| {
        let key = (data.as_ptr() as usize, data.len());
        let fp = crate::gpu::fp_bytes(data);
        let mut cb = c.const_bufs.lock().unwrap();
        if let Some((b, f)) = cb.get_mut(&key) {
            if *f != fp {
                c.queue.write_buffer(b, 0, data);
                *f = fp;
            }
            return b.clone();
        }
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: data.len() as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        c.queue.write_buffer(&b, 0, data);
        cb.insert(key, (b.clone(), fp));
        b
    };
    let unif = |d: &[u32]| {
        c.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(d),
                usage: wgpu::BufferUsages::UNIFORM,
            })
    };
    let bg = |layout: &wgpu::BindGroupLayout, bufs: &[&wgpu::Buffer]| {
        let e: Vec<_> = bufs
            .iter()
            .enumerate()
            .map(|(i, b)| bind_buf(i as u32, b))
            .collect();
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout,
            entries: &e,
        })
    };
    // Buffers usable both as compute storage and copy src/dst (K-loop slicing).
    let rwc = |n: usize| {
        c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (n.max(1) * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_SRC
                | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        })
    };
    let h_buf = rwc(k * hidden);
    c.queue
        .write_buffer(&h_buf, 0, bytemuck::cast_slice(&h[..k * hidden]));
    // Narrow, opt-in parity tap for the batch-vs-token bring-up.  It copies
    // the post-residual hidden of one layer without changing the normal
    // graph or readback contract; the result is printed only after the final
    // submit has completed.  Keep this diagnostic off the hot path.
    let tap_layer = std::env::var("CMF_GRAPH_TAP_LAYER")
        .ok()
        .and_then(|v| v.parse::<usize>().ok())
        .filter(|&li| li < layers.len());
    let tap_stage = tap_layer.map(|_| c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("batch-graph-tap"),
        size: (k * hidden * 4) as u64,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    }));
    let n1 = rwc(k * hidden);
    // The token graph's Prism transform is one 1024-block per activation
    // row.  Reuse that shader for the batched graph by carrying an explicit
    // row count in FwhtP; the sign table remains one validated vector per
    // supported width, not k duplicated copies.
    let prism_signs = model
        .header
        .arch
        .prism_hadamard
        .as_ref()
        .map(|cfg| stor(bytemuck::cast_slice(&cfg.signs)));
    let prism_sign_offset = |width: usize| -> Option<usize> {
        let cfg = model.header.arch.prism_hadamard.as_ref()?;
        let mut off = 0usize;
        for &w in &cfg.widths {
            if w == width {
                return Some(off);
            }
            off = off.checked_add(w)?;
        }
        None
    };
    let prism_round16 = model
        .header
        .arch
        .prism_hadamard
        .as_ref()
        .is_some_and(|cfg| cfg.activation_f16);
    let prism_rot_width = hidden.max(inter).max(nh * hd).max(nkv * hd);
    let prism_rot = rwc(k * prism_rot_width);
    let prism_input_b =
        |enc: &mut wgpu::CommandEncoder,
         mats: &[&GMat],
         src: &wgpu::Buffer,
         width: usize|
         -> Option<wgpu::Buffer> {
            let mut op = crate::gpu::GraphPrismOp::None;
            for m in mats {
                if m.prism != crate::gpu::GraphPrismOp::None {
                    if op != crate::gpu::GraphPrismOp::None && op != m.prism {
                        return None;
                    }
                    op = m.prism;
                }
            }
            if op == crate::gpu::GraphPrismOp::None {
                return Some(src.clone());
            }
            if op != crate::gpu::GraphPrismOp::Forward
                || width == 0
                || width > prism_rot_width
                || k == 0
            {
                return None;
            }
            let cfg = model.header.arch.prism_hadamard.as_ref()?;
            let block = cfg.block_size;
            let sign_offset = prism_sign_offset(width)?;
            let signs = prism_signs.as_ref()?;
            if block != 1024 || width % block != 0 {
                return None;
            }
            let pipe = c.fwht.as_ref()?;
            let p = unif(&[
                width as u32,
                block as u32,
                sign_offset as u32,
                0,
                u32::from(prism_round16),
                k as u32,
                0,
                0,
            ]);
            let layout = pipe.get_bind_group_layout(0);
            let bg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("prism-fwht-batch"),
                layout: &layout,
                entries: &[
                    bind_buf(0, src),
                    bind_buf(1, &prism_rot),
                    bind_buf(2, signs),
                    bind_buf(3, &p),
                ],
            });
            let active_bytes = (k * width * std::mem::size_of::<f32>() * 2) as u64;
            let tsw = batch_kernel_ts_pair(c, 1, active_bytes);
            let mut pass = begin_pass_with(enc, Some("prism-fwht-batch"), tsw);
            pass.set_pipeline(pipe);
            pass.set_bind_group(0, &bg, &[]);
            pass.dispatch_workgroups((k * width / block) as u32, 1, 1);
            Some(prism_rot.clone())
        };
    let any_gate = layers.iter().any(|l| {
        matches!(
            &l.attn,
            crate::gpu::GraphAttn::Full {
                output_gate: true,
                ..
            }
        )
    });
    let qdim = nh * hd * (1 + any_gate as usize);
    let (gnv, _gnk, gdk, gdv, _gkk, gcdim) = gdn_dims.unwrap_or((1, 1, 1, 1, 1, 1));
    // batched GEMM outputs
    let qraw_b = rwc(k * qdim);
    let kb_b = rwc(k * nkv * hd);
    let vb_b = rwc(k * nkv * hd);
    let attn_bb = rwc(k * nh * hd);
    let qkv_b = rwc(k * gcdim);
    let z_b = rwc(k * gnv * gdv);
    let gdo_b = rwc(k * gnv * gdv);
    let ob = rwc(k * hidden);
    let gbuf = rwc(k * inter);
    let ubuf = rwc(k * inter);
    let abuf = rwc(k * inter);
    // per-position scratch
    let _n1_s = rwc(hidden);
    let _qraw_s = rwc(qdim);
    let _kb_s = rwc(nkv * hd);
    let _vb_s = rwc(nkv * hd);
    let qout_s = rwc(nh * hd);
    let gout_s = rwc(nh * hd);
    let attn_s = rwc(nh * hd);
    // Split-K attend partials for the per-row attention above 256 positions.
    let pacc_s = rwc(nh * cap.div_ceil(ATTEND_GCK) * hd);
    let pml_s = rwc(nh * cap.div_ceil(ATTEND_GCK) * 2);
    let _qkv_s = rwc(gcdim);
    // k rows for the k-looped conv/step twins — row i at i*cdim.
    let cq_s = rwc(k * gcdim);
    let _z_s = rwc(gnv * gdv);
    let _a_s = rwc(gnv);
    let _b_s = rwc(gnv);
    // Whole-batch a/b planes: one token-axis matvec per layer fills them,
    // and gdn_step reads its token's row via GdnP.tok.
    let a_bb = rwc(k * gnv);
    let b_bb = rwc(k * gnv);
    let _gdo_s = rwc(gnv * gdv);
    let invf_b = stor(bytemuck::cast_slice(invf));
    let dummy_hd = stor(bytemuck::cast_slice(&vec![0f32; hd]));
    // KV mirror + GDN state (fresh; batch appends positions pos0..pos0+k).
    let mut kvbufs: Vec<Option<(wgpu::Buffer, wgpu::Buffer)>> = Vec::with_capacity(layers.len());
    let mut gdnbufs: Vec<Option<(wgpu::Buffer, wgpu::Buffer)>> = Vec::with_capacity(layers.len());
    {
        let mut kvm = c.attn_kv.lock().unwrap();
        let mut gsm = c.gdn_state.lock().unwrap();
        let mut gcm = c.gdn_cursor.lock().unwrap();
        for (li, l) in layers.iter().enumerate() {
            match &l.attn {
                crate::gpu::GraphAttn::Full { cpu_k, cpu_v, .. } => {
                    if o1.get(li).is_some_and(|v| v.is_some()) {
                        // Sealed O(1) owns this layer's attention state. Do
                        // not allocate or advance an exact-KV mirror that a
                        // later fallback could accidentally read.
                        kvbufs.push(None);
                        gdnbufs.push(None);
                        continue;
                    }
                    let e = kv_mirror_ensure(c, &mut kvm, (kv_id, li), nkv, hd, cap);
                    if e.synced > pos0 {
                        bgraph_refused("KV mirror is ahead of batch position");
                        return batch_outcome(true, false);
                    }
                    if e.synced < pos0 {
                        if cpu_k.len() < nkv
                            || cpu_v.len() < nkv
                            || cpu_k
                                .iter()
                                .zip(cpu_v.iter())
                                .any(|(kh, vh)| kh.len() / hd < pos0 || vh.len() / hd < pos0)
                        {
                            bgraph_refused("CPU KV seed missing for nonzero batch position");
                            return batch_outcome(o1_started || state_started, false);
                        }
                        for hh in 0..nkv {
                            let off = ((hh * cap + e.synced) * hd * 4) as u64;
                            c.queue.write_buffer(
                                &e.k,
                                off,
                                bytemuck::cast_slice(&cpu_k[hh][e.synced * hd..pos0 * hd]),
                            );
                            c.queue.write_buffer(
                                &e.v,
                                off,
                                bytemuck::cast_slice(&cpu_v[hh][e.synced * hd..pos0 * hd]),
                            );
                        }
                        e.synced = pos0;
                    }
                    kvbufs.push(Some((e.k.clone(), e.v.clone())));
                    gdnbufs.push(None);
                }
                crate::gpu::GraphAttn::Gdn {
                    cpu_state,
                    nv,
                    nk,
                    dk,
                    dv,
                    kk,
                    ..
                } => {
                    let key = (kv_id, li);
                    let e = gsm.entry(key).or_insert_with(|| {
                        let ring_sz = (gcdim * (_gkk.max(1).saturating_sub(1)) * 4) as u64;
                        let s_sz = (gnv * gdk * gdv * 4) as u64;
                        let mk = |sz: u64| {
                            let bf = c.device.create_buffer(&wgpu::BufferDescriptor {
                                label: Some("gdn-state"),
                                size: sz.max(4),
                                usage: wgpu::BufferUsages::STORAGE
                                    | wgpu::BufferUsages::COPY_DST
                                    | wgpu::BufferUsages::COPY_SRC,
                                mapped_at_creation: false,
                            });
                            c.queue.write_buffer(&bf, 0, &vec![0u8; sz.max(4) as usize]);
                            bf
                        };
                        let (ring, sbuf) = (mk(ring_sz), mk(s_sz));
                        // A CPU prefix (including bounded O(1) calibration)
                        // may have already advanced the recurrent state. Seed
                        // a newly-created device entry from that exact layout
                        // instead of silently starting the batch from zero.
                        let want = (ring_sz + s_sz) as usize / 4;
                        if cpu_state.len() == want && want > 0 {
                            let ring_n = ring_sz as usize / 4;
                            c.queue.write_buffer(
                                &ring,
                                0,
                                bytemuck::cast_slice(&cpu_state[..ring_n]),
                            );
                            c.queue.write_buffer(
                                &sbuf,
                                0,
                                bytemuck::cast_slice(&cpu_state[ring_n..]),
                            );
                        }
                        (ring, sbuf)
                    });
                    gcm.entry(key).or_insert(GdnCursor {
                        dims: (*nv, *nk, *dk, *dv, *kk, 2 * nk * dk + nv * dv),
                        next_pos: pos0,
                    });
                    gdnbufs.push(Some((e.0.clone(), e.1.clone())));
                    kvbufs.push(None);
                }
                crate::gpu::GraphAttn::ShortConv { .. } => {
                    // Unreachable: the batch resolve declined this mixer
                    // above. Kept explicit so a future prefill kernel has
                    // to think about the ring, not inherit a None.
                    kvbufs.push(None);
                    gdnbufs.push(None);
                }
            }
        }
    }
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("batch-graph"),
        });
    // Pass merging for this encoder (and its chunk successors — same
    // slot, same key): every begin_pass below hands back one open pass;
    // copies, timestamps, swaps and finishes flush it first.
    let _merge_guard = PassMergeGuard::new(&enc);
    let go =
        |enc: &mut wgpu::CommandEncoder, p: &wgpu::ComputePipeline, b: &wgpu::BindGroup, g: u32| {
            let mut pass = begin_pass(enc);
            pass.set_pipeline(p);
            pass.set_bind_group(0, b, &[]);
            pass.dispatch_workgroups(g, 1, 1);
        };
    let flags = |qn: bool, kn: bool, late: bool| {
        (if qn { 2u32 } else { 0 })
            | (if kn { 4 } else { 0 })
            | (if gemma { 8 } else { 0 })
            | (if late { 32 } else { 0 })
    };
    let rms_u = unif(&[hidden as u32, if gemma { 1 } else { 0 }, eps.to_bits(), 0]);
    let silu_u = unif(&[(k * inter) as u32, 0, 0, 0]);
    // Batched GEMM matvec (q8_row / q1) into a [k·rows] output.
    let ematb = |enc: &mut wgpu::CommandEncoder,
                 m: &GMat,
                 xs: &wgpu::Buffer,
                 y: &wgpu::Buffer,
                 rows: usize,
                 cols: usize| {
        match m.kind {
            0 => encode_q8_mm(c, enc, &m.buf, m.rs.as_ref().unwrap(), xs, y, rows, cols, k),
            5 => encode_q4_tile_mm(c, enc, &c.q4t_mm, &m.buf, xs, y, rows, cols, k),
            // The 2-bit plane: the tile GEMM handles any k (the MoE prefill's
            // kernel).  Its fourth uniform word is the descriptor-aware
            // center bit, not the q4 cooperative activation scale: affine
            // q2tp must stay `(code-1)·s` on the batched path just as it is
            // on the resident single-token matvec.  Passing the generic
            // q4 helper here used zero for every source and silently applied
            // the ordinary `(code-1.5)·s` center.
            9 => {
                // The coop arm is only for a validated Prism forward
                // activation_f16 boundary.  Ordinary Q2TP and non-Prism
                // affine-looking descriptors remain on the scalar decoder.
                let use_q2_coop = m.affine
                    && m.prism == crate::gpu::GraphPrismOp::Forward
                    && prism_round16
                    && c.q2tp_mm_coop.is_some()
                    && k >= 16
                    && cols % 32 == 0;
                encode_q2_tile_mm(
                    c,
                    enc,
                    &m.buf,
                    xs,
                    y,
                    rows,
                    cols,
                    k,
                    m.affine,
                    use_q2_coop,
                );
            }
            6 if std::env::var("CMF_BATCH_Q4_SCALAR").is_ok_and(|v| v != "0") => {
                // Diagnostic parity arm for q4tp: use the established
                // one-row matvec over offset slices, keeping batch state and
                // attention unchanged while isolating tiled-GEMM behavior.
                let p_buf = uniform_u32x4(c, [(cols / 32) as u32, rows as u32, cols as u32, 0]);
                let layout = c.q4tp_mv.get_bind_group_layout(0);
                for i in 0..k {
                    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: Some("batch-q4-scalar"),
                        layout: &layout,
                        entries: &[
                            bind_buf(0, &m.buf),
                            bind_buf_off(1, xs, (i * cols * 4) as u64, (cols * 4) as u64),
                            bind_buf_off(2, y, (i * rows * 4) as u64, (rows * 4) as u64),
                            bind_buf(3, &p_buf),
                        ],
                    });
                    let mut pass = begin_pass(enc);
                    pass.set_pipeline(&c.q4tp_mv);
                    pass.set_bind_group(0, &bind, &[]);
                    pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
                }
            }
            6 => {
                // The 64-wide GEMM tile wastes a small batch (a k=3 verify
                // keeps 3 rows of 64 busy). The batched matvec kernel
                // streams the weight once per batch element with the batch
                // as the fast dispatch axis — its other callers proved it;
                // the tile GEMM keeps the big-chunk prefill.
                if k <= 4 && cols / 32 <= 64 && c.use_mv4 {
                    // Narrow rows only: with eight x vec4 fetches PER BATCH
                    // ELEMENT in flight the register file overflows at the
                    // dense 17408-wide shapes and the amortized weight read
                    // buys nothing back (measured 470 us/layer against the
                    // pair kernel's 3x167). MoE-width experts keep the win.
                    let gpr = cols / 32;
                    let p_buf = q4tp_mv_params(c, gpr, rows, k);
                    let layout = c.q4tp_mv_k.get_bind_group_layout(0);
                    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &layout,
                        entries: &[
                            bind_buf(0, &m.buf),
                            bind_buf(2, y),
                            bind_buf(3, &p_buf),
                            bind_buf(4, &m.buf),
                            bind_buf(5, xs),
                        ],
                    });
                    let mut pass = begin_pass(enc);
                    pass.set_pipeline(&c.q4tp_mv_k);
                    pass.set_bind_group(0, &bind, &[]);
                    pass.dispatch_workgroups((rows as u32).div_ceil(8).min(MAX_WG), 1, 1);
                } else if k <= 16 && c.use_mv4 {
                    let _ = encode_q4tp_mv4_b(c, enc, &m.buf, xs, y, rows, cols, k);
                } else if batch_coop_on() && c.q4tp_mm_coop_s.is_some() && cols % 32 == 0 {
                    // Big-chunk prefill on the matrix units: the scalar
                    // tile GEMM holds ~5 TFLOP/s, the cooperative one ~50.
                    // f16 operands, f32 accumulate; the activation scale is
                    // computed ON THE DEVICE (this panel never reaches the
                    // host) and read by the kernel at both ends.
                    let asc = {
                        let mut sc = c.scratch.lock().unwrap();
                        Scratch::ensure(
                            &c.device,
                            &mut sc.amax1,
                            4,
                            wgpu::BufferUsages::STORAGE,
                            "batch-ascale",
                        )
                    };
                    if encode_act_absmax(c, enc, xs, k * cols, &asc) {
                        encode_q4_tile_mm_full(
                            c,
                            enc,
                            c.q4tp_mm_coop_s.as_ref().unwrap(),
                            &m.buf,
                            xs,
                            y,
                            rows,
                            cols,
                            k,
                            0.0,
                            Some(&asc),
                        );
                    } else {
                        encode_q4_tile_mm(c, enc, &c.q4tp_mm, &m.buf, xs, y, rows, cols, k)
                    }
                } else {
                    encode_q4_tile_mm(c, enc, &c.q4tp_mm, &m.buf, xs, y, rows, cols, k)
                }
            }
            _ => encode_q1_mm(c, enc, &m.buf, xs, y, rows, cols, k),
        }
    };
    // Two batched projections of one input in one dispatch when both are
    // wide q4tp and the batch fits the bku kernel; otherwise two `ematb`.
    let ematb2 = |enc: &mut wgpu::CommandEncoder,
                  a: &GMat,
                  b: &GMat,
                  xs: &wgpu::Buffer,
                  ya: &wgpu::Buffer,
                  yb: &wgpu::Buffer,
                  rows_a: usize,
                  rows_b: usize,
                  cols: usize| {
        // The int8 arm first: two dp4a dispatches (each quantizes the
        // shared x once, cheaply) beat the fused f32 pair on the verify's
        // shapes — the pair kernel is the f32 batched matvec's arithmetic
        // twice over. `ematb` routes to it when the switch is on.
        if a.kind == 6
            && b.kind == 6
            && c.use_mv4
            && verify_i8_on()
            && (2..=8).contains(&k)
            && cols / 32 > 64
            && cols % 32 == 0
        {
            encode_q4tp_mv4_b_i8(c, enc, &a.buf, xs, ya, rows_a, cols, k);
            encode_q4tp_mv4_b_i8(c, enc, &b.buf, xs, yb, rows_b, cols, k);
            return;
        }
        if a.kind == 6
            && b.kind == 6
            && c.use_mv4
            && encode_q4tp_mv4_b_x2(c, enc, &a.buf, &b.buf, xs, ya, yb, rows_a, rows_b, cols, k)
        {
            return;
        }
        ematb(enc, a, xs, ya, rows_a, cols);
        ematb(enc, b, xs, yb, rows_b, cols);
    };
    // SINGLE-row matvec for the per-token stretches inside the batch (the MoE
    // router/gate run once per token). `ematb` bakes nb=k into the GEMM: fed a
    // one-row buffer it reads k rows past the end and writes k rows into a
    // one-row output — and it has no f32 arm at all, so a kind-4 router fell
    // into the q1 decoder. Both were enough to turn the answer into noise.
    let emat1 = |enc: &mut wgpu::CommandEncoder,
                 m: &GMat,
                 xs: &wgpu::Buffer,
                 y: &wgpu::Buffer,
                 rows: usize,
                 cols: usize| {
        match m.kind {
            0 => encode_matvec(c, enc, &m.buf, xs, m.rs.as_ref().unwrap(), y, rows, cols),
            1 => encode_matvec_q1(c, enc, &m.buf, xs, y, rows, cols),
            5 => {
                if c.use_mv4 {
                    let gpr = cols / 32;
                    let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, cols as u32, 0]);
                    let layout = c.q4t_mv8.get_bind_group_layout(0);
                    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &layout,
                        entries: &[
                            bind_buf(0, &m.buf),
                            bind_buf(2, y),
                            bind_buf(3, &p_buf),
                            bind_buf(4, &m.buf),
                            bind_buf(5, xs),
                        ],
                    });
                    let mut pass = begin_pass(enc);
                    pass.set_pipeline(&c.q4t_mv8);
                    pass.set_bind_group(0, &bind, &[]);
                    pass.dispatch_workgroups((rows as u32).div_ceil(8).min(MAX_WG), 1, 1);
                } else {
                    encode_q1t_like(c, enc, &c.q4t_mv, &m.buf, xs, y, rows, cols)
                }
            }
            6 => {
                if c.use_mv4 {
                    encode_q4tp_mv4(c, enc, &m.buf, xs, y, rows, cols)
                } else {
                    encode_q1t_like(c, enc, &c.q4tp_mv, &m.buf, xs, y, rows, cols)
                }
            }
            7 => {
                // Both scale planes live inside the buffer, so the kernel
                // needs the true `cols` alongside the word count.
                let p_buf = uniform_u32x4(c, [(cols / 4) as u32, rows as u32, cols as u32, 0]);
                let layout = c.q8_2f_mv.get_bind_group_layout(0);
                let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &layout,
                    entries: &[
                        bind_buf(0, &m.buf),
                        bind_buf(1, xs),
                        bind_buf(2, y),
                        bind_buf(3, &p_buf),
                    ],
                });
                let mut pass = begin_pass(enc);
                pass.set_pipeline(&c.q8_2f_mv);
                pass.set_bind_group(0, &bind, &[]);
                pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
            }
            9 if std::env::var("CMF_BATCH_Q2_SCALAR").is_ok_and(|v| v != "0") => {
                // Diagnostic/correctness arm: run the established single-row
                // q2tp kernel over offset slices.  It preserves the same
                // descriptor center and activation layout while isolating a
                // possible tiled-GEMM mismatch from the recurrent graph.
                let gpr = cols / 32;
                let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, 1, m.affine as u32]);
                let layout = q2tp_pipeline(c).get_bind_group_layout(0);
                for i in 0..k {
                    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: Some("batch-q2-scalar"),
                        layout: &layout,
                        entries: &[
                            bind_buf(0, &m.buf),
                            bind_buf_off(2, y, (i * rows * 4) as u64, (rows * 4) as u64),
                            bind_buf(3, &p_buf),
                            bind_buf_off(5, xs, (i * cols * 4) as u64, (cols * 4) as u64),
                        ],
                    });
                    let mut pass = begin_pass(enc);
                    pass.set_pipeline(q2tp_pipeline(c));
                    pass.set_bind_group(0, &bind, &[]);
                    pass.dispatch_workgroups(mv_grid((rows as u32).div_ceil(16)), 1, 1);
                }
            }
            9 => {
                let gpr = cols / 32;
                let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, 1, m.affine as u32]);
                let layout = q2tp_pipeline(c).get_bind_group_layout(0);
                let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &layout,
                    entries: &[
                        bind_buf(0, &m.buf),
                        bind_buf(2, y),
                        bind_buf(3, &p_buf),
                        bind_buf(5, xs),
                    ],
                });
                let mut pass = begin_pass(enc);
                pass.set_pipeline(q2tp_pipeline(c));
                pass.set_bind_group(0, &bind, &[]);
                pass.dispatch_workgroups(mv_grid((rows as u32).div_ceil(16)), 1, 1);
            }
            _ => encode_f32matvec(c, enc, &m.buf, xs, y, rows, cols),
        }
    };
    let cp =
        |enc: &mut wgpu::CommandEncoder,
         src: &wgpu::Buffer,
         so: usize,
         dst: &wgpu::Buffer,
         n: usize| enc.copy_buffer_to_buffer(src, (so * 4) as u64, dst, 0, (n * 4) as u64);
    // Однострочные срезы батча для MoE: его ядра написаны на ОДИН токен,
    // поэтому i-я строка копируется сюда, считается и уезжает обратно.
    let row_in = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("bg-row-in"),
        size: (hidden * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE
            | wgpu::BufferUsages::COPY_DST
            | wgpu::BufferUsages::COPY_SRC,
        mapped_at_creation: false,
    });
    let row_out = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("bg-row-out"),
        size: (hidden * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE
            | wgpu::BufferUsages::COPY_DST
            | wgpu::BufferUsages::COPY_SRC,
        mapped_at_creation: false,
    });
    let moe_bufs = lws.iter().find_map(|w| match &w.ffn {
        BFfn::Moe {
            n_exp,
            top_k,
            inter: mi,
            ..
        } => Some((*n_exp, *top_k + 1, *mi)),
        _ => None,
    });
    let moe_bufs = moe_bufs.map(|(mn, ms, mi)| {
        let mk = |n: usize, label: &str| {
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some(label),
                size: (n * 4).max(4) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            })
        };
        (
            mk(k * mn, "bg-mlogit"),
            mk(1, "bg-mslog"),
            mk(k * ms, "bg-msel"),
            mk(k * ms, "bg-mwt"),
            mk(k * ms * mi, "bg-mact"),
        )
    });
    let cpo =
        |enc: &mut wgpu::CommandEncoder,
         src: &wgpu::Buffer,
         dst: &wgpu::Buffer,
         dof: usize,
         n: usize| enc.copy_buffer_to_buffer(src, 0, dst, (dof * 4) as u64, (n * 4) as u64);
    // Bootstrap first layer's input norm over all k rows.
    let inw0 = stor(bytemuck::cast_slice(layers[0].input_norm));
    go(
        &mut enc,
        &c.rmsnorm_b,
        &bg(&c.layout_rmsnorm_b, &[&h_buf, &inw0, &n1, &rms_u]),
        k as u32,
    );
    // `CMF_BATCH_TS=1`: coarse GPU stage stamps over the batch — the
    // k-independent fixed cost lives somewhere in here and host timers
    // cannot see past the submit boundary.
    let bts_on = std::env::var("CMF_BATCH_TS").is_ok() && c.ts_query.is_some();
    let mut bts_lbl: Vec<u8> = Vec::new();
    macro_rules! bts {
        ($enc:expr, $lbl:expr) => {
            if bts_on {
                if let Some((qs, _, _)) = c.ts_query.as_ref() {
                    let n = bts_lbl.len() as u32;
                    let cap = if std::env::var("CMF_BATCH_KERNEL_TS").as_deref() == Ok("1") {
                        1024
                    } else {
                        250
                    };
                    if n < cap {
                        flush_pass(&$enc);
                        $enc.write_timestamp(qs, n);
                        bts_lbl.push($lbl);
                    }
                }
            }
        };
    }
    bts!(enc, 0);
    // The token graph's pipelined submission, ported: the card starts the
    // first layers of the verify/prefill batch while the host still
    // encodes the rest. Timestamps must stay inside ONE submission-window
    // accounting, so the stamps simply ride whichever encoder is current.
    let bchunk = if graph_split_n() > 1 {
        layers.len().div_ceil(graph_split_n()).max(4)
    } else {
        usize::MAX
    };
    for (li, l) in layers.iter().enumerate() {
        if li > 0 && bchunk != usize::MAX && li % bchunk == 0 {
            flush_pass(&enc);
            let full = std::mem::replace(
                &mut enc,
                c.device
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                        label: Some("batch-graph"),
                    }),
            );
            submit(c, finish_enc(full));
        }
        let lw = &lws[li];
        let pnw = stor(bytemuck::cast_slice(l.post_norm));
        match (&lw.attn, &l.attn) {
            (
                LAttn::Full { wq, wk, wv, wo },
                crate::gpu::GraphAttn::Full {
                    q_norm,
                    k_norm,
                    late_qk_norm,
                    output_gate,
                    ..
                },
            ) => {
                let o1_here = o1.get(li).and_then(|v| v.as_ref());
                let qnw = stor(bytemuck::cast_slice(q_norm.unwrap_or(&vec![0f32; hd])));
                let knw = stor(bytemuck::cast_slice(k_norm.unwrap_or(&vec![0f32; hd])));
                let qrows = nh * hd * (1 + *output_gate as usize);
                let Some(n1_p) = prism_input_b(&mut enc, &[wq, wk, wv], &n1, hidden) else {
                    bgraph_refused("Prism transform unavailable for batched attention input");
                    return batch_outcome(o1_started || state_started, false);
                };
                ematb(&mut enc, wq, &n1_p, &qraw_b, qrows, hidden);
                ematb2(
                    &mut enc,
                    wk,
                    wv,
                    &n1_p,
                    &kb_b,
                    &vb_b,
                    nkv * hd,
                    nkv * hd,
                    hidden,
                );
                if let Some(views) = o1_here {
                    // Sealed O(1) attention uses the same rotated Q/K/V
                    // projections as the exact path, then runs the existing
                    // far/push/attend kernels for each batch row in causal
                    // order. `goff` is the row's KV-group offset in kb_b/vb_b.
                    let (
                        dmeta,
                        drk,
                        drv,
                        dsk,
                        dsv,
                        dkt,
                        dqt,
                        dmu,
                        dmz,
                        dth,
                        gg,
                        hh_,
                        mm,
                        ww,
                        nns,
                        sc,
                    ) = {
                        let map = c.o1m.lock().unwrap();
                        let Some(d) = map.get(&(kv_id, li)) else {
                            graph_refused("o1 batch mirror missing after admission");
                            return batch_outcome(o1_started || state_started, false);
                        };
                        (
                            d.meta.clone(),
                            d.ring_k.clone(),
                            d.ring_v.clone(),
                            d.sink_k.clone(),
                            d.sink_v.clone(),
                            d.k_tilde.clone(),
                            d.qt.clone(),
                            d.mu.clone(),
                            d.mz.clone(),
                            d.that.clone(),
                            d.g,
                            d.h,
                            d.m,
                            d.w,
                            d.ns,
                            d.scale,
                        )
                    };
                    let rect_fm = views
                        .first()
                        .and_then(|v| v.heads.first())
                        .is_some_and(|h| h.rect_fm);
                    let mut pass = begin_pass(&mut enc);
                    for i in 0..k {
                        let p = positions[i];
                        let rope_u = uniform_u32x8(
                            c,
                            [
                                nh as u32,
                                nkv as u32,
                                hd as u32,
                                rd as u32,
                                p as u32,
                                flags(q_norm.is_some(), k_norm.is_some(), *late_qk_norm)
                                    | if *output_gate { 1 } else { 0 },
                                eps.to_bits(),
                                i as u32,
                            ],
                        );
                        let o1_u = uniform_u32x8(
                            c,
                            [
                                hh_ as u32,
                                mm as u32,
                                ww as u32,
                                (nns as u32) | (u32::from(rect_fm) << 8),
                                hd as u32,
                                hd as u32,
                                sc.to_bits(),
                                (i * nh * hd) as u32,
                            ],
                        );
                        let bg_rope = bg(
                            &c.layout_attn_rope,
                            &[
                                &qraw_b, &kb_b, &qout_s, &gout_s, &qnw, &knw, &invf_b, &rope_u,
                            ],
                        );
                        let bg_far = bg(
                            &c.layout_o1_far,
                            &[&dmeta, &drk, &drv, &dqt, &dmz, &dth, &o1_u],
                        );
                        let bg_push = bg(
                            &c.layout_o1_push,
                            &[&dmeta, &kb_b, &vb_b, &drk, &drv, &o1_u],
                        );
                        let bg_att = bg(
                            &c.layout_o1_attend,
                            &[
                                &dmeta, &qout_s, &drk, &drv, &dsk, &dsv, &dkt, &dmu, &dmz, &dth,
                                &attn_s, &o1_u,
                            ],
                        );
                        pass.set_pipeline(&c.attn_rope);
                        pass.set_bind_group(0, &bg_rope, &[]);
                        pass.dispatch_workgroups((nh + nkv) as u32, 1, 1);
                        pass.set_pipeline(&c.o1_far);
                        pass.set_bind_group(0, &bg_far, &[]);
                        pass.dispatch_workgroups((gg * hh_ * mm) as u32, 1, 1);
                        pass.set_pipeline(&c.o1_push);
                        pass.set_bind_group(0, &bg_push, &[]);
                        pass.dispatch_workgroups(gg as u32, 1, 1);
                        pass.set_pipeline(&c.o1_attend);
                        pass.set_bind_group(0, &bg_att, &[]);
                        pass.dispatch_workgroups((gg * hh_) as u32, 1, 1);
                        if *output_gate {
                            let gm_u = unif(&[(nh * hd) as u32, 0, 0, 0]);
                            pass.set_pipeline(&c.gate_mul);
                            pass.set_bind_group(
                                0,
                                &bg(&c.layout_gate_mul, &[&gout_s, &attn_s, &gm_u]),
                                &[],
                            );
                            pass.dispatch_workgroups(((nh * hd) as u32).div_ceil(256), 1, 1);
                        }
                        encode_blit_p(
                            &mut pass,
                            c,
                            &attn_s,
                            &attn_bb,
                            nh * hd,
                            0,
                            i * nh * hd,
                            None,
                        );
                    }
                } else {
                    let (kbuf, vbuf) = kvbufs[li].as_ref().unwrap();
                    // ONE compute pass for every position: the loop's four
                    // dispatches per token each carried their own pass, and
                    // pass boundaries — not the math — were 4.3 of this
                    // stage's 5.4 ms. In-pass dispatch ordering already
                    // guarantees each sees the previous one's writes.
                    let mut pass = begin_pass(&mut enc);
                    for i in 0..k {
                        let p = positions[i];
                        let gate_flag = if *output_gate { 1u32 } else { 0 };
                        let rope_u = uniform_u32x8(
                            c,
                            [
                                nh as u32,
                                nkv as u32,
                                hd as u32,
                                rd as u32,
                                p as u32,
                                flags(q_norm.is_some(), k_norm.is_some(), *late_qk_norm) | gate_flag,
                                eps.to_bits(),
                                i as u32,
                            ],
                        );
                        let kv_u = uniform_u32x4(
                            c,
                            [nkv as u32, hd as u32, cap as u32, (p | (i << 20)) as u32],
                        );
                        let at_u = unif(&[
                            nh as u32,
                            (nh / nkv) as u32,
                            hd as u32,
                            cap as u32,
                            (p + 1) as u32,
                            attn_scale.to_bits(),
                            0,
                            0,
                        ]);
                        pass.set_pipeline(&c.attn_rope);
                        pass.set_bind_group(
                            0,
                            &bg(
                                &c.layout_attn_rope,
                                &[
                                    &qraw_b, &kb_b, &qout_s, &gout_s, &qnw, &knw, &invf_b, &rope_u,
                                ],
                            ),
                            &[],
                        );
                        pass.dispatch_workgroups((nh + nkv) as u32, 1, 1);
                        pass.set_pipeline(&c.kv_append);
                        // The first Full-attention K/V append is the
                        // persistent-state admission boundary.  A failed
                        // submit/readback after this point must be surfaced
                        // as Failed so the caller clears state instead of
                        // falling back to a stale CPU cache.
                        state_started = true;
                        pass.set_bind_group(
                            0,
                            &bg(&c.layout_kv, &[&kb_b, &vb_b, kbuf, vbuf, &kv_u]),
                            &[],
                        );
                        pass.dispatch_workgroups(((nkv * hd) as u32).div_ceil(256), 1, 1);
                        // The same attend arms as the token graph: the
                        // 256-lane decode kernel at short context, the
                        // GQA-shared split-K past ATTEND_SPLIT_MIN. The
                        // 32-lane per-head kernel this loop used to call
                        // walked every cached position with a barrier
                        // pair — at a 2.3k-token prompt a k=4 verify spent
                        // more in it than in all its matvecs (spec decoded
                        // 26 tok/s against a plain 44 on that prompt).
                        let n_ctx = p + 1;
                        let hpk = nh / nkv;
                        let split_ok = n_ctx > ATTEND_SPLIT_MIN
                            && c.attend_gpart.is_some()
                            && hpk <= 8
                            && hd <= 256
                            && hd % 4 == 0
                            && nh % nkv == 0;
                        if split_ok {
                            let nc = cap.div_ceil(ATTEND_GCK);
                            let nc_used = n_ctx.div_ceil(ATTEND_GCK);
                            let ap_u = unif(&[
                                nh as u32,
                                hpk as u32,
                                hd as u32,
                                cap as u32,
                                n_ctx as u32,
                                ATTEND_GCK as u32,
                                nc as u32,
                                attn_scale.to_bits(),
                            ]);
                            let bg_part = bg(
                                c.layout_attend_gpart.as_ref().unwrap(),
                                &[&qout_s, kbuf, vbuf, &pacc_s, &pml_s, &ap_u],
                            );
                            let bg_merge = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                label: None,
                                layout: &c.layout_attend_merge,
                                entries: &[
                                    bind_buf(3, &pacc_s),
                                    bind_buf(4, &pml_s),
                                    bind_buf(5, &ap_u),
                                    bind_buf(6, &attn_s),
                                ],
                            });
                            pass.set_pipeline(c.attend_gpart.as_ref().unwrap());
                            pass.set_bind_group(0, &bg_part, &[]);
                            pass.dispatch_workgroups(nkv as u32, nc_used as u32, 1);
                            pass.set_pipeline(&c.attend_merge);
                            pass.set_bind_group(0, &bg_merge, &[]);
                            pass.dispatch_workgroups(nh as u32, 1, 1);
                        } else if c.attend_dec && hd <= 256 {
                            let dec_l = c.gqa_attend_dec.get_bind_group_layout(0);
                            pass.set_pipeline(&c.gqa_attend_dec);
                            pass.set_bind_group(
                                0,
                                &bg(&dec_l, &[&qout_s, kbuf, vbuf, &attn_s, &at_u]),
                                &[],
                            );
                            pass.dispatch_workgroups(nh as u32, 1, 1);
                        } else {
                            let (ap, al) = attend_pipes(c, hd);
                            pass.set_pipeline(ap);
                            pass.set_bind_group(
                                0,
                                &bg(al, &[&qout_s, kbuf, vbuf, &attn_s, &at_u]),
                                &[],
                            );
                            pass.dispatch_workgroups(nh as u32, 1, 1);
                        }
                        if *output_gate {
                            let gm_u = unif(&[(nh * hd) as u32, 0, 0, 0]);
                            pass.set_pipeline(&c.gate_mul);
                            pass.set_bind_group(
                                0,
                                &bg(&c.layout_gate_mul, &[&gout_s, &attn_s, &gm_u]),
                                &[],
                            );
                            pass.dispatch_workgroups(((nh * hd) as u32).div_ceil(256), 1, 1);
                        }
                        encode_blit_p(
                            &mut pass,
                            c,
                            &attn_s,
                            &attn_bb,
                            nh * hd,
                            0,
                            i * nh * hd,
                            None,
                        );
                    }
                }
                let Some(attn_p) = prism_input_b(&mut enc, &[wo], &attn_bb, nh * hd) else {
                    bgraph_refused("Prism transform unavailable for batched output projection");
                    return batch_outcome(o1_started || state_started, false);
                };
                ematb(&mut enc, wo, &attn_p, &ob, hidden, nh * hd);
                bts!(enc, 1);
            }
            (
                LAttn::Gdn {
                    qkv,
                    z,
                    a,
                    b,
                    out,
                    nv,
                    nk,
                    dk,
                    dv,
                    kk,
                    cdim,
                },
                crate::gpu::GraphAttn::Gdn {
                    conv1d,
                    a_log,
                    dt_bias,
                    norm,
                    ..
                },
            ) => {
                state_started = true;
                let (ring, s) = gdnbufs[li].as_ref().unwrap();
                // Speculative rounds: a snapshot slot per position, taken
                // right after this position's state advance. The buffer
                // lives per (kv_id, layer) and regrows if k does.
                let snap = if spec.is_some() && std::env::var("CMF_SPEC_NOSNAP").is_err() {
                    let ring_sz = (cdim * kk.saturating_sub(1) * 4) as u64;
                    let s_sz = (nv * dk * dv * 4) as u64;
                    let mut m = c.gdn_snap.lock().unwrap();
                    let e = m.entry((kv_id, li)).or_insert_with(|| {
                        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
                            label: Some("gdn-snap"),
                            size: (k as u64 * (ring_sz + s_sz)).max(4),
                            usage: wgpu::BufferUsages::STORAGE
                                | wgpu::BufferUsages::COPY_DST
                                | wgpu::BufferUsages::COPY_SRC,
                            mapped_at_creation: false,
                        });
                        (b, ring_sz, s_sz, k, pos0)
                    });
                    if e.3 < k {
                        e.0 = c.device.create_buffer(&wgpu::BufferDescriptor {
                            label: Some("gdn-snap"),
                            size: (k as u64 * (ring_sz + s_sz)).max(4),
                            usage: wgpu::BufferUsages::STORAGE
                                | wgpu::BufferUsages::COPY_DST
                                | wgpu::BufferUsages::COPY_SRC,
                            mapped_at_creation: false,
                        });
                        e.3 = k;
                    }
                    // A snapshot buffer is reused across rounds.  Its slot
                    // zero belongs to this batch's absolute base position,
                    // which restore needs to put the recurrent cursor back.
                    e.4 = pos0;
                    Some((e.0.clone(), ring_sz, s_sz))
                } else {
                    None
                };
                let taps = stor(bytemuck::cast_slice(conv1d));
                let alog = stor(bytemuck::cast_slice(a_log));
                let dtb = stor(bytemuck::cast_slice(dt_bias));
                let gnorm = stor(bytemuck::cast_slice(norm));
                bts!(enc, 6);
                if a.prism != crate::gpu::GraphPrismOp::None
                    || b.prism != crate::gpu::GraphPrismOp::None
                    || a.affine
                    || b.affine
                {
                    // This specialized F32 auxiliary lane has no transform
                    // slot.  Refuse a descriptor that would otherwise make
                    // the original normalized-input binding ambiguous.
                    bgraph_refused("batched GDN a/b carries unsupported transform");
                    return batch_outcome(o1_started || state_started, false);
                }
                // Prism rotates only the wide learned projections.  The
                // GDN a/b auxiliaries are plain f32 vectors in source space;
                // unlike qkv/z they must consume the original RMS-normalized
                // input, not the FWHT/F16 view used by the tiled projections.
                // Keeping them in this transform group silently changed the
                // recurrent decay/update coefficients for every batch row.
                let Some(n1_p) = prism_input_b(&mut enc, &[qkv, z], &n1, hidden) else {
                    bgraph_refused("Prism transform unavailable for batched GDN input");
                    return batch_outcome(o1_started || state_started, false);
                };
                ematb2(&mut enc, qkv, z, &n1_p, &qkv_b, &z_b, *cdim, nv * dv, hidden);
                let _gc_p = unif(&[*cdim as u32, *kk as u32, 0, 0]);
                let _gd_p = unif(&[
                    *nv as u32,
                    *dk as u32,
                    *dv as u32,
                    (nk * dk) as u32,
                    (nv / nk) as u32,
                    *cdim as u32,
                    eps.to_bits(),
                    0,
                ]);
                // Token offsets ride in the kernels' spare uniform words:
                // conv reads its token's qkv slice, step reads/writes its
                // token's z/output rows in the BATCH buffers. The staging
                // copies this replaces were 4 of the 8 commands per token
                // per GDN layer of a chunk.
                // a/b for EVERY token in one dispatch each — the per-token
                // matvecs were 1920 of the chunk's ~4800 remaining commands.
                let fb_u = uniform_u32x4(c, [hidden as u32, *nv as u32, 0, 0]);
                {
                    let mut pass = begin_pass(&mut enc);
                    for (w, y) in [(&a.buf, &a_bb), (&b.buf, &b_bb)] {
                        pass.set_pipeline(&c.f32_matvec_b);
                        // a/b are untransformed auxiliary projections.  They
                        // intentionally use the original normalized input,
                        // matching the token graph's GDN path.
                        pass.set_bind_group(0, &bg(&c.layout_f32b, &[w, &n1, y, &fb_u]), &[]);
                        pass.dispatch_workgroups((*nv as u32).min(MAX_WG), k as u32, 1);
                    }
                }
                bts!(enc, 7);
                {
                    // The position recurrence lives INSIDE two k-looped
                    // kernels: one dispatch of conv (columns independent —
                    // no barriers at all), one of step (heads independent),
                    // with each position's (ring, S) written straight into
                    // the snapshot buffer by the kernels themselves. The
                    // per-dispatch chain this replaces spent 7-8 ms of a
                    // 3-position verify on barrier drains.
                    let (snap_stride, ring_els, snap_buf) = match snap.as_ref() {
                        Some((b, ring_sz, s_sz)) => (
                            ((*ring_sz + *s_sz) / 4) as u32,
                            (*ring_sz / 4) as u32,
                            b.clone(),
                        ),
                        None => (0u32, 0u32, row_in.clone()),
                    };
                    let gc_pt = unif(&[
                        *cdim as u32,
                        *kk as u32,
                        k as u32,
                        snap_stride.min(1),
                        snap_stride,
                        0,
                        0,
                        0,
                    ]);
                    let gd_pt = unif(&[
                        *nv as u32,
                        *dk as u32,
                        *dv as u32,
                        (nk * dk) as u32,
                        (nv / nk) as u32,
                        *cdim as u32,
                        eps.to_bits(),
                        k as u32,
                        snap_stride,
                        ring_els,
                        0,
                        0,
                    ]);
                    let mut pass = begin_pass(&mut enc);
                    pass.set_pipeline(&c.gdn_conv_k);
                    pass.set_bind_group(
                        0,
                        &{
                            let l = c.gdn_conv_k.get_bind_group_layout(0);
                            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                label: None,
                                layout: &l,
                                entries: &[
                                    bind_buf(0, &qkv_b),
                                    bind_buf(1, &taps),
                                    bind_buf(2, ring),
                                    bind_buf(3, &cq_s),
                                    bind_buf(4, &gc_pt),
                                    bind_buf(5, &snap_buf),
                                ],
                            })
                        },
                        &[],
                    );
                    pass.dispatch_workgroups((*cdim as u32).div_ceil(256), 1, 1);
                    // Keep the proven scalar k-loop available as an explicit
                    // correctness arm.  The vec4 parallel kernel is useful
                    // only after its state layout has been independently
                    // matched; CMF_BATCH_GDN_SAFE=1 intentionally exercises
                    // the same reduction/order as the token graph and does
                    // not silently trade a wrong recurrent state for speed.
                    let token_gdn = std::env::var("CMF_BATCH_GDN_TOKEN")
                        .map(|v| v != "0")
                        .unwrap_or(false)
                        && snap.is_none();
                    let safe_gdn = std::env::var("CMF_BATCH_GDN_SAFE")
                        .map(|v| v != "0")
                        .unwrap_or(false)
                        && snap.is_none();
                    let needs_norm_k;
                    if token_gdn {
                        // Exact token-kernel A/B: keep the batched
                        // projections/conv, but run the already validated
                        // per-position GDN step and norm over offset views.
                        // This isolates the k-looped vec4 recurrence without
                        // changing state ownership or the batch handoff.
                        for i in 0..k {
                            let gd_pi = unif(&[
                                *nv as u32,
                                *dk as u32,
                                *dv as u32,
                                (nk * dk) as u32,
                                (nv / nk) as u32,
                                *cdim as u32,
                                eps.to_bits(),
                                i as u32,
                            ]);
                            let bg_i = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                label: None,
                                layout: &c.gdn_step_par.get_bind_group_layout(0),
                                entries: &[
                                    bind_buf_off(0, &cq_s, (i * gcdim * 4) as u64, (gcdim * 4) as u64),
                                    bind_buf(2, &a_bb),
                                    bind_buf(3, &b_bb),
                                    bind_buf(4, &alog),
                                    bind_buf(5, &dtb),
                                    bind_buf(7, s),
                                    bind_buf(8, &gdo_b),
                                    bind_buf(9, &gd_pi),
                                ],
                            });
                            pass.set_pipeline(&c.gdn_step_par);
                            pass.set_bind_group(0, &bg_i, &[]);
                            pass.dispatch_workgroups(*nv as u32, (*dv as u32).div_ceil(4), 1);
                            let bn_i = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                label: None,
                                layout: &c.gdn_step_norm.get_bind_group_layout(0),
                                entries: &[
                                    bind_buf(1, &z_b),
                                    bind_buf(6, &gnorm),
                                    bind_buf(8, &gdo_b),
                                    bind_buf(9, &gd_pi),
                                ],
                            });
                            pass.set_pipeline(&c.gdn_step_norm);
                            pass.set_bind_group(0, &bn_i, &[]);
                            pass.dispatch_workgroups(*nv as u32, 1, 1);
                        }
                        // gdn_step_par + gdn_step_norm already emits the
                        // gated/normalized output for every row.
                        needs_norm_k = false;
                    } else if safe_gdn {
                        pass.set_pipeline(&c.gdn_step_k);
                        pass.set_bind_group(
                            0,
                            &{
                                let l = c.gdn_step_k.get_bind_group_layout(0);
                                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                    label: None,
                                    layout: &l,
                                    entries: &[
                                        bind_buf(0, &cq_s),
                                        bind_buf(1, &z_b),
                                        bind_buf(2, &a_bb),
                                        bind_buf(3, &b_bb),
                                        bind_buf(4, &alog),
                                        bind_buf(5, &dtb),
                                        bind_buf(6, &gnorm),
                                        bind_buf(7, s),
                                        bind_buf(8, &gdo_b),
                                        bind_buf(9, &gd_pt),
                                        bind_buf(10, &snap_buf),
                                    ],
                                })
                            },
                            &[],
                        );
                        pass.dispatch_workgroups(*nv as u32, 1, 1);
                        // gdn_step_k includes the gated RMS normalization.
                        needs_norm_k = false;
                    } else {
                        pass.set_pipeline(&c.gdn_step_par_k);
                        pass.set_bind_group(
                            0,
                            &{
                                let l = c.gdn_step_par_k.get_bind_group_layout(0);
                                // The auto layout keeps only what the entry
                                // point touches: no z, no norm weight here.
                                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                    label: None,
                                    layout: &l,
                                    entries: &[
                                        bind_buf(0, &cq_s),
                                        bind_buf(2, &a_bb),
                                        bind_buf(3, &b_bb),
                                        bind_buf(4, &alog),
                                        bind_buf(5, &dtb),
                                        bind_buf(7, s),
                                        bind_buf(8, &gdo_b),
                                        bind_buf(9, &gd_pt),
                                        bind_buf(10, &snap_buf),
                                    ],
                                })
                            },
                            &[],
                        );
                        pass.dispatch_workgroups(*nv as u32, (*dv as u32).div_ceil(4), 1);
                        // The parallel raw step leaves normalization to its
                        // separate k-wide pass below.
                        needs_norm_k = true;
                    }
                    if needs_norm_k {
                        pass.set_pipeline(&c.gdn_step_norm_k);
                        pass.set_bind_group(
                            0,
                            &{
                                let l = c.gdn_step_norm_k.get_bind_group_layout(0);
                                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                    label: None,
                                    layout: &l,
                                    entries: &[
                                        bind_buf(1, &z_b),
                                        bind_buf(6, &gnorm),
                                        bind_buf(8, &gdo_b),
                                        bind_buf(9, &gd_pt),
                                    ],
                                })
                            },
                            &[],
                        );
                        pass.dispatch_workgroups(*nv as u32, 1, 1);
                    }
                }
                bts!(enc, 5);
                let Some(gdo_p) = prism_input_b(&mut enc, &[out], &gdo_b, nv * dv) else {
                    bgraph_refused("Prism transform unavailable for batched GDN output");
                    return batch_outcome(o1_started || state_started, false);
                };
                ematb(&mut enc, out, &gdo_p, &ob, hidden, nv * dv);
                bts!(enc, 2);
            }
            _ => return batch_outcome(o1_started || state_started, false),
        }
        go(
            &mut enc,
            &c.add_rmsnorm_b,
            &bg(&c.layout_add_rmsnorm_b, &[&h_buf, &ob, &pnw, &n1, &rms_u]),
            k as u32,
        );
        match &lw.ffn {
            BFfn::Dense {
                gate,
                up,
                down,
                width,
            } => {
                let inter = *width; // this layer's, not the model's
                let Some(n1_p) = prism_input_b(&mut enc, &[gate, up], &n1, hidden) else {
                    bgraph_refused("Prism transform unavailable for batched FFN input");
                    return batch_outcome(o1_started || state_started, false);
                };
                ematb2(&mut enc, gate, up, &n1_p, &gbuf, &ubuf, inter, inter, hidden);
                go(
                    &mut enc,
                    &c.silu,
                    &bg(&c.layout_silu, &[&gbuf, &ubuf, &dummy_hd, &abuf, &silu_u]),
                    ((k * inter) as u32).div_ceil(256),
                );
                let Some(abuf_p) = prism_input_b(&mut enc, &[down], &abuf, inter) else {
                    bgraph_refused("Prism transform unavailable for batched FFN output");
                    return batch_outcome(o1_started || state_started, false);
                };
                ematb(&mut enc, down, &abuf_p, &ob, hidden, inter);
            }
            // Routing is per token, so the experts run token by token —
            // but inside THIS submit, next to the batched attention and
            // projections. Same four kernels the token graph uses, fed a
            // one-row slice of the batch and writing one row back.
            BFfn::Moe {
                router,
                sgate,
                gate_all,
                up_all,
                down_all,
                n_exp,
                top_k,
                inter: mi,
                norm_topk,
                q4tp,
                gu_q2,
                sigmoid,
                bias,
                shared_gated,
                route_scale,
            } => {
                let (mlogit, mslog, msel, mwt, mact) = moe_bufs.as_ref().unwrap();
                let mut continue_ffn = true;
                let bias_buf = bias.clone().unwrap_or_else(|| {
                    c.device
                        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                            label: Some("bmoe-sel-bias0"),
                            contents: &[0u8; 4],
                            usage: wgpu::BufferUsages::STORAGE,
                        })
                });
                let slots = *top_k + 1;
                let mat16 = |rows: usize, cols: usize| -> u32 {
                    let n = if *q4tp {
                        cortiq_core::quant::expected_nbytes(
                            cortiq_core::TensorDtype::Q4TiledP,
                            &[rows, cols],
                        )
                        .unwrap_or(0)
                    } else {
                        rows * (cols / 32) * 18
                    };
                    (n / 2) as u32
                };
                let sg_fold = sgate.kind == 4;
                // Same flags word as the token graph's select kernel (the
                // batch shares the layer's ungated-shared / sigmoid / bias
                // routing; the shared slot is always present here).
                let sel_u = uniform_u32x8(
                    c,
                    [
                        *n_exp as u32,
                        *top_k as u32,
                        u32::from(*norm_topk)
                            | (u32::from(*sigmoid) << 1)
                            | (u32::from(bias.is_some()) << 2)
                            | (1u32 << 3)
                            | (u32::from(!*shared_gated) << 4),
                        ((hidden as u32) << 8) | (u32::from(sg_fold) * 4),
                        route_scale.to_bits(),
                        0,
                        0,
                        0,
                    ],
                );
                // Gate/up stride follows the GU dtype: the mixed profile
                // packs them q2tp while down stays q4tp (see mat16 for
                // the down side).
                let gu_stride = if *gu_q2 {
                    (cortiq_core::quant::expected_nbytes(
                        cortiq_core::TensorDtype::Q2TiledP,
                        &[*mi, hidden],
                    )
                    .unwrap_or(0)
                        / 2) as u32
                } else {
                    mat16(*mi, hidden)
                };
                let gu_u = uniform_u32x8(
                    c,
                    [
                        (hidden / 32) as u32,
                        *mi as u32,
                        slots as u32,
                        gu_stride,
                        0,
                        0,
                        0,
                        0,
                    ],
                );
                let dn_u = uniform_u32x4(
                    c,
                    [
                        (*mi / 32) as u32,
                        hidden as u32,
                        slots as u32,
                        mat16(hidden, *mi),
                    ],
                );
                let (p_gu, p_dn, l_gu, l_dn) = if *gu_q2 {
                    // Mixed profile: 2-bit gate/up kernel, exact q4tp down
                    // (same pair the DSV4 chain path runs).
                    (
                        &c.moe_gate_up_q2tp,
                        &c.moe_down_q4tp,
                        &c.layout_moe_gu_q2tp,
                        &c.layout_moe_dn_q4tp,
                    )
                } else if *q4tp {
                    (
                        &c.moe_gate_up_q4tp,
                        &c.moe_down_q4tp,
                        &c.layout_moe_gu_q4tp,
                        &c.layout_moe_dn_q4tp,
                    )
                } else {
                    (
                        &c.moe_gate_up,
                        &c.moe_down,
                        &c.layout_moe_gu,
                        &c.layout_moe_dn,
                    )
                };
                {
                    use std::sync::atomic::{AtomicBool, Ordering};
                    static SAID: AtomicBool = AtomicBool::new(false);
                    if !SAID.swap(true, Ordering::Relaxed)
                        && std::env::var("CMF_GRAPH_SPEC_TIME").is_ok()
                    {
                        eprintln!(
                            "batch-moe path: q4tp={} router.kind={} sgate.kind={} n_exp={}",
                            q4tp, router.kind, sgate.kind, n_exp
                        );
                    }
                }
                if *q4tp && !*gu_q2 && router.kind == 4 && sgate.kind == 4 && *n_exp <= 256 {
                    // Uniform q4tp experts + f32 router/gate: k router
                    // matvecs (offset bindings, no staging rows) plus THREE
                    // token-axis dispatches for select/experts/down. The
                    // loop below is ~7 commands per token per layer and
                    // clocks the chunk at per-position speed.
                    // Router for every token in ONE dispatch; per-row math is
                    // f32_matvec verbatim, so the logits stay bit-identical.
                    let fr_u = uniform_u32x4(c, [hidden as u32, *n_exp as u32, 0, 0]);
                    {
                        let mut pass = begin_pass(&mut enc);
                        pass.set_pipeline(&c.f32_matvec_b);
                        pass.set_bind_group(
                            0,
                            &bg(&c.layout_f32b, &[&router.buf, &n1, mlogit, &fr_u]),
                            &[],
                        );
                        pass.dispatch_workgroups((*n_exp as u32).min(MAX_WG), k as u32, 1);
                    }
                    let bg_sel = bg(
                        &c.layout_moe_sel_b,
                        &[mlogit, &n1, msel, mwt, &sel_u, &sgate.buf, &bias_buf],
                    );
                    let bg_gu = bg(
                        &c.layout_moe_gu_b,
                        &[gate_all, up_all, &n1, msel, mact, &gu_u],
                    );
                    let bg_dn = bg(&c.layout_moe_dn_b, &[down_all, mact, msel, mwt, &ob, &dn_u]);
                    let mut pass = begin_pass(&mut enc);
                    pass.set_pipeline(&c.moe_select_b);
                    pass.set_bind_group(0, &bg_sel, &[]);
                    pass.dispatch_workgroups(k as u32, 1, 1);
                    pass.set_pipeline(&c.moe_gate_up_q4tp_b);
                    pass.set_bind_group(0, &bg_gu, &[]);
                    pass.dispatch_workgroups(*mi as u32, slots as u32, k as u32);
                    pass.set_pipeline(&c.moe_down_q4tp_b);
                    pass.set_bind_group(0, &bg_dn, &[]);
                    pass.dispatch_workgroups(hidden as u32, k as u32, 1);
                    drop(pass);
                    continue_ffn = false;
                }
                if continue_ffn {
                    for i in 0..k {
                        cp(&mut enc, &n1, i * hidden, &row_in, hidden);
                        let bg_sel = bg(
                            &c.layout_moe_sel,
                            &[mlogit, mslog, msel, mwt, &sel_u, &sgate.buf, &row_in, &bias_buf],
                        );
                        let bg_gu = bg(l_gu, &[gate_all, up_all, &row_in, msel, mact, &gu_u]);
                        let bg_dn = bg(l_dn, &[down_all, mact, msel, mwt, &row_out, &dn_u]);
                        emat1(&mut enc, router, &row_in, mlogit, *n_exp, hidden);
                        if !sg_fold {
                            emat1(&mut enc, sgate, &row_in, mslog, 1, hidden);
                        }
                        let mut pass = begin_pass(&mut enc);
                        pass.set_pipeline(&c.moe_select);
                        pass.set_bind_group(0, &bg_sel, &[]);
                        pass.dispatch_workgroups(1, 1, 1);
                        pass.set_pipeline(p_gu);
                        pass.set_bind_group(0, &bg_gu, &[]);
                        pass.dispatch_workgroups(*mi as u32, slots as u32, 1);
                        pass.set_pipeline(p_dn);
                        pass.set_bind_group(0, &bg_dn, &[]);
                        pass.dispatch_workgroups(hidden as u32, 1, 1);
                        drop(pass);
                        cpo(&mut enc, &row_out, &ob, i * hidden, hidden);
                    }
                }
            }
        }
        bts!(enc, 3);
        if li + 1 < layers.len() {
            let inw_next = stor(bytemuck::cast_slice(layers[li + 1].input_norm));
            go(
                &mut enc,
                &c.add_rmsnorm_b,
                &bg(
                    &c.layout_add_rmsnorm_b,
                    &[&h_buf, &ob, &inw_next, &n1, &rms_u],
                ),
                k as u32,
            );
        } else {
            let ax_u = unif(&[1.0f32.to_bits(), (k * hidden) as u32, 0, 0]);
            go(
                &mut enc,
                &c.axpy,
                &bg(&c.layout_axpy, &[&ob, &h_buf, &ax_u]),
                ((k * hidden) as u32).div_ceil(256),
            );
        }
        bts!(enc, 4);
        if tap_layer == Some(li) {
            if let Some(tap) = tap_stage.as_ref() {
                flush_pass(&enc);
                enc.copy_buffer_to_buffer(&h_buf, 0, tap, 0, (k * hidden * 4) as u64);
            }
        }
    }
    let size = (k * hidden * 4) as u64;
    let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("bg-stage"),
        size,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    let t_enc_done = std::time::Instant::now();
    let ok = if let Some(sp) = spec.as_mut() {
        // Speculative tail: final-norm each row, one batched lm_head GEMM,
        // read every position's logits back beside the hiddens. The CPU
        // argmaxes them — a verify wants the last accepted row's WHOLE
        // logits anyway (the sampler's contract at the loop top).
        let Some(lm) = resolve(&sp.lm, sp.lm_rows, hidden) else {
            bgraph_refused("spec tail: lm head weight did not resolve");
            return batch_outcome(o1_started || state_started, false);
        };
        let fnw = stor(bytemuck::cast_slice(sp.final_norm));
        let n1b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("spec-n1"),
            size,
            usage: wgpu::BufferUsages::STORAGE,
            mapped_at_creation: false,
        });
        go(
            &mut enc,
            &c.rmsnorm_b,
            &bg(&c.layout_rmsnorm_b, &[&h_buf, &fnw, &n1b, &rms_u]),
            k as u32,
        );
        let lsize = (k * sp.lm_rows * 4) as u64;
        let lbuf = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("spec-logits"),
            size: lsize,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let Some(n1b_p) = prism_input_b(&mut enc, &[&lm], &n1b, hidden) else {
            bgraph_refused("Prism transform unavailable for batched lm_head");
            return batch_outcome(o1_started || state_started, false);
        };
        ematb(&mut enc, &lm, &n1b_p, &lbuf, sp.lm_rows, hidden);
        bts!(enc, 5);
        batch_kernel_ts_resolve(c, &mut enc);
        if bts_on && bts_lbl.len() > 1 {
            if let Some((qs, resolve, tstage)) = c.ts_query.as_ref() {
                flush_pass(&enc);
                enc.resolve_query_set(qs, 0..bts_lbl.len() as u32, resolve, 0);
                flush_pass(&enc);
                enc.copy_buffer_to_buffer(resolve, 0, tstage, 0, bts_lbl.len() as u64 * 8);
            }
        }
        sp.logits_out.resize(k * sp.lm_rows, 0.0);
        readback2(
            c,
            enc,
            (&h_buf, &mut h[..k * hidden]),
            (&lbuf, &mut sp.logits_out[..]),
        )
    } else {
        batch_kernel_ts_resolve(c, &mut enc);
        if bts_on && bts_lbl.len() > 1 {
            if let Some((qs, resolve, tstage)) = c.ts_query.as_ref() {
                flush_pass(&enc);
                enc.resolve_query_set(qs, 0..bts_lbl.len() as u32, resolve, 0);
                flush_pass(&enc);
                enc.copy_buffer_to_buffer(resolve, 0, tstage, 0, bts_lbl.len() as u64 * 8);
            }
        }
        readback(c, enc, &h_buf, &stage, size, &mut h[..k * hidden])
    };
    if std::env::var("CMF_GRAPH_SPEC_TIME").is_ok() {
        let post = t_enc_done.elapsed().as_secs_f64() * 1e3;
        eprintln!(
            "batch-graph: encode {:.1} ms | gpu+readback {:.1} ms",
            t_bfn.elapsed().as_secs_f64() * 1e3 - post,
            post,
        );
    }
    if ok {
        if let Some(tap) = tap_stage.as_ref() {
            let bytes = (k * hidden * 4) as u64;
            let (tx, rx) = std::sync::mpsc::channel();
            tap.slice(..bytes).map_async(wgpu::MapMode::Read, move |r| {
                let _ = tx.send(r);
            });
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            if rx.recv().map(|r| r.is_ok()).unwrap_or(false) {
                if let Ok(raw) = tap.get_mapped_range(..bytes) {
                    let vals: &[f32] = bytemuck::cast_slice(&raw);
                    let last = &vals[(k - 1) * hidden..k * hidden];
                    let mut norm = 0.0f64;
                    for &v in last { norm += (v as f64) * (v as f64); }
                    eprintln!(
                        "batch-tap layer={} rows={} row0={:?} rowlast={:?} normlast={:.9e}",
                        tap_layer.unwrap_or(usize::MAX),
                        k,
                        &vals[..hidden.min(4)],
                        &vals[(k - 1) * hidden..(k - 1) * hidden + hidden.min(4)],
                        norm.sqrt()
                    );
                    drop(raw);
                }
            }
            tap.unmap();
        }
    }
    if ok && bts_on && bts_lbl.len() > 1 {
        if let Some((_, _, tstage)) = c.ts_query.as_ref() {
            let bytes = bts_lbl.len() as u64 * 8;
            let (tx, rx) = std::sync::mpsc::channel();
            tstage.map_async(wgpu::MapMode::Read, ..bytes, move |r| {
                let _ = tx.send(r);
            });
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            if rx.recv().map(|r| r.is_ok()).unwrap_or(false) {
                if let Ok(raw) = tstage.get_mapped_range(..bytes) {
                    let t: Vec<u64> = bytemuck::cast_slice::<u8, u64>(&raw).to_vec();
                    drop(raw);
                    let mut agg = [(0f64, 0u32); 8];
                    for i in 1..bts_lbl.len() {
                        let d = t[i].saturating_sub(t[i - 1]) as f64 * c.ts_period as f64 / 1e6;
                        let e = &mut agg[bts_lbl[i] as usize];
                        e.0 += d;
                        e.1 += 1;
                    }
                    let names = [
                        "start", "attn", "gdn-out", "ffn", "norm", "recur", "gdn-in", "gdn-proj",
                    ];
                    let line: Vec<String> = agg
                        .iter()
                        .enumerate()
                        .filter(|(_, e)| e.1 > 0)
                        .map(|(i, e)| format!("{}={:.2}ms/{}", names[i], e.0, e.1))
                        .collect();
                    eprintln!("batch-ts: {}", line.join(" "));
                }
            }
            tstage.unmap();
        }
    }
    if ok {
        batch_kernel_ts_report(c, k);
    }
    if ok {
        let mut kvm = c.attn_kv.lock().unwrap();
        for li in 0..layers.len() {
            if let Some(m) = kvm.get_mut(&(kv_id, li)) {
                m.synced = pos0 + k;
            }
        }
        let next = pos0 + k;
        let mut gcm = c.gdn_cursor.lock().unwrap();
        for li in 0..layers.len() {
            if let Some(cur) = gcm.get_mut(&(kv_id, li)) {
                cur.next_pos = next;
            }
        }
        drop(gcm);
        let mut om = c.o1m.lock().unwrap();
        for (li, views) in o1.iter().enumerate() {
            if views.is_some() {
                if let Some(d) = om.get_mut(&(kv_id, li)) {
                    d.next_pos = Some(next);
                }
            }
        }
    }
    batch_outcome(o1_started || state_started, ok)
}

/// After a partial speculative acceptance: put every GDN layer's (ring, S)
/// back to the snapshot taken after batch position `slot` — the last
/// position whose input token was real. `base_pos` and `expected_layers` are
/// part of the call contract: a stale snapshot from an earlier round, or a
/// partial snapshot set, must fail closed instead of moving only some
/// recurrent layers and letting the next graph consume mixed state.
pub fn gdn_spec_restore(kv_id: u64, slot: usize, base_pos: usize, expected_layers: usize) -> bool {
    let Some(c) = ctx() else { return false };
    let snaps = c.gdn_snap.lock().unwrap();
    let entries: Vec<(usize, &(wgpu::Buffer, u64, u64, usize, usize))> = snaps
        .iter()
        .filter_map(|((id, li), value)| (*id == kv_id).then_some((*li, value)))
        .collect();
    if entries.len() != expected_layers || entries.is_empty() {
        return false;
    }
    let Some(next) = slot
        .checked_add(1)
        .and_then(|accepted_rows| base_pos.checked_add(accepted_rows))
    else {
        return false;
    };
    let states = c.gdn_state.lock().unwrap();
    let mut common_dims: Option<(usize, usize, usize, usize, usize, usize)> = None;
    let mut cursors = c.gdn_cursor.lock().unwrap();
    // Every layer must describe this same speculative batch anchor and have
    // a matching device state/cursor.  Derive the byte geometry from that
    // cursor as well as trusting the snapshot metadata: otherwise a reused
    // `(kv_id, layer)` entry from another model could copy a valid-looking
    // prefix with the wrong stride. The batch admission path currently
    // requires homogeneous GDN geometry, so reject a mixed set here too.
    for (li, (_, ring_sz, s_sz, slots, snap_base)) in &entries {
        let Some(cur) = cursors.get(&(kv_id, *li)) else {
            return false;
        };
        if !states.contains_key(&(kv_id, *li))
            || slot >= *slots
            || *snap_base != base_pos
            || cur.next_pos < next
        {
            return false;
        }
        let (nv, _nk, dk, dv, kk, cdim) = cur.dims;
        let Some(ring_words) = cdim.checked_mul(kk.saturating_sub(1)) else {
            return false;
        };
        let Some(state_words) = nv.checked_mul(dk).and_then(|v| v.checked_mul(dv)) else {
            return false;
        };
        let want_ring = (ring_words as u64).saturating_mul(4);
        let want_state = (state_words as u64).saturating_mul(4);
        if *ring_sz != want_ring || *s_sz != want_state {
            return false;
        }
        if let Some(prev) = common_dims {
            if prev != cur.dims {
                return false;
            }
        } else {
            common_dims = Some(cur.dims);
        }
    }
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("gdn-restore"),
        });
    for (li, (snap, ring_sz, s_sz, _, _)) in &entries {
        let Some((ring, s)) = states.get(&(kv_id, *li)) else {
            return false;
        };
        let off = slot as u64 * (*ring_sz + *s_sz);
        if *ring_sz > 0 {
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(snap, off, ring, 0, *ring_sz);
        }
        if *s_sz > 0 {
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(snap, off + *ring_sz, s, 0, *s_sz);
        }
    }
    submit(c, finish_enc(enc));
    for (li, _) in &entries {
        let Some(cur) = cursors.get_mut(&(kv_id, *li)) else {
            // All cursors were validated before submit; this is defensive
            // against an impossible concurrent deletion and still fails
            // closed for the caller.
            return false;
        };
        cur.next_pos = next;
    }
    true
}

/// Drop the device K/V mirror for a pipeline (called on cache clear).
pub fn kv_mirror_reset(kv_id: u64) {
    if let Some(c) = ctx() {
        c.attn_kv.lock().unwrap().retain(|(id, _), _| *id != kv_id);
        c.gdn_state
            .lock()
            .unwrap()
            .retain(|(id, _), _| *id != kv_id);
        c.gdn_cursor
            .lock()
            .unwrap()
            .retain(|(id, _), _| *id != kv_id);
        c.gdn_snap.lock().unwrap().retain(|(id, _), _| *id != kv_id);
        c.o1m.lock().unwrap().retain(|(id, _), _| *id != kv_id);
        c.graph_bgs
            .lock()
            .unwrap()
            .retain(|(_, _, id), _| *id != kv_id);
    }
}

/// Rewind the logical row count of one exact-attention mirror. Speculative
/// MTP drafts append rows directly on the device while the CPU owner retains
/// only the real anchor; after verification, accepted warm rows overwrite from
/// that anchor rather than being rejected as ahead of the next graph position.
/// The storage itself is retained for the next append.
pub fn kv_mirror_set_stored(kv_id: u64, layer: usize, stored: usize) -> bool {
    let Some(c) = ctx() else { return false };
    let mut mirrors = c.attn_kv.lock().unwrap();
    let Some(m) = mirrors.get_mut(&(kv_id, layer)) else {
        return false;
    };
    if stored > m.cap {
        return false;
    }
    m.synced = stored;
    true
}

/// GDN depthwise conv step (bring-up / parity): updates cq [cdim] and shifts
/// the ring [(kk-1)·cdim] in place.
pub fn gdn_conv_gpu(
    qkv: &[f32],
    taps: &[f32],
    ring: &mut [f32],
    cdim: usize,
    kk: usize,
    cq: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    let qb = storage_bytes(c, bytemuck::cast_slice(qkv));
    let tb = storage_bytes(c, bytemuck::cast_slice(taps));
    let rb = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(ring),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        });
    let cb = rw_f32(c, cdim, true);
    let p = uniform_u32x4(c, [cdim as u32, kk as u32, 0, 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &c.layout_gdn_conv,
        entries: &[
            bind_buf(0, &qb),
            bind_buf(1, &tb),
            bind_buf(2, &rb),
            bind_buf(3, &cb),
            bind_buf(4, &p),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
    {
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.gdn_conv);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((cdim as u32).div_ceil(256), 1, 1);
    }
    let rsz = (ring.len() * 4) as u64;
    let csz = (cdim * 4) as u64;
    let sr = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: None,
        size: rsz,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    let scq = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: None,
        size: csz,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&rb, 0, &sr, 0, rsz);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&cb, 0, &scq, 0, csz);
    submit(c, finish_enc(enc));
    sr.slice(..).map_async(wgpu::MapMode::Read, |_| {});
    scq.slice(..).map_async(wgpu::MapMode::Read, |_| {});
    if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        return false;
    }
    let (Ok(dr), Ok(dc)) = (
        sr.slice(..).get_mapped_range(),
        scq.slice(..).get_mapped_range(),
    ) else {
        return false;
    };
    ring.copy_from_slice(bytemuck::cast_slice(&dr[..ring.len() * 4]));
    cq[..cdim].copy_from_slice(bytemuck::cast_slice(&dc[..cdim * 4]));
    true
}

/// GDN decode step (bring-up / parity): one workgroup per v-head. `s` is the
/// [nv·dk·dv] recurrent state, updated in place; writes `o` [nv·dv].
#[allow(clippy::too_many_arguments)]
pub fn gdn_step_gpu(
    cq: &[f32],
    z: &[f32],
    a: &[f32],
    b: &[f32],
    alog: &[f32],
    dtb: &[f32],
    norm: &[f32],
    s: &mut [f32],
    nv: usize,
    dk: usize,
    dv: usize,
    kd: usize,
    rep: usize,
    cdim: usize,
    eps: f32,
    o: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    let sb = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("gdn-s"),
            contents: bytemuck::cast_slice(s),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        });
    let ob = rw_f32(c, nv * dv, true);
    let p = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("gdn-p"),
            contents: bytemuck::cast_slice(&[
                nv as u32,
                dk as u32,
                dv as u32,
                kd as u32,
                rep as u32,
                cdim as u32,
                eps.to_bits(),
                0u32,
            ]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let sbuf = |d: &[f32]| storage_bytes(c, bytemuck::cast_slice(d));
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("gdn-bg"),
        layout: &c.layout_gdn,
        entries: &[
            bind_buf(0, &sbuf(cq)),
            bind_buf(1, &sbuf(z)),
            bind_buf(2, &sbuf(a)),
            bind_buf(3, &sbuf(b)),
            bind_buf(4, &sbuf(alog)),
            bind_buf(5, &sbuf(dtb)),
            bind_buf(6, &sbuf(norm)),
            bind_buf(7, &sb),
            bind_buf(8, &ob),
            bind_buf(9, &p),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("gdn") });
    {
        let mut pass = begin_pass_with(&mut enc, Some("gdn"), None);
        pass.set_pipeline(&c.gdn_step);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(nv as u32, 1, 1);
    }
    // read back updated S and o
    let ssz = (s.len() * 4) as u64;
    let osz = (nv * dv * 4) as u64;
    let stage_s = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: None,
        size: ssz,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    let stage_o = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: None,
        size: osz,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&sb, 0, &stage_s, 0, ssz);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&ob, 0, &stage_o, 0, osz);
    submit(c, finish_enc(enc));
    stage_s.slice(..).map_async(wgpu::MapMode::Read, |_| {});
    stage_o.slice(..).map_async(wgpu::MapMode::Read, |_| {});
    if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        return false;
    }
    let (Ok(ds), Ok(dobuf)) = (
        stage_s.slice(..).get_mapped_range(),
        stage_o.slice(..).get_mapped_range(),
    ) else {
        return false;
    };
    s.copy_from_slice(bytemuck::cast_slice(&ds[..s.len() * 4]));
    o[..nv * dv].copy_from_slice(bytemuck::cast_slice(&dobuf[..nv * dv * 4]));
    true
}

/// One full attention sub-block resident on the GPU in a SINGLE command
/// encoder: rmsnorm → QKV (q1) → rope/qk-norm → kv_append → attend → O (q1)
/// → residual. The K/V cache lives on the device ([nkv,cap,hd]) and persists
/// across tokens; only the updated hidden is read back. This is the token
/// graph's attention half — it collapses ~6 per-op submits into one.
/// `flags` follows attn_rope_qkn (2=qnorm 4=knorm 8=gemma; gate unsupported
/// here). Weights are raw q1 payloads (bring-up path; production keys the
/// resident VRAM cache). Returns false without a GPU context.
#[allow(clippy::too_many_arguments)]
pub fn attn_block_gpu(
    h_in: &[f32],
    attn_norm_w: &[f32],
    wq: &[u8],
    wk: &[u8],
    wv: &[u8],
    wo: &[u8],
    qnw: &[f32],
    knw: &[f32],
    invf: &[f32],
    kbuf: &wgpu::Buffer,
    vbuf: &wgpu::Buffer,
    nh: usize,
    nkv: usize,
    hd: usize,
    rd: usize,
    hidden: usize,
    cap: usize,
    stored: usize,
    flags: u32,
    eps: f32,
    h_out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    let unif = |data: &[u32]| {
        c.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("blk-u"),
                contents: bytemuck::cast_slice(data),
                usage: wgpu::BufferUsages::UNIFORM,
            })
    };
    let stor = |data: &[u8]| {
        c.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("blk-w"),
                contents: data,
                usage: wgpu::BufferUsages::STORAGE,
            })
    };
    // Resident buffers.
    let h_buf = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("blk-h"),
            contents: bytemuck::cast_slice(&h_in[..hidden]),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        });
    let normw_b = stor(bytemuck::cast_slice(&attn_norm_w[..hidden]));
    let normed_b = rw_f32(c, hidden, false);
    let wq_b = stor(wq);
    let wk_b = stor(wk);
    let wv_b = stor(wv);
    let wo_b = stor(wo);
    let qraw_b = rw_f32(c, nh * hd, false);
    let k_b = rw_f32(c, nkv * hd, false);
    let v_b = rw_f32(c, nkv * hd, false);
    let qout_b = rw_f32(c, nh * hd, false);
    let gout_b = rw_f32(c, nh * hd, false);
    let qnw_b = stor(bytemuck::cast_slice(qnw));
    let knw_b = stor(bytemuck::cast_slice(knw));
    let invf_b = stor(bytemuck::cast_slice(invf));
    let attn_b = rw_f32(c, nh * hd, false);
    let o_b = rw_f32(c, hidden, false);
    let bg = |layout: &wgpu::BindGroupLayout, bufs: &[&wgpu::Buffer]| {
        let entries: Vec<wgpu::BindGroupEntry> = bufs
            .iter()
            .enumerate()
            .map(|(i, b)| bind_buf(i as u32, b))
            .collect();
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout,
            entries: &entries,
        })
    };
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("attn-block"),
        });
    let dispatch = |enc: &mut wgpu::CommandEncoder,
                    pipe: &wgpu::ComputePipeline,
                    bind: &wgpu::BindGroup,
                    groups: u32| {
        let mut pass = begin_pass(enc);
        pass.set_pipeline(pipe);
        pass.set_bind_group(0, bind, &[]);
        pass.dispatch_workgroups(groups, 1, 1);
    };
    // 1. rmsnorm(h) -> normed
    let rms_p = unif(&[hidden as u32, 0, eps.to_bits(), 0]);
    dispatch(
        &mut enc,
        &c.rmsnorm,
        &bg(&c.layout_rmsnorm, &[&h_buf, &normw_b, &normed_b, &rms_p]),
        1,
    );
    // 2. QKV (q1) from normed
    encode_matvec_q1(c, &mut enc, &wq_b, &normed_b, &qraw_b, nh * hd, hidden);
    encode_matvec_q1(c, &mut enc, &wk_b, &normed_b, &k_b, nkv * hd, hidden);
    encode_matvec_q1(c, &mut enc, &wv_b, &normed_b, &v_b, nkv * hd, hidden);
    // 3. rope + qk-norm
    let rq_p = unif(&[
        nh as u32,
        nkv as u32,
        hd as u32,
        rd as u32,
        stored as u32,
        flags,
        eps.to_bits(),
        0,
    ]);
    dispatch(
        &mut enc,
        &c.attn_rope,
        &bg(
            &c.layout_attn_rope,
            &[
                &qraw_b, &k_b, &qout_b, &gout_b, &qnw_b, &knw_b, &invf_b, &rq_p,
            ],
        ),
        (nh + nkv) as u32,
    );
    // 4. kv_append
    let kv_p = unif(&[nkv as u32, hd as u32, cap as u32, stored as u32]);
    let kv_groups = ((nkv * hd) as u32).div_ceil(256);
    dispatch(
        &mut enc,
        &c.kv_append,
        &bg(&c.layout_kv, &[&k_b, &v_b, kbuf, vbuf, &kv_p]),
        kv_groups,
    );
    // 5. attend
    let at_p = unif(&[
        nh as u32,
        (nh / nkv) as u32,
        hd as u32,
        cap as u32,
        (stored + 1) as u32,
        (1.0 / (hd as f32).sqrt()).to_bits(),
        0,
        0,
    ]);
    {
        let (ap, al) = attend_pipes(c, hd);
        dispatch(
            &mut enc,
            ap,
            &bg(al, &[&qout_b, kbuf, vbuf, &attn_b, &at_p]),
            nh as u32,
        );
    }
    // 6. O (q1)
    encode_matvec_q1(c, &mut enc, &wo_b, &attn_b, &o_b, hidden, nh * hd);
    // 7. residual h += o
    let ax_p = unif(&[1.0f32.to_bits(), hidden as u32, 0, 0]);
    dispatch(
        &mut enc,
        &c.axpy,
        &bg(&c.layout_axpy, &[&o_b, &h_buf, &ax_p]),
        (hidden as u32).div_ceil(256),
    );
    // readback updated hidden
    let size = (hidden * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "blk-stage",
    );
    let ok = readback(c, enc, &h_buf, &stage, size, &mut h_out[..hidden]);
    drop(sc);
    ok
}

/// q1 kernel body (weight_key = None — no residency cache; test path).
fn dispatch_q1(
    c: &Ctx,
    weight_key: Option<(usize, usize)>,
    payload: &[u8],
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    let gpr = cols / 32;
    let q_buf = match weight_key {
        Some(k) => match weight_buffer(c, k, payload) {
            Some(b) => b,
            None => return false, // over VRAM budget — honest CPU path
        },
        None => c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("q1-weights"),
                contents: payload,
                usage: wgpu::BufferUsages::STORAGE,
            }),
    };
    let mut sc = c.scratch.lock().unwrap();
    let xs_buf = Scratch::ensure(
        &c.device,
        &mut sc.xs,
        (cols * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "q1-xs",
    );
    c.queue
        .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&xs[..cols]));
    let y_size = (rows * 4) as u64;
    let y_buf = Scratch::ensure(
        &c.device,
        &mut sc.y,
        y_size,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "q1-y",
    );
    let params = [(gpr / 2) as u32, rows as u32, 0u32, 0u32];
    let p_buf = match &sc.params {
        Some(b) => b.clone(),
        None => {
            let b = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("q1-params"),
                size: 16,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            sc.params = Some(b.clone());
            b
        }
    };
    c.queue
        .write_buffer(&p_buf, 0, bytemuck::cast_slice(&params));
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        y_size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "q1-stage",
    );
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q1-bg"),
        layout: &c.layout_q1,
        entries: &[
            bind_buf(0, &q_buf),
            bind_buf(1, &xs_buf),
            bind_buf(2, &y_buf),
            bind_buf(3, &p_buf),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("q1") });
    {
        let mut pass = begin_pass_with(&mut enc, Some("q1"), None);
        pass.set_pipeline(&c.q1);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((rows as u32).div_ceil(c.q1_rows).min(MAX_WG), 1, 1);
    }
    let ok = readback(c, enc, &y_buf, &stage_buf, y_size, &mut out[..rows]);
    drop(sc);
    ok
}

/// GEMM of the prefill batch: `pre` are prescaled inputs row-major [b, cols],
/// out — row-major [b, rows]. Weights are resident in VRAM. false = CPU path.
#[allow(clippy::too_many_arguments)]
/// The two-field int8 GEMM with the column field handed to the device.
///
/// `q8_matmat` takes an activation the caller has already multiplied by that
/// field — a full copy of the panel, per call, on the host. When the weight
/// goes to the matrix units the field is folded into the plane instead, and
/// nothing touches the activation at all.
#[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 {
    let Some(c) = ctx() else { return false };
    if cols % 4 != 0 || rows == 0 || b == 0 || col_field.len() < cols {
        return false;
    }
    let entry = &model.tensors[idx];
    if entry.shape.first().copied().unwrap_or(0) < rows {
        return false;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return false;
    };
    let bytes = model.primary_bytes();
    let payload_len = if entry.dtype == cortiq_core::TensorDtype::Q8_2f {
        entry.nbytes as usize
    } else {
        rows * cols
    };
    if abs + payload_len > bytes.len()
        || row_scale.len() < rows
        || xs.len() < b * cols
        || out.len() < b * rows
    {
        return false;
    }
    dispatch_matmat_keep(
        c,
        {
            note_layer((model.uid() as usize, idx), &model.tensors[idx].name);
            Some((model.uid() as usize, idx))
        },
        &bytes[abs..abs + payload_len],
        row_scale,
        Some(&col_field[..cols]),
        xs,
        b,
        rows,
        cols,
        Some(out),
        None,
    )
    .is_some()
}

pub fn q8_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    row_scale: &[f32],
    pre: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if cols % 4 != 0 || rows == 0 || b == 0 {
        return false;
    }
    let entry = &model.tensors[idx];
    if entry.shape.first().copied().unwrap_or(0) < rows {
        return false;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return false;
    };
    let bytes = model.primary_bytes();
    let payload_len = if entry.dtype == cortiq_core::TensorDtype::Q8_2f {
        entry.nbytes as usize
    } else {
        rows * cols
    };
    if abs + payload_len > bytes.len()
        || row_scale.len() < rows
        || pre.len() < b * cols
        || out.len() < b * rows
    {
        return false;
    }
    let full_quant = &bytes[abs..abs + payload_len];
    dispatch_matmat(
        c,
        {
            note_layer((model.uid() as usize, idx), &model.tensors[idx].name);
            Some((model.uid() as usize, idx))
        },
        full_quant,
        row_scale,
        pre,
        b,
        rows,
        cols,
        out,
    )
}

/// The per-row scales and the column field of a `q8_row`/`q8_2f` weight,
/// read straight out of the mapped file: `[int8 body][f16 row scales][f16
/// column field]`.
fn q8_fields(
    model: &Arc<CmfModel>,
    idx: usize,
    rows: usize,
    cols: usize,
) -> Option<(Vec<f32>, Vec<f32>)> {
    use cortiq_core::TensorDtype as D;
    let entry = &model.tensors[idx];
    let abs = model.entry_abs_offset(entry)?;
    let bytes = model.primary_bytes();
    let n = rows * cols;
    let need = n + rows * 2 + if entry.dtype == D::Q8_2f { cols * 2 } else { 0 };
    if abs + need > bytes.len() {
        return None;
    }
    let b = &bytes[abs..abs + need];
    let f16 = |o: usize| cortiq_core::quant::f16_to_f32(u16::from_le_bytes([b[o], b[o + 1]]));
    let rs: Vec<f32> = (0..rows).map(|o| f16(n + o * 2)).collect();
    let cf: Vec<f32> = if entry.dtype == D::Q8_2f {
        (0..cols).map(|i| f16(n + rows * 2 + i * 2)).collect()
    } else {
        vec![1.0; cols]
    };
    Some((rs, cf))
}

/// The panel a projection produces, left on the card, for WHICHEVER codec the
/// weight is in. The fused DiT paths were written against `q4tp` and asked for
/// it by name; a container packed any other way then fell back to per-op GEMMs
/// with a readback between every one.
///
/// `xs` is the host activation. For the two-field codec the column field is
/// folded into it here — that is what leaves a plain per-row int8 product.
pub(crate) fn fused_panel_keep(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
) -> Option<wgpu::Buffer> {
    use cortiq_core::TensorDtype as D;
    let entry = &model.tensors[idx];
    match entry.dtype {
        D::Q4TiledP => {
            let mut unused = Vec::new();
            tp_matmat_keep(model, idx, xs, b, rows, cols, &mut unused, false)
        }
        D::Q8Row | D::Q8_2f if fused_any() => {
            let (rs, cf) = q8_fields(model, idx, rows, cols)?;
            let col = (entry.dtype == D::Q8_2f).then_some(&cf[..]);
            q8_matmat_keep(model, idx, &rs, col, xs, b, rows, cols)
        }
        _ => None,
    }
}

/// `y = x · col` on a panel already on the card, into a fresh buffer. The
/// two-field codec's column field has to meet the activation somewhere; when
/// the activation is the previous kernel's output, it meets it here instead
/// of on a round trip home.
fn colscale_keep(c: &Ctx, src: &wgpu::Buffer, col: &[f32], n_total: usize) -> Option<wgpu::Buffer> {
    let cols = col.len();
    if cols == 0 || n_total == 0 || n_total % cols != 0 {
        return None;
    }
    let col_buf = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("cs-col"),
            contents: bytemuck::cast_slice(col),
            usage: wgpu::BufferUsages::STORAGE,
        });
    let y = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("cs-y"),
        size: (n_total * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE,
        mapped_at_creation: false,
    });
    let pbuf = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("cs-p"),
            contents: bytemuck::cast_slice(&[n_total as u32, cols as u32, 0u32, 0u32]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("cs"),
        layout: &c.layout_colscale,
        entries: &[
            bind_buf(0, src),
            bind_buf(1, &col_buf),
            bind_buf(2, &y),
            bind_buf(3, &pbuf),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("cs") });
    {
        let mut pass = begin_pass_with(&mut enc, Some("cs"), None);
        pass.set_pipeline(&c.colscale);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((n_total as u32).div_ceil(256).min(MAX_WG), 1, 1);
    }
    c.queue.submit(Some(enc.finish()));
    Some(y)
}

/// A projection whose input is ALREADY on the card, for whichever codec —
/// the second half of every fused pair (attention → out, SwiGLU → fc2).
pub(crate) fn fused_gemm_from_device(
    model: &Arc<CmfModel>,
    idx: usize,
    src: &wgpu::Buffer,
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    use cortiq_core::TensorDtype as D;
    match model.tensors[idx].dtype {
        D::Q4TiledP => {
            tp_matmat_impl(
                model,
                idx,
                &[],
                b,
                rows,
                cols,
                Some(out),
                Some(src),
                0,
                false,
                None,
            )
                .is_some()
        }
        dt @ (D::Q8Row | D::Q8_2f) if fused_any() => {
            let Some(c) = ctx() else { return false };
            let Some((rs, cf)) = q8_fields(model, idx, rows, cols) else {
                return false;
            };
            let col = (dt == D::Q8_2f).then_some(&cf[..]);
            let entry = &model.tensors[idx];
            let Some(abs) = model.entry_abs_offset(entry) else {
                return false;
            };
            let bytes = model.primary_bytes();
            let payload_len = if dt == D::Q8_2f {
                entry.nbytes as usize
            } else {
                rows * cols
            };
            if abs + payload_len > bytes.len() {
                return false;
            }
            dispatch_matmat_keep(
                c,
                {
                    note_layer((model.uid() as usize, idx), &model.tensors[idx].name);
                    Some((model.uid() as usize, idx))
                },
                &bytes[abs..abs + payload_len],
                &rs,
                col,
                &[],
                b,
                rows,
                cols,
                Some(out),
                Some(src),
            )
            .is_some()
        }
        _ => false,
    }
}

/// The int8 GEMM with its result left on the card — the entry the fused DiT
/// paths use for a `q8_row`/`q8_2f` weight. `pre` is the activation already
/// carrying the column field, which is what turns the two-field codec into a
/// plain per-row int8 product.
pub fn q8_matmat_keep(
    model: &Arc<CmfModel>,
    idx: usize,
    row_scale: &[f32],
    col_scale: Option<&[f32]>,
    pre: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
) -> Option<wgpu::Buffer> {
    let c = ctx()?;
    if cols % 4 != 0 || rows == 0 || b == 0 {
        return None;
    }
    let entry = &model.tensors[idx];
    if entry.shape.first().copied().unwrap_or(0) < rows {
        return None;
    }
    let abs = model.entry_abs_offset(entry)?;
    let bytes = model.primary_bytes();
    let payload_len = if entry.dtype == cortiq_core::TensorDtype::Q8_2f {
        entry.nbytes as usize
    } else {
        rows * cols
    };
    if abs + payload_len > bytes.len() || row_scale.len() < rows || pre.len() < b * cols {
        return None;
    }
    let full_quant = &bytes[abs..abs + payload_len];
    dispatch_matmat_keep(
        c,
        {
            note_layer((model.uid() as usize, idx), &model.tensors[idx].name);
            Some((model.uid() as usize, idx))
        },
        full_quant,
        row_scale,
        col_scale,
        pre,
        b,
        rows,
        cols,
        None,
        None,
    )
}

/// Batched q1 GEMM (prefill): resident 1-bit weight, batch of raw-f32 inputs,
/// one 2D dispatch of q1_mul_mm, one readback. cols must be a 64-multiple (the
/// q1 format packs whole tile-pairs). Weights resident + cached; x through the
/// pooled scratch.
pub fn q1_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    pre: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if cols % 64 != 0 || rows == 0 || b == 0 || pre.len() < b * cols || out.len() < b * rows {
        return false;
    }
    let entry = &model.tensors[idx];
    if entry.shape.first().copied().unwrap_or(0) < rows {
        return false;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return false;
    };
    let bytes = model.primary_bytes();
    let plen = entry.nbytes as usize;
    if abs + plen > bytes.len() {
        return false;
    }
    let Some(w) = weight_buffer_l(
        c,
        (model.uid() as usize, idx),
        &bytes[abs..abs + plen],
        layer_of_name(&model.tensors[idx].name),
    ) else {
        return false; // over VRAM budget → CPU path
    };
    let mut sc = c.scratch.lock().unwrap();
    let xs_buf = Scratch::ensure(
        &c.device,
        &mut sc.xs,
        (b * cols * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "q1mm-xs",
    );
    c.queue
        .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&pre[..b * cols]));
    let y_size = (b * rows * 4) as u64;
    let y_buf = Scratch::ensure(
        &c.device,
        &mut sc.y,
        y_size,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "q1mm-y",
    );
    let p_buf = uniform_u32x4(c, [(cols / 4) as u32, rows as u32, b as u32, 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q1mm-bg"),
        layout: &c.layout_q1mm, // q1_mul_mm omits binding 2 (no row-scale)
        entries: &[
            bind_buf(0, &w),
            bind_buf(1, &xs_buf),
            bind_buf(3, &y_buf),
            bind_buf(4, &p_buf),
        ],
    });
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        y_size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "q1mm-stage",
    );
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("q1mm"),
        });
    {
        let mut pass = begin_pass_with(&mut enc, Some("q1mm"), None);
        pass.set_pipeline(&c.q1_mm);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(
            (rows as u32).div_ceil(64).min(MAX_WG),
            (b as u32).div_ceil(64),
            1,
        );
    }
    let ok = readback(c, enc, &y_buf, &stage_buf, y_size, &mut out[..b * rows]);
    drop(sc);
    ok
}

/// matmat kernel: resident weights + rs + batch of inputs, 2D dispatch, readback.
#[allow(clippy::too_many_arguments)]
fn dispatch_matmat(
    c: &Ctx,
    weight_key: Option<(usize, usize)>,
    full_quant: &[u8],
    row_scale: &[f32],
    pre: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    dispatch_matmat_keep(
        c,
        weight_key,
        full_quant,
        row_scale,
        None,
        pre,
        b,
        rows,
        cols,
        Some(out),
        None,
    )
    .is_some()
}

/// The same int8 GEMM, with the result LEFT on the card when `out` is None.
///
/// This is what lets the two-field codec ride the fused DiT paths: they hand
/// one kernel's output straight to the next without a readback, and a
/// readback drains the queue. `q4tp` had such a variant from the start
/// (`tp_matmat_keep`) and `q8_2f` did not, which is the whole reason the
/// eight-bit build went through per-op GEMMs while the four-bit one fused.
#[allow(clippy::too_many_arguments)]
fn dispatch_matmat_keep(
    c: &Ctx,
    weight_key: Option<(usize, usize)>,
    full_quant: &[u8],
    row_scale: &[f32],
    // The two-field codec's column field, when the weight has one. On the
    // matrix-unit arm it is folded into the weight plane, which is where it
    // belongs: the activation is then untouched, and a 40 MB host multiply
    // per projection — 400 of them in a render — disappears.
    col_scale: Option<&[f32]>,
    pre: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    mut out: Option<&mut [f32]>,
    src: Option<&wgpu::Buffer>,
) -> Option<wgpu::Buffer> {
    if full_quant.len() < rows * cols
        || row_scale.len() < rows
        || (src.is_none() && pre.len() < b * cols)
        || out.as_ref().is_some_and(|o| o.len() < b * rows)
    {
        return None;
    }
    let q_buf = match weight_key {
        Some(k) => match weight_buffer(c, k, full_quant) {
            Some(b) => b,
            None => return None, // over VRAM budget — honest CPU path
        },
        None => c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("mm-weights"),
                contents: full_quant,
                usage: wgpu::BufferUsages::STORAGE,
            }),
    };
    // rs cached per tensor (row0 sentinel = full-tensor scales).
    let rs_buf = match weight_key {
        Some((base, idx)) => c
            .rs_bufs
            .lock()
            .unwrap()
            .entry((base, (idx, usize::MAX)))
            .or_insert_with(|| {
                crate::gpu::probe_note_cold();
                c.device
                    .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: Some("mm-rs"),
                        contents: bytemuck::cast_slice(&row_scale[..rows]),
                        usage: wgpu::BufferUsages::STORAGE,
                    })
            })
            .clone(),
        None => c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("mm-rs"),
                contents: bytemuck::cast_slice(&row_scale[..rows]),
                usage: wgpu::BufferUsages::STORAGE,
            }),
    };
    // Pooled scratch for the whole op (encode → submit → poll).
    let mut sc = c.scratch.lock().unwrap();
    let xs_buf = match src {
        // The operand is already on the card: the kernel before us left it
        // there, and taking delivery just to hand it back is the round trip
        // this whole path exists to remove.
        Some(bf) => bf.clone(),
        None => Scratch::ensure(
            &c.device,
            &mut sc.xs,
            (b * cols * 4) as u64,
            wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            "mm-xs",
        ),
    };
    // Only a host operand is uploaded. Writing here unconditionally would
    // overwrite the resident panel the previous kernel just produced — and
    // with `pre` empty, as the device-side callers pass it, index past the
    // end of an empty slice.
    if src.is_none() {
        c.queue
            .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&pre[..b * cols]));
    }
    let y_size = (b * rows * 4) as u64;
    let y_buf = Scratch::ensure(
        &c.device,
        &mut sc.y,
        y_size,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "mm-y",
    );
    let params = [(cols / 4) as u32, rows as u32, b as u32, 0u32];
    let p_buf = match &sc.params {
        Some(bf) => bf.clone(),
        None => {
            let bf = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("mm-params"),
                size: 16,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            sc.params = Some(bf.clone());
            bf
        }
    };
    c.queue
        .write_buffer(&p_buf, 0, bytemuck::cast_slice(&params));
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        y_size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "mm-stage",
    );
    // The matrix units, for an int8 weight. The four-bit path unpacks its
    // weight into an f16 plane and hands that to the cooperative GEMM; int8
    // had no such arm and ran the scalar `mul_mm` instead, which is most of
    // what separated the two codecs on the clock. The plane is the same, the
    // GEMM after it is the same — only the unpacker differs, and int8's is
    // one multiply. Worth its pass only when the batch amortizes it, hence
    // the same b >= 64 gate the four-bit arm uses.
    let coop = if b >= 64 && cols % 2 == 0 && !std::env::var("CMF_Q8_COOP").is_ok_and(|v| v == "0")
    {
        c.q4tp_mm_coop_f16
            .as_ref()
            .and_then(|_| dq8_f16_plane(c, &mut sc, &q_buf, &rs_buf, col_scale, rows, cols))
    } else {
        None
    };
    // A resident operand never passed through the host, so max|x| has to be
    // taken on the card; without that reduction the f16 operands overflow.
    let can_dev_scale =
        (c.act_amax_part.is_some() && c.act_amax_fold.is_some()) || c.act_absmax.is_some();
    let coop = if src.is_some() && !can_dev_scale {
        None
    } else {
        coop
    };
    if let (Some((plane, bind_dq)), Some(mm_pipe)) = (&coop, c.q4tp_mm_coop_f16.as_ref()) {
        if std::env::var("CMF_GPU_DEBUG").is_ok() {
            static ONCE: std::sync::Once = std::sync::Once::new();
            ONCE.call_once(|| eprintln!("int8 GEMM: matrix units (plane {rows}x{cols})"));
        }
        let dev_scale = src.is_some();
        let ascale: f32 = if dev_scale {
            0.0
        } else {
            let mx =
                pre[..b * cols]
                    .iter()
                    .fold(0f32, |m, v| if v.is_finite() { m.max(v.abs()) } else { m });
            if mx > 1000.0 {
                1000.0 / mx
            } else {
                1.0
            }
        };
        let cp = [
            (cols / 4) as u32,
            rows as u32,
            b as u32,
            if dev_scale {
                0xFFFF_FFFFu32
            } else {
                ascale.to_bits()
            },
        ];
        let cp_buf = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("q8coop-params"),
                contents: bytemuck::cast_slice(&cp),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let asc_buf = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("q8coop-ascale"),
            size: 4,
            usage: wgpu::BufferUsages::STORAGE,
            mapped_at_creation: false,
        });
        // The partials buffer for the device-side max|x| is taken HERE, under
        // the guard this function already holds — `encode_act_absmax` would
        // otherwise take that same lock a second time, and std's Mutex is not
        // reentrant. That deadlock looks from outside like a card at 0% with
        // the process idling, which is exactly how it presented.
        let amax_parts = Scratch::ensure(
            &c.device,
            &mut sc.amaxp,
            512 * 4,
            wgpu::BufferUsages::STORAGE,
            "q8coop-amax-parts",
        );
        let bind_mm = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("q8coop-bg"),
            layout: &mm_pipe.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, plane),
                bind_buf(1, &xs_buf),
                bind_buf(2, &y_buf),
                bind_buf(3, &cp_buf),
                bind_buf(4, &asc_buf),
            ],
        });
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("q8coop"),
            });
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(c.q8_dq_f16.as_ref().unwrap());
            pass.set_bind_group(0, bind_dq, &[]);
            let wgs = ((rows * cols / 2) as u32).div_ceil(256);
            pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
        }
        if dev_scale {
            encode_act_absmax_with(c, &mut enc, &xs_buf, b * cols, &asc_buf, Some(&amax_parts));
        }
        {
            let mut pass = begin_pass_with(&mut enc, Some("q8coop"), None);
            pass.set_pipeline(mm_pipe);
            pass.set_bind_group(0, &bind_mm, &[]);
            pass.dispatch_workgroups(
                (rows as u32).div_ceil(64).min(MAX_WG),
                (b as u32).div_ceil(64),
                1,
            );
        }
        return match out.as_mut() {
            Some(o) => {
                let ok = readback(c, enc, &y_buf, &stage_buf, y_size, &mut o[..b * rows]);
                drop(sc);
                ok.then_some(y_buf)
            }
            None => {
                c.queue.submit(Some(enc.finish()));
                drop(sc);
                Some(y_buf)
            }
        };
    }
    // The scalar arm has no plane to fold into, so the field meets the
    // activation the old way — on the host for an operand that came from
    // there, through the colscale kernel for one already on the card.
    let scaled_xs;
    let xs_buf = match col_scale {
        Some(cf) if src.is_none() => {
            let pres: Vec<f32> = pre[..b * cols]
                .chunks_exact(cols)
                .flat_map(|r| r.iter().zip(cf).map(|(a, c)| a * c))
                .collect();
            c.queue
                .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&pres));
            xs_buf
        }
        Some(cf) => match colscale_keep(c, &xs_buf, cf, b * cols) {
            Some(bf) => {
                scaled_xs = bf;
                scaled_xs
            }
            None => return None,
        },
        None => xs_buf,
    };
    let use_mm = b >= 32;
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("mm-bg"),
        // Auto bind-group layouts are pipeline-exclusive in wgpu — pick
        // the layout of the pipeline this dispatch actually uses.
        layout: if use_mm { &c.layout_mmm } else { &c.layout_mm },
        entries: &[
            bind_buf(0, &q_buf),
            bind_buf(1, &xs_buf),
            bind_buf(2, &rs_buf),
            bind_buf(3, &y_buf),
            bind_buf(4, &p_buf),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("mm") });
    {
        let mut pass = begin_pass_with(&mut enc, Some("mm"), None);
        if use_mm {
            pass.set_pipeline(&c.mul_mm);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(
                (rows as u32).div_ceil(64).min(MAX_WG),
                (b as u32).div_ceil(64),
                1,
            );
        } else {
            pass.set_pipeline(&c.matmat);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups((rows as u32).min(MAX_WG), b as u32, 1);
        }
    }
    match out.as_mut() {
        Some(o) => {
            let ok = readback(c, enc, &y_buf, &stage_buf, y_size, &mut o[..b * rows]);
            drop(sc);
            ok.then_some(y_buf)
        }
        None => {
            // Nothing to bring home: submit and hand the buffer on.
            c.queue.submit(Some(enc.finish()));
            drop(sc);
            Some(y_buf)
        }
    }
}

/// q1t batched GEMM (prefill) on wgpu — register-blocked base GEMM then the
/// sparse overlay, two passes in one encoder. Raw f32 x, scales in the tiles.
/// Batched q4t GEMM (imagegen DiT prefill shapes) — the wgpu twin of
/// the Metal q4t_matmat: one q4t_mul_mm dispatch reading the 18-byte
/// tiles from the cached weight buffer. The CPU/GPU probe arbitrates
/// per process exactly as on Metal.
/// q4tp twin of `q4t_matmat` — the batched GEMM the wide-batch arm of
/// `QTensor::matmat` reaches for (DiT prefill, MoE experts, dense FFN
/// batches). Without it a q4tp model kept that arm on the CPU while q4t
/// went to the device.
/// Per-phase microseconds of the DiT attention under
/// `CMF_DIT_ATTN_PROF=1`: 0 = QK, 1 = softmax, 2 = PV.
pub static DIT_PHASE: [std::sync::atomic::AtomicU64; 3] = [
    std::sync::atomic::AtomicU64::new(0),
    std::sync::atomic::AtomicU64::new(0),
    std::sync::atomic::AtomicU64::new(0),
];

/// The three walls, or None when the mode is off.
pub fn dit_phase_report() -> Option<String> {
    if std::env::var("CMF_DIT_ATTN_PROF").is_err() {
        return None;
    }
    let v: Vec<f64> = DIT_PHASE
        .iter()
        .map(|a| a.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6)
        .collect();
    Some(format!(
        "dit attention: qk {:.2} s · softmax {:.2} s · pv {:.2} s",
        v[0], v[1], v[2]
    ))
}

/// Is this tensor's weight buffer already on the card? The probe asks
/// before it decides which arm a still-cold call should take.
pub fn weight_is_resident(model: &Arc<CmfModel>, idx: usize) -> bool {
    let Some(c) = ctx() else { return false };
    c.weight_bufs
        .lock()
        .unwrap()
        .contains_key(&(model.uid() as usize, idx))
}

/// q4tp GEMM whose result STAYS on the device: the caller gets the
/// buffer, not a host copy. The DiT computes qkv on the card, reads
/// 160 MB back per block, and the attention split uploads the same
/// bytes again — 320 MB a block of pure round trip.
///
/// Deliberately NOT on the backend facade: a `wgpu::Buffer` cannot cross
/// it, and a facade that returns `Option<()>` (as the first version did)
/// throws away the only thing the caller wanted. The user of this is the
/// attention path in this same module, once qk-norm and RoPE move to the
/// device — until then the panel has to come home anyway.
pub(crate) fn q4tp_matmat_dev(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
) -> Option<wgpu::Buffer> {
    let c = ctx()?;
    let mut out = Vec::new();
    // Reuse the whole validated path; `keep` swaps the readback for a
    // handle to the same buffer.
    let buf = tp_matmat_keep(model, idx, xs, b, rows, cols, &mut out, false)?;
    let _ = c;
    Some(buf)
}

/// Qwen's resident chain uses independent panels for the two streams and
/// therefore cannot use the historical shared `Scratch::y` result slot. The
/// GEMM implementation already has the complete codec/scale validation; this
/// narrow wrapper only supplies a caller-owned storage destination and keeps
/// the result on the device.
fn qwen_q4tp_gemm_keep(
    model: &Arc<CmfModel>,
    idx: usize,
    src: &wgpu::Buffer,
    dst: &wgpu::Buffer,
    batch: usize,
    rows: usize,
    cols: usize,
) -> bool {
    qwen_q4tp_gemm_keep_offset(model, idx, src, 0, dst, batch, rows, cols)
}

/// The output projection can bind the image or text window directly from
/// the joint token-major attention panel.  Keeping the byte offset here
/// avoids a device-side split copy before the two distinct output weights.
fn qwen_q4tp_gemm_keep_offset(
    model: &Arc<CmfModel>,
    idx: usize,
    src: &wgpu::Buffer,
    src_offset: u64,
    dst: &wgpu::Buffer,
    batch: usize,
    rows: usize,
    cols: usize,
) -> bool {
    tp_matmat_impl(
        model,
        idx,
        &[],
        batch,
        rows,
        cols,
        None,
        Some(src),
        src_offset,
        false,
        Some(dst),
    )
    .is_some()
}

/// Encode one Q4TP projection into a caller-owned command buffer.  The
/// resident plane cache is keyed by model/tensor identity and capped by the
/// existing VRAM policy; a miss falls back to the established in-kernel
/// Q4TP GEMM.  No submission or host readback occurs here.
#[allow(clippy::too_many_arguments)]
fn qwen_q4tp_gemm_encode(
    c: &Ctx,
    model: &Arc<CmfModel>,
    idx: usize,
    src: &wgpu::Buffer,
    src_offset: u64,
    dst: &wgpu::Buffer,
    batch: usize,
    rows: usize,
    cols: usize,
    enc: &mut wgpu::CommandEncoder,
) -> bool {
    if batch == 0 || rows == 0 || cols == 0 || cols % 32 != 0 || rows % 32 != 0 {
        return false;
    }
    let Some(src_bytes) = batch
        .checked_mul(cols)
        .and_then(|n| n.checked_mul(4))
        .and_then(|n| u64::try_from(n).ok())
    else {
        return false;
    };
    if src_offset % 256 != 0
        || src_offset
            .checked_add(src_bytes)
            .is_none_or(|end| end > src.size())
        || !qwen_storage_binding_fit(c, src_bytes as usize, "chain activation")
    {
        return false;
    }
    let Some(dst_bytes) = batch
        .checked_mul(rows)
        .and_then(|n| n.checked_mul(4))
        .and_then(|n| u64::try_from(n).ok())
    else {
        return false;
    };
    if dst_bytes > dst.size() || !qwen_storage_binding_fit(c, dst_bytes as usize, "chain output") {
        return false;
    }
    let Some(entry) = model.tensors.get(idx) else {
        return false;
    };
    if entry.dtype != cortiq_core::TensorDtype::Q4TiledP
        || entry.shape.as_slice() != [rows, cols]
    {
        return false;
    }
    let Some(q_buf) = qwen_q4tp_weight(c, model, idx, rows, cols) else {
        return false;
    };

    let can_scale = (c.act_amax_part.is_some() && c.act_amax_fold.is_some())
        || c.act_absmax.is_some();
    let coop = c.q4tp_mm_coop_f16.is_some() && c.q4tp_dq_f16.is_some() && can_scale;
    if coop {
        if let Some((plane, fresh)) = plane_cached(
            c,
            (model.uid() as usize, idx),
            &q_buf,
            rows,
            cols,
            8192,
        ) {
            if let (Some(dq), Some(mm)) = (c.q4tp_dq_f16.as_ref(), c.q4tp_mm_coop_f16.as_ref()) {
                if let Some(bind_dq) = fresh {
                    let mut pass = begin_pass_with(enc, Some("qwen-chain-dequant"), None);
                    pass.set_pipeline(dq);
                    pass.set_bind_group(0, &bind_dq, &[]);
                    let pairs = (rows * cols / 2) as u32;
                    let groups = pairs.div_ceil(256);
                    pass.dispatch_workgroups(groups.min(MAX_WG), groups.div_ceil(MAX_WG), 1);
                }
                let asc = c.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("qwen-chain-ascale"),
                    size: 4,
                    usage: wgpu::BufferUsages::STORAGE,
                    mapped_at_creation: false,
                });
                if encode_act_absmax_offset(c, enc, src, src_offset, batch * cols, &asc) {
                    encode_q4_tile_mm_full_offset(
                        c,
                        enc,
                        mm,
                        &plane,
                        src,
                        src_offset,
                        dst,
                        rows,
                        cols,
                        batch,
                        0.0,
                        Some(&asc),
                    );
                    qwen_chain_debug_path(model.uid(), idx, rows, cols, batch, "cached-f16");
                    return true;
                }
            }
        }
    }
    // Once the bounded f16 plane bank is full, keep using the existing
    // cooperative Q4TP kernel instead of silently dropping to the scalar
    // tile decoder. It dequantizes in the GEMM's tile loop, so it does not
    // retain another model-wide plane, while the matrix units still handle
    // the multiply. The device scale is the same one used by the cached-f16
    // arm, and the branch remains opt-out for isolating a fallback run.
    if c.q4tp_mm_coop_s.is_some()
        && can_scale
        && std::env::var("CMF_QWEN_IMAGE_FUSED_MLP_COOP").as_deref() != Ok("0")
    {
        let asc = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("qwen-chain-ascale-direct"),
            size: 4,
            usage: wgpu::BufferUsages::STORAGE,
            mapped_at_creation: false,
        });
        if encode_act_absmax_offset(c, enc, src, src_offset, batch * cols, &asc) {
            encode_q4_tile_mm_full_offset(
                c,
                enc,
                c.q4tp_mm_coop_s.as_ref().unwrap(),
                &q_buf,
                src,
                src_offset,
                dst,
                rows,
                cols,
                batch,
                0.0,
                Some(&asc),
            );
            qwen_chain_debug_path(model.uid(), idx, rows, cols, batch, "direct-coop");
            return true;
        }
    }
    // The scalar shader is the exact portable GPU fallback. It uses the
    // same source window and leaves every caller's CPU/Metal fallback intact.
    qwen_chain_debug_path(model.uid(), idx, rows, cols, batch, "scalar");
    encode_q4_tile_mm_full_offset(
        c,
        enc,
        &c.q4tp_mm,
        &q_buf,
        src,
        src_offset,
        dst,
        rows,
        cols,
        batch,
        0.0,
        None,
    );
    true
}

/// Emit one bounded diagnostic per Qwen projection/path. The full model has
/// thousands of block invocations, so logging every call would perturb the
/// timing and swamp the useful evidence. `CMF_GPU_DEBUG=1` enables this
/// summary; normal inference does no set allocation or formatting.
fn qwen_chain_debug_path(
    uid: u64,
    idx: usize,
    rows: usize,
    cols: usize,
    batch: usize,
    path: &'static str,
) {
    if std::env::var("CMF_GPU_DEBUG").is_err() {
        return;
    }
    use std::collections::HashSet;
    use std::sync::{Mutex, OnceLock};
    static SEEN: OnceLock<Mutex<HashSet<(u64, usize, &'static str)>>> = OnceLock::new();
    let seen = SEEN.get_or_init(|| Mutex::new(HashSet::new()));
    if seen.lock().unwrap().insert((uid, idx, path)) {
        eprintln!(
            "qwen-chain GEMM path={path} tensor={idx} shape={rows}x{cols} batch={batch} uid={uid}"
        );
    }
}

/// Music-3's FFN with the activations held on the device: ff_in's
/// GEMM keeps its result, the GLU runs where that result lives, and
/// ff_out reads it in place — only `h` goes up (5.6 MB) and only the
/// block's output comes back (5.6 MB). The host arm of the same chain
/// moved 68 MB more per block-step, and on this class of stand the
/// split timer put transfers at 82% of the device arm's total; the
/// GEMMs themselves run at ~5 TFLOP/s and were never the cost.
///
/// `false` = refused (no ctx, no pipeline, or a GEMM declined) — the
/// caller's host arm is the fallback and produces identical audio.
#[allow(clippy::too_many_arguments)]
pub fn music3_ffn(
    model: &Arc<CmfModel>,
    idx_in: usize,
    idx_out: usize,
    h: &[f32],
    bias_in: &[f32],
    n: usize,
    hs: usize,
    inter: usize,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    let Some(glu) = c.music3_glu.as_ref() else {
        return false;
    };
    if h.len() < n * hs || bias_in.len() < 2 * inter || out.len() < n * hs || n == 0 {
        return false;
    }
    // ff_in: [n, 2·inter] stays on the card.
    let mut unused = Vec::new();
    let Some(gu) = tp_matmat_keep(model, idx_in, h, n, 2 * inter, hs, &mut unused, false) else {
        return false;
    };
    // GLU in place. The bias buffer is cached by pointer+fingerprint,
    // so it crosses the bus once per process, not once per call.
    let bias = bake_weight(c, &bias_in[..2 * inter], "m3-glu-b");
    let act = {
        let mut sc = c.scratch.lock().unwrap();
        Scratch::ensure(
            &c.device,
            &mut sc.m3act,
            (n * inter * 4) as u64,
            wgpu::BufferUsages::STORAGE,
            "m3-act",
        )
    };
    let u = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[n as u32, inter as u32, 0u32, 0u32]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &glu.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, &gu),
            bind_buf(1, &act),
            bind_buf(2, &bias),
            bind_buf(3, &u),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("m3-glu"),
        });
    {
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(glu);
        pass.set_bind_group(0, &bind, &[]);
        let groups = ((n * inter) as u32).div_ceil(256);
        pass.dispatch_workgroups(groups.min(65_535), groups.div_ceil(65_535), 1);
    }
    submit(c, finish_enc(enc));
    // ff_out reads the resident activations; the queue orders the three
    // submissions, so nothing here waits until the final readback.
    tp_matmat_impl(
        model,
        idx_out,
        &[],
        n,
        hs,
        inter,
        Some(out),
        Some(&act),
        0,
        false,
        None,
    )
    .is_some()
}

pub fn q4tp_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    tp_matmat(model, idx, xs, b, rows, cols, out, false)
}

/// The same, over a two-bit weight plane.
pub fn q2tp_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    tp_matmat(model, idx, xs, b, rows, cols, out, true)
}

/// Batched q2tp GEMM with the explicit descriptor center correction.
pub fn q2tp_affine_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    if !crate::prism::is_affine_target(model, &model.tensors[idx].name) {
        return false;
    }
    tp_matmat(model, idx, xs, b, rows, cols, out, true)
}

/// Single-token q2tp matvec through the descriptor-aware WGSL kernel. The
/// payload is the ordinary dtype16 plane in both modes; `affine` only selects
/// the explicit center correction required by the model descriptor.
pub fn q2tp_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    q2tp_matvec_impl(model, idx, xs, rows, cols, out, false)
}

/// The q2tp_affine twin of [`q2tp_matvec`]. It refuses a tensor that is not
/// explicitly listed by the affine descriptor, so a caller cannot
/// accidentally apply the center shift to an ordinary q2tp payload.
pub fn q2tp_affine_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    q2tp_matvec_impl(model, idx, xs, rows, cols, out, true)
}

fn q2tp_matvec_impl(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
    affine: bool,
) -> bool {
    let Some(c) = ctx() else { return false };
    if cols % 32 != 0 || rows == 0 || xs.len() < cols || out.len() < rows {
        return false;
    }
    let entry = &model.tensors[idx];
    if entry.dtype != cortiq_core::TensorDtype::Q2TiledP
        || entry.shape.first().copied().unwrap_or(0) < rows
        || affine != crate::prism::is_affine_target(model, &entry.name)
    {
        return false;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return false;
    };
    let bytes = model.primary_bytes();
    let plen = entry.nbytes as usize;
    let Some(need) =
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q2TiledP, &[rows, cols])
    else {
        return false;
    };
    if plen < need || abs + plen > bytes.len() {
        return false;
    }
    if std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1") {
        use std::sync::atomic::{AtomicUsize, Ordering};
        static N: AtomicUsize = AtomicUsize::new(0);
        let n = N.fetch_add(1, Ordering::Relaxed);
        if n < 512 {
            eprintln!(
                "q2tp-gpu matvec #{n} name={} shape={}x{} affine={affine}",
                entry.name, rows, cols
            );
        }
    }
    let Some(q_buf) = weight_buffer_l(
        c,
        (model.uid() as usize, idx),
        &bytes[abs..abs + plen],
        layer_of_name(&entry.name),
    ) else {
        return false;
    };
    // The ladder cache is a run-owned component experiment. It is built
    // only for the explicit affine target and only when both optional
    // pipelines were admitted; every failure falls back to the validated
    // row-local shader below.
    let ladder_cache = if affine
        && std::env::var("CMF_Q2_LADDER_CACHE").as_deref() == Ok("1")
    {
        ensure_q2_ladder_cache(c, model, idx, rows, cols)
    } else {
        None
    };
    let mut sc = c.scratch.lock().unwrap();
    let xs_buf = Scratch::ensure(
        &c.device,
        &mut sc.xs,
        (cols * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "q2tp-xs",
    );
    c.queue
        .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&xs[..cols]));
    let y_buf = Scratch::ensure(
        &c.device,
        &mut sc.y,
        (rows * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "q2tp-y",
    );
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        (rows * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "q2tp-stage",
    );
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("q2tp-mv"),
        });
    if let Some((ladders, row_ids, row_id_base)) = ladder_cache {
        encode_q2tp_ladder_cache(
            c,
            &mut enc,
            &q_buf,
            &xs_buf,
            &y_buf,
            &ladders,
            &row_ids,
            row_id_base,
            rows,
            cols,
            affine,
        );
    } else if q2tp_dp4a_on() && affine {
        encode_q2tp_mv1_i8(c, &mut enc, &q_buf, &xs_buf, &y_buf, rows, cols);
    } else {
        encode_q2tp_mv16w(c, &mut enc, &q_buf, &xs_buf, &y_buf, rows, cols, affine);
    }
    let ok = readback(
        c,
        enc,
        &y_buf,
        &stage_buf,
        (rows * 4) as u64,
        &mut out[..rows],
    );
    ok
}

/// `q4tp` and `q2tp` differ in their weight plane and their ladder, and
/// in nothing this function does: same buffers, same bind group, same
/// dispatch. Only the pipeline and the expected byte count follow the
/// width.
#[allow(clippy::too_many_arguments)]
/// Dequantize a q4tp plane into f16 in a REUSED scratch buffer, once
/// per GEMM call. Caching whole planes was the obvious idea and the
/// wrong one: this model's DiT would want 38 GB of them. The unpack
/// pass costs a fraction of a millisecond; what it buys is a GEMM whose
/// matrix units are not waiting on a nibble unpacker that re-runs for
/// every 64-row tile of activations.
/// THE NAME LIES, AND IT IS THE BIGGEST LEVER LEFT ON EVERY DiT: this
/// "dequantize once" plane is ONE SCRATCH SLOT shared by every weight,
/// so each GEMM overwrites the previous one's plane and the unpack has
/// to be re-encoded before every call — every weight, every block,
/// every step. Nothing is cached but the allocation.
///
/// That is why the f16 arm loses on Lumina (11.23 s against 11.02 for
/// the scalar arm, 2304×2304 weights) and wins on the video DiT only
/// because its 21504×5376 GEMM is large enough to swallow the unpack
/// it repeats. Both would be faster with a real cache.
///
/// The fix is a plane cache keyed by (model uid, tensor index) — the
/// caller has that identity, this function does not, so it belongs one
/// level up. Bound it by VRAM: an f16 plane is 2× the q4 bytes, so
/// Lumina's DiT is ~6 GB of planes (fits a 32 GB card and pays for
/// itself over 30 steps) while MiniMax's 25.7 GB does not — the cache
/// has to evict, and the honest policy is largest-first by call count.
/// An f16 plane per weight, kept for the life of the context. Returns
/// the plane and — only when it was just created — the bind group its
/// unpack pass needs, so a warm plane costs the caller one hash lookup
/// and no dispatch at all.
///
/// Bounded, because a plane is 2× the q4 bytes: `CMF_PLANE_CACHE_MB`
/// (default 8192) is the ceiling, and past it this refuses rather than
/// evicting — a DiT that does not fit keeps the old behaviour instead
/// of thrashing.
fn plane_cached(
    c: &Ctx,
    key: (usize, usize),
    q_buf: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    default_cap_mb: u64,
) -> Option<(wgpu::Buffer, Option<wgpu::BindGroup>)> {
    if cols % 2 != 0 {
        return None;
    }
    let pipe = c.q4tp_dq_f16.as_ref()?;
    let bytes = (rows * cols * 2) as u64;
    let cap = std::env::var("CMF_PLANE_CACHE_MB")
        .ok()
        .and_then(|v| v.parse::<u64>().ok())
        .unwrap_or(default_cap_mb)
        * 1024
        * 1024;
    let mut m = c.planes.lock().unwrap();
    if let Some((p, _)) = m.get(&key) {
        return Some((p.clone(), None));
    }
    // Half the cap is reserved for small planes. The DiT runs before the 3D
    // VAE and its weights are twice the size, so first-come-first-served gave
    // the whole cache to the DiT and the VAE — where the same tensor is
    // unpacked dozens of times per render rather than four — kept missing.
    // Measured at 768x448: caching the DiT's alone took the denoise from
    // 72.9 s to 60.6 s while the VAE's FFN phase stayed at 24 s.
    const BIG: u64 = 96 * 1024 * 1024;
    let used: u64 = m.values().map(|(_, b)| *b).sum();
    let used_big: u64 = m.values().map(|(_, b)| *b).filter(|b| *b > BIG).sum();
    if used + bytes > cap || (bytes > BIG && used_big + bytes > cap / 2) {
        return None;
    }
    let plane = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("q4tp-plane-cached"),
        size: bytes,
        usage: wgpu::BufferUsages::STORAGE,
        mapped_at_creation: false,
    });
    let u = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[cols as u32, rows as u32, 0u32, 0u32]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &pipe.get_bind_group_layout(0),
        entries: &[bind_buf(0, q_buf), bind_buf(1, &plane), bind_buf(2, &u)],
    });
    m.insert(key, (plane.clone(), bytes));
    Some((plane, Some(bind)))
}

fn dq_f16_plane(
    c: &Ctx,
    sc: &mut Scratch,
    q_buf: &wgpu::Buffer,
    rows: usize,
    cols: usize,
) -> Option<(wgpu::Buffer, wgpu::BindGroup)> {
    dq_f16_plane_slot(c, sc, q_buf, rows, cols, false)
}

fn dq_f16_plane_slot(
    c: &Ctx,
    sc: &mut Scratch,
    q_buf: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    second: bool,
) -> Option<(wgpu::Buffer, wgpu::BindGroup)> {
    if cols % 2 != 0 {
        return None;
    }
    let pipe = c.q4tp_dq_f16.as_ref()?;
    let bytes = (rows * cols * 2) as u64;
    // The caller already holds the scratch guard; std's Mutex is not
    // reentrant, and taking it again here hung the device at 0%.
    let slot = if second { &mut sc.dqw2 } else { &mut sc.dqw };
    let plane = Scratch::ensure(
        &c.device,
        slot,
        bytes,
        wgpu::BufferUsages::STORAGE,
        "q4tp-dq-plane",
    );
    let u = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[cols as u32, rows as u32, 0u32, 0u32]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &pipe.get_bind_group_layout(0),
        entries: &[bind_buf(0, q_buf), bind_buf(1, &plane), bind_buf(2, &u)],
    });
    Some((plane, bind))
}

/// The int8 payload unpacked into the packed-f16 plane the coop GEMM reads.
/// Same slot and same shape as the four-bit `dq_f16_plane`; only the unpacker
/// differs, because there is nothing to unpack — one multiply per weight.
fn dq8_f16_plane(
    c: &Ctx,
    sc: &mut Scratch,
    q_buf: &wgpu::Buffer,
    rs_buf: &wgpu::Buffer,
    col: Option<&[f32]>,
    rows: usize,
    cols: usize,
) -> Option<(wgpu::Buffer, wgpu::BindGroup)> {
    if cols % 2 != 0 {
        return None;
    }
    let pipe = c.q8_dq_f16.as_ref()?;
    let plane = Scratch::ensure(
        &c.device,
        &mut sc.dqw,
        (rows * cols * 2) as u64,
        wgpu::BufferUsages::STORAGE,
        "q8-dq-plane",
    );
    let u = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[
                cols as u32,
                rows as u32,
                u32::from(col.is_some()),
                0u32,
            ]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let cbuf = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("q8-col"),
            contents: bytemuck::cast_slice(col.unwrap_or(&[1.0f32])),
            usage: wgpu::BufferUsages::STORAGE,
        });
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &pipe.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, q_buf),
            bind_buf(1, &plane),
            bind_buf(2, &u),
            bind_buf(3, rs_buf),
            bind_buf(4, &cbuf),
        ],
    });
    Some((plane, bind))
}

/// max|x| over an activation panel that lives on the card, folded into
/// the reciprocal scale the f16 GEMM reads from binding 4. Two stages
/// when both halves exist: one workgroup cannot saturate a card's
/// bandwidth, and this panel is hundreds of megabytes.
fn encode_act_absmax(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    act: &wgpu::Buffer,
    n: usize,
    asc: &wgpu::Buffer,
) -> bool {
    encode_act_absmax_with_offset(c, enc, act, 0, n, asc, None)
}

/// `encode_act_absmax` for a resident window inside a larger token panel.
/// The reduction sees exactly `n` f32 values beginning at `offset`; this is
/// used by Qwen's image/text output projections after joint attention.
fn encode_act_absmax_offset(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    act: &wgpu::Buffer,
    offset: u64,
    n: usize,
    asc: &wgpu::Buffer,
) -> bool {
    encode_act_absmax_with_offset(c, enc, act, offset, n, asc, None)
}

/// The same, for a caller that already holds the scratch guard: it passes the
/// partials buffer in rather than making this take the lock a second time.
/// std's Mutex is not reentrant, and that deadlock looks from outside like a
/// card sitting at 0% with the process idle.
fn encode_act_absmax_with(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    act: &wgpu::Buffer,
    n: usize,
    asc: &wgpu::Buffer,
    parts: Option<&wgpu::Buffer>,
) -> bool {
    encode_act_absmax_with_offset(c, enc, act, 0, n, asc, parts)
}

/// Internal activation reduction with an optional source window.  Keeping
/// the existing zero-offset wrapper avoids changing the many established
/// GEMM callers while the Qwen chain can bind each stream without a copy.
fn encode_act_absmax_with_offset(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    act: &wgpu::Buffer,
    act_offset: u64,
    n: usize,
    asc: &wgpu::Buffer,
    parts: Option<&wgpu::Buffer>,
) -> bool {
    let act_bytes = (n as u64).saturating_mul(4);
    let act_entry = if act_offset == 0 {
        bind_buf(0, act)
    } else {
        bind_buf_off(0, act, act_offset, act_bytes)
    };
    let ap = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[n as u32, 0u32, 0u32, 0u32]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    match (c.act_amax_part.as_ref(), c.act_amax_fold.as_ref()) {
        (Some(part), Some(fold)) => {
            let nparts = 512usize;
            let pbuf = match parts {
                Some(b) => b.clone(),
                None => {
                    let mut sc = c.scratch.lock().unwrap();
                    Scratch::ensure(
                        &c.device,
                        &mut sc.amaxp,
                        (nparts * 4) as u64,
                        wgpu::BufferUsages::STORAGE,
                        "q4tpmm-amax-parts",
                    )
                }
            };
            let bg1 = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &part.get_bind_group_layout(0),
                entries: &[act_entry, bind_buf(1, &pbuf), bind_buf(2, &ap)],
            });
            let fp = c
                .device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytemuck::cast_slice(&[nparts as u32, 0u32, 0u32, 0u32]),
                    usage: wgpu::BufferUsages::UNIFORM,
                });
            let bg2 = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &fold.get_bind_group_layout(0),
                entries: &[bind_buf(0, &pbuf), bind_buf(1, asc), bind_buf(2, &fp)],
            });
            let mut pass = begin_pass(enc);
            pass.set_pipeline(part);
            pass.set_bind_group(0, &bg1, &[]);
            pass.dispatch_workgroups(nparts as u32, 1, 1);
            drop(pass);
            let mut pass = begin_pass(enc);
            pass.set_pipeline(fold);
            pass.set_bind_group(0, &bg2, &[]);
            pass.dispatch_workgroups(1, 1, 1);
            true
        }
        _ => match c.act_absmax.as_ref() {
            Some(amax) => {
                let bg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &amax.get_bind_group_layout(0),
                    entries: &[act_entry, bind_buf(1, asc), bind_buf(2, &ap)],
                });
                let mut pass = begin_pass(enc);
                pass.set_pipeline(amax);
                pass.set_bind_group(0, &bg, &[]);
                pass.dispatch_workgroups(1, 1, 1);
                true
            }
            None => false,
        },
    }
}

fn tp_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
    two_bit: bool,
) -> bool {
    tp_matmat_impl(model, idx, xs, b, rows, cols, Some(out), None, 0, two_bit, None).is_some()
}

/// Result stays on the device; the caller owns the returned handle.
fn tp_matmat_keep(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    _unused: &mut Vec<f32>,
    two_bit: bool,
) -> Option<wgpu::Buffer> {
    tp_matmat_impl(model, idx, xs, b, rows, cols, None, None, 0, two_bit, None)
}

#[allow(clippy::too_many_arguments)]
fn tp_matmat_impl(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: Option<&mut [f32]>,
    src: Option<&wgpu::Buffer>,
    src_offset: u64,
    two_bit: bool,
    dst: Option<&wgpu::Buffer>,
) -> Option<wgpu::Buffer> {
    let c = ctx()?;
    let _gate = c.mm_gate.lock().unwrap();
    let _gpr = cols / 32;
    if cols % 32 != 0 || rows == 0 || b == 0 {
        return None;
    }
    let entry = &model.tensors[idx];
    // This kernel reads the four-bit tiled layout and nothing else. Handed a
    // tensor in another codec it would read int8 as tiles and return plausible
    // garbage — which is how an eight-bit VAE decoded to a flat grey frame for
    // an hour this afternoon. Refusing sends the caller to a path that can.
    let expected_dtype = if two_bit {
        cortiq_core::TensorDtype::Q2TiledP
    } else {
        cortiq_core::TensorDtype::Q4TiledP
    };
    if entry.dtype != expected_dtype {
        return None;
    }
    if entry.shape.first().copied().unwrap_or(0) < rows {
        return None;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return None;
    };
    let bytes = model.primary_bytes();
    let plen = entry.nbytes as usize;
    let dt = if two_bit {
        cortiq_core::TensorDtype::Q2TiledP
    } else {
        cortiq_core::TensorDtype::Q4TiledP
    };
    let Some(need) = cortiq_core::quant::expected_nbytes(dt, &[rows, cols]) else {
        return None;
    };
    if std::env::var("CMF_GPU_DEBUG").is_ok() && b >= 512 {
        use std::collections::HashSet;
        use std::sync::Mutex;
        static SEEN: Mutex<Option<HashSet<(usize, usize)>>> = Mutex::new(None);
        let mut g = SEEN.lock().unwrap();
        if g.get_or_insert_with(HashSet::new).insert((rows, cols)) {
            eprintln!(
                "tp_matmat entry: {rows}x{cols} plen={plen} need={need} xs={} b*cols={} out={} b*rows={}",
                xs.len(),
                b * cols,
                out.as_ref().map_or(0, |o| o.len()),
                b * rows
            );
        }
    }
    if plen < need
        || abs + plen > bytes.len()
        || (src.is_none() && xs.len() < b * cols)
        || out.as_ref().is_some_and(|o| o.len() < b * rows)
    {
        return None;
    }
    let q_buf = weight_buffer_l(
        c,
        (model.uid() as usize, idx),
        &bytes[abs..abs + plen],
        layer_of_name(&model.tensors[idx].name),
    )?;
    let mut sc = c.scratch.lock().unwrap();
    // A resident A operand: the kernel before us left it on the card, so
    // there is nothing to upload and — the point of the exercise —
    // nothing to read back first. A readback drains the whole queue, and
    // that stall was the last one left inside a DiT block.
    let xs_buf = match src {
        Some(bf) => bf.clone(),
        None => {
            let bf = Scratch::ensure(
                &c.device,
                &mut sc.xs,
                (b * cols * 4) as u64,
                wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                "q4tpmm-xs",
            );
            c.queue
                .write_buffer(&bf, 0, bytemuck::cast_slice(&xs[..b * cols]));
            bf
        }
    };
    let src_bytes = (b * cols * 4) as u64;
    if src_offset % 256 != 0
        || src_offset
            .checked_add(src_bytes)
            .is_none_or(|end| end > xs_buf.size())
    {
        return None;
    }
    let y_size = (b * rows * 4) as u64;
    let y_buf = match dst {
        Some(buf) => buf.clone(),
        None => Scratch::ensure(
            &c.device,
            &mut sc.y,
            y_size,
            wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            "q4tpmm-y",
        ),
    };
    // Activation scale for the coop arm: bring max|x| to ~1000 so the
    // f16 operands cannot overflow (the DiT's modulated activations run
    // past 65504 — that overflow was this kernel's NaN). 0 = no scaling,
    // so the scalar arm and every other caller keep their numerics.
    let coop_arm = !two_bit && c.q4tp_mm_coop.is_some();
    if std::env::var("CMF_GPU_DEBUG").is_ok() && b >= 512 {
        use std::collections::HashSet;
        use std::sync::Mutex;
        static SEEN: Mutex<Option<HashSet<(usize, usize)>>> = Mutex::new(None);
        let mut g = SEEN.lock().unwrap();
        let set = g.get_or_insert_with(HashSet::new);
        if set.insert((rows, cols)) {
            eprintln!(
                "tp_matmat(render): b={b} rows={rows} cols={cols} two_bit={two_bit} coop={coop_arm}"
            );
        }
    }
    // The scan happens only if the card is not going to do it (see dev_scale
    // below): scanning and then ignoring the result is what this path did for
    // an afternoon.
    let can_reduce_here =
        (c.act_amax_part.is_some() && c.act_amax_fold.is_some()) || c.act_absmax.is_some();
    let host_scan = std::env::var("CMF_FFN_HOST_SCAN").as_deref() == Ok("1")
        || !can_reduce_here
        || std::env::var("CMF_PROJ_DEV_SCAN").as_deref() != Ok("1");
    let ascale: f32 = if coop_arm && src.is_none() && host_scan {
        let mx = xs[..b * cols]
            .iter()
            .fold(0f32, |m, v| if v.is_finite() { m.max(v.abs()) } else { m });
        if mx > 1000.0 {
            1000.0 / mx
        } else {
            1.0
        }
    } else {
        0.0
    };
    let mut params = [(cols / 4) as u32, rows as u32, b as u32, ascale.to_bits()];
    // The q2tp MM shader shares this fourth word with q4tp's activation
    // scale. For q2tp it is instead a validated descriptor bit selecting
    // `(c-1)·s`; ordinary q2tp remains the raw `(c-1.5)·s` decode.
    if two_bit && crate::prism::is_affine_target(model, &entry.name) {
        params[3] = 1;
    }
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        y_size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "q4tpmm-stage",
    );
    // Dequantize-once: unpack the weight plane into f16 (cached per
    // tensor, so repeat calls — every layer, every step — pay nothing)
    // and run the pure f16 GEMM. The in-kernel unpacker spends about as
    // many scalar ops per weight tile as the matrix units spend MACs on
    // it, and repeats them for every 64-row tile of activations; this
    // pays that once.
    // Worth unpacking the plane only when the batch amortizes it: the
    // pass touches rows×cols regardless of b, so a narrow GEMM pays for
    // a dequantizer it barely uses (measured: nanbeige prefill chunks
    // lost 7× to it before this gate).
    // The batch at which unpacking the whole plane pays for itself. 64 was
    // measured on a language model's prefill; the video VAE calls this with
    // far narrower panels, and its four-bit decode costs four times what the
    // same decode costs an eight-bit container that always takes the plane.
    // CMF_Q4TP_PLANE_MIN moves the line so the answer can be measured.
    let plane_min: usize = std::env::var("CMF_Q4TP_PLANE_MIN")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(64);
    let dq = if coop_arm && c.q4tp_mm_coop_f16.is_some() && b >= plane_min {
        dq_f16_plane(c, &mut sc, &q_buf, rows, cols)
    } else {
        None
    };
    drop(sc);
    // A resident operand never passed through the host, so nothing here
    // has seen max|x| — the f16 arm needs the card to compute it. If
    // that reduction is missing, take the scalar arm instead of running
    // f16 unscaled: unscaled f16 on DiT activations overflows to NaN,
    // which is exactly how this kernel failed the first time.
    let can_dev_scale =
        (c.act_amax_part.is_some() && c.act_amax_fold.is_some()) || c.act_absmax.is_some();
    let dq = if src.is_some() && !can_dev_scale {
        None
    } else {
        dq
    };
    let dq_plane = dq.as_ref().map(|(p, _)| p.clone());
    let (mm_pipe, w_bind) = match (&dq_plane, c.q4tp_mm_coop_f16.as_ref()) {
        (Some(dq), Some(pipe)) => (pipe, dq),
        _ => (mm_pipeline(c, true, two_bit), &q_buf),
    };
    // The f16 twin declares a fifth binding (the device-scale buffer);
    // it takes a one-element dummy here because this path computes the
    // scale on the host and passes it in `pad`.
    let f16_arm = dq_plane.is_some() && c.q4tp_mm_coop_f16.is_some();
    // A host operand's scale could be taken on the card too, not only a
    // resident one: the scan is 13.2 M floats at DiT shapes, on one core,
    // before anything is submitted. The same change in the packed FFN moved
    // its host time per call from 116.4 ms to 77.0 ms and a 768×448 denoise
    // from 69.6 s to 62.4 s with byte-identical output.
    //
    // Here it is OPT-IN (`CMF_PROJ_DEV_SCAN=1`) for one reason only: the
    // machine it would have been measured on went away before a single frame
    // came out of this path. Shipping it on by default would be shipping an
    // untested render, and the FFN half of the same idea is what this release
    // actually stands on.
    let dev_scale = f16_arm
        && (src.is_some()
            || (can_dev_scale
                && !host_scan
                && std::env::var("CMF_PROJ_DEV_SCAN").as_deref() == Ok("1")));
    if dev_scale {
        params[3] = 0xFFFF_FFFFu32;
    }
    // Per call, NOT the shared `sc.params` slot. That slot is written
    // outside the scratch lock and read at submit time, so two threads
    // in this function raced: one GEMM ran with the other's rows/cols.
    // The batch parity test caught it the moment the write moved a few
    // lines later — the window had always been there.
    let p_buf = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("q4tpmm-params"),
            contents: bytemuck::cast_slice(&params),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let asc_buf = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("q4tpmm-ascale"),
        size: 4,
        usage: wgpu::BufferUsages::STORAGE,
        mapped_at_creation: false,
    });
    let dummy_scale = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[0.0f32]),
            usage: wgpu::BufferUsages::STORAGE,
        });
    let src_entry = if src_offset == 0 {
        bind_buf(1, &xs_buf)
    } else {
        bind_buf_off(1, &xs_buf, src_offset, src_bytes)
    };
    let mut entries = vec![
        bind_buf(0, w_bind),
        src_entry,
        bind_buf(2, &y_buf),
        bind_buf(3, &p_buf),
    ];
    // Both cooperative kernels declare the fifth (scale) binding now —
    // the in-kernel one grew it for the batched prefill. The scalar tile
    // GEMMs (q4t, q2tp, and q4tp without cooperative matrices) do not.
    if f16_arm {
        entries.push(bind_buf(4, if dev_scale { &asc_buf } else { &dummy_scale }));
    }
    let bind_mm = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q4tpmm-bg"),
        layout: &mm_pipe.get_bind_group_layout(0),
        entries: &entries,
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("q4tpmm"),
        });
    if let (Some((_, bind_dq)), Some(pipe_dq)) = (&dq, c.q4tp_dq_f16.as_ref()) {
        // Same encoder as the GEMM: a submit of its own made the driver
        // serialize two round trips where one belongs.
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(pipe_dq);
        pass.set_bind_group(0, bind_dq, &[]);
        let wgs = ((rows * cols / 2) as u32).div_ceil(256);
        pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
    }
    if dev_scale {
        encode_act_absmax(c, &mut enc, &xs_buf, b * cols, &asc_buf);
    }
    {
        let mut pass = begin_pass_with(&mut enc, Some("q4tpmm"), None);
        pass.set_pipeline(mm_pipe);
        pass.set_bind_group(0, &bind_mm, &[]);
        pass.dispatch_workgroups(
            (rows as u32).div_ceil(64).min(MAX_WG),
            (b as u32).div_ceil(64),
            1,
        );
    }
    match out {
        Some(o) => {
            // CMF_MM_SPLIT=1: land the kernel alone first, then time the
            // readback separately — the one number that says whether the
            // next work is a faster GEMM or a resident chain.
            if crate::mm_ab::split_on() {
                let t0 = std::time::Instant::now();
                submit(c, finish_enc(enc));
                if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
                    return None;
                }
                let t_k = t0.elapsed();
                let enc2 = c
                    .device
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                        label: Some("mm-split-rb"),
                    });
                let t0 = std::time::Instant::now();
                let ok = readback(c, enc2, &y_buf, &stage_buf, y_size, &mut o[..b * rows]);
                crate::mm_ab::split_note(b, rows, cols, t_k, t0.elapsed());
                return if ok { Some(y_buf) } else { None };
            }
            if readback(c, enc, &y_buf, &stage_buf, y_size, &mut o[..b * rows]) {
                Some(y_buf)
            } else {
                None
            }
        }
        None => {
            // No host copy: submit and hand the buffer over.
            submit(c, finish_enc(enc));
            Some(y_buf)
        }
    }
}

pub fn q4t_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    let _gate = c.mm_gate.lock().unwrap();
    let gpr = cols / 32;
    if cols % 32 != 0 || rows == 0 || b == 0 {
        return false;
    }
    let entry = &model.tensors[idx];
    if entry.shape.first().copied().unwrap_or(0) < rows {
        return false;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return false;
    };
    let bytes = model.primary_bytes();
    let plen = entry.nbytes as usize;
    if plen < rows * gpr * 18
        || abs + plen > bytes.len()
        || xs.len() < b * cols
        || out.len() < b * rows
    {
        return false;
    }
    let q_buf = match weight_buffer_l(
        c,
        (model.uid() as usize, idx),
        &bytes[abs..abs + plen],
        layer_of_name(&model.tensors[idx].name),
    ) {
        Some(bf) => bf,
        None => return false,
    };
    let mut sc = c.scratch.lock().unwrap();
    let xs_buf = Scratch::ensure(
        &c.device,
        &mut sc.xs,
        (b * cols * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "q4tmm-xs",
    );
    c.queue
        .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&xs[..b * cols]));
    let y_size = (b * rows * 4) as u64;
    let y_buf = Scratch::ensure(
        &c.device,
        &mut sc.y,
        y_size,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "q4tmm-y",
    );
    let params = [(cols / 4) as u32, rows as u32, b as u32, 0u32];
    let p_buf = match &sc.params {
        Some(bf) => bf.clone(),
        None => {
            let bf = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("q4tmm-params"),
                size: 16,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            sc.params = Some(bf.clone());
            bf
        }
    };
    c.queue
        .write_buffer(&p_buf, 0, bytemuck::cast_slice(&params));
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        y_size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "q4tmm-stage",
    );
    let entries = [
        bind_buf(0, &q_buf),
        bind_buf(1, &xs_buf),
        bind_buf(2, &y_buf),
        bind_buf(3, &p_buf),
    ];
    let bind_mm = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q4tmm-bg"),
        layout: &c.q4t_mm.get_bind_group_layout(0),
        entries: &entries,
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("q4tmm"),
        });
    {
        let mut pass = begin_pass_with(&mut enc, Some("q4tmm"), None);
        pass.set_pipeline(&c.q4t_mm);
        pass.set_bind_group(0, &bind_mm, &[]);
        pass.dispatch_workgroups(
            (rows as u32).div_ceil(64).min(MAX_WG),
            (b as u32).div_ceil(64),
            1,
        );
    }
    readback(c, enc, &y_buf, &stage_buf, y_size, &mut out[..b * rows])
}

/// One q4tp matvec through the WGSL kernel, weight fetched from the model.
/// Exists so the shader can be pinned to `dequant_q4tp` in a test: the token
/// graph is the only other caller, and a wrong kernel there still produces
/// fluent text.
#[doc(hidden)]
/// Single-token q4tp matvec — the lm_head class — through the DEDICATED
/// matvec kernel rather than the batched GEMM at b=1. The GEMM measured
/// 11.73 ms against the host's 9.51 on a 129280x4096 head, which is how a
/// route that should have been a rout ended up losing its own probe.
pub fn q4tp_matvec(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    q4tp_matvec_for_test(model, idx, xs, rows, cols, out)
}

pub fn q4tp_matvec_for_test(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    q4tp_matvec_batch_for_test(model, idx, xs, 1, rows, cols, out)
}

/// `batch` activation vectors end to end against one weight, `batch * rows`
/// outputs. Exposed so the dequant-pinned test can hold the batched form to
/// the same definition as the single one — the batch dimension shares the
/// kernel with the matvec, so a bug in it is a bug in decode too.
#[allow(clippy::too_many_arguments)]
pub fn q4tp_matvec_batch_for_test(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    batch: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if cols % 32 != 0
        || rows == 0
        || batch == 0
        || xs.len() < batch * cols
        || out.len() < batch * rows
    {
        return false;
    }
    let entry = &model.tensors[idx];
    if entry.dtype != cortiq_core::TensorDtype::Q4TiledP {
        return false;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return false;
    };
    let bytes = model.primary_bytes();
    let plen = entry.nbytes as usize;
    let Some(need) =
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
    else {
        return false;
    };
    if plen < need || abs + plen > bytes.len() {
        return false;
    }
    let Some(q_buf) = weight_buffer_l(
        c,
        (model.uid() as usize, idx),
        &bytes[abs..abs + plen],
        layer_of_name(&model.tensors[idx].name),
    ) else {
        return false;
    };
    let mut sc = c.scratch.lock().unwrap();
    let xs_buf = Scratch::ensure(
        &c.device,
        &mut sc.xs,
        (batch * cols * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "q4tp-xs",
    );
    c.queue
        .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&xs[..batch * cols]));
    let y_size = (batch * rows * 4) as u64;
    let y_buf = Scratch::ensure(
        &c.device,
        &mut sc.y,
        y_size,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "q4tp-y",
    );
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        y_size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "q4tp-stage",
    );
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("q4tp-mv"),
        });
    // The f32 kernels of the runtime's arm (bku / bk / the singles): this
    // helper backs the 1e-4 dequant-reference test; the int8-activation
    // arm the verify runs by default has its own bounded parity test.
    if c.use_mv4 {
        if !encode_q4tp_mv4_b_with(
            c, &mut enc, &q_buf, &xs_buf, &y_buf, rows, cols, batch, false,
        ) {
            return false;
        }
    } else if batch == 1 {
        encode_q1t_like(c, &mut enc, &c.q4tp_mv, &q_buf, &xs_buf, &y_buf, rows, cols);
    } else {
        return false;
    }
    readback(c, enc, &y_buf, &stage_buf, y_size, &mut out[..batch * rows])
}

/// DiT attention on wgpu: per head, scores = scale·Q·Kᵀ → row softmax →
/// P·V, then one unstack of the [nh][n][hd] panel into [n][nh·hd]. All
/// of it in ONE submission with the scores and the panel resident on the
/// device — the CPU only ships Q/K/V in and the result out. Head-major
/// inputs, matching `gpu_metal::dit_attention`.
#[allow(clippy::too_many_arguments)]
/// Attention straight from an interleaved qkv panel: the split into
/// head-major planes happens on the device. `false` means the split
/// kernel is unavailable and the caller should repack on the host.
#[allow(clippy::too_many_arguments)]
#[allow(clippy::too_many_arguments)]
/// qkv GEMM + attention with the panel never leaving the card: the
/// projection writes into a device buffer and the split reads it there.
/// Requires the device to own qk-norm and RoPE too (`nr`), because the
/// host loop that used to apply them is the only other reason the panel
/// came home. `false` = anything refused, and the caller runs its host
/// chain unchanged.
#[allow(clippy::too_many_arguments)]
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 {
    let inner = nh * hd;
    let Some(panel) = fused_panel_keep(model, qkv_idx, xn, n, 3 * inner, hidden) else {
        return false;
    };
    dit_attention_packed_src(
        &[],
        Some(&panel),
        nh,
        n,
        hd,
        scale,
        Some(nr),
        out,
        None,
        0,
        None,
    )
}

/// qkv GEMM, attention, AND the output projection with nothing crossing
/// the bus in between: the projection reads the attention panel where
/// the unstack left it, and only its own result (n×hidden) comes home.
/// This was the last round trip inside a DiT block — and a readback is
/// not just its own bytes, it drains the queue, so the card idled once
/// per block waiting for the host to take delivery.
/// `false` = refused at the door (every condition is checked BEFORE any
/// work, so the caller's fallback never repeats work this already did).
#[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 {
    let Some(c) = ctx() else { return false };
    let inner = nh * hd;
    let can_dev_scale =
        (c.act_amax_part.is_some() && c.act_amax_fold.is_some()) || c.act_absmax.is_some();
    // The f16 dequant and the coop GEMM are the FOUR-BIT path's tools; an
    // int8 weight needs neither, and demanding them here is what kept a
    // q8_2f container out of the fused path entirely.
    let four_bit = model.tensors[qkv_idx].dtype == cortiq_core::TensorDtype::Q4TiledP;
    let refuse = if !can_dev_scale {
        "no device absmax"
    } else if four_bit && c.q4tp_mm_coop_f16.is_none() {
        "no coop f16 gemm"
    } else if four_bit && c.q4tp_dq_f16.is_none() {
        "no f16 dequant"
    } else if c.dit_qkv_split.is_none() {
        "no qkv split"
    } else if inner % 32 != 0 {
        "inner not 32-aligned"
    } else if n < 64 {
        "batch too small"
    } else if proj.len() < n * hidden {
        "proj too small"
    } else {
        ""
    };
    if !refuse.is_empty() {
        if std::env::var("CMF_GPU_DEBUG").is_ok() {
            static ONCE: std::sync::Once = std::sync::Once::new();
            ONCE.call_once(|| eprintln!("dit_qkv_attn_out refused: {refuse}"));
        }
        return false;
    }
    let Some(panel) = fused_panel_keep(model, qkv_idx, xn, n, 3 * inner, hidden) else {
        if std::env::var("CMF_GPU_DEBUG").is_ok() {
            static ONCE: std::sync::Once = std::sync::Once::new();
            ONCE.call_once(|| eprintln!("dit_qkv_attn_out refused: qkv gemm"));
        }
        return false;
    };
    let mut ab = None;
    if !dit_attention_packed_src(
        &[],
        Some(&panel),
        nh,
        n,
        hd,
        scale,
        Some(nr),
        &mut [],
        Some(&mut ab),
        0,
        None,
    ) {
        if std::env::var("CMF_GPU_DEBUG").is_ok() {
            static ONCE: std::sync::Once = std::sync::Once::new();
            ONCE.call_once(|| eprintln!("dit_qkv_attn_out refused: attention"));
        }
        return false;
    }
    let Some(ab) = ab else { return false };
    let ok = fused_gemm_from_device(model, out_idx, &ab, n, hidden, inner, proj);
    if !ok && std::env::var("CMF_GPU_DEBUG").is_ok() {
        static ONCE: std::sync::Once = std::sync::Once::new();
        ONCE.call_once(|| eprintln!("dit_qkv_attn_out refused: out gemm {hidden}x{inner} b={n}"));
    }
    ok
}

/// The VAE decoder's attention half on the card: qkv GEMM, bias, plain
/// RMS norm on q and k, RoPE, attention, output projection. Only the
/// projection's result comes home. Its panel is head-interleaved and
/// its norm carries no weight — a ones vector makes `dit_qknorm_rope`
/// compute exactly `rms_norm_plain`, so no second kernel exists for it.
#[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 {
    let Some(c) = ctx() else { return false };
    let inner = nh * hd;
    let can_dev_scale =
        (c.act_amax_part.is_some() && c.act_amax_fold.is_some()) || c.act_absmax.is_some();
    let four_bit = model
        .tensors
        .get(qkv_idx)
        .is_some_and(|e| e.dtype == cortiq_core::TensorDtype::Q4TiledP);
    if !can_dev_scale
        || (four_bit && c.q4tp_mm_coop_f16.is_none())
        || (four_bit && c.q4tp_dq_f16.is_none())
        || c.dit_qkv_split.is_none()
        || c.dit_qknorm.is_none()
        || inner != dim
        || inner % 32 != 0
        || n < 64
        || angles.is_empty()
        || qkv_bias.len() < 3 * dim
        || proj.len() < n * dim
    {
        return false;
    }
    // Whichever codec the container is packed in. This used to call the
    // four-bit entry by name while the caller's gate had been widened to any
    // weight with a device GEMM, so an eight-bit VAE read its int8 payload as
    // four-bit tiles: a decode that ran fast and returned a flat grey frame.
    let Some(panel) = fused_panel_keep(model, qkv_idx, xn, n, 3 * inner, dim) else {
        return false;
    };
    let ones = vec![1.0f32; hd];
    let mut ab = None;
    if !dit_attention_packed_src(
        &[],
        Some(&panel),
        nh,
        n,
        hd,
        scale,
        Some((angles, &ones[..], &ones[..], eps)),
        &mut [],
        Some(&mut ab),
        1,
        Some(qkv_bias),
    ) {
        return false;
    }
    let Some(ab) = ab else { return false };
    fused_gemm_from_device(model, out_idx, &ab, n, dim, inner, proj)
}

/// Bisect handle: the VAE's head-interleaved split plus the weightless
/// q/k norm and RoPE, fed a HOST panel that already carries its bias.
/// Isolates those two kernels from the resident-GEMM handoff.
#[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 {
    let ones = vec![1.0f32; hd];
    dit_attention_packed_src(
        qkv,
        None,
        nh,
        n,
        hd,
        scale,
        Some((angles, &ones[..], &ones[..], eps)),
        out,
        None,
        layout,
        None,
    )
}

/// Diagnostic: run ONLY the split and read the q plane back. No norm,
/// no RoPE, no attention — so a mismatch here is the hand-off itself.
#[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 {
    let Some(c) = ctx() else { return false };
    let Some(split) = c.dit_qkv_split.as_ref() else {
        return false;
    };
    let _gate = c.mm_gate.lock().unwrap();
    let inner = nh * hd;
    if qkv.len() < n * 3 * inner || out_q.len() < nh * n * hd {
        return false;
    }
    let plane = (nh * n * hd * 4) as u64;
    let st = wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST;
    // Own buffers, NOT the shared slots: a grow-only slot keeps the
    // usage flags it was FIRST created with, so asking an existing
    // `dit-q` for COPY_SRC is silently not granted and the copy fails
    // validation. That trap is documented two functions up; it caught
    // this probe anyway.
    let mk = |size: u64, usage: wgpu::BufferUsages, label: &str| {
        c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size,
            usage,
            mapped_at_creation: false,
        })
    };
    let src = mk((n * 3 * inner * 4) as u64, st, "dso-qkv");
    let qb = mk(plane, st | wgpu::BufferUsages::COPY_SRC, "dso-q");
    let kb = mk(plane, st, "dso-k");
    let vb = mk(plane, st, "dso-v");
    let stage = mk(
        plane,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dso-stage",
    );
    c.queue
        .write_buffer(&src, 0, bytemuck::cast_slice(&qkv[..n * 3 * inner]));
    let u = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[
                n as u32, nh as u32, hd as u32, layout, 0u32, 0u32, 0u32, 0u32,
            ]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let dummy = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[0.0f32]),
            usage: wgpu::BufferUsages::STORAGE,
        });
    let bg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &split.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, &src),
            bind_buf(1, &qb),
            bind_buf(2, &kb),
            bind_buf(3, &vb),
            bind_buf(4, &u),
            bind_buf(5, &dummy),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("dso") });
    {
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(split);
        pass.set_bind_group(0, &bg, &[]);
        let wgs = ((n * inner) as u32).div_ceil(256);
        pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
    }
    // Stage two, same treatment: qk-norm + RoPE over q only.
    if let (Some((angles, eps)), Some(pipe)) = (norm, c.dit_qknorm.as_ref()) {
        let pairs = if angles.is_empty() {
            0
        } else {
            angles.len() / n
        };
        let ang = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(if angles.is_empty() {
                    &[0.0f32][..]
                } else {
                    angles
                }),
                usage: wgpu::BufferUsages::STORAGE,
            });
        let ones = vec![1.0f32; hd];
        let w = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&ones),
                usage: wgpu::BufferUsages::STORAGE,
            });
        let up = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&[
                    n as u32,
                    nh as u32,
                    hd as u32,
                    pairs as u32,
                    eps.to_bits(),
                    0u32,
                    layout,
                    0u32,
                ]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let bgn = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &pipe.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &src),
                bind_buf(1, &qb),
                bind_buf(2, &ang),
                bind_buf(3, &w),
                bind_buf(4, &up),
                bind_buf(5, &dummy),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(pipe);
        pass.set_bind_group(0, &bgn, &[]);
        let jobs = (n * nh) as u32;
        pass.dispatch_workgroups(jobs.min(65535), jobs.div_ceil(65535), 1);
    }
    readback(c, enc, &qb, &stage, plane, &mut out_q[..nh * n * hd])
}

pub fn dit_attention_packed(
    qkv: &[f32],
    nh: usize,
    n: usize,
    hd: usize,
    scale: f32,
    nr: Option<(&[f32], &[f32], &[f32], f32)>,
    out: &mut [f32],
) -> bool {
    dit_attention_packed_src(qkv, None, nh, n, hd, scale, nr, out, None, 0, None)
}

/// The two things `dit_attention_packed_src` checks before it can do
/// anything: a live context and the split pipeline. A caller that is
/// about to SKIP work on the strength of this arm has to know now — a
/// refusal discovered later leaves q/k unnormalized with no way back.
pub fn dit_attention_packed_ready() -> bool {
    ctx().is_some_and(|c| c.dit_qkv_split.is_some())
}

/// The same, with the panel possibly ALREADY on the card (`pre`): its
/// GEMM kept it there, so nothing is uploaded and nothing was read back
/// — 320 MB a block that used to cross the bus twice.
#[allow(clippy::too_many_arguments)]
pub fn dit_attention_packed_src(
    qkv: &[f32],
    pre: Option<&wgpu::Buffer>,
    nh: usize,
    n: usize,
    hd: usize,
    scale: f32,
    nr: Option<(&[f32], &[f32], &[f32], f32)>,
    out: &mut [f32],
    keep: Option<&mut Option<wgpu::Buffer>>,
    layout: u32,
    qkv_bias: Option<&[f32]>,
) -> bool {
    let Some(c) = ctx() else { return false };
    let Some(split) = c.dit_qkv_split.as_ref() else {
        return false;
    };
    let inner = nh * hd;
    if (pre.is_none() && qkv.len() < n * 3 * inner) || (keep.is_none() && out.len() < n * inner) {
        return false;
    }
    let plane = (nh * n * hd * 4) as u64;
    let (qb, kb, vb, src) = {
        let mut sc = c.scratch.lock().unwrap();
        // COPY_DST too: these are the SAME slots the host-repack path
        // uploads into, and a grow-only slot keeps whatever usage it was
        // first created with — leaving it out made a later host call
        // fail on a buffer this one had allocated.
        let st = wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST;
        let src = Scratch::ensure(
            &c.device,
            &mut sc.dqkv,
            (n * 3 * inner * 4) as u64,
            st | wgpu::BufferUsages::COPY_DST,
            "dit-qkv",
        );
        (
            Scratch::ensure(&c.device, &mut sc.dq, plane, st, "dit-q"),
            Scratch::ensure(&c.device, &mut sc.dk, plane, st, "dit-k"),
            Scratch::ensure(&c.device, &mut sc.dv, plane, st, "dit-v"),
            src,
        )
    };
    let src = match pre {
        Some(b) => b.clone(),
        None => {
            c.queue
                .write_buffer(&src, 0, bytemuck::cast_slice(&qkv[..n * 3 * inner]));
            src
        }
    };
    let u = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[
                n as u32,
                nh as u32,
                hd as u32,
                layout,
                u32::from(qkv_bias.is_some()),
                0u32,
                0u32,
                0u32,
            ]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    // The bias binding always exists; one element stands in when the
    // panel has none, because a declared binding must be bound.
    let bias_buf = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("dit-qkv-bias"),
            contents: bytemuck::cast_slice(qkv_bias.unwrap_or(&[0.0f32][..])),
            usage: wgpu::BufferUsages::STORAGE,
        });
    let bg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &split.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, &src),
            bind_buf(1, &qb),
            bind_buf(2, &kb),
            bind_buf(3, &vb),
            bind_buf(4, &u),
            bind_buf(5, &bias_buf),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dit-split"),
        });
    {
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(split);
        pass.set_bind_group(0, &bg, &[]);
        let wgs = ((n * inner) as u32).div_ceil(256);
        pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
    }
    // SUBMIT IT. This encoder was built, filled and dropped — the split
    // never ran, and only the qk-norm encoder below reached the queue.
    // q and k survived that (qk-norm reads the packed panel itself and
    // rewrites both planes), so the DiT could not see it; v is written
    // by nobody else, so it stayed zero and P·V came out exactly zero.
    // That is the VAE's all-zero attention, and the reason the layout
    // experiments all agreed: zero does not depend on addressing.
    submit(c, finish_enc(enc));
    // qk-norm + RoPE on the card: two dispatches over the SAME kernel,
    // q and k differing only in the output plane, the weight buffer and
    // src_off. When `nr` is None the caller already did this on the host.
    if std::env::var("CMF_GPU_DEBUG").is_ok() {
        // Once per LAYOUT, not once per process: a single `Once` is
        // spent by the DiT's first call and the VAE's never prints.
        static SEEN: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(0);
        let bit = 1u32 << layout.min(1);
        if SEEN.fetch_or(bit, std::sync::atomic::Ordering::Relaxed) & bit == 0 {
            eprintln!(
                "packed_src: layout={layout} bias={} qknorm={} nr={} pre={} n={n} nh={nh} hd={hd} pairs={}",
                qkv_bias.is_some(),
                c.dit_qknorm.is_some(),
                nr.is_some(),
                pre.is_some(),
                nr.map_or(0, |(a, _, _, _)| if a.is_empty() { 0 } else { a.len() / n }),
            );
        }
    }
    if let (Some(pipe), Some((angles, qw, kw, eps))) = (c.dit_qknorm.as_ref(), nr) {
        let pairs = if angles.is_empty() {
            0
        } else {
            angles.len() / n
        };
        let ang_b = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("dit-ang"),
                contents: bytemuck::cast_slice(if angles.is_empty() {
                    &[0.0f32][..]
                } else {
                    angles
                }),
                usage: wgpu::BufferUsages::STORAGE,
            });
        for (half, wts, dst) in [(0usize, qw, &qb), (1usize, kw, &kb)] {
            let w_b = c
                .device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some("dit-qkw"),
                    contents: bytemuck::cast_slice(wts),
                    usage: wgpu::BufferUsages::STORAGE,
                });
            let u = c
                .device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytemuck::cast_slice(&[
                        n as u32,
                        nh as u32,
                        hd as u32,
                        pairs as u32,
                        eps.to_bits(),
                        (half * inner) as u32,
                        layout,
                        u32::from(qkv_bias.is_some()),
                    ]),
                    usage: wgpu::BufferUsages::UNIFORM,
                });
            let bg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &pipe.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, &src),
                    bind_buf(1, dst),
                    bind_buf(2, &ang_b),
                    bind_buf(3, &w_b),
                    bind_buf(4, &u),
                    bind_buf(5, &bias_buf),
                ],
            });
            let mut e = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("dit-qkn"),
                });
            {
                let mut pass = begin_pass(&mut e);
                pass.set_pipeline(pipe);
                pass.set_bind_group(0, &bg, &[]);
                let jobs = (n * nh) as u32;
                pass.dispatch_workgroups(jobs.min(65535), jobs.div_ceil(65535), 1);
            }
            submit(c, finish_enc(e));
        }
    }
    // The planes are on the card now; the shared path skips its uploads.
    dit_attention_inner(&[], &[], &[], nh, nh, n, hd, scale, out, true, keep)
}

pub fn dit_attention(
    qh: &[f32],
    kh: &[f32],
    vh: &[f32],
    nh: usize,
    nkv: usize,
    n: usize,
    hd: usize,
    scale: f32,
    out: &mut [f32],
) -> bool {
    dit_attention_inner(qh, kh, vh, nh, nkv, n, hd, scale, out, false, None)
}

/// `resident` = the q/k/v planes are ALREADY in the scratch slots (the
/// device-side split just wrote them), so the host slices are empty and
/// the uploads are skipped.
#[allow(clippy::too_many_arguments)]
fn dit_attention_inner(
    qh: &[f32],
    kh: &[f32],
    vh: &[f32],
    nh: usize,
    nkv: usize,
    n: usize,
    hd: usize,
    scale: f32,
    out: &mut [f32],
    resident: bool,
    keep: Option<&mut Option<wgpu::Buffer>>,
) -> bool {
    let Some(c) = ctx() else { return false };
    if nh == 0 || nkv == 0 || n == 0 || hd == 0 || nh % nkv != 0 {
        return false;
    }
    if !resident && (qh.len() < nh * n * hd || kh.len() < nkv * n * hd || vh.len() < nkv * n * hd) {
        return false;
    }
    // Empty in keep mode: the panel is handed over as a buffer,
    // so there is no host slice to size-check.
    if keep.is_none() && out.len() < n * nh * hd {
        return false;
    }
    let dev = &c.device;
    // Grow-only slots, not fresh allocations per call: a render calls
    // this 26 times per forward and the driver's allocator is not free.
    // `Scratch::ensure` also flags the cold call so the contention
    // tripwire does not read a one-off buffer creation as a busy device.
    let mut sc = c.scratch.lock().unwrap();
    let st = wgpu::BufferUsages::STORAGE;
    let up = |slot: &mut Option<(wgpu::Buffer, u64)>, data: &[f32], label: &str| -> wgpu::Buffer {
        let b = Scratch::ensure(
            dev,
            slot,
            (data.len() * 4) as u64,
            st | wgpu::BufferUsages::COPY_DST,
            label,
        );
        // Empty slice = the plane is already resident (the device-side
        // split wrote it); allocating the slot is all that is needed.
        if !data.is_empty() {
            c.queue.write_buffer(&b, 0, bytemuck::cast_slice(data));
        }
        b
    };
    let qb = up(
        &mut sc.dq,
        if resident { &[] } else { &qh[..nh * n * hd] },
        "dit-q",
    );
    let kb = up(
        &mut sc.dk,
        if resident { &[] } else { &kh[..nkv * n * hd] },
        "dit-k",
    );
    let vb = up(
        &mut sc.dv,
        if resident { &[] } else { &vh[..nkv * n * hd] },
        "dit-v",
    );
    let scb = Scratch::ensure(dev, &mut sc.dsc, (n * n * 4) as u64, st, "dit-scores");
    let pb = Scratch::ensure(dev, &mut sc.dpan, (nh * n * hd * 4) as u64, st, "dit-panel");
    let ab = Scratch::ensure(
        dev,
        &mut sc.dout,
        (n * nh * hd * 4) as u64,
        st | wgpu::BufferUsages::COPY_SRC,
        "dit-out",
    );
    let stage = Scratch::ensure(
        dev,
        &mut sc.dstage,
        (n * nh * hd * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dit-stage",
    );
    drop(sc);

    // One uniform per distinct (m, k, n, scale) shape; the head offset
    // rides in the bound slice, not the params.
    let params = |m: u32, k: u32, nn: u32, sc: f32| -> wgpu::Buffer {
        let raw = [m, k, nn, sc.to_bits(), 0u32, 0u32, 0u32, 0u32];
        let b = dev.create_buffer(&wgpu::BufferDescriptor {
            label: Some("dit-params"),
            size: 32,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        c.queue.write_buffer(&b, 0, bytemuck::cast_slice(&raw));
        b
    };
    let p_qk = params(n as u32, hd as u32, n as u32, scale);
    let p_sm = params(n as u32, hd as u32, n as u32, 1.0);
    let p_pv = params(n as u32, n as u32, hd as u32, 1.0);
    let p_un = params(nh as u32, n as u32, hd as u32, 1.0);

    let bind = |pipe: &wgpu::ComputePipeline,
                a: &wgpu::Buffer,
                ao: u64,
                al: u64,
                b: &wgpu::Buffer,
                bo: u64,
                bl: u64,
                cc: &wgpu::Buffer,
                co: u64,
                cl: u64,
                pp: &wgpu::Buffer|
     -> wgpu::BindGroup {
        dev.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("dit-bg"),
            layout: &pipe.get_bind_group_layout(0),
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                        buffer: a,
                        offset: ao,
                        size: std::num::NonZeroU64::new(al),
                    }),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                        buffer: b,
                        offset: bo,
                        size: std::num::NonZeroU64::new(bl),
                    }),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                        buffer: cc,
                        offset: co,
                        size: std::num::NonZeroU64::new(cl),
                    }),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: pp.as_entire_binding(),
                },
            ],
        })
    };

    let hpk = nh / nkv;
    let head = (n * hd * 4) as u64;
    let sc_len = (n * n * 4) as u64;
    let mut enc = dev.create_command_encoder(&wgpu::CommandEncoderDescriptor {
        label: Some("dit-attn"),
    });
    // Default now, on the numbers: QK and PV both ride the matrix units
    // and PV's operand is transposed first (4.76 → 2.91 s a step, frames
    // within 0.1% of the scalar path). `CMF_DIT_ATTN_COOP=0` opts out.
    let coop_dit = c
        .dit_gemm_coop
        .as_ref()
        .filter(|_| std::env::var("CMF_DIT_ATTN_COOP").as_deref() != Ok("0"));
    // PV's right operand is read down columns in v's [n][hd] layout —
    // 4.76 s a step against QK's 1.65 at the same FLOPs. Transpose it
    // once per block and PV becomes the NT product QK already is.
    // PV keeps its own switch for bisecting, but defaults on with the
    // rest: `CMF_DIT_PV_COOP=0` leaves PV on the scalar tile kernel.
    let pv_coop_on = std::env::var("CMF_DIT_PV_COOP").as_deref() != Ok("0");
    let vtb = match (coop_dit.filter(|_| pv_coop_on), c.dit_v_transpose.as_ref()) {
        (Some(_), Some(tp)) => {
            let vt = {
                let mut sc = c.scratch.lock().unwrap();
                Scratch::ensure(
                    &c.device,
                    &mut sc.dvt,
                    (nh * n * hd * 4) as u64,
                    wgpu::BufferUsages::STORAGE,
                    "dit-vt",
                )
            };
            let u = c
                .device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytemuck::cast_slice(&[n as u32, nh as u32, hd as u32, 0u32]),
                    usage: wgpu::BufferUsages::UNIFORM,
                });
            let bg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &tp.get_bind_group_layout(0),
                entries: &[bind_buf(0, &vb), bind_buf(1, &vt), bind_buf(2, &u)],
            });
            let mut e = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("dit-vt"),
                });
            {
                let mut pass = begin_pass(&mut e);
                pass.set_pipeline(tp);
                pass.set_bind_group(0, &bg, &[]);
                let wgs = ((n * nh * hd) as u32).div_ceil(256);
                pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
            }
            submit(c, finish_enc(e));
            Some(vt)
        }
        _ => None,
    };
    for h in 0..nh {
        let kv = (h / hpk) as u64;
        // QK on the matrix units: x = this head's q rows, w = its k
        // rows, y = its score plane; offsets in ELEMENTS, as the kernel
        // reads them.
        let bg_qk_coop = coop_dit.map(|p| {
            let u = c
                .device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytemuck::cast_slice(&[
                        (hd / 4) as u32,
                        n as u32,
                        n as u32,
                        scale.to_bits(),
                        (h * n * hd) as u32,
                        (kv as usize * n * hd) as u32,
                        0u32,
                        hd as u32,
                    ]),
                    usage: wgpu::BufferUsages::UNIFORM,
                });
            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &p.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, &kb),
                    bind_buf(1, &qb),
                    bind_buf(2, &scb),
                    bind_buf(3, &u),
                ],
            })
        });
        let bg_pv_coop = match (coop_dit, &vtb) {
            (Some(p), Some(vt)) => {
                let u = c
                    .device
                    .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: None,
                        contents: bytemuck::cast_slice(&[
                            (n / 4) as u32,
                            hd as u32,
                            n as u32,
                            1.0f32.to_bits(),
                            0u32,
                            (kv as usize * n * hd) as u32,
                            (h * n * hd) as u32,
                            n as u32,
                        ]),
                        usage: wgpu::BufferUsages::UNIFORM,
                    });
                Some(c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &p.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, vt),
                        bind_buf(1, &scb),
                        bind_buf(2, &pb),
                        bind_buf(3, &u),
                    ],
                }))
            }
            _ => None,
        };
        let bg_qk = bind(
            &c.dit_qk,
            &qb,
            h as u64 * head,
            head,
            &kb,
            kv * head,
            head,
            &scb,
            0,
            sc_len,
            &p_qk,
        );
        // Naga derives each pipeline's layout from the bindings it
        // actually uses: softmax touches only the scores and the params,
        // so its group has two entries, not four.
        let bg_sm = dev.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("dit-sm-bg"),
            layout: &c.dit_softmax.get_bind_group_layout(0),
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: scb.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: p_sm.as_entire_binding(),
                },
            ],
        });
        let bg_pv = bind(
            &c.dit_pv,
            &scb,
            0,
            sc_len,
            &vb,
            kv * head,
            head,
            &pb,
            h as u64 * head,
            head,
            &p_pv,
        );
        // Each stage reads what the previous one wrote to the SAME
        // scores buffer, so each gets its own pass: wgpu inserts the
        // memory barrier at pass boundaries, and three dispatches inside
        // one pass raced (max pixel error 38/255 against the CPU path).
        // CMF_DIT_ATTN_PROF=1 splits the three phases into their own
        // submits and waits between them, so each wall is separable.
        // A measurement mode, not a path: the syncs cost real time.
        let prof = std::env::var("CMF_DIT_ATTN_PROF").is_ok();
        let split_now = |enc: &mut wgpu::CommandEncoder, slot: usize| {
            if !prof {
                return;
            }
            flush_pass(&enc);
            let e = std::mem::replace(
                enc,
                c.device
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }),
            );
            let t = std::time::Instant::now();
            submit(c, finish_enc(e));
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            DIT_PHASE[slot].fetch_add(
                t.elapsed().as_micros() as u64,
                std::sync::atomic::Ordering::Relaxed,
            );
        };
        {
            let mut pass = begin_pass_with(&mut enc, Some("dit-qk"), None);
            // Matrix units when the device has them and the caller opted
            // in: same NT product, f16 operands, f32 accumulator. The
            // scalar tile kernel below runs this attention at 0.5
            // TFLOP/s where the card's q4tp GEMM holds 51.9.
            match (&coop_dit, &bg_qk_coop) {
                (Some(p), Some(bgc)) => {
                    pass.set_pipeline(p);
                    pass.set_bind_group(0, bgc, &[]);
                    pass.dispatch_workgroups((n as u32).div_ceil(64), (n as u32).div_ceil(64), 1);
                }
                _ => {
                    pass.set_pipeline(&c.dit_qk);
                    pass.set_bind_group(0, &bg_qk, &[]);
                    pass.dispatch_workgroups((n as u32).div_ceil(64), (n as u32).div_ceil(64), 1);
                }
            }
        }
        split_now(&mut enc, 0);
        {
            let mut pass = begin_pass_with(&mut enc, Some("dit-sm"), None);
            pass.set_pipeline(&c.dit_softmax);
            pass.set_bind_group(0, &bg_sm, &[]);
            pass.dispatch_workgroups(n as u32, 1, 1);
        }
        split_now(&mut enc, 1);
        {
            let mut pass = begin_pass_with(&mut enc, Some("dit-pv"), None);
            match (&coop_dit, &bg_pv_coop) {
                (Some(p), Some(bgc)) => {
                    pass.set_pipeline(p);
                    pass.set_bind_group(0, bgc, &[]);
                    pass.dispatch_workgroups((hd as u32).div_ceil(64), (n as u32).div_ceil(64), 1);
                }
                _ => {
                    pass.set_pipeline(&c.dit_pv);
                    pass.set_bind_group(0, &bg_pv, &[]);
                    pass.dispatch_workgroups((hd as u32).div_ceil(64), (n as u32).div_ceil(64), 1);
                }
            }
        }
        split_now(&mut enc, 2);
    }
    {
        let total = (nh * n * hd) as u32;
        // unstack reads the panel (0) and writes the output (2).
        let bg_un = dev.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("dit-un-bg"),
            layout: &c.dit_unstack.get_bind_group_layout(0),
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: pb.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: ab.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: p_un.as_entire_binding(),
                },
            ],
        });
        let mut pass = begin_pass_with(&mut enc, Some("dit-unstack"), None);
        pass.set_pipeline(&c.dit_unstack);
        pass.set_bind_group(0, &bg_un, &[]);
        pass.dispatch_workgroups_flat(total.div_ceil(256));
    }
    match keep {
        // The caller is the next GEMM in the same block: hand it the
        // panel where it already is. Nothing crosses the bus, and the
        // queue is never drained mid-block.
        Some(slot) => {
            submit(c, finish_enc(enc));
            *slot = Some(ab);
            true
        }
        None => readback(
            c,
            enc,
            &ab,
            &stage,
            (n * nh * hd * 4) as u64,
            &mut out[..n * nh * hd],
        ),
    }
}

/// Causal chunk attention on wgpu: `b` new queries against `s0 + b`
/// cached keys, per head, with the causal bound applied in the softmax.
/// Same three kernels as the DiT path, rectangular this time.
///
/// This is the prefill attention the CPU path only has on aarch64 — its
/// batched attend needs Accelerate or the NEON micro-GEMM, so x86 fell
/// back to a per-position scalar loop. Measured on a 256-core EPYC that
/// loop was 30% of a 512-token prefill and 46% of a 1024-token one.
///
/// `q` is head-major [nh][b][hd] (post-RoPE); `k`/`v` are per-kv-head
/// contiguous [s0+b][hd] — the cache's own layout. `out` is
/// [b][nh·hd].
#[allow(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 {
    let Some(c) = ctx() else { return false };
    let n = s0 + b;
    if nh == 0 || nkv == 0 || b == 0 || hd == 0 || nh % nkv != 0 || n == 0 {
        return false;
    }
    if q.len() < nh * b * hd || k.len() != nkv || v.len() != nkv {
        return false;
    }
    for h in 0..nkv {
        if k[h].len() < n * hd || v[h].len() < n * hd {
            return false;
        }
    }
    if out.len() < b * nh * hd {
        return false;
    }
    let dev = &c.device;
    let mut sc = c.scratch.lock().unwrap();
    let st = wgpu::BufferUsages::STORAGE;
    let qb = Scratch::ensure(
        dev,
        &mut sc.dq,
        (nh * b * hd * 4) as u64,
        st | wgpu::BufferUsages::COPY_DST,
        "ca-q",
    );
    c.queue
        .write_buffer(&qb, 0, bytemuck::cast_slice(&q[..nh * b * hd]));
    // K/V are per-head slices of the CPU cache: pack them back to back
    // so one buffer serves every head at a known stride.
    let kvsz = (nkv * n * hd * 4) as u64;
    let kb = Scratch::ensure(
        dev,
        &mut sc.dk,
        kvsz,
        st | wgpu::BufferUsages::COPY_DST,
        "ca-k",
    );
    let vb = Scratch::ensure(
        dev,
        &mut sc.dv,
        kvsz,
        st | wgpu::BufferUsages::COPY_DST,
        "ca-v",
    );
    for h in 0..nkv {
        let off = (h * n * hd * 4) as u64;
        c.queue
            .write_buffer(&kb, off, bytemuck::cast_slice(&k[h][..n * hd]));
        c.queue
            .write_buffer(&vb, off, bytemuck::cast_slice(&v[h][..n * hd]));
    }
    let scb = Scratch::ensure(dev, &mut sc.dsc, (b * n * 4) as u64, st, "ca-scores");
    let pb = Scratch::ensure(dev, &mut sc.dpan, (nh * b * hd * 4) as u64, st, "ca-panel");
    let ab = Scratch::ensure(
        dev,
        &mut sc.dout,
        (b * nh * hd * 4) as u64,
        st | wgpu::BufferUsages::COPY_SRC,
        "ca-out",
    );
    let stage = Scratch::ensure(
        dev,
        &mut sc.dstage,
        (b * nh * hd * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "ca-stage",
    );
    drop(sc);

    let params = |m: u32, kk: u32, nn: u32, s: f32, s0v: u32, caus: u32| -> wgpu::Buffer {
        let raw = [m, kk, nn, s.to_bits(), s0v, caus, 0u32, 0u32];
        let bf = dev.create_buffer(&wgpu::BufferDescriptor {
            label: Some("ca-params"),
            size: 32,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        c.queue.write_buffer(&bf, 0, bytemuck::cast_slice(&raw));
        bf
    };
    let p_qk = params(b as u32, hd as u32, n as u32, scale, s0 as u32, 0);
    let p_sm = params(b as u32, hd as u32, n as u32, 1.0, s0 as u32, 1);
    let p_pv = params(b as u32, n as u32, hd as u32, 1.0, s0 as u32, 0);
    let p_un = params(nh as u32, b as u32, hd as u32, 1.0, 0, 0);

    fn slot(bf: &wgpu::Buffer, off: u64, len: u64, bind: u32) -> wgpu::BindGroupEntry<'_> {
        wgpu::BindGroupEntry {
            binding: bind,
            resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                buffer: bf,
                offset: off,
                size: std::num::NonZeroU64::new(len),
            }),
        }
    }
    let qhead = (b * hd * 4) as u64;
    let khead = (n * hd * 4) as u64;
    let sc_len = (b * n * 4) as u64;
    let hpk = nh / nkv;
    let mut enc = dev.create_command_encoder(&wgpu::CommandEncoderDescriptor {
        label: Some("chunk-attend"),
    });
    for h in 0..nh {
        let kv = (h / hpk) as u64;
        let bg_qk = dev.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("ca-qk"),
            layout: &c.dit_qk.get_bind_group_layout(0),
            entries: &[
                slot(&qb, h as u64 * qhead, qhead, 0),
                slot(&kb, kv * khead, khead, 1),
                slot(&scb, 0, sc_len, 2),
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: p_qk.as_entire_binding(),
                },
            ],
        });
        let bg_sm = dev.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("ca-sm"),
            layout: &c.dit_softmax.get_bind_group_layout(0),
            entries: &[
                slot(&scb, 0, sc_len, 2),
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: p_sm.as_entire_binding(),
                },
            ],
        });
        let bg_pv = dev.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("ca-pv"),
            layout: &c.dit_pv.get_bind_group_layout(0),
            entries: &[
                slot(&scb, 0, sc_len, 0),
                slot(&vb, kv * khead, khead, 1),
                slot(&pb, h as u64 * qhead, qhead, 2),
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: p_pv.as_entire_binding(),
                },
            ],
        });
        // One pass per stage: each reads what the previous wrote to the
        // same scores buffer, and dispatches inside one pass do not
        // order against each other.
        {
            let mut pass = begin_pass_with(&mut enc, Some("ca-qk"), None);
            pass.set_pipeline(&c.dit_qk);
            pass.set_bind_group(0, &bg_qk, &[]);
            pass.dispatch_workgroups((n as u32).div_ceil(64), (b as u32).div_ceil(64), 1);
        }
        {
            let mut pass = begin_pass_with(&mut enc, Some("ca-sm"), None);
            pass.set_pipeline(&c.dit_softmax);
            pass.set_bind_group(0, &bg_sm, &[]);
            pass.dispatch_workgroups(b as u32, 1, 1);
        }
        {
            let mut pass = begin_pass_with(&mut enc, Some("ca-pv"), None);
            pass.set_pipeline(&c.dit_pv);
            pass.set_bind_group(0, &bg_pv, &[]);
            pass.dispatch_workgroups((hd as u32).div_ceil(64), (b as u32).div_ceil(64), 1);
        }
    }
    {
        let bg_un = dev.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("ca-un"),
            layout: &c.dit_unstack.get_bind_group_layout(0),
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: pb.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: ab.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: p_un.as_entire_binding(),
                },
            ],
        });
        let mut pass = begin_pass_with(&mut enc, Some("ca-un"), None);
        pass.set_pipeline(&c.dit_unstack);
        pass.set_bind_group(0, &bg_un, &[]);
        pass.dispatch_workgroups_flat(((nh * b * hd) as u32).div_ceil(256));
    }
    readback(
        c,
        enc,
        &ab,
        &stage,
        (b * nh * hd * 4) as u64,
        &mut out[..b * nh * hd],
    )
}

/// Fused QKV on wgpu: one upload of the normed chunk, three GEMMs, one
/// readback of Q|K|V laid out back to back. The unfused route pays three
/// submits and three uploads of the same X — at a 512-token chunk that
/// is the same 6 MB shipped three times, 44 times per prefill.
/// Weights stay cached in VRAM. `out` receives q (b·rq), then k (b·rk),
/// then v (b·rv).
#[allow(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 {
    let Some(c) = ctx() else { return false };
    if cols % 32 != 0 || b == 0 {
        return false;
    }
    let need = b * (rq + rk + rv);
    if xs.len() < b * cols || out.len() < need {
        return false;
    }
    let bytes = model.primary_bytes();
    let gpr = cols / 32;
    let wbuf = |idx: usize, rows: usize| -> Option<wgpu::Buffer> {
        let entry = &model.tensors[idx];
        if entry.shape.first().copied().unwrap_or(0) < rows {
            return None;
        }
        let abs = model.entry_abs_offset(entry)?;
        let plen = entry.nbytes as usize;
        if plen < rows * gpr * 18 || abs + plen > bytes.len() {
            return None;
        }
        weight_buffer_l(
            c,
            (model.uid() as usize, idx),
            &bytes[abs..abs + plen],
            layer_of_name(&model.tensors[idx].name),
        )
    };
    let (Some(bq), Some(bk), Some(bv)) = (wbuf(wq, rq), wbuf(wk, rk), wbuf(wv, rv)) else {
        return false;
    };

    let dev = &c.device;
    let mut sc = c.scratch.lock().unwrap();
    let st = wgpu::BufferUsages::STORAGE;
    let xs_buf = Scratch::ensure(
        dev,
        &mut sc.xs,
        (b * cols * 4) as u64,
        st | wgpu::BufferUsages::COPY_DST,
        "qkv-xs",
    );
    c.queue
        .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&xs[..b * cols]));
    let y_size = (need * 4) as u64;
    let y_buf = Scratch::ensure(
        dev,
        &mut sc.y,
        y_size,
        st | wgpu::BufferUsages::COPY_SRC,
        "qkv-y",
    );
    let stage = Scratch::ensure(
        dev,
        &mut sc.stage,
        y_size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "qkv-stage",
    );
    drop(sc);

    let params = |rows: usize| -> wgpu::Buffer {
        let raw = [(cols / 4) as u32, rows as u32, b as u32, 0u32];
        let bf = dev.create_buffer(&wgpu::BufferDescriptor {
            label: Some("qkv-params"),
            size: 16,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        c.queue.write_buffer(&bf, 0, bytemuck::cast_slice(&raw));
        bf
    };
    let layout = c.q4t_mm.get_bind_group_layout(0);
    let mut enc =
        dev.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("qkv") });
    let mut off = 0u64;
    for (wbf, rows) in [(&bq, rq), (&bk, rk), (&bv, rv)] {
        let pbf = params(rows);
        let bg = dev.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("qkv-bg"),
            layout: &layout,
            entries: &[
                bind_buf(0, wbf),
                bind_buf(1, &xs_buf),
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                        buffer: &y_buf,
                        offset: off,
                        size: std::num::NonZeroU64::new((b * rows * 4) as u64),
                    }),
                },
                bind_buf(3, &pbf),
            ],
        });
        let mut pass = begin_pass_with(&mut enc, Some("qkv-mm"), None);
        pass.set_pipeline(&c.q4t_mm);
        pass.set_bind_group(0, &bg, &[]);
        pass.dispatch_workgroups(
            (rows as u32).div_ceil(64).min(MAX_WG),
            (b as u32).div_ceil(64),
            1,
        );
        drop(pass);
        off += (b * rows * 4) as u64;
    }
    readback(c, enc, &y_buf, &stage, y_size, &mut out[..need])
}

/// Fused DiT SwiGLU FFN on wgpu: g=X·W1ᵀ, u=X·W3ᵀ, silu(g)·u,
/// y=·W2ᵀ — four passes, ONE submission, one readback. The unfused
/// per-op route pays 3 submits and ships the [b, inter] intermediates
/// across PCIe twice; on discrete cards that overhead dominates the
/// GEMM itself. Weights stay cached in VRAM.
#[allow(clippy::too_many_arguments)]
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 {
    ffn_q4(model, w1, w3, w2, xs, b, hidden, inter, false, out)
}

/// The q4tp twin. Same three GEMMs, same scratch, the ladder layout's
/// kernel — without it every Vulkan/DX12/Android box ran a q4tp image
/// model's diffusion transformer entirely on the CPU: the facade had no
/// wgpu arm to offer and said so by returning false.
#[allow(clippy::too_many_arguments)]
/// SwiGLU FFN whose fc1 emits gate and up PACKED in one row — the
/// MiniMax-H3 DiT's layout — with everything resident between the two
/// GEMMs. Before this, the intermediate crossed the bus twice per block
/// (at render size the gate/up panel alone is ~660 MB down and the
/// activation ~330 MB back up); now one upload of x and one readback of
/// the result. `false` = shapes or dtypes outside the contract, and the
/// host loop runs as before.
#[allow(clippy::too_many_arguments)]
/// The SwiGLU fold over a packed `[gate|up]` panel that is already on the
/// card, into a fresh compact `[b][inter]` panel. The four-bit FFN does this
/// inside its own encoder; the int8 pair needs it as a step of its own.
fn silu_packed_keep(
    c: &Ctx,
    gu: &wgpu::Buffer,
    b: usize,
    inter: usize,
    bias: Option<&[f32]>,
) -> Option<wgpu::Buffer> {
    let n = b.checked_mul(inter)?;
    if n == 0 {
        return None;
    }
    let act = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("pffn8-act"),
        size: (n * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE,
        mapped_at_creation: false,
    });
    let bias = bias.filter(|v| v.len() >= 2 * inter);
    let ps = uniform_u32x4(c, [n as u32, inter as u32, u32::from(bias.is_some()), 0]);
    let bbuf = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("pffn8-bias"),
            contents: bytemuck::cast_slice(bias.unwrap_or(&[0.0f32])),
            usage: wgpu::BufferUsages::STORAGE,
        });
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("pffn8-silu"),
        layout: &c.ffn_silu_packed.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, gu),
            bind_buf(1, &act),
            bind_buf(2, &ps),
            bind_buf(3, &bbuf),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("pffn8-silu"),
        });
    {
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.ffn_silu_packed);
        pass.set_bind_group(0, &bind, &[]);
        let wgs = (n as u32).div_ceil(256);
        pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
    }
    c.queue.submit(Some(enc.finish()));
    Some(act)
}

/// The FFN pair for an int8 container: fc1 from the host activation, the
/// SwiGLU fold, and fc2 — with the panel between them left on the card.
///
/// The four-bit twin below is a single kernel because a four-bit weight has to
/// be unpacked before the matrix units can touch it, and that unpack is worth
/// hiding inside the GEMM. An int8 weight needs no unpack, so the same work is
/// just two GEMMs; what mattered was never reading the intermediate home.
pub fn q8_ffn_packed(
    model: &Arc<CmfModel>,
    w1: usize,
    w2: usize,
    xs: &[f32],
    b: usize,
    hidden: usize,
    inter: usize,
    bias: Option<&[f32]>,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if b == 0 || xs.len() < b * hidden || out.len() < b * hidden {
        return false;
    }
    let Some(gu) = fused_panel_keep(model, w1, xs, b, 2 * inter, hidden) else {
        return false;
    };
    let Some(act) = silu_packed_keep(c, &gu, b, inter, bias) else {
        return false;
    };
    fused_gemm_from_device(model, w2, &act, b, hidden, inter, out)
}

/// The FFN pair for whichever codec the container is packed in.
#[allow(clippy::too_many_arguments)]
pub fn ffn_packed(
    model: &Arc<CmfModel>,
    w1: usize,
    w2: usize,
    xs: &[f32],
    b: usize,
    hidden: usize,
    inter: usize,
    bias: Option<&[f32]>,
    out: &mut [f32],
) -> bool {
    use cortiq_core::TensorDtype as D;
    // The [gate|up] panel is b·2·inter floats and a storage binding stops at
    // 2 GiB. A 321-frame clip rendered bidirectionally asks for exactly that
    // and the render died at a validation error rather than at anything
    // physical, so the rows are split into passes that fit. The FFN is
    // row-wise; splitting it changes nothing but the buffer sizes.
    //
    // The budget is a gigabyte, not the limit itself: the scratch allocator
    // rounds capacity up to the next power of two, so a panel of 2 GiB minus
    // a row still gets a 2 GiB buffer — and the bind group binds the whole
    // buffer. Anything above 2^30 rounds into the wall.
    let per = (PANEL_BUDGET_BYTES / (2 * inter * 4)).max(1);
    if b > per {
        for r0 in (0..b).step_by(per) {
            let r1 = (r0 + per).min(b);
            let ok = ffn_packed(
                model,
                w1,
                w2,
                &xs[r0 * hidden..r1 * hidden],
                r1 - r0,
                hidden,
                inter,
                bias,
                &mut out[r0 * hidden..r1 * hidden],
            );
            if !ok {
                return false;
            }
        }
        return true;
    }
    let t_ffn = std::time::Instant::now();
    let took = match model.tensors.get(w1).map(|e| e.dtype) {
        Some(D::Q4TiledP) => q4tp_ffn_packed(model, w1, w2, xs, b, hidden, inter, bias, out),
        Some(D::Q8Row | D::Q8_2f) if fused_any() => {
            q8_ffn_packed(model, w1, w2, xs, b, hidden, inter, bias, out)
        }
        _ => false,
    };
    // Calls and microseconds per shape. The 3D VAE's FFN phase costs the same
    // 24 s at 512x288 as at 768x448, which means the cost does not scale with
    // the work — so the question is how many times this is called and what
    // each call costs, not what the kernel does inside.
    if std::env::var("CMF_FFN_COUNT").is_ok() {
        use std::sync::atomic::{AtomicU64, Ordering};
        static CALLS: AtomicU64 = AtomicU64::new(0);
        static MICROS: AtomicU64 = AtomicU64::new(0);
        let n = CALLS.fetch_add(1, Ordering::Relaxed) + 1;
        let us = MICROS.fetch_add(t_ffn.elapsed().as_micros() as u64, Ordering::Relaxed)
            + t_ffn.elapsed().as_micros() as u64;
        if n % 200 == 0 {
            eprintln!(
                "ffn_packed: {n} calls, {:.1} s total, {:.2} ms each (last {hidden}x{inter}, b={b})",
                us as f64 / 1e6,
                us as f64 / 1e3 / n as f64
            );
        }
    }
    // A fused FFN that quietly refuses sends the caller to a path that costs
    // several times as much, and the only trace is a slower render. Say it
    // once per shape, so a profile can be read against what actually ran.
    if std::env::var("CMF_GPU_DEBUG").is_ok() {
        use std::collections::HashSet;
        use std::sync::Mutex;
        static SEEN: Mutex<Option<HashSet<(usize, usize, bool)>>> = Mutex::new(None);
        let mut g = SEEN.lock().unwrap();
        if g.get_or_insert_with(HashSet::new)
            .insert((hidden, inter, took))
        {
            eprintln!(
                "ffn_packed: {} for {hidden}x{inter} (b={b}, dtype={:?})",
                if took { "fused" } else { "REFUSED" },
                model.tensors.get(w1).map(|e| e.dtype)
            );
        }
    }
    took
}

/// What one FFN pass may allocate. A storage binding stops at 2 GiB minus four
/// bytes on this class of card, and the scratch allocator rounds capacity up
/// to the next power of two — so the largest panel that does not round into
/// that wall is a gigabyte.
const PANEL_BUDGET_BYTES: usize = 1024 * 1024 * 1024;

/// `CMF_FUSED_ANY=0` puts a non-four-bit container back on the per-op path —
/// the A/B switch for what the codec-agnostic fusion is worth.
pub(crate) fn fused_any() -> bool {
    std::env::var("CMF_FUSED_ANY").as_deref() != Ok("0")
}

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 {
    let Some(c) = ctx() else { return false };
    let _gate = c.mm_gate.lock().unwrap();
    let t_stage = std::time::Instant::now();
    if hidden % 32 != 0 || inter % 32 != 0 || b == 0 {
        return false;
    }
    if xs.len() < b * hidden || out.len() < b * hidden {
        return false;
    }
    let bytes = model.primary_bytes();
    let wbuf = |idx: usize, rows: usize, cols: usize| -> Option<wgpu::Buffer> {
        let e = model.tensors.get(idx)?;
        if e.dtype != cortiq_core::TensorDtype::Q4TiledP
            || e.shape.first().copied()? != rows
            || e.shape.get(1).copied()? != cols
        {
            return None;
        }
        let abs = model.entry_abs_offset(e)?;
        let plen = e.nbytes as usize;
        if abs + plen > bytes.len() {
            return None;
        }
        weight_buffer_l(
            c,
            (model.uid() as usize, idx),
            &bytes[abs..abs + plen],
            layer_of_name(&model.tensors[idx].name),
        )
    };
    let (Some(w1b), Some(w2b)) = (wbuf(w1, 2 * inter, hidden), wbuf(w2, hidden, inter)) else {
        return false;
    };
    let mut sc = c.scratch.lock().unwrap();
    let st = wgpu::BufferUsages::STORAGE;
    let xs_buf = Scratch::ensure(
        &c.device,
        &mut sc.xs,
        (b * hidden * 4) as u64,
        st | wgpu::BufferUsages::COPY_DST,
        "pffn-x",
    );
    let gu_buf = Scratch::ensure(
        &c.device,
        &mut sc.g,
        (b * 2 * inter * 4) as u64,
        st,
        "pffn-gu",
    );
    let act_buf = Scratch::ensure(&c.device, &mut sc.u, (b * inter * 4) as u64, st, "pffn-act");
    let y_size = (b * hidden * 4) as u64;
    let y_buf = Scratch::ensure(
        &c.device,
        &mut sc.y,
        y_size,
        st | wgpu::BufferUsages::COPY_SRC,
        "pffn-y",
    );
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        y_size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "pffn-stage",
    );
    // MEASURED AND REJECTED, so opt-in (`CMF_PLANE_CACHE_MB=<n>`): keeping the
    // unpacked plane across calls looked like the obvious fix for the 3D VAE,
    // which calls this with the same 33 M-weight tensor over and over while
    // the four-bit unpack — a 5-bit rung ladder per group — is paid every
    // time. At 768×448 its FFN phase is 24-26 s against 5.4 s for the same
    // FFN in an eight-bit container.
    //
    // Eight runs say the cache is not it. Denoise, alternating arms:
    // off 72.9 / 71.1 / 69.5 / 76.6, on 60.6 / 71.1 / 73.6 / 76.2. One fast
    // run out of four, and the spread WITHIN an arm is 10% — larger than
    // anything being claimed. The VAE's FFN phase also costs the same 24 s
    // at 512×288 as at 768×448, which says its cost does not scale with the
    // work: what is left to chase is per-call host overhead, not the unpack.
    let key1 = (model.uid() as usize, w1);
    let dq1 = plane_cached(c, key1, &w1b, 2 * inter, hidden, 0)
        .map(|(p, b)| (p, b))
        .or_else(|| dq_f16_plane(c, &mut sc, &w1b, 2 * inter, hidden).map(|(p, b)| (p, Some(b))));
    // Taken under the guard this function still holds, for the device-side
    // max|x| below — `encode_act_absmax` would take the same lock again.
    let amax_parts = Scratch::ensure(
        &c.device,
        &mut sc.amaxp,
        512 * 4,
        wgpu::BufferUsages::STORAGE,
        "pffn-amax-parts",
    );
    drop(sc);
    c.queue
        .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&xs[..b * hidden]));
    let bias = bias.filter(|v| v.len() >= 2 * inter);
    let ps2 = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[
                (b * inter) as u32,
                inter as u32,
                u32::from(bias.is_some()),
                0u32,
            ]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let bbuf = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("pffn-bias"),
            contents: bytemuck::cast_slice(bias.unwrap_or(&[0.0f32])),
            usage: wgpu::BufferUsages::STORAGE,
        });
    let bg_silu = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("pffn-silu-bg"),
        layout: &c.ffn_silu_packed.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, &gu_buf),
            bind_buf(1, &act_buf),
            bind_buf(2, &ps2),
            bind_buf(3, &bbuf),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("pffn"),
        });
    // fc1 over the dequantized plane when the tensor cores are up. A plane
    // that came from the cache is already unpacked — encoding the pass again
    // would redo the very work the cache exists to skip.
    if let (Some((_, Some(bind_dq))), Some(pipe_dq)) = (&dq1, c.q4tp_dq_f16.as_ref()) {
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(pipe_dq);
        pass.set_bind_group(0, bind_dq, &[]);
        let wgs = ((2 * inter * hidden / 2) as u32).div_ceil(256);
        pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
    }
    let (p1, w1_bind) = match (&dq1, c.q4tp_mm_coop_f16.as_ref()) {
        (Some((plane, _)), Some(pipe)) => (pipe, plane.clone()),
        _ => (mm_pipeline(c, true, false), w1b.clone()),
    };
    // f16 operands cap at 65504 and a DiT's modulated activations pass
    // it. The scale is computed from the input we just uploaded; the
    // kernel divides it back out at the store.
    let f16_arm = dq1.is_some() && c.q4tp_mm_coop_f16.is_some();
    // The activation scale is taken on the card, where the resident-operand
    // path already computes it, instead of scanning the panel on the host —
    // 13.2 M floats per call at DiT shapes, on one core, before anything is
    // submitted. Three quarters of a four-bit FFN call was host time.
    //
    // MEASURED, RTX 5090, 768×448, two runs each way: host time per call
    // 116.4 / 118.6 ms with the host scan against 77.0 ms with this, and the
    // render's denoise 68.1 / 71.1 s against 62.0 / 62.8 s. Ten percent, with
    // the output byte-identical across all four. `CMF_FFN_HOST_SCAN=1` keeps
    // the old arm.
    let dev_scan = f16_arm
        && std::env::var("CMF_FFN_HOST_SCAN").as_deref() != Ok("1")
        && (c.act_absmax.is_some() || (c.act_amax_part.is_some() && c.act_amax_fold.is_some()));
    let ascale = if dev_scan {
        0.0
    } else {
        let mx = xs[..b * hidden]
            .iter()
            .fold(0f32, |m, v| if v.is_finite() { m.max(v.abs()) } else { m });
        if mx > 1000.0 {
            1000.0 / mx
        } else {
            1.0
        }
    };
    let asc_dev = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("pffn-ascale-1"),
        size: 4,
        usage: wgpu::BufferUsages::STORAGE,
        mapped_at_creation: false,
    });
    let host_scale = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[ascale]),
            usage: wgpu::BufferUsages::STORAGE,
        });
    if dev_scan {
        encode_act_absmax_with(
            c,
            &mut enc,
            &xs_buf,
            b * hidden,
            &asc_dev,
            Some(&amax_parts),
        );
    }
    encode_q4_tile_mm_full(
        c,
        &mut enc,
        p1,
        &w1_bind,
        &xs_buf,
        &gu_buf,
        2 * inter,
        hidden,
        b,
        ascale,
        if dev_scan {
            Some(&asc_dev)
        } else {
            f16_arm.then_some(&host_scale)
        },
    );
    {
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.ffn_silu_packed);
        pass.set_bind_group(0, &bg_silu, &[]);
        // 2-D grid: b·inter/256 passes 65 535 workgroups at render batch
        // sizes, and clamping x leaves the tail of the activations as
        // whatever the buffer held (measured: frames 92% smaller, i.e.
        // nearly blank, at 512×288 while 256×160 looked fine).
        let wgs = ((b * inter) as u32).div_ceil(256);
        pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
    }
    // fc2's input is the SwiGLU panel, which lives on the device — so
    // its activation scale is computed there too (`act_absmax` writes
    // 1/scale into one float) and the GEMM reads it from that buffer.
    // Without a scale the f16 operands overflow: measured as nearly
    // blank frames at 512×288 while 256×160 looked perfect. When the
    // reduction or the f16 twin is unavailable, the scalar arm runs —
    // it reads f32 and cannot overflow.
    let asc_buf = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("pffn-ascale"),
        size: 4,
        usage: wgpu::BufferUsages::STORAGE,
        mapped_at_creation: false,
    });
    // MEASURED NEUTRAL AT BOTH SIZES, so opt-in: putting fc2 on the matrix
    // units costs a second plane unpack plus a max-reduction over the
    // activation panel. On MiniMax-H3, 22 frames, four steps:
    //
    //   512×288 (b ≈ 1.1k rows): 137.4 s without, 139.8 s with.
    //   768×448 (b = 2449 rows): one alternating pair said 75.3 s of denoise
    //     without against 70.9 s with — and the next pair, back to back,
    //     said 71.3 s without against 70.9 s with. The first pair was the
    //     stand drifting (its container pages had gone cold), not the
    //     kernel. A default set from that would have been noise shipped.
    //
    // `CMF_FFN_FC2_COOP=1` takes it; a multi-workgroup reduction is what
    // would tip it. Row counts are from the render itself — CMF_GPU_DEBUG=1
    // prints them — not from arithmetic on the frame size.
    let want_fc2_coop = std::env::var("CMF_FFN_FC2_COOP").as_deref() == Ok("1");
    let dq2 = if want_fc2_coop {
        let mut sc = c.scratch.lock().unwrap();
        dq_f16_plane_slot(c, &mut sc, &w2b, hidden, inter, true)
    } else {
        None
    };
    // Putting fc2 on the matrix units costs a second plane unpack plus a
    // max-reduction over the activation panel, and whether that pays depends
    // The reduction that feeds the f16 arm its activation scale comes in two
    // shapes on this backend — a single-workgroup `act_absmax` and a
    // two-stage `act_amax_part`/`act_amax_fold`. The gate below asked for the
    // first one by name, so on a card that has only the second the fc2 plane
    // never ran however the switch was set. That is why it kept measuring
    // neutral: the arm under test was not the arm being taken.
    let have_reduction =
        c.act_absmax.is_some() || (c.act_amax_part.is_some() && c.act_amax_fold.is_some());
    match (
        (have_reduction && want_fc2_coop).then_some(()),
        &dq2,
        c.q4tp_mm_coop_f16.as_ref(),
    ) {
        (Some(()), Some((plane2, bind_dq2)), Some(pipe)) => {
            {
                let mut pass = begin_pass(&mut enc);
                pass.set_pipeline(c.q4tp_dq_f16.as_ref().unwrap());
                pass.set_bind_group(0, bind_dq2, &[]);
                let wgs = ((hidden * inter / 2) as u32).div_ceil(256);
                pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
            }
            let ap = c
                .device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytemuck::cast_slice(&[(b * inter) as u32, 0u32, 0u32, 0u32]),
                    usage: wgpu::BufferUsages::UNIFORM,
                });
            // Two stages when both halves are available: N workgroups
            // read the panel in parallel (one SM cannot saturate a card's
            // bandwidth), then a single fold turns the partials into the
            // scale. Falls back to the one-workgroup form otherwise.
            match (c.act_amax_part.as_ref(), c.act_amax_fold.as_ref()) {
                (Some(part), Some(fold)) => {
                    let nparts = 512usize;
                    let pbuf = {
                        let mut sc = c.scratch.lock().unwrap();
                        Scratch::ensure(
                            &c.device,
                            &mut sc.amaxp,
                            (nparts * 4) as u64,
                            wgpu::BufferUsages::STORAGE,
                            "pffn-amax-parts",
                        )
                    };
                    let bg1 = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &part.get_bind_group_layout(0),
                        entries: &[bind_buf(0, &act_buf), bind_buf(1, &pbuf), bind_buf(2, &ap)],
                    });
                    let fp = c
                        .device
                        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                            label: None,
                            contents: bytemuck::cast_slice(&[nparts as u32, 0u32, 0u32, 0u32]),
                            usage: wgpu::BufferUsages::UNIFORM,
                        });
                    let bg2 = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &fold.get_bind_group_layout(0),
                        entries: &[bind_buf(0, &pbuf), bind_buf(1, &asc_buf), bind_buf(2, &fp)],
                    });
                    let mut pass = begin_pass(&mut enc);
                    pass.set_pipeline(part);
                    pass.set_bind_group(0, &bg1, &[]);
                    pass.dispatch_workgroups(nparts as u32, 1, 1);
                    drop(pass);
                    let mut pass = begin_pass(&mut enc);
                    pass.set_pipeline(fold);
                    pass.set_bind_group(0, &bg2, &[]);
                    pass.dispatch_workgroups(1, 1, 1);
                }
                _ => {
                    // The one-workgroup form, for a card without the split
                    // reduction. `have_reduction` guarantees one of the two.
                    let Some(amax) = c.act_absmax.as_ref() else {
                        return false;
                    };
                    let bg_amax = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &amax.get_bind_group_layout(0),
                        entries: &[
                            bind_buf(0, &act_buf),
                            bind_buf(1, &asc_buf),
                            bind_buf(2, &ap),
                        ],
                    });
                    let mut pass = begin_pass(&mut enc);
                    pass.set_pipeline(amax);
                    pass.set_bind_group(0, &bg_amax, &[]);
                    pass.dispatch_workgroups(1, 1, 1);
                }
            }
            encode_q4_tile_mm_full(
                c,
                &mut enc,
                pipe,
                plane2,
                &act_buf,
                &y_buf,
                hidden,
                inter,
                b,
                0.0,
                Some(&asc_buf),
            );
        }
        _ => encode_q4_tile_mm(
            c, &mut enc, &c.q4tp_mm, &w2b, &act_buf, &y_buf, hidden, inter, b,
        ),
    }
    // Where one call spends itself, under CMF_FFN_COUNT: everything before the
    // readback is host work — buffer creation, bind groups, the host scan for
    // max|x| — and the readback is the wait for the card.
    //
    // MEASURED, RTX 5090, 768×448: the DiT shape (5376×14336, b=2447) costs
    // 126.2 ms of host against 41.4 ms of card, and the 3D VAE's
    // (2048×8192, b=1797) 100.2 ms of host against 12.6 ms of card. The
    // eight-bit container's whole call at that same VAE shape is 14 ms — so
    // the four-bit GPU half is already as fast as the int8 path, and the gap
    // everyone has been chasing is host-side preparation, three quarters of
    // it. That is the thread to pull: fewer buffers and bind groups per call,
    // and the max|x| scan moved to the card where the resident-operand path
    // already computes it.
    let t_host = t_stage.elapsed();
    let ok = readback(c, enc, &y_buf, &stage_buf, y_size, &mut out[..b * hidden]);
    if std::env::var("CMF_FFN_COUNT").is_ok() {
        use std::sync::atomic::{AtomicU64, Ordering};
        static HOST: AtomicU64 = AtomicU64::new(0);
        static WAIT: AtomicU64 = AtomicU64::new(0);
        static N: AtomicU64 = AtomicU64::new(0);
        let total = t_stage.elapsed();
        let h = HOST.fetch_add(t_host.as_micros() as u64, Ordering::Relaxed)
            + t_host.as_micros() as u64;
        let w = WAIT.fetch_add((total - t_host).as_micros() as u64, Ordering::Relaxed)
            + (total - t_host).as_micros() as u64;
        let n = N.fetch_add(1, Ordering::Relaxed) + 1;
        if n % 200 == 0 {
            eprintln!(
                "q4tp_ffn split: {n} calls, host {:.1} ms each, card {:.1} ms each ({hidden}x{inter}, b={b})",
                h as f64 / 1e3 / n as f64,
                w as f64 / 1e3 / n as f64
            );
        }
    }
    ok
}

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 {
    ffn_q4(model, w1, w3, w2, xs, b, hidden, inter, true, out)
}

#[allow(clippy::too_many_arguments)]
fn ffn_q4(
    model: &Arc<CmfModel>,
    w1: usize,
    w3: usize,
    w2: usize,
    xs: &[f32],
    b: usize,
    hidden: usize,
    inter: usize,
    q4tp: bool,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if hidden % 32 != 0 || inter % 32 != 0 || b == 0 {
        return false;
    }
    let bytes = model.primary_bytes();
    let wbuf = |idx: usize, rows: usize, cols: usize| -> Option<wgpu::Buffer> {
        let entry = &model.tensors[idx];
        let abs = model.entry_abs_offset(entry)?;
        let plen = entry.nbytes as usize;
        let want = if q4tp {
            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
                .unwrap_or(usize::MAX)
        } else {
            rows * (cols / 32) * 18
        };
        if plen < want || abs + plen > bytes.len() {
            return None;
        }
        weight_buffer_l(
            c,
            (model.uid() as usize, idx),
            &bytes[abs..abs + plen],
            layer_of_name(&model.tensors[idx].name),
        )
    };
    let (Some(q1), Some(q3), Some(q2)) = (
        wbuf(w1, inter, hidden),
        wbuf(w3, inter, hidden),
        wbuf(w2, hidden, inter),
    ) else {
        return false;
    };
    if xs.len() < b * hidden || out.len() < b * hidden {
        return false;
    }
    let mut sc = c.scratch.lock().unwrap();
    let xs_buf = Scratch::ensure(
        &c.device,
        &mut sc.xs,
        (b * hidden * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "q4tffn-xs",
    );
    c.queue
        .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&xs[..b * hidden]));
    let panel = (b * inter * 4) as u64;
    let g_buf = Scratch::ensure(
        &c.device,
        &mut sc.g,
        panel,
        wgpu::BufferUsages::STORAGE,
        "q4tffn-g",
    );
    let u_buf = Scratch::ensure(
        &c.device,
        &mut sc.u,
        panel,
        wgpu::BufferUsages::STORAGE,
        "q4tffn-u",
    );
    let y_size = (b * hidden * 4) as u64;
    let y_buf = Scratch::ensure(
        &c.device,
        &mut sc.y,
        y_size,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "q4tffn-y",
    );
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        y_size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "q4tffn-stage",
    );
    // Content-keyed uniforms: three shapes live in one submission, so
    // one rewritable buffer cannot serve them.
    let p13 = uniform_u32x4(c, [(hidden / 4) as u32, inter as u32, b as u32, 0]);
    let p2 = uniform_u32x4(c, [(inter / 4) as u32, hidden as u32, b as u32, 0]);
    let psilu = uniform_u32x4(c, [(b * inter) as u32, 0, 0, 0]);
    let mm = mm_pipeline(c, q4tp, false);
    let mm_layout = mm.get_bind_group_layout(0);
    let bind_mm = |q: &wgpu::Buffer, x: &wgpu::Buffer, y: &wgpu::Buffer, p: &wgpu::Buffer| {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("q4tffn-bg"),
            layout: &mm_layout,
            entries: &[
                bind_buf(0, q),
                bind_buf(1, x),
                bind_buf(2, y),
                bind_buf(3, p),
            ],
        })
    };
    let bg1 = bind_mm(&q1, &xs_buf, &g_buf, &p13);
    let bg3 = bind_mm(&q3, &xs_buf, &u_buf, &p13);
    let bg2 = bind_mm(&q2, &g_buf, &y_buf, &p2);
    let bg_silu = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q4tffn-silu-bg"),
        layout: &c.ffn_silu.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, &g_buf),
            bind_buf(1, &u_buf),
            bind_buf(2, &psilu),
        ],
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("q4tffn"),
        });
    let mm_pass = |enc: &mut wgpu::CommandEncoder, bg: &wgpu::BindGroup, rows: usize| {
        let mut pass = begin_pass_with(enc, Some("q4tffn-mm"), None);
        pass.set_pipeline(mm);
        pass.set_bind_group(0, bg, &[]);
        pass.dispatch_workgroups(
            (rows as u32).div_ceil(64).min(MAX_WG),
            (b as u32).div_ceil(64),
            1,
        );
    };
    mm_pass(&mut enc, &bg1, inter);
    mm_pass(&mut enc, &bg3, inter);
    {
        let mut pass = begin_pass_with(&mut enc, Some("q4tffn-silu"), None);
        pass.set_pipeline(&c.ffn_silu);
        pass.set_bind_group(0, &bg_silu, &[]);
        // 2-D grid: b·inter/256 passes 65 535 workgroups at render batch
        // sizes, and clamping x leaves the tail of the activations as
        // whatever the buffer held (measured: frames 92% smaller, i.e.
        // nearly blank, at 512×288 while 256×160 looked fine).
        let wgs = ((b * inter) as u32).div_ceil(256);
        pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
    }
    mm_pass(&mut enc, &bg2, hidden);
    readback(c, enc, &y_buf, &stage_buf, y_size, &mut out[..b * hidden])
}

pub fn q1t_matmat(
    model: &Arc<CmfModel>,
    idx: usize,
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    let gpr = cols / 32;
    if cols % 32 != 0 || rows == 0 || b == 0 {
        return false;
    }
    let entry = &model.tensors[idx];
    if entry.shape.first().copied().unwrap_or(0) < rows {
        return false;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return false;
    };
    let bytes = model.primary_bytes();
    let plen = entry.nbytes as usize;
    if plen < rows * gpr * 9
        || abs + plen > bytes.len()
        || xs.len() < b * cols
        || out.len() < b * rows
    {
        return false;
    }
    dispatch_q1t_mm(
        c,
        {
            note_layer((model.uid() as usize, idx), &model.tensors[idx].name);
            Some((model.uid() as usize, idx))
        },
        &bytes[abs..abs + plen],
        xs,
        b,
        rows,
        cols,
        out,
    )
}

#[allow(clippy::too_many_arguments)]
fn dispatch_q1t_mm(
    c: &Ctx,
    weight_key: Option<(usize, usize)>,
    payload: &[u8],
    xs: &[f32],
    b: usize,
    rows: usize,
    cols: usize,
    out: &mut [f32],
) -> bool {
    let q_buf = match weight_key {
        Some(k) => match weight_buffer(c, k, payload) {
            Some(bf) => bf,
            None => return false,
        },
        None => c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("q1tmm-weights"),
                contents: payload,
                usage: wgpu::BufferUsages::STORAGE,
            }),
    };
    let mut sc = c.scratch.lock().unwrap();
    let xs_buf = Scratch::ensure(
        &c.device,
        &mut sc.xs,
        (b * cols * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "q1tmm-xs",
    );
    c.queue
        .write_buffer(&xs_buf, 0, bytemuck::cast_slice(&xs[..b * cols]));
    let y_size = (b * rows * 4) as u64;
    let y_buf = Scratch::ensure(
        &c.device,
        &mut sc.y,
        y_size,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "q1tmm-y",
    );
    let params = [(cols / 4) as u32, rows as u32, b as u32, 0u32];
    let p_buf = match &sc.params {
        Some(bf) => bf.clone(),
        None => {
            let bf = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("q1tmm-params"),
                size: 16,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            sc.params = Some(bf.clone());
            bf
        }
    };
    c.queue
        .write_buffer(&p_buf, 0, bytemuck::cast_slice(&params));
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        y_size,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "q1tmm-stage",
    );
    let entries = [
        bind_buf(0, &q_buf),
        bind_buf(1, &xs_buf),
        bind_buf(2, &y_buf),
        bind_buf(3, &p_buf),
    ];
    let bind_mm = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q1tmm-bg"),
        layout: &c.q1t_mm.get_bind_group_layout(0),
        entries: &entries,
    });
    let bind_ov = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q1tov-bg"),
        layout: &c.q1t_ovmm.get_bind_group_layout(0),
        entries: &entries,
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("q1tmm"),
        });
    {
        let mut pass = begin_pass_with(&mut enc, Some("q1tmm"), None);
        pass.set_pipeline(&c.q1t_mm);
        pass.set_bind_group(0, &bind_mm, &[]);
        pass.dispatch_workgroups(
            (rows as u32).div_ceil(64).min(MAX_WG),
            (b as u32).div_ceil(64),
            1,
        );
    }
    {
        // Separate pass = a barrier, so the overlay reads the finished base.
        let mut pass = begin_pass_with(&mut enc, Some("q1tov"), None);
        pass.set_pipeline(&c.q1t_ovmm);
        pass.set_bind_group(0, &bind_ov, &[]);
        pass.dispatch_workgroups((rows as u32).div_ceil(64).min(MAX_WG), 1, 1);
    }
    let ok = readback(c, enc, &y_buf, &stage_buf, y_size, &mut out[..b * rows]);
    drop(sc);
    ok
}

/// Copy the output buffer GPU→staging→CPU (map+poll). Single readback path
/// for matvec/matmat.
/// Spin briefly for a submission to land instead of sleeping on it.
fn spin_wait() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| {
        std::env::var("CMF_GPU_SPIN")
            .map(|v| v != "0")
            .unwrap_or(true)
    })
}

// An adaptive version of the above was written and reverted. The idea:
// spinning before the block is a bet that the submission lands in
// microseconds, and on an Adreno 642L a readback waits 46.5 ms, so watch
// the waits and stop spinning when they are long. It is worth NOTHING
// here — four runs, alternating: 0.897 / 0.901 / 0.877 / 0.896 tok/s
// with the spin on, off, on, off. The 9% that started the chase was one
// outlier. And the self-switching version tripped on the laptop too,
// which would have taken the spin away from the platform where it was
// measured to pay. `CMF_GPU_SPIN=0` remains for anyone who wants to try.

/// Submissions to the device, all sites. A round trip costs a fence
/// whatever it carries, so "how many a token" is the number that decides
/// whether a decode step is compute-bound or latency-bound — and it is not
/// derivable from anything else the profile prints.
pub static SUBMITS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// Aggregate readback cost for the ordinary per-op path.  Unlike the
/// opt-in line trace, these counters make the bounded benchmark's fence and
/// mapped-copy cost joinable with its submit/pass counters without printing
/// once per projection.
pub static READBACK_CALLS: std::sync::atomic::AtomicU64 =
    std::sync::atomic::AtomicU64::new(0);
pub static READBACK_BYTES: std::sync::atomic::AtomicU64 =
    std::sync::atomic::AtomicU64::new(0);
pub static READBACK_WAIT_NS: std::sync::atomic::AtomicU64 =
    std::sync::atomic::AtomicU64::new(0);
pub static READBACK_COPY_NS: std::sync::atomic::AtomicU64 =
    std::sync::atomic::AtomicU64::new(0);
pub static READBACK_TOTAL_NS: std::sync::atomic::AtomicU64 =
    std::sync::atomic::AtomicU64::new(0);

/// Print the minimum per-op GPU accounting for a bounded decode.  GPU
/// timestamp queries remain a separate graph-only facility; on this path the
/// empty/one-dispatch round trip and fence/readback counters are the honest
/// measurements available without changing command ordering.
pub fn perf_report() {
    if std::env::var("CMF_PERF_PROFILE").as_deref() != Ok("1") {
        return;
    }
    let n = READBACK_CALLS.load(std::sync::atomic::Ordering::Relaxed);
    let submits = SUBMITS.load(std::sync::atomic::Ordering::Relaxed);
    let passes = PASSES.load(std::sync::atomic::Ordering::Relaxed);
    eprintln!(
        "[perf-wgpu] submits={} passes={} readback_calls={} readback_mb={:.3} readback_wait_ms={:.3} readback_copy_ms={:.3} readback_total_ms={:.3} weight_upload_ms={:.3} weight_upload_mb={:.3} dispatched_weight_mb={:.3} cache_hits={} cache_misses={}",
        submits,
        passes,
        n,
        READBACK_BYTES.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6,
        READBACK_WAIT_NS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6,
        READBACK_COPY_NS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6,
        READBACK_TOTAL_NS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6,
        UPLOAD_NS.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6,
        UPLOAD_BYTES.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6,
        WEIGHT_BYTES.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6,
        RES_HITS.load(std::sync::atomic::Ordering::Relaxed),
        RES_MISSES.load(std::sync::atomic::Ordering::Relaxed),
    );
}

/// What a round trip to this device costs, with nothing in it.
///
/// The number the whole mobile-GPU investigation needed and did not
/// have: if an EMPTY submit-and-wait already costs milliseconds, no
/// kernel work can be blamed for a slow token, and no amount of kernel
/// tuning can help. Returns (empty submit+fence, submit+dispatch+readback)
/// in milliseconds per iteration.
pub fn roundtrip_bench(n: usize) -> Option<(f64, f64)> {
    let c = ctx()?;
    let n = n.max(1);
    // 1. Submit an empty encoder and wait for the queue to drain.
    let t0 = std::time::Instant::now();
    for _ in 0..n {
        let enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("rt-empty"),
            });
        submit(c, finish_enc(enc));
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
    }
    let empty_ms = t0.elapsed().as_secs_f64() * 1e3 / n as f64;

    // 2. The shape the per-op path actually pays: one trivial dispatch,
    //    then a staged readback of its output.
    let y = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("rt-y"),
        size: 1024,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        mapped_at_creation: false,
    });
    let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("rt-stage"),
        size: 1024,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    let p_buf = uniform_u32x4(c, [256, 0, 0, 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &c.zero.get_bind_group_layout(0),
        entries: &[bind_buf(0, &y), bind_buf(1, &p_buf)],
    });
    let mut out = vec![0f32; 256];
    let t1 = std::time::Instant::now();
    for _ in 0..n {
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("rt-one"),
            });
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(&c.zero);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(1, 1, 1);
        }
        if !readback(c, enc, &y, &stage, 1024, &mut out) {
            return None;
        }
    }
    let one_ms = t1.elapsed().as_secs_f64() * 1e3 / n as f64;
    Some((empty_ms, one_ms))
}

/// Weight residency: how much went to the card and how long it took. The
/// first token pays all of it, and on a 92 GB expert stack that was half an
/// hour — worth knowing as a rate rather than as "the bench is slow to start".
pub static UPLOAD_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static UPLOAD_BYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// Staged upload (default): map a staging buffer and issue the copy here,
/// instead of handing the bytes to `queue.write_buffer`. `CMF_GPU_UPLOAD=map`
/// restores the historical write_buffer path.
fn upload_staged() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| {
        std::env::var("CMF_GPU_UPLOAD")
            .map(|v| v == "staged")
            .unwrap_or(true)
    })
}

/// Bound transient upload residency. Creating one mapped staging buffer as
/// large as every dense tensor made Vulkan's allocator retain a second,
/// size-fragmented copy of most of a model: the 15.6 GB Granite 30B Q4TP graph
/// occupied 33.5 GB, and its 29.3 GB Q8_2F twin OOMed a 46 GB A40 before the
/// fallback could run. Fixed-size chunks let the allocator reuse one small
/// class while the final device-local buffers remain budget-accounted.
fn upload_chunk_bytes() -> usize {
    static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
    *N.get_or_init(|| {
        std::env::var("CMF_GPU_UPLOAD_CHUNK_MB")
            .ok()
            .and_then(|v| v.parse::<usize>().ok())
            .unwrap_or(64)
            .max(1)
            * 1024
            * 1024
    })
}

#[inline]
fn note_submit(c: &Ctx) {
    SUBMITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    if slot_check() {
        c.slot_writes.lock().unwrap().clear();
    }
}

#[inline]
fn submit(c: &Ctx, buf: wgpu::CommandBuffer) {
    note_submit(c);
    c.queue.submit(Some(buf));
}

#[inline]
fn submit_empty(c: &Ctx) {
    note_submit(c);
    c.queue.submit(std::iter::empty());
}

/// `CMF_DSV4_SLOT_CHECK=1`: panic if a per-layer slot is written twice
/// before its submission. Off by default — it costs a lock per slot write —
/// but the toy gate runs with it on, because the invariant it guards is a
/// convention and conventions are what the 50.280 was made of.
fn slot_check() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| std::env::var("CMF_DSV4_SLOT_CHECK").is_ok_and(|v| v != "0"))
}

fn note_slot_write(c: &Ctx, kind: &str, key: (u8, u64, usize)) {
    if !slot_check() {
        return;
    }
    let mut m = c.slot_writes.lock().unwrap();
    let n = m.entry(key).or_insert(0);
    *n += 1;
    assert!(
        *n == 1,
        "{kind}-слот {key:?} записан {n} раз до одной отправки — \
         последняя запись достанется ВСЕМ проходам, которые его читают \
         (столкновение тегов или два слоя на один ключ)",
    );
}

/// Two device buffers, one staging buffer, one fence. The chain used to
/// read its folded vector and then call `dsv4_state_read` for the
/// hyper-connection state — two submissions and two map-waits for a token
/// that is otherwise a single submission.
/// Compute passes opened, all sites. On a decode step the dsv4 chain opens
/// about thirty a layer, and a pass costs the driver a fixed amount whatever
/// it dispatches — which is why a 32-layer TOY with 128-wide tensors still
/// waits 35 ms a token. Reported per token next to the submissions.
pub static PASSES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// Split a flat element dispatch across x and y: a dimension caps at
/// 65535 workgroups, and Lumina at 512x512 asks for 75060 of them for one
/// SwiGLU (2085 tokens x 9216 intermediate / 256). Shaders paired with
/// this read their index as `gid.x + gid.y * nwg.x * WG`, so the x-extent
/// stays whatever is chosen here, and a y of 1 leaves the old arithmetic
/// untouched for every caller still under the cap.
fn flat_groups(groups: u32) -> (u32, u32) {
    const MAX: u32 = 65535;
    if groups <= MAX {
        (groups, 1)
    } else {
        (MAX, groups.div_ceil(MAX))
    }
}

trait FlatDispatch {
    /// `dispatch_workgroups` for a flat element count, folded into two
    /// dimensions when it overflows one.
    fn dispatch_workgroups_flat(&mut self, groups: u32);
}

impl FlatDispatch for wgpu::ComputePass<'_> {
    fn dispatch_workgroups_flat(&mut self, groups: u32) {
        let (x, y) = flat_groups(groups);
        self.dispatch_workgroups(x, y, 1);
    }
}

/// The q4tp GEMM on cooperative matrices — tensor cores on NVIDIA,
/// simdgroup matrices on Apple. Kept in its own module because the
/// `enable` directive below is a parse error on a device that does not
/// offer the feature, so it must never reach one.
///
/// The shape: a workgroup of 256 threads is eight subgroups; the output
/// tile is 64 tokens by 64 weight rows; subgroup `s` owns rows
/// `s*8 .. s*8+8` and holds eight 8x8 accumulators across the 64 columns.
/// Weights are staged dequantized, as in the scalar kernel — the matrix
/// units want f32 planes, not nibbles.
#[cfg(feature = "gpu")]
/// The q4tp GEMM on the card's matrix units — tensor cores on NVIDIA —
/// through wgpu's own cooperative matrices. On the shape this is written
/// against, f16 operands with an f32 accumulator at 16x16x16:
///
/// ```text
/// scalar WGSL      25 077 GFLOP/s    3.53 ms
/// cooperative      50 730 GFLOP/s    1.75 ms    2.02x
/// ```
///
/// and a 512x512x30-step render goes 20.2 -> 14.7 s. The image moves in
/// the last bits — f16 multiplies where the scalar kernel used f32 — but
/// not away from the truth: against the diffusers reference it measures
/// 22.5 dB where the scalar path measures 22.3.
///
/// Two things about the WGSL surface, both settled by disassembling naga's
/// SPIR-V beside glslang's rather than by reading documentation:
///
/// * `coopLoad` emits **ColumnMajor** and `coopLoadT` **RowMajor** — the
///   opposite way round from what the names suggest. A square probe cannot
///   see this: transposing both operands and the result cancels out, and
///   the arithmetic only breaks once the two tiles have different strides,
///   which is to say once it is a real GEMM.
/// * wgpu's own documentation says the implementation "currently only
///   supports 8x8 f32 matrices". It compiles and runs 16x16 f16 — but the
///   feature flag is raised on the weaker configuration too, so the shape
///   is checked against `adapter.cooperative_matrix_properties()` before
///   this pipeline is built. `CMF_COOP=0` opts out.
/// Dequantize a q4tp weight plane into packed f16 pairs ONCE per GEMM,
/// so the matrix units stop waiting on nibble unpacking. Measured on an
/// RTX 5090: the in-kernel path spends about as many scalar ops
/// unpacking a 64×32 weight tile as it spends MACs on it, and it repeats
/// that work for every 64-row tile of activations — tensor cores idle
/// through a dequantizer. One pass, then a pure f16 GEMM.
const COOP_DQ_SRC: &str = r#"
enable f16;

struct DqP { cols: u32, rows: u32, pad0: u32, pad1: u32 };
@group(0) @binding(0) var<storage, read> qsrc: array<u32>;
@group(0) @binding(1) var<storage, read_write> dst: array<u32>;
@group(0) @binding(2) var<uniform> dp: DqP;

fn dq_byte(off: u32) -> u32 {
    return (qsrc[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
}

// One thread per PAIR of weights (a packed u32 of two f16).
@compute @workgroup_size(256)
fn q4tp_dq_f16(@builtin(global_invocation_id) gid: vec3<u32>) {
    let pairs_per_row = dp.cols / 2u;
    // 2-D grid: a plane of 77M pairs needs 300k workgroups and the
    // per-dimension limit is 65 535. Clamping the x dimension (the
    // obvious shortcut) silently dequantizes a fifth of the weight and
    // leaves the rest as garbage the GEMM will happily multiply.
    let idx = gid.y * (65535u * 256u) + gid.x;
    if (idx >= dp.rows * pairs_per_row) { return; }
    let row = idx / pairs_per_row;
    let pair = idx % pairs_per_row;
    let col0 = pair * 2u;
    let gpr = dp.cols >> 5u;
    let params_b = dp.rows * gpr * 16u;
    let codes_b = params_b + dp.rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;
    let g = col0 >> 5u;
    let bit = g * 5u;
    let cb = codes_b + row * cstride + (bit >> 3u);
    let sh = bit & 7u;
    var cv = dq_byte(cb);
    if (sh > 3u) { cv = cv | (dq_byte(cb + 1u) << 8u); }
    let pr = unpack2x16float(qsrc[(params_b >> 2u) + row]);
    let scale = exp2(pr.x + f32((cv >> sh) & 31u) * pr.y);
    let toff = (row * gpr + g) * 16u;
    let pp = col0 - g * 32u;
    let bo = toff + pp / 2u;
    let b0 = dq_byte(bo);
    // Two consecutive columns share a byte when pp is even, which it is:
    // col0 is even by construction (one thread per pair).
    let w0 = (f32(b0 & 0xFu) - 8.0) * scale;
    let w1 = (f32(b0 >> 4u) - 8.0) * scale;
    dst[idx] = pack2x16float(vec2<f32>(w0, w1));
}
"#;

/// int8 → the same packed-f16 plane the coop GEMM eats. The two-field codec's
/// column field is already in the activation by the time a GEMM runs, so what
/// is left on the weight side is one scale per row — `w = q·row[o]`.
const COOP_DQ8_SRC: &str = r#"
enable f16;

struct DqP { cols: u32, rows: u32, pad0: u32, pad1: u32 };
@group(0) @binding(0) var<storage, read> qsrc: array<u32>;
@group(0) @binding(1) var<storage, read_write> dst: array<u32>;
@group(0) @binding(2) var<uniform> dp: DqP;
@group(0) @binding(3) var<storage, read> rsc: array<f32>;
// The two-field codec's column field, or a one-element dummy when the
// weight has none (dp.pad0 says which).
@group(0) @binding(4) var<storage, read> csc: array<f32>;

fn s8(off: u32) -> f32 {
    let b = (qsrc[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
    if (b > 127u) { return f32(b) - 256.0; }
    return f32(b);
}

// One thread per PAIR of weights, as the four-bit twin does — the plane is
// packed u32s of two f16 and the GEMM reads it that way.
@compute @workgroup_size(256)
fn q8_dq_f16(@builtin(global_invocation_id) gid: vec3<u32>) {
    let pairs_per_row = dp.cols / 2u;
    let idx = gid.y * (65535u * 256u) + gid.x;
    if (idx >= dp.rows * pairs_per_row) { return; }
    let row = idx / pairs_per_row;
    let pair = idx % pairs_per_row;
    let base = row * dp.cols + pair * 2u;
    let sc = rsc[row];
    var c0 = 1.0;
    var c1 = 1.0;
    if (dp.pad0 != 0u) {
        c0 = csc[pair * 2u];
        c1 = csc[pair * 2u + 1u];
    }
    dst[idx] = pack2x16float(vec2<f32>(s8(base) * sc * c0, s8(base + 1u) * sc * c1));
}
"#;

/// The same tiled coop GEMM as `q4tp_mm_coop`, reading an ALREADY
/// dequantized f16 plane (packed pairs). Everything else — tile shape,
/// load choreography, the activation scale in `pmm.pad` — is verbatim.
const COOP_MM_F16_SRC: &str = r#"
enable wgpu_cooperative_matrix;
enable f16;

struct MmP { cols4: u32, rows: u32, nb: u32, pad: u32 };
@group(0) @binding(0) var<storage, read> wmm: array<u32>;
@group(0) @binding(1) var<storage, read> xmm: array<f32>;
@group(0) @binding(2) var<storage, read_write> ymm: array<f32>;
@group(0) @binding(3) var<uniform> pmm: MmP;
// Device-computed activation scale (1 element). Used when pmm.pad is
// the sentinel 0xFFFFFFFF — a host-side scale still rides in pmm.pad.
@group(0) @binding(4) var<storage, read> pmm_s: array<f32>;

const KS: u32 = 32u;
var<workgroup> cm_a: array<f16, 64 * 32>;
var<workgroup> cm_b: array<f16, 64 * 32>;
var<workgroup> cm_c: array<f32, 64 * 64>;

@compute @workgroup_size(128)
fn q4tp_mm_coop_f16(@builtin(workgroup_id) wid: vec3<u32>,
                    @builtin(local_invocation_index) tid: u32,
                    @builtin(subgroup_id) sg: u32) {
    let cols = pmm.cols4 * 4u;
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;
    var c0: coop_mat16x16<f32, C>;
    var c1: coop_mat16x16<f32, C>;
    var c2: coop_mat16x16<f32, C>;
    var c3: coop_mat16x16<f32, C>;
    // ONE scale for both ends of the kernel. It used to read the uniform
    // here and the buffer at the store: with the device-computed scale
    // (pad = sentinel) the operand went in UNSCALED and the result came
    // out divided by 1000/max|x| — off by max|x|/1000, on exactly the
    // inputs large enough to have needed the scale. Caught by
    // wgpu_q4tp_mm_coop_f16_device_scale_matches_scalar (rel rms 2.0
    // before this line, 3e-4 after).
    let asc_in = select(bitcast<f32>(pmm.pad), pmm_s[0], pmm.pad == 0xFFFFFFFFu);
    let ainv = select(1.0, asc_in, asc_in > 0.0);

    var k0 = 0u;
    loop {
        if (k0 >= cols) { break; }
        for (var t = tid; t < 64u * 8u; t = t + 128u) {
            let m = t / 8u;
            let k4 = (t % 8u) * 4u;
            let col0 = k0 + k4;
            let dst = m * KS + k4;
            var v = vec4<f32>(0.0);
            if (m0 + m < pmm.nb && col0 < cols) {
                let base = (m0 + m) * cols + col0;
                v = vec4<f32>(xmm[base], xmm[base + 1u], xmm[base + 2u], xmm[base + 3u]);
            }
            cm_a[dst] = f16(v.x * ainv); cm_a[dst + 1u] = f16(v.y * ainv);
            cm_a[dst + 2u] = f16(v.z * ainv); cm_a[dst + 3u] = f16(v.w * ainv);
        }
        for (var t = tid; t < 64u * 8u; t = t + 128u) {
            let n = t / 8u;
            let k4 = (t % 8u) * 4u;
            let col0 = k0 + k4;
            let bd = n * KS + k4;
            var p0 = vec2<f32>(0.0);
            var p1 = vec2<f32>(0.0);
            if (n0 + n < pmm.rows && col0 < cols) {
                let base = ((n0 + n) * cols + col0) >> 1u;
                p0 = unpack2x16float(wmm[base]);
                p1 = unpack2x16float(wmm[base + 1u]);
            }
            cm_b[bd] = f16(p0.x); cm_b[bd + 1u] = f16(p0.y);
            cm_b[bd + 2u] = f16(p1.x); cm_b[bd + 3u] = f16(p1.y);
        }
        workgroupBarrier();
        {
            let a = coopLoadT<coop_mat16x16<f16, A>>(&cm_a[sg * 512u + 0u], 32u);
            let b0 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[0u + 0u], 32u);
            let b1 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[512u + 0u], 32u);
            let b2 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1024u + 0u], 32u);
            let b3 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1536u + 0u], 32u);
            c0 = coopMultiplyAdd(a, b0, c0);
            c1 = coopMultiplyAdd(a, b1, c1);
            c2 = coopMultiplyAdd(a, b2, c2);
            c3 = coopMultiplyAdd(a, b3, c3);
        }
        {
            let a = coopLoadT<coop_mat16x16<f16, A>>(&cm_a[sg * 512u + 16u], 32u);
            let b0 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[0u + 16u], 32u);
            let b1 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[512u + 16u], 32u);
            let b2 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1024u + 16u], 32u);
            let b3 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1536u + 16u], 32u);
            c0 = coopMultiplyAdd(a, b0, c0);
            c1 = coopMultiplyAdd(a, b1, c1);
            c2 = coopMultiplyAdd(a, b2, c2);
            c3 = coopMultiplyAdd(a, b3, c3);
        }
        workgroupBarrier();
        k0 = k0 + KS;
    }
    workgroupBarrier();
    coopStoreT(c0, &cm_c[sg * 16u * 64u + 0u], 64u);
    coopStoreT(c1, &cm_c[sg * 16u * 64u + 16u], 64u);
    coopStoreT(c2, &cm_c[sg * 16u * 64u + 32u], 64u);
    coopStoreT(c3, &cm_c[sg * 16u * 64u + 48u], 64u);
    workgroupBarrier();
    let aback = select(1.0, 1.0 / asc_in, asc_in > 0.0);
    for (var t = tid; t < 64u * 64u; t = t + 128u) {
        let m = t / 64u;
        let n = t % 64u;
        if (m0 + m < pmm.nb && n0 + n < pmm.rows) {
            ymm[(m0 + m) * pmm.rows + n0 + n] = cm_c[m * 64u + n] * aback;
        }
    }
}
"#;

/// max|x| over a device panel, into one f32. The fused FFN's second GEMM
/// reads its input from a buffer the host never sees, so the activation
/// scale its f16 operands need has to be computed HERE — reading that
/// panel back to find it would be the round trip the fusion exists to
/// avoid. One workgroup: 256 partial maxima, then a tree reduce.
const COOP_AMAX_SRC: &str = r#"
struct AmP { n: u32, _a: u32, _b: u32, _c: u32 };
@group(0) @binding(0) var<storage, read>       amx : array<f32>;
@group(0) @binding(1) var<storage, read_write> amo : array<f32>;
@group(0) @binding(2) var<uniform>             amp : AmP;
var<workgroup> am_red: array<f32, 256>;

// Stage 1: one workgroup per slice, each writing its own partial max.
// A single workgroup walking the whole panel (330 MB at render size)
// was what made the device-side scale a wash — one SM reading what the
// card can read with hundreds.
@compute @workgroup_size(256)
fn act_absmax_part(@builtin(local_invocation_index) lid: u32,
                   @builtin(workgroup_id) wid: vec3<u32>,
                   @builtin(num_workgroups) nwg: vec3<u32>) {
    var m = 0.0;
    var i = wid.x * 256u + lid;
    let stride = nwg.x * 256u;
    loop {
        if (i >= amp.n) { break; }
        let v = abs(amx[i]);
        if (v > m && v < 3.0e38) { m = v; }
        i = i + stride;
    }
    am_red[lid] = m;
    workgroupBarrier();
    var st = 128u;
    loop {
        if (st == 0u) { break; }
        if (lid < st) { am_red[lid] = max(am_red[lid], am_red[lid + st]); }
        workgroupBarrier();
        st = st >> 1u;
    }
    if (lid == 0u) { amo[wid.x] = am_red[0]; }
}

// Stage 2: fold the partials into the scale the GEMM multiplies by.
@compute @workgroup_size(256)
fn act_absmax_fold(@builtin(local_invocation_index) lid: u32) {
    var m = 0.0;
    var i = lid;
    loop {
        if (i >= amp.n) { break; }
        let v = amx[i];
        if (v > m) { m = v; }
        i = i + 256u;
    }
    am_red[lid] = m;
    workgroupBarrier();
    var st = 128u;
    loop {
        if (st == 0u) { break; }
        if (lid < st) { am_red[lid] = max(am_red[lid], am_red[lid + st]); }
        workgroupBarrier();
        st = st >> 1u;
    }
    if (lid == 0u) {
        let mx = am_red[0];
        amo[0] = select(1.0, 1000.0 / mx, mx > 1000.0);
    }
}

@compute @workgroup_size(256)
fn act_absmax(@builtin(local_invocation_index) lid: u32) {
    var m = 0.0;
    var i = lid;
    loop {
        if (i >= amp.n) { break; }
        let v = abs(amx[i]);
        if (v > m && v < 3.0e38) { m = v; }
        i = i + 256u;
    }
    am_red[lid] = m;
    workgroupBarrier();
    var stride = 128u;
    loop {
        if (stride == 0u) { break; }
        if (lid < stride) {
            am_red[lid] = max(am_red[lid], am_red[lid + stride]);
        }
        workgroupBarrier();
        stride = stride >> 1u;
    }
    if (lid == 0u) {
        // Store the SCALE the GEMM multiplies by on load: bring the peak
        // to ~1000, well inside f16's range, or leave the data alone.
        let mx = am_red[0];
        amo[0] = select(1.0, 1000.0 / mx, mx > 1000.0);
    }
}
"#;

const COOP_MM_SRC: &str = r#"
enable wgpu_cooperative_matrix;
enable f16;

struct MmP { cols4: u32, rows: u32, nb: u32, pad: u32 };
@group(0) @binding(0) var<storage, read> qmm: array<u32>;
@group(0) @binding(1) var<storage, read> xmm: array<f32>;
@group(0) @binding(2) var<storage, read_write> ymm: array<f32>;
@group(0) @binding(3) var<uniform> pmm: MmP;
// Activation scale: the value in `pmm.pad` (0 = none), or, when pad is
// the sentinel 0xFFFFFFFF, the DEVICE-computed one in this buffer — the
// batched-prefill GEMMs feed on activations the host never sees, and
// f16 operands cap at 65504.
@group(0) @binding(4) var<storage, read> qmm_s: array<f32>;

const KS: u32 = 32u;

// Operands are f16 because that is what the matrix units take; the
// accumulator stays f32, which is the configuration the hardware reports
// (M16 N16 K16, f16 x f16 -> f32).
var<workgroup> cm_a: array<f16, 64 * 32>;
// k-major: the B role is K x N read row by row, and the layout probe
// settled that `coopLoad` means row-major for both operands. Staging the
// weights the other way round and asking for a transposed load reads
// something else again — measured, twice, both wrong.
var<workgroup> cm_b: array<f16, 64 * 32>;
// The result plane must be f32: storing an f32 accumulator into an f16
// array compiles and writes nothing usable.
var<workgroup> cm_c: array<f32, 64 * 64>;

fn cm_byte(off: u32) -> u32 {
    return (qmm[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
}

// Two entry points over one body: `q4tp_mm_coop` (host scale in
// `pmm.pad`, four bindings — every existing caller) and
// `q4tp_mm_coop_s` (the scale from the device buffer at binding 4 when
// pad is the sentinel — the batched prefill's operands never reach the
// host). wgpu builds the auto layout PER entry point from the bindings
// it touches, so the first keeps its four-slot layout.
fn coop_body(wid: vec3<u32>, tid: u32, sg: u32, asc_in: f32) {
    let cols = pmm.cols4 * 4u;
    let gpr = cols >> 5u;
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;
    // Four subgroups, and which one this is comes from the builtin rather
    // than from `tid / 32`: nothing promises the driver hands out lanes to
    // subgroups in that order, and a cooperative matrix belongs to a
    // subgroup, not to a range of local indices.

    var c0: coop_mat16x16<f32, C>;
    var c1: coop_mat16x16<f32, C>;
    var c2: coop_mat16x16<f32, C>;
    var c3: coop_mat16x16<f32, C>;
    let ainv = select(1.0, asc_in, asc_in > 0.0);
    let aback = select(1.0, 1.0 / asc_in, asc_in > 0.0);

    let params_b = pmm.rows * gpr * 16u;
    let codes_b = params_b + pmm.rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;

    var k0 = 0u;
    loop {
        if (k0 >= cols) { break; }
        // 64 rows x 32 k of each side; 128 threads take sixteen apiece.
        for (var t = tid; t < 64u * 8u; t = t + 128u) {
            let m = t / 8u;
            let k4 = (t % 8u) * 4u;
            let col0 = k0 + k4;
            let dst = m * KS + k4;
            var v = vec4<f32>(0.0);
            if (m0 + m < pmm.nb && col0 < cols) {
                let base = (m0 + m) * cols + col0;
                v = vec4<f32>(xmm[base], xmm[base + 1u], xmm[base + 2u], xmm[base + 3u]);
            }
            cm_a[dst] = f16(v.x * ainv); cm_a[dst + 1u] = f16(v.y * ainv);
            cm_a[dst + 2u] = f16(v.z * ainv); cm_a[dst + 3u] = f16(v.w * ainv);
        }
        for (var t = tid; t < 64u * 8u; t = t + 128u) {
            let n = t / 8u;
            let k4 = (t % 8u) * 4u;
            let col0 = k0 + k4;
            var wv = vec4<f32>(0.0);
            if (n0 + n < pmm.rows && col0 < cols) {
                let g = col0 >> 5u;
                let wrow = n0 + n;
                let bit = g * 5u;
                let cb = codes_b + wrow * cstride + (bit >> 3u);
                let sh = bit & 7u;
                var cv = cm_byte(cb);
                if (sh > 3u) { cv = cv | (cm_byte(cb + 1u) << 8u); }
                let pr = unpack2x16float(qmm[(params_b >> 2u) + wrow]);
                let scale = exp2(pr.x + f32((cv >> sh) & 31u) * pr.y);
                let toff = (wrow * gpr + g) * 16u;
                let pp = col0 - g * 32u;
                let bo = toff + pp / 2u;
                let b0 = cm_byte(bo);
                let b1 = cm_byte(bo + 1u);
                wv[0u] = (f32(b0 & 0xFu) - 8.0) * scale;
                wv[1u] = (f32(b0 >> 4u) - 8.0) * scale;
                wv[2u] = (f32(b1 & 0xFu) - 8.0) * scale;
                wv[3u] = (f32(b1 >> 4u) - 8.0) * scale;
            }
            // n-major, four consecutive k in a row: `coopLoad` reads it
            // ColumnMajor, which turns [n][k] into the K x N the B role
            // wants without a transposing load.
            let bd = n * KS + k4;
            cm_b[bd] = f16(wv.x); cm_b[bd + 1u] = f16(wv.y);
            cm_b[bd + 2u] = f16(wv.z); cm_b[bd + 3u] = f16(wv.w);
        }
        workgroupBarrier();
        {
            let a = coopLoadT<coop_mat16x16<f16, A>>(&cm_a[sg * 512u + 0u], 32u);
            let b0 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[0u + 0u], 32u);
            let b1 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[512u + 0u], 32u);
            let b2 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1024u + 0u], 32u);
            let b3 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1536u + 0u], 32u);
            c0 = coopMultiplyAdd(a, b0, c0);
            c1 = coopMultiplyAdd(a, b1, c1);
            c2 = coopMultiplyAdd(a, b2, c2);
            c3 = coopMultiplyAdd(a, b3, c3);
        }
        {
            let a = coopLoadT<coop_mat16x16<f16, A>>(&cm_a[sg * 512u + 16u], 32u);
            let b0 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[0u + 16u], 32u);
            let b1 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[512u + 16u], 32u);
            let b2 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1024u + 16u], 32u);
            let b3 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1536u + 16u], 32u);
            c0 = coopMultiplyAdd(a, b0, c0);
            c1 = coopMultiplyAdd(a, b1, c1);
            c2 = coopMultiplyAdd(a, b2, c2);
            c3 = coopMultiplyAdd(a, b3, c3);
        }
        workgroupBarrier();
        k0 = k0 + KS;
    }
    workgroupBarrier();
    coopStoreT(c0, &cm_c[sg * 16u * 64u + 0u], 64u);
    coopStoreT(c1, &cm_c[sg * 16u * 64u + 16u], 64u);
    coopStoreT(c2, &cm_c[sg * 16u * 64u + 32u], 64u);
    coopStoreT(c3, &cm_c[sg * 16u * 64u + 48u], 64u);
    workgroupBarrier();
    for (var t = tid; t < 64u * 64u; t = t + 128u) {
        let m = t / 64u;
        let n = t % 64u;
        if (m0 + m < pmm.nb && n0 + n < pmm.rows) {
            ymm[(m0 + m) * pmm.rows + n0 + n] = cm_c[m * 64u + n] * aback;
        }
    }
}

@compute @workgroup_size(128)
fn q4tp_mm_coop(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(local_invocation_index) tid: u32,
                @builtin(subgroup_id) sg: u32) {
    coop_body(wid, tid, sg, bitcast<f32>(pmm.pad));
}

@compute @workgroup_size(128)
fn q4tp_mm_coop_s(@builtin(workgroup_id) wid: vec3<u32>,
                  @builtin(local_invocation_index) tid: u32,
                  @builtin(subgroup_id) sg: u32) {
    let asc = select(bitcast<f32>(pmm.pad), qmm_s[0], pmm.pad == 0xFFFFFFFFu);
    coop_body(wid, tid, sg, asc);
}
"#;

const COOP_Q2_MM_SRC: &str = r#"
enable wgpu_cooperative_matrix;
enable f16;

struct MmP { cols4: u32, rows: u32, nb: u32, pad: u32 };
@group(0) @binding(0) var<storage, read> qmm: array<u32>;
@group(0) @binding(1) var<storage, read> xmm: array<f32>;
@group(0) @binding(2) var<storage, read_write> ymm: array<f32>;
@group(0) @binding(3) var<uniform> pmm: MmP;
const KS: u32 = 32u;

// Operands are f16 because that is what the matrix units take; the
// accumulator stays f32, which is the configuration the hardware reports
// (M16 N16 K16, f16 x f16 -> f32).
var<workgroup> cm_a: array<f16, 64 * 32>;
// k-major: the B role is K x N read row by row, and the layout probe
// settled that `coopLoad` means row-major for both operands. Staging the
// weights the other way round and asking for a transposed load reads
// something else again — measured, twice, both wrong.
var<workgroup> cm_b: array<f16, 64 * 32>;
// The result plane must be f32: storing an f32 accumulator into an f16
// array compiles and writes nothing usable.
var<workgroup> cm_c: array<f32, 64 * 64>;

fn cm_byte(off: u32) -> u32 {
    return (qmm[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu;
}

// The fourth uniform word is the explicit Q2 affine descriptor bit. It is
// not an activation scale/sentinel: the Q2 candidate deliberately starts
// with the validated Prism activation_f16 boundary and scale 1.0.
fn coop_body(wid: vec3<u32>, tid: u32, sg: u32) {
    let cols = pmm.cols4 * 4u;
    let gpr = cols >> 5u;
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;
    // Four subgroups, and which one this is comes from the builtin rather
    // than from `tid / 32`: nothing promises the driver hands out lanes to
    // subgroups in that order, and a cooperative matrix belongs to a
    // subgroup, not to a range of local indices.

    var c0: coop_mat16x16<f32, C>;
    var c1: coop_mat16x16<f32, C>;
    var c2: coop_mat16x16<f32, C>;
    var c3: coop_mat16x16<f32, C>;
    let params_b = pmm.rows * gpr * 8u;
    let codes_b = params_b + pmm.rows * 4u;
    let cstride = (gpr * 5u + 7u) / 8u;

    var k0 = 0u;
    loop {
        if (k0 >= cols) { break; }
        // 64 rows x 32 k of each side; 128 threads take sixteen apiece.
        for (var t = tid; t < 64u * 8u; t = t + 128u) {
            let m = t / 8u;
            let k4 = (t % 8u) * 4u;
            let col0 = k0 + k4;
            let dst = m * KS + k4;
            var v = vec4<f32>(0.0);
            if (m0 + m < pmm.nb && col0 < cols) {
                let base = (m0 + m) * cols + col0;
                v = vec4<f32>(xmm[base], xmm[base + 1u], xmm[base + 2u], xmm[base + 3u]);
            }
            cm_a[dst] = f16(v.x); cm_a[dst + 1u] = f16(v.y);
            cm_a[dst + 2u] = f16(v.z); cm_a[dst + 3u] = f16(v.w);
        }
        for (var t = tid; t < 64u * 8u; t = t + 128u) {
            let n = t / 8u;
            let k4 = (t % 8u) * 4u;
            let col0 = k0 + k4;
            var wv = vec4<f32>(0.0);
            if (n0 + n < pmm.rows && col0 < cols) {
                let g = col0 >> 5u;
                let wrow = n0 + n;
                let bit = g * 5u;
                let cb = codes_b + wrow * cstride + (bit >> 3u);
                let sh = bit & 7u;
                var cv = cm_byte(cb);
                if (sh > 3u) { cv = cv | (cm_byte(cb + 1u) << 8u); }
                let pr = unpack2x16float(qmm[(params_b >> 2u) + wrow]);
                let rung = (cv >> sh) & 31u;
                var scale = 0.0;
                if (rung != 0u) {
                    scale = exp2(pr.x + f32(rung - 1u) * pr.y);
                }
                let toff = (wrow * gpr + g) * 8u;
                let pp = col0 - g * 32u;
                let bo = toff + pp / 4u;
                let by = cm_byte(bo);
                let center = select(1.5, 1.0, pmm.pad != 0u);
                wv[0u] = (f32(by & 3u) - center) * scale;
                wv[1u] = (f32((by >> 2u) & 3u) - center) * scale;
                wv[2u] = (f32((by >> 4u) & 3u) - center) * scale;
                wv[3u] = (f32((by >> 6u) & 3u) - center) * scale;
            }
            // n-major, four consecutive k in a row: `coopLoad` reads it
            // ColumnMajor, which turns [n][k] into the K x N the B role
            // wants without a transposing load.
            let bd = n * KS + k4;
            cm_b[bd] = f16(wv.x); cm_b[bd + 1u] = f16(wv.y);
            cm_b[bd + 2u] = f16(wv.z); cm_b[bd + 3u] = f16(wv.w);
        }
        workgroupBarrier();
        {
            let a = coopLoadT<coop_mat16x16<f16, A>>(&cm_a[sg * 512u + 0u], 32u);
            let b0 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[0u + 0u], 32u);
            let b1 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[512u + 0u], 32u);
            let b2 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1024u + 0u], 32u);
            let b3 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1536u + 0u], 32u);
            c0 = coopMultiplyAdd(a, b0, c0);
            c1 = coopMultiplyAdd(a, b1, c1);
            c2 = coopMultiplyAdd(a, b2, c2);
            c3 = coopMultiplyAdd(a, b3, c3);
        }
        {
            let a = coopLoadT<coop_mat16x16<f16, A>>(&cm_a[sg * 512u + 16u], 32u);
            let b0 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[0u + 16u], 32u);
            let b1 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[512u + 16u], 32u);
            let b2 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1024u + 16u], 32u);
            let b3 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1536u + 16u], 32u);
            c0 = coopMultiplyAdd(a, b0, c0);
            c1 = coopMultiplyAdd(a, b1, c1);
            c2 = coopMultiplyAdd(a, b2, c2);
            c3 = coopMultiplyAdd(a, b3, c3);
        }
        workgroupBarrier();
        k0 = k0 + KS;
    }
    workgroupBarrier();
    coopStoreT(c0, &cm_c[sg * 16u * 64u + 0u], 64u);
    coopStoreT(c1, &cm_c[sg * 16u * 64u + 16u], 64u);
    coopStoreT(c2, &cm_c[sg * 16u * 64u + 32u], 64u);
    coopStoreT(c3, &cm_c[sg * 16u * 64u + 48u], 64u);
    workgroupBarrier();
    for (var t = tid; t < 64u * 64u; t = t + 128u) {
        let m = t / 64u;
        let n = t % 64u;
        if (m0 + m < pmm.nb && n0 + n < pmm.rows) {
            ymm[(m0 + m) * pmm.rows + n0 + n] = cm_c[m * 64u + n];
        }
    }
}

@compute @workgroup_size(128)
fn q2tp_mm_coop(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(local_invocation_index) tid: u32,
                @builtin(subgroup_id) sg: u32) {
    coop_body(wid, tid, sg);
}

"#;

/// The bake's forward GEMM on the same matrix units: y[nb,rows] =
/// x[nb,cols] · w[rows,cols]ᵀ with both operands ALREADY f32 in memory —
/// frozen weights the replica dequantized, or trained masters. The
/// staging, tile shape and load/store choreography are `q4tp_mm_coop`'s
/// verbatim (every hard-won layout fact above applies unchanged); the
/// only difference is that the B tile reads f32 planes instead of
/// dequantizing nibbles. f16 operands, f32 accumulator — the same
/// numerics contract the DiT ships with.
/// The DiT's attention GEMMs on the matrix units. Same tiling and load
/// choreography as `COOP_NT_SRC`, plus what attention needs: per-head
/// offsets into the packed planes, a score scale, and GQA's head
/// mapping. Measured motive: `dit_qk`/`dit_pv` are scalar f32 and run
/// the attention at 0.5 TFLOP/s where this card's q4tp GEMM holds 51.9.
/// Split an interleaved qkv panel ([token][q|k|v][head][dim]) into the
/// head-major planes attention wants. Measured motive: doing this on
/// the host cost 4.4 s of a 7.3 s attention phase — more than the
/// device work it was feeding.
/// v per head from [n][hd] to [hd][n], in its OWN module. It lived in
/// the split kernel's module first, and could not be bound: wgpu builds
/// an auto layout PER ENTRY POINT, so a kernel touching bindings 0, 1
/// and 4 of a five-binding module gets a three-entry layout whose
/// numbering is its own ("no declaration for binding 0").
/// qk-norm + RoPE for ONE of q/k, scattered head-major in the same pass.
/// The draft that tried to do q and k inside one dispatch could not:
/// WGSL cannot take a storage binding as a parameter, so the helper had
/// no way to write two different planes. Two dispatches with different
/// bindings solve it without duplicating a line of shader.
///
/// One workgroup per (token, head): the norm needs a reduction over the
/// head dimension. Host semantics (`mmh3::norm_rope_w`): mean of squares
/// + eps, then the weight, then rotate x[j] against x[j+pairs].
const DIT_QKNORM_SRC: &str = r#"
// This kernel reads the PACKED PANEL, not the split planes — so it
// needs the layout and the bias just as the split does. It re-derives
// q and k itself and overwrites what the split wrote for them.
struct QnP { n: u32, nh: u32, hd: u32, pairs: u32, eps: f32, src_off: u32, mode: u32, bias: u32 };
@group(0) @binding(0) var<storage, read>       qn_src: array<f32>;
@group(0) @binding(1) var<storage, read_write> qn_dst: array<f32>;
@group(0) @binding(2) var<storage, read>       qn_ang: array<f32>;
@group(0) @binding(3) var<storage, read>       qn_w: array<f32>;
@group(0) @binding(4) var<uniform>             qn_p: QnP;
@group(0) @binding(5) var<storage, read>       qn_bias: array<f32>;

var<workgroup> qn_buf: array<f32, 256>;
var<workgroup> qn_red: array<f32, 64>;

@compute @workgroup_size(64)
fn dit_qknorm_rope(@builtin(workgroup_id) wid: vec3<u32>,
                   @builtin(local_invocation_index) lid: u32) {
    let job = wid.y * 65535u + wid.x;
    if (job >= qn_p.n * qn_p.nh) { return; }
    let p = job / qn_p.nh;
    let h = job - p * qn_p.nh;
    let inner = qn_p.nh * qn_p.hd;
    // qn_p.src_off selects q (0) or k (inner) inside the packed panel.
    let row = p * 3u * inner;
    var src = row + qn_p.src_off + h * qn_p.hd;
    if (qn_p.mode == 1u) {
        // Head-interleaved: [q_h | k_h | v_h] per head. src_off still
        // says which of the two this dispatch is doing.
        src = row + h * 3u * qn_p.hd + select(0u, qn_p.hd, qn_p.src_off > 0u);
    }
    let dst = (h * qn_p.n + p) * qn_p.hd;

    var acc = 0.0;
    var i = lid;
    loop {
        if (i >= qn_p.hd) { break; }
        var x = qn_src[src + i];
        if (qn_p.bias == 1u) {
            x = x + qn_bias[src + i - row];
        }
        qn_buf[i] = x;
        acc = acc + x * x;
        i = i + 64u;
    }
    qn_red[lid] = acc;
    workgroupBarrier();
    var st = 32u;
    loop {
        if (st == 0u) { break; }
        if (lid < st) { qn_red[lid] = qn_red[lid] + qn_red[lid + st]; }
        workgroupBarrier();
        st = st >> 1u;
    }
    // eps < 0 is the rope-only sentinel: a model with no qk-norm
    // (Music-3) still needs the rotation done here, because this pass
    // is the only reason the packed panel ever came home. The branch is
    // on a uniform, so the barriers around it stay legal.
    var inv = 1.0;
    if (qn_p.eps >= 0.0) {
        inv = 1.0 / sqrt(qn_red[0] / f32(qn_p.hd) + qn_p.eps);
    }
    workgroupBarrier();
    i = lid;
    loop {
        if (i >= qn_p.hd) { break; }
        if (qn_p.eps >= 0.0) {
            qn_buf[i] = qn_buf[i] * inv * qn_w[i];
        }
        i = i + 64u;
    }
    workgroupBarrier();
    i = lid;
    loop {
        if (i >= qn_p.hd) { break; }
        var out = qn_buf[i];
        if (i < qn_p.pairs) {
            let a = qn_ang[p * qn_p.pairs + i];
            out = qn_buf[i] * cos(a) - qn_buf[i + qn_p.pairs] * sin(a);
        } else if (i < 2u * qn_p.pairs) {
            let j = i - qn_p.pairs;
            let a = qn_ang[p * qn_p.pairs + j];
            out = qn_buf[j] * sin(a) + qn_buf[i] * cos(a);
        }
        qn_dst[dst + i] = out;
        i = i + 64u;
    }
}
"#;

const DIT_VT_SRC: &str = r#"
struct VtP { n: u32, nh: u32, hd: u32, _p: u32 };
@group(0) @binding(0) var<storage, read>       vt_src: array<f32>;
@group(0) @binding(1) var<storage, read_write> vt_dst: array<f32>;
@group(0) @binding(2) var<uniform>             vt_p: VtP;

@compute @workgroup_size(256)
fn dit_v_transpose(@builtin(global_invocation_id) gid: vec3<u32>) {
    let total = vt_p.n * vt_p.nh * vt_p.hd;
    let i = gid.y * (65535u * 256u) + gid.x;
    if (i >= total) { return; }
    let h = i / (vt_p.n * vt_p.hd);
    let r = i - h * vt_p.n * vt_p.hd;
    let p = r / vt_p.hd;
    let d = r - p * vt_p.hd;
    vt_dst[h * vt_p.n * vt_p.hd + d * vt_p.n + p] = vt_src[i];
}
"#;

const MUSIC3_GLU_SRC: &str = r#"
// The elementwise middle of Music-3's FFN, on the device so the two
// GEMMs around it never leave. The host arm of this chain read back
// 45 MB of ff_in output and re-uploaded 23 MB of activations per
// block-step, on a stand whose DMA measures 6-33 ms per 8.5 MB — the
// transfers were 82% of the device arm's time. GLU order is VALUE
// first, gate second, bias added to BOTH halves before they meet,
// exactly as the host loop spells it.
struct GluP { n: u32, inter: u32, _a: u32, _b: u32 };
@group(0) @binding(0) var<storage, read>       glu_gu: array<f32>;
@group(0) @binding(1) var<storage, read_write> glu_act: array<f32>;
@group(0) @binding(2) var<storage, read>       glu_bias: array<f32>;
@group(0) @binding(3) var<uniform>             glup: GluP;

@compute @workgroup_size(256)
fn music3_glu(@builtin(global_invocation_id) gid: vec3<u32>) {
    let total = glup.n * glup.inter;
    let i = gid.y * (65535u * 256u) + gid.x;
    if (i >= total) { return; }
    let p = i / glup.inter;
    let j = i - p * glup.inter;
    let row = p * 2u * glup.inter;
    let v = glu_gu[row + j] + glu_bias[j];
    let g = glu_gu[row + glup.inter + j] + glu_bias[glup.inter + j];
    glu_act[i] = v * g / (1.0 + exp(-g));
}
"#;

const CONV1D_IM2COL_SRC: &str = r#"
// The 1D twin of `vae_im2col`, and it writes the TRANSPOSED layout the
// NT GEMM wants directly: col[t·(ic·k) + (i·k + j)] = x[i, t + j·dil - pad].
// The host arm built this buffer, then built a second one to transpose
// it, then uploaded the result — three passes over as much as 2.37 GB
// for a 20-second song, when the source `x` is k times smaller. Here
// only `x` crosses the bus and the columns are expanded where they are
// consumed. `p0`/`tile` slice the time axis so one binding stays under
// Vulkan's 2 GiB, which is the limit a 20-second render died on.
struct C1P { ic: u32, n: u32, k: u32, dil: u32, pad: u32, p0: u32, tile: u32, _a: u32 };
@group(0) @binding(0) var<storage, read>       c1_x: array<f32>;
@group(0) @binding(1) var<storage, read_write> c1_col: array<f32>;
@group(0) @binding(2) var<uniform>             c1p: C1P;

@compute @workgroup_size(256)
fn conv1d_im2col(@builtin(global_invocation_id) gid: vec3<u32>) {
    let kdim = c1p.ic * c1p.k;
    let total = c1p.tile * kdim;
    let idx = gid.y * (65535u * 256u) + gid.x;
    if (idx >= total) { return; }
    let t = idx / kdim;
    let r = idx - t * kdim;
    let i = r / c1p.k;
    let j = r - i * c1p.k;
    let p = i32(c1p.p0 + t) + i32(j * c1p.dil) - i32(c1p.pad);
    var v = 0.0;
    if (p >= 0 && p < i32(c1p.n)) { v = c1_x[i * c1p.n + u32(p)]; }
    c1_col[idx] = v;
}
"#;

const VAE_IM2COL_SRC: &str = r#"
// One patch column per pixel: col[t][r] = x[i, y+dy-pad, x+dx-pad] with
// r = (i·k + dy)·k + dx. That is exactly A for an NT GEMM against a
// weight already stored as [oc, ic·k·k], so nothing is repacked.
struct IcP { ic: u32, h: u32, w: u32, k: u32, pad: u32, p0: u32, tile: u32, _a: u32 };
@group(0) @binding(0) var<storage, read>       ic_x: array<f32>;
@group(0) @binding(1) var<storage, read_write> ic_col: array<f32>;
@group(0) @binding(2) var<uniform>             icp: IcP;

@compute @workgroup_size(256)
fn vae_im2col(@builtin(global_invocation_id) gid: vec3<u32>) {
    let kk = icp.k * icp.k;
    let kdim = icp.ic * kk;
    let total = icp.tile * kdim;
    let idx = gid.y * (65535u * 256u) + gid.x;
    if (idx >= total) { return; }
    let t = idx / kdim;
    let r = idx - t * kdim;
    let p = icp.p0 + t;
    let y = p / icp.w;
    let xx = p - y * icp.w;
    let i = r / kk;
    let rem = r - i * kk;
    let dy = rem / icp.k;
    let dx = rem - dy * icp.k;
    let sy = i32(y) + i32(dy) - i32(icp.pad);
    let sx = i32(xx) + i32(dx) - i32(icp.pad);
    var v = 0.0;
    if (sy >= 0 && sy < i32(icp.h) && sx >= 0 && sx < i32(icp.w)) {
        v = ic_x[(i * icp.h + u32(sy)) * icp.w + u32(sx)];
    }
    ic_col[idx] = v;
}
"#;

const DIT_SPLIT_SRC: &str = r#"
// `mode` 0: a token row is [q(inner) | k | v] — the DiT's panel.
// `mode` 1: it is head-interleaved, [q_h | k_h | v_h] per head — the
// VAE's. Same kernel, one branch, because the only other difference
// between the two attentions is a bias this now adds in place.
struct SpP { n: u32, nh: u32, hd: u32, mode: u32, bias: u32, _a: u32, _b: u32, _c: u32 };
@group(0) @binding(0) var<storage, read>       sqkv: array<f32>;
@group(0) @binding(1) var<storage, read_write> sq: array<f32>;
@group(0) @binding(2) var<storage, read_write> sk: array<f32>;
@group(0) @binding(3) var<storage, read_write> sv: array<f32>;
@group(0) @binding(4) var<uniform>             sp: SpP;
@group(0) @binding(5) var<storage, read>       sbias: array<f32>;

@compute @workgroup_size(256)
fn dit_qkv_split(@builtin(global_invocation_id) gid: vec3<u32>) {
    let inner = sp.nh * sp.hd;
    let total = sp.n * inner;
    let i = gid.y * (65535u * 256u) + gid.x;
    if (i >= total) { return; }
    let p = i / inner;          // token
    let r = i - p * inner;      // head*hd + d
    let h = r / sp.hd;
    let d = r - h * sp.hd;
    let row = p * 3u * inner;
    var qo = row + h * sp.hd + d;
    var ko = qo + inner;
    var vo = qo + 2u * inner;
    if (sp.mode == 1u) {
        qo = row + h * 3u * sp.hd + d;
        ko = qo + sp.hd;
        vo = qo + 2u * sp.hd;
    }
    let dst = (h * sp.n + p) * sp.hd + d;
    var qv = sqkv[qo];
    var kv = sqkv[ko];
    var vv = sqkv[vo];
    if (sp.bias == 1u) {
        qv = qv + sbias[qo - row];
        kv = kv + sbias[ko - row];
        vv = vv + sbias[vo - row];
    }
    sq[dst] = qv;
    sk[dst] = kv;
    sv[dst] = vv;
}
"#;

const DIT_COOP_SRC: &str = r#"
enable wgpu_cooperative_matrix;
enable f16;

struct DcP {
    cols4: u32,  // k/4 (the reduction width)
    rows: u32,   // output columns of the NT product
    nb: u32,     // output rows
    scale: f32,  // applied at the store (QK only; 1.0 for PV)
    a_off: u32,  // element offset of this head's x plane
    b_off: u32,  // element offset of this head's w plane
    c_off: u32,  // element offset of this head's y plane
    kk: u32,     // TRUE reduction width; 0 = use cols4*4 for both
};
@group(0) @binding(0) var<storage, read> dcw: array<f32>;
@group(0) @binding(1) var<storage, read> dcx: array<f32>;
@group(0) @binding(2) var<storage, read_write> dcy: array<f32>;
@group(0) @binding(3) var<uniform> dcp: DcP;

const KS: u32 = 32u;
var<workgroup> dm_a: array<f16, 64 * 32>;
var<workgroup> dm_b: array<f16, 64 * 32>;
var<workgroup> dm_c: array<f32, 64 * 64>;

@compute @workgroup_size(128)
fn dit_gemm_coop(@builtin(workgroup_id) wid: vec3<u32>,
                 @builtin(local_invocation_index) tid: u32,
                 @builtin(subgroup_id) sg: u32) {
    // PV reduces over the token count, which is not a multiple of four
    // at render size (1859 → 1856 truncated, three columns of every
    // score row silently dropped, frames 59% off). The row STRIDE and
    // the loop BOUND are the same number and it is this one.
    let cols = select(dcp.cols4 * 4u, dcp.kk, dcp.kk > 0u);
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;
    var c0: coop_mat16x16<f32, C>;
    var c1: coop_mat16x16<f32, C>;
    var c2: coop_mat16x16<f32, C>;
    var c3: coop_mat16x16<f32, C>;

    var k0 = 0u;
    loop {
        if (k0 >= cols) { break; }
        for (var t = tid; t < 64u * 8u; t = t + 128u) {
            let m = t / 8u;
            let k4 = (t % 8u) * 4u;
            let col0 = k0 + k4;
            let dst = m * KS + k4;
            var v = vec4<f32>(0.0);
            if (m0 + m < dcp.nb && col0 < cols) {
                let base = dcp.a_off + (m0 + m) * cols + col0;
                if (col0 + 3u < cols) {
                    v = vec4<f32>(dcx[base], dcx[base + 1u], dcx[base + 2u], dcx[base + 3u]);
                } else {
                    v.x = dcx[base];
                    if (col0 + 1u < cols) { v.y = dcx[base + 1u]; }
                    if (col0 + 2u < cols) { v.z = dcx[base + 2u]; }
                    if (col0 + 3u < cols) { v.w = dcx[base + 3u]; }
                }
            }
            dm_a[dst] = f16(v.x); dm_a[dst + 1u] = f16(v.y);
            dm_a[dst + 2u] = f16(v.z); dm_a[dst + 3u] = f16(v.w);
        }
        for (var t = tid; t < 64u * 8u; t = t + 128u) {
            let nn = t / 8u;
            let k4 = (t % 8u) * 4u;
            let col0 = k0 + k4;
            let bd = nn * KS + k4;
            var wv = vec4<f32>(0.0);
            if (n0 + nn < dcp.rows && col0 < cols) {
                let base = dcp.b_off + (n0 + nn) * cols + col0;
                if (col0 + 3u < cols) {
                    wv = vec4<f32>(dcw[base], dcw[base + 1u], dcw[base + 2u], dcw[base + 3u]);
                } else {
                    wv.x = dcw[base];
                    if (col0 + 1u < cols) { wv.y = dcw[base + 1u]; }
                    if (col0 + 2u < cols) { wv.z = dcw[base + 2u]; }
                    if (col0 + 3u < cols) { wv.w = dcw[base + 3u]; }
                }
            }
            dm_b[bd] = f16(wv.x); dm_b[bd + 1u] = f16(wv.y);
            dm_b[bd + 2u] = f16(wv.z); dm_b[bd + 3u] = f16(wv.w);
        }
        workgroupBarrier();
        {
            let a = coopLoadT<coop_mat16x16<f16, A>>(&dm_a[sg * 512u + 0u], 32u);
            let b0 = coopLoad<coop_mat16x16<f16, B>>(&dm_b[0u], 32u);
            let b1 = coopLoad<coop_mat16x16<f16, B>>(&dm_b[512u], 32u);
            let b2 = coopLoad<coop_mat16x16<f16, B>>(&dm_b[1024u], 32u);
            let b3 = coopLoad<coop_mat16x16<f16, B>>(&dm_b[1536u], 32u);
            c0 = coopMultiplyAdd(a, b0, c0);
            c1 = coopMultiplyAdd(a, b1, c1);
            c2 = coopMultiplyAdd(a, b2, c2);
            c3 = coopMultiplyAdd(a, b3, c3);
        }
        {
            let a = coopLoadT<coop_mat16x16<f16, A>>(&dm_a[sg * 512u + 16u], 32u);
            let b0 = coopLoad<coop_mat16x16<f16, B>>(&dm_b[16u], 32u);
            let b1 = coopLoad<coop_mat16x16<f16, B>>(&dm_b[528u], 32u);
            let b2 = coopLoad<coop_mat16x16<f16, B>>(&dm_b[1040u], 32u);
            let b3 = coopLoad<coop_mat16x16<f16, B>>(&dm_b[1552u], 32u);
            c0 = coopMultiplyAdd(a, b0, c0);
            c1 = coopMultiplyAdd(a, b1, c1);
            c2 = coopMultiplyAdd(a, b2, c2);
            c3 = coopMultiplyAdd(a, b3, c3);
        }
        workgroupBarrier();
        k0 = k0 + KS;
    }
    workgroupBarrier();
    coopStoreT(c0, &dm_c[sg * 16u * 64u + 0u], 64u);
    coopStoreT(c1, &dm_c[sg * 16u * 64u + 16u], 64u);
    coopStoreT(c2, &dm_c[sg * 16u * 64u + 32u], 64u);
    coopStoreT(c3, &dm_c[sg * 16u * 64u + 48u], 64u);
    workgroupBarrier();
    for (var t = tid; t < 64u * 64u; t = t + 128u) {
        let m = t / 64u;
        let nn = t % 64u;
        if (m0 + m < dcp.nb && n0 + nn < dcp.rows) {
            dcy[dcp.c_off + (m0 + m) * dcp.rows + n0 + nn] = dm_c[m * 64u + nn] * dcp.scale;
        }
    }
}
"#;

const COOP_NT_SRC: &str = r#"
enable wgpu_cooperative_matrix;
enable f16;

struct MmP { cols4: u32, rows: u32, nb: u32, pad: u32 };
@group(0) @binding(0) var<storage, read> wmm: array<vec4<f32>>;
@group(0) @binding(1) var<storage, read> xmm: array<vec4<f32>>;
@group(0) @binding(2) var<storage, read_write> ymm: array<f32>;
@group(0) @binding(3) var<uniform> pmm: MmP;

const KS: u32 = 32u;

// Two K-tiles of A and B, not one: tile k+1 stages while the matrix
// units chew tile k, so the global-memory latency hides behind the mma
// work and each K step costs ONE barrier where the old kernel paid two.
// That kernel measured 614-1537 GFLOP/s on a 3090 across the DiT's
// shapes — flat against an 8x spread in work, which is the signature of
// paying per step, not per flop. Global loads are vec4 for the same
// reason: four scalar f32 reads per thread per tile were four trips.
var<workgroup> cm_a: array<f16, 2 * 64 * 32>;
var<workgroup> cm_b: array<f16, 2 * 64 * 32>;
var<workgroup> cm_c: array<f32, 64 * 64>;

fn stage_nt(buf: u32, k0: u32, m0: u32, n0: u32, tid: u32) {
    let cols = pmm.cols4 * 4u;
    let ab = buf * 2048u;
    for (var t = tid; t < 512u; t = t + 128u) {
        let m = t / 8u;
        let kv = (t % 8u) * 4u;
        let col0 = k0 + kv;
        var v = vec4<f32>();
        if (m0 + m < pmm.nb && col0 < cols) {
            v = xmm[((m0 + m) * cols + col0) >> 2u];
        }
        let dst = ab + m * KS + kv;
        cm_a[dst] = f16(v.x); cm_a[dst + 1u] = f16(v.y);
        cm_a[dst + 2u] = f16(v.z); cm_a[dst + 3u] = f16(v.w);
    }
    for (var t = tid; t < 512u; t = t + 128u) {
        let n = t / 8u;
        let kv = (t % 8u) * 4u;
        let col0 = k0 + kv;
        var v = vec4<f32>();
        if (n0 + n < pmm.rows && col0 < cols) {
            v = wmm[((n0 + n) * cols + col0) >> 2u];
        }
        let dst = ab + n * KS + kv;
        cm_b[dst] = f16(v.x); cm_b[dst + 1u] = f16(v.y);
        cm_b[dst + 2u] = f16(v.z); cm_b[dst + 3u] = f16(v.w);
    }
}

@compute @workgroup_size(128)
fn gemm_nt_coop(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(local_invocation_index) tid: u32,
                @builtin(subgroup_id) sg: u32) {
    let cols = pmm.cols4 * 4u;
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;

    var c0: coop_mat16x16<f32, C>;
    var c1: coop_mat16x16<f32, C>;
    var c2: coop_mat16x16<f32, C>;
    var c3: coop_mat16x16<f32, C>;

    stage_nt(0u, 0u, m0, n0, tid);
    workgroupBarrier();
    var k0 = 0u;
    var buf = 0u;
    loop {
        let nk = k0 + KS;
        // Issue the next tile's global loads BEFORE the mma below — by
        // the barrier they are in flight behind the math, not after it.
        if (nk < cols) { stage_nt(buf ^ 1u, nk, m0, n0, tid); }
        let ab = buf * 2048u;
        {
            let a = coopLoadT<coop_mat16x16<f16, A>>(&cm_a[ab + sg * 512u + 0u], 32u);
            let b0 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[ab + 0u + 0u], 32u);
            let b1 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[ab + 512u + 0u], 32u);
            let b2 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[ab + 1024u + 0u], 32u);
            let b3 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[ab + 1536u + 0u], 32u);
            c0 = coopMultiplyAdd(a, b0, c0);
            c1 = coopMultiplyAdd(a, b1, c1);
            c2 = coopMultiplyAdd(a, b2, c2);
            c3 = coopMultiplyAdd(a, b3, c3);
        }
        {
            let a = coopLoadT<coop_mat16x16<f16, A>>(&cm_a[ab + sg * 512u + 16u], 32u);
            let b0 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[ab + 0u + 16u], 32u);
            let b1 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[ab + 512u + 16u], 32u);
            let b2 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[ab + 1024u + 16u], 32u);
            let b3 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[ab + 1536u + 16u], 32u);
            c0 = coopMultiplyAdd(a, b0, c0);
            c1 = coopMultiplyAdd(a, b1, c1);
            c2 = coopMultiplyAdd(a, b2, c2);
            c3 = coopMultiplyAdd(a, b3, c3);
        }
        workgroupBarrier();
        k0 = nk;
        buf = buf ^ 1u;
        if (k0 >= cols) { break; }
    }
    coopStoreT(c0, &cm_c[sg * 16u * 64u + 0u], 64u);
    coopStoreT(c1, &cm_c[sg * 16u * 64u + 16u], 64u);
    coopStoreT(c2, &cm_c[sg * 16u * 64u + 32u], 64u);
    coopStoreT(c3, &cm_c[sg * 16u * 64u + 48u], 64u);
    workgroupBarrier();
    for (var t = tid; t < 64u * 64u; t = t + 128u) {
        let m = t / 64u;
        let n = t % 64u;
        if (m0 + m < pmm.nb && n0 + n < pmm.rows) {
            ymm[(m0 + m) * pmm.rows + n0 + n] = cm_c[m * 64u + n];
        }
    }
}
"#;

/// The backward twin on the same units: dx[nb,rows] = dy[nb,cols] ·
/// w[cols,rows] — the reduction runs over w's ROWS (its leading axis), so
/// the B staging reads a column of w per output: four stride-`rows` loads
/// where the forward kernel read four consecutive. Everything else —
/// tiles, loads, the store — is the forward kernel unchanged.
const COOP_NN_SRC: &str = r#"
enable wgpu_cooperative_matrix;
enable f16;

struct MmP { cols4: u32, rows: u32, nb: u32, pad: u32 };
@group(0) @binding(0) var<storage, read> wmm: array<f32>;
@group(0) @binding(1) var<storage, read> xmm: array<f32>;
@group(0) @binding(2) var<storage, read_write> ymm: array<f32>;
@group(0) @binding(3) var<uniform> pmm: MmP;

const KS: u32 = 32u;

var<workgroup> cm_a: array<f16, 64 * 32>;
var<workgroup> cm_b: array<f16, 64 * 32>;
var<workgroup> cm_c: array<f32, 64 * 64>;

@compute @workgroup_size(128)
fn gemm_nn_coop(@builtin(workgroup_id) wid: vec3<u32>,
                @builtin(local_invocation_index) tid: u32,
                @builtin(subgroup_id) sg: u32) {
    let cols = pmm.cols4 * 4u;
    let m0 = wid.y * 64u;
    let n0 = wid.x * 64u;

    var c0: coop_mat16x16<f32, C>;
    var c1: coop_mat16x16<f32, C>;
    var c2: coop_mat16x16<f32, C>;
    var c3: coop_mat16x16<f32, C>;

    var k0 = 0u;
    loop {
        if (k0 >= cols) { break; }
        for (var t = tid; t < 64u * 8u; t = t + 128u) {
            let m = t / 8u;
            let k4 = (t % 8u) * 4u;
            let col0 = k0 + k4;
            let dst = m * KS + k4;
            var v = vec4<f32>(0.0);
            if (m0 + m < pmm.nb && col0 < cols) {
                let base = (m0 + m) * cols + col0;
                v = vec4<f32>(xmm[base], xmm[base + 1u], xmm[base + 2u], xmm[base + 3u]);
            }
            cm_a[dst] = f16(v.x); cm_a[dst + 1u] = f16(v.y);
            cm_a[dst + 2u] = f16(v.z); cm_a[dst + 3u] = f16(v.w);
        }
        for (var t = tid; t < 64u * 8u; t = t + 128u) {
            let n = t / 8u;
            let k4 = (t % 8u) * 4u;
            let col0 = k0 + k4;
            var wv = vec4<f32>(0.0);
            if (n0 + n < pmm.rows && col0 < cols) {
                // w is [cols][rows] row-major here: the four reduction
                // neighbours live a full row apart.
                let base = col0 * pmm.rows + n0 + n;
                let lim = cols - col0;
                wv.x = wmm[base];
                if (lim > 1u) { wv.y = wmm[base + pmm.rows]; }
                if (lim > 2u) { wv.z = wmm[base + 2u * pmm.rows]; }
                if (lim > 3u) { wv.w = wmm[base + 3u * pmm.rows]; }
            }
            let bd = n * KS + k4;
            cm_b[bd] = f16(wv.x); cm_b[bd + 1u] = f16(wv.y);
            cm_b[bd + 2u] = f16(wv.z); cm_b[bd + 3u] = f16(wv.w);
        }
        workgroupBarrier();
        {
            let a = coopLoadT<coop_mat16x16<f16, A>>(&cm_a[sg * 512u + 0u], 32u);
            let b0 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[0u + 0u], 32u);
            let b1 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[512u + 0u], 32u);
            let b2 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1024u + 0u], 32u);
            let b3 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1536u + 0u], 32u);
            c0 = coopMultiplyAdd(a, b0, c0);
            c1 = coopMultiplyAdd(a, b1, c1);
            c2 = coopMultiplyAdd(a, b2, c2);
            c3 = coopMultiplyAdd(a, b3, c3);
        }
        {
            let a = coopLoadT<coop_mat16x16<f16, A>>(&cm_a[sg * 512u + 16u], 32u);
            let b0 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[0u + 16u], 32u);
            let b1 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[512u + 16u], 32u);
            let b2 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1024u + 16u], 32u);
            let b3 = coopLoad<coop_mat16x16<f16, B>>(&cm_b[1536u + 16u], 32u);
            c0 = coopMultiplyAdd(a, b0, c0);
            c1 = coopMultiplyAdd(a, b1, c1);
            c2 = coopMultiplyAdd(a, b2, c2);
            c3 = coopMultiplyAdd(a, b3, c3);
        }
        workgroupBarrier();
        k0 = k0 + KS;
    }
    workgroupBarrier();
    coopStoreT(c0, &cm_c[sg * 16u * 64u + 0u], 64u);
    coopStoreT(c1, &cm_c[sg * 16u * 64u + 16u], 64u);
    coopStoreT(c2, &cm_c[sg * 16u * 64u + 32u], 64u);
    coopStoreT(c3, &cm_c[sg * 16u * 64u + 48u], 64u);
    workgroupBarrier();
    for (var t = tid; t < 64u * 64u; t = t + 128u) {
        let m = t / 64u;
        let n = t % 64u;
        if (m0 + m < pmm.nb && n0 + n < pmm.rows) {
            ymm[(m0 + m) * pmm.rows + n0 + n] = cm_c[m * 64u + n];
        }
    }
}
"#;

/// Did the device come up with cooperative matrices? Read by the GEMM to
/// pick its pipeline, and by the shader-module builder to decide whether
/// the tensor-core source can be compiled at all.
static COOP_OK: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);

pub fn coop_matrix_active() -> bool {
    COOP_OK.load(std::sync::atomic::Ordering::Relaxed)
}

#[inline]
fn begin_pass(enc: &mut wgpu::CommandEncoder) -> PassHandle {
    // A stage label set by the BT frame turns this pass into a timestamped
    // one; everything else stays on the cold path of one atomic load.
    let which = BT_TS_STAGE.load(std::sync::atomic::Ordering::Relaxed);
    let timestamp_writes = if which != 0 {
        ctx().and_then(|c| ts_pair(c, which))
    } else {
        None
    };
    begin_pass_with(enc, None, timestamp_writes)
}

// ── Pass merging. A compute-pass boundary costs ~9 µs on this Vulkan
// stack (measured: 700 one-dispatch passes 9.7 ms against 3.4 ms for
// the same 700 dispatches in ONE pass), and a token issued 147 of them,
// a batched verify ~500. Dispatches inside a pass are serialized with
// memory visibility, so a graph needs no boundary at all except where
// the ENCODER itself is used — a copy, a timestamp, a submit. An
// encoder that opts in (`PassMergeGuard`, the two graph functions) gets
// ONE open pass per thread that every `begin_pass` on it hands back
// (`forget_lifetime` detaches it from the borrow), and every direct
// encoder use flushes it first (`flush_pass`, `finish_enc`). Encoders
// that did not opt in behave exactly as before. `CMF_PASS_MERGE=0`
// turns the merge off.
thread_local! {
    static MERGE_KEYS: std::cell::RefCell<std::collections::HashSet<usize>> =
        std::cell::RefCell::new(std::collections::HashSet::new());
    static OPEN_PASS: std::cell::RefCell<Option<(usize, wgpu::ComputePass<'static>)>> =
        const { std::cell::RefCell::new(None) };
}

fn pass_merge_on() -> bool {
    static N: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *N.get_or_init(|| std::env::var("CMF_PASS_MERGE").as_deref() != Ok("0"))
}

/// The address of the encoder object, whether the caller holds it by
/// value or behind a reference — the merge state is keyed on it.
trait EncAddr {
    fn enc_addr(&self) -> usize;
}
impl EncAddr for wgpu::CommandEncoder {
    fn enc_addr(&self) -> usize {
        self as *const wgpu::CommandEncoder as usize
    }
}
impl EncAddr for &mut wgpu::CommandEncoder {
    fn enc_addr(&self) -> usize {
        (&**self) as *const wgpu::CommandEncoder as usize
    }
}
impl EncAddr for &wgpu::CommandEncoder {
    fn enc_addr(&self) -> usize {
        (*self) as *const wgpu::CommandEncoder as usize
    }
}

/// End the merged pass open on this thread, if any. Call before any
/// direct use of an encoder (copy, timestamp, resolve, finish, swap).
/// Deliberately NOT keyed on the encoder: a helper that takes the graph's
/// encoder BY VALUE (readback) sees it at a new address, and a local
/// encoder finishing while the graph's pass is open only splits that
/// pass early — harmless, since dispatch order inside the graph's own
/// encoder is unchanged.
fn flush_pass<E: EncAddr>(enc: &E) {
    let _ = enc.enc_addr();
    OPEN_PASS.with(|o| {
        let mut o = o.borrow_mut();
        if o.is_some() {
            *o = None; // drop ends the pass into its encoder
        }
    });
}

/// `finish_enc(enc)` with the merged pass flushed first.
fn finish_enc(enc: wgpu::CommandEncoder) -> wgpu::CommandBuffer {
    flush_pass(&enc);
    enc.finish()
}

/// Opt an encoder into pass merging for the guard's lifetime; drop
/// flushes and unregisters.
struct PassMergeGuard {
    key: usize,
}
impl PassMergeGuard {
    fn new<E: EncAddr>(enc: &E) -> Self {
        let key = enc.enc_addr();
        if pass_merge_on() {
            MERGE_KEYS.with(|m| {
                m.borrow_mut().insert(key);
            });
        }
        Self { key }
    }
}
impl Drop for PassMergeGuard {
    fn drop(&mut self) {
        let key = self.key;
        OPEN_PASS.with(|o| {
            let mut o = o.borrow_mut();
            if matches!(&*o, Some((k, _)) if *k == key) {
                *o = None;
            }
        });
        MERGE_KEYS.with(|m| {
            m.borrow_mut().remove(&key);
        });
    }
}

/// What `begin_pass` hands out: a compute pass that is either this
/// encoder's merged open pass (given back to the thread-local on drop)
/// or a plain pass ended on drop. Derefs to the pass, so call sites are
/// unchanged.
pub struct PassHandle {
    pass: Option<wgpu::ComputePass<'static>>,
    key: usize,
    merged: bool,
}
impl std::ops::Deref for PassHandle {
    type Target = wgpu::ComputePass<'static>;
    fn deref(&self) -> &Self::Target {
        self.pass.as_ref().expect("pass handle")
    }
}
impl std::ops::DerefMut for PassHandle {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.pass.as_mut().expect("pass handle")
    }
}
impl Drop for PassHandle {
    fn drop(&mut self) {
        if let Some(p) = self.pass.take() {
            if self.merged {
                OPEN_PASS.with(|o| {
                    *o.borrow_mut() = Some((self.key, p));
                });
            }
            // else: `p` drops here and the pass ends.
        }
    }
}

fn begin_pass_with<'a>(
    enc: &'a mut wgpu::CommandEncoder,
    label: Option<&'a str>,
    timestamp_writes: Option<wgpu::ComputePassTimestampWrites<'a>>,
) -> PassHandle {
    let key = (*enc).enc_addr();
    let merge = timestamp_writes.is_none() && MERGE_KEYS.with(|m| m.borrow().contains(&key));
    if merge {
        // Hand back the open pass if it is this encoder's; otherwise
        // flush whatever is open (another encoder's) and open ours.
        let taken = OPEN_PASS.with(|o| {
            let mut o = o.borrow_mut();
            match o.take() {
                Some((k, p)) if k == key => Some(p),
                Some((_, p)) => {
                    drop(p);
                    None
                }
                None => None,
            }
        });
        if let Some(p) = taken {
            return PassHandle {
                pass: Some(p),
                key,
                merged: true,
            };
        }
        PASSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        let p = enc
            .begin_compute_pass(&wgpu::ComputePassDescriptor {
                label,
                timestamp_writes: None,
            })
            .forget_lifetime();
        return PassHandle {
            pass: Some(p),
            key,
            merged: true,
        };
    }
    // A plain pass: make sure no merged pass is open on this encoder
    // (a second pass on a locked encoder is a validation error).
    flush_pass(&*enc);
    PASSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    let p = enc
        .begin_compute_pass(&wgpu::ComputePassDescriptor {
            label,
            timestamp_writes,
        })
        .forget_lifetime();
    PassHandle {
        pass: Some(p),
        key,
        merged: false,
    }
}

fn readback2(
    c: &Ctx,
    mut enc: wgpu::CommandEncoder,
    a: (&wgpu::Buffer, &mut [f32]),
    b: (&wgpu::Buffer, &mut [f32]),
) -> bool {
    let (a_buf, a_out) = a;
    let (b_buf, b_out) = b;
    let a_bytes = (a_out.len() * 4) as u64;
    let b_bytes = (b_out.len() * 4) as u64;
    // COPY_BUFFER_ALIGNMENT is 4; the second slice starts on a 16-byte
    // boundary so the map range stays comfortably aligned on every backend.
    let off = a_bytes.div_ceil(16) * 16;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage2,
        off + b_bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dsv4-pair-stage",
    );
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(a_buf, 0, &stage, 0, a_bytes);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(b_buf, 0, &stage, off, b_bytes);
    submit(c, finish_enc(enc));
    let slice = stage.slice(..off + b_bytes);
    let done = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
    let d2 = done.clone();
    slice.map_async(wgpu::MapMode::Read, move |_| {
        d2.store(true, std::sync::atomic::Ordering::Release);
    });
    if spin_wait() {
        let t0 = std::time::Instant::now();
        loop {
            let _ = c.device.poll(wgpu::PollType::Poll);
            if done.load(std::sync::atomic::Ordering::Acquire) {
                break;
            }
            if t0.elapsed() > std::time::Duration::from_millis(2) {
                if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
                    stage.unmap();
                    return false;
                }
                break;
            }
            std::hint::spin_loop();
        }
    } else if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        stage.unmap();
        return false;
    }
    {
        let Ok(data) = slice.get_mapped_range() else {
            stage.unmap();
            return false;
        };
        a_out.copy_from_slice(bytemuck::cast_slice(&data[..a_bytes as usize]));
        b_out.copy_from_slice(bytemuck::cast_slice(
            &data[off as usize..(off + b_bytes) as usize],
        ));
    }
    stage.unmap();
    true
}

/// Two buffers in one submit and one map: `a` lands at staging[0..a_size),
/// `b` right after it. The MTP draft's graph step reads its logits AND
/// its block hidden back this way instead of paying a second fence.
#[allow(clippy::too_many_arguments)]
fn readback_two(
    c: &Ctx,
    mut enc: wgpu::CommandEncoder,
    a_buf: &wgpu::Buffer,
    a_size: u64,
    b_buf: &wgpu::Buffer,
    b_size: u64,
    staging: &wgpu::Buffer,
    out_a: &mut [f32],
    out_b: &mut [f32],
) -> bool {
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(a_buf, 0, staging, 0, a_size);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(b_buf, 0, staging, a_size, b_size);
    submit(c, finish_enc(enc));
    let slice = staging.slice(..a_size + b_size);
    let done = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
    let d2 = done.clone();
    slice.map_async(wgpu::MapMode::Read, move |_| {
        d2.store(true, std::sync::atomic::Ordering::Release);
    });
    if spin_wait() {
        let t0 = std::time::Instant::now();
        loop {
            let _ = c.device.poll(wgpu::PollType::Poll);
            if done.load(std::sync::atomic::Ordering::Acquire) {
                break;
            }
            if t0.elapsed() > std::time::Duration::from_millis(2) {
                if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
                    staging.unmap();
                    return false;
                }
                break;
            }
            std::hint::spin_loop();
        }
    } else if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        staging.unmap();
        return false;
    }
    {
        let Ok(data) = slice.get_mapped_range() else {
            staging.unmap();
            return false;
        };
        let a_len = out_a.len() * 4;
        par_copy(out_a, bytemuck::cast_slice(&data[..a_len]));
        let b_off = a_size as usize;
        par_copy(
            out_b,
            bytemuck::cast_slice(&data[b_off..b_off + out_b.len() * 4]),
        );
    }
    staging.unmap();
    true
}

fn readback(
    c: &Ctx,
    mut enc: wgpu::CommandEncoder,
    y_buf: &wgpu::Buffer,
    staging: &wgpu::Buffer,
    y_size: u64,
    out: &mut [f32],
) -> bool {
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(y_buf, 0, staging, 0, y_size);
    let t_dma = std::time::Instant::now();
    submit(c, finish_enc(enc));
    let slice = staging.slice(..y_size);
    let done = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
    let d2 = done.clone();
    slice.map_async(wgpu::MapMode::Read, move |_| {
        d2.store(true, std::sync::atomic::Ordering::Release);
    });
    // A blocking wait hands the thread to the OS scheduler, and getting it
    // back costs more than the work did: these submissions finish in tens of
    // microseconds and there are 86 of them a token. Spin on the queue for a
    // short while first, then block — a decode that stalls for a real reason
    // must not burn a core forever. CMF_GPU_SPIN=0 reverts.

    if spin_wait() {
        let t0 = std::time::Instant::now();
        loop {
            let _ = c.device.poll(wgpu::PollType::Poll);
            if done.load(std::sync::atomic::Ordering::Acquire) {
                break;
            }
            if t0.elapsed() > std::time::Duration::from_millis(2) {
                if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
                    staging.unmap();
                    return false;
                }
                break;
            }
            std::hint::spin_loop();
        }
    } else if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        staging.unmap();
        return false;
    }

    let t_cp = std::time::Instant::now();
    {
        let Ok(data) = slice.get_mapped_range() else {
            staging.unmap();
            return false;
        };
        par_copy(out, bytemuck::cast_slice(&data[..out.len() * 4]));
        // One line per big readback under CMF_RB_TRACE=1: how much of it
        // is the DMA+fence and how much the mapped-memory copy. 45 MB at
        // 0.8 GB/s effective needed this split to know which half to fix.
        if y_size > (8 << 20) && std::env::var("CMF_RB_TRACE").as_deref() == Ok("1") {
            let full = t_dma.elapsed().as_secs_f64() * 1e3;
            let cp = t_cp.elapsed().as_secs_f64() * 1e3;
            eprintln!(
                "rb {:.1} MB: dma+fence {:.1} ms, copy {:.1} ms",
                y_size as f64 / 1e6,
                full - cp,
                cp
            );
        }
    }
    staging.unmap();
    READBACK_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    READBACK_BYTES.fetch_add(y_size, std::sync::atomic::Ordering::Relaxed);
    READBACK_WAIT_NS.fetch_add(
        t_cp.duration_since(t_dma).as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    READBACK_COPY_NS.fetch_add(
        t_cp.elapsed().as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    READBACK_TOTAL_NS.fetch_add(
        t_dma.elapsed().as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    true
}

/// A mapped range is uncached host memory: one thread streams it at
/// ~1-2 GB/s, and the bake's big planes (a 264 MB dw readback) turn that
/// memcpy into the call's dominant cost. Split copies above a few MB
/// across threads — read bandwidth scales nearly linearly.
fn par_copy(out: &mut [f32], src: &[f32]) {
    // 4 MB, and it stays there. The tempting change is to lower it so the
    // decode's own per-token copy qualifies — a folded lm_head reads back
    // the whole logit row, 993 KB at Qwen3.6's 248320 vocab, which is
    // real single-threaded time out of a 20 ms token. MEASURED, and it
    // LOSES: at 256 KB the production loop goes 44.4 -> 42.5 tok/s and
    // greedy 49.0 -> 47.3 (medians of three, RTX 5090). `thread::scope`
    // spawns fresh OS threads on every call, and for a megabyte once a
    // token that costs more than the copy. It pays for the bake's 264 MB
    // planes, which is what the threshold was chosen for. A pooled
    // version would be a different measurement; this one is not it.
    // `CMF_PARCOPY_MIN` (bytes) is there to re-run the experiment.
    let par_min: usize = std::env::var("CMF_PARCOPY_MIN")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(4 << 20);
    if out.len() * 4 >= par_min {
        let lanes = std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(4)
            .min(16);
        let chunk = out.len().div_ceil(lanes);
        std::thread::scope(|s| {
            for (dst, sr) in out.chunks_mut(chunk).zip(src.chunks(chunk)) {
                s.spawn(move || dst.copy_from_slice(sr));
            }
        });
    } else {
        out.copy_from_slice(src);
    }
}

fn bind_buf(binding: u32, buf: &wgpu::Buffer) -> wgpu::BindGroupEntry<'_> {
    wgpu::BindGroupEntry {
        binding,
        resource: buf.as_entire_binding(),
    }
}

/// A window into a strided per-token buffer: how the batched frame hands a
/// single-token kernel one token's slice without changing the kernel. The
/// offset must respect `min_storage_buffer_offset_alignment` (256 on every
/// card this runs on), which every per-token stride here does.
fn bind_buf_off(binding: u32, buf: &wgpu::Buffer, off: u64, size: u64) -> wgpu::BindGroupEntry<'_> {
    wgpu::BindGroupEntry {
        binding,
        resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
            buffer: buf,
            offset: off,
            size: std::num::NonZeroU64::new(size),
        }),
    }
}

fn storage_bytes(c: &Ctx, data: &[u8]) -> wgpu::Buffer {
    c.device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: data,
            usage: wgpu::BufferUsages::STORAGE,
        })
}

fn uniform_u32x4(c: &Ctx, v: [u32; 4]) -> wgpu::Buffer {
    // Content-keyed cache: these params (rows/cols/flags) repeat every token,
    // so build each once and clone the handle thereafter.
    let mut u = c.uniforms.lock().unwrap();
    if let Some(b) = u.get(&v) {
        return b.clone();
    }
    let b = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&v),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    u.insert(v, b.clone());
    b
}

/// Three words and a float, in one uniform. The cache is keyed on the bits,
/// which is exactly right: two params differ iff their bytes differ.
fn uniform_mixed(c: &Ctx, v: [u32; 3], f: f32) -> wgpu::Buffer {
    uniform_u32x4(c, [v[0], v[1], v[2], f.to_bits()])
}

fn rw_f32(c: &Ctx, n: usize, copy_src: bool) -> wgpu::Buffer {
    let usage = if copy_src {
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC
    } else {
        wgpu::BufferUsages::STORAGE
    };
    c.device.create_buffer(&wgpu::BufferDescriptor {
        label: None,
        size: (n * 4) as u64,
        usage,
        mapped_at_creation: false,
    })
}

/// Resident quant weights of tensor `idx` (the whole tensor, cached by (file,idx)).
fn tensor_weight(
    c: &Ctx,
    model: &Arc<CmfModel>,
    idx: usize,
    rows: usize,
    cols: usize,
) -> Option<wgpu::Buffer> {
    let entry = &model.tensors[idx];
    if entry.shape.first().copied().unwrap_or(0) < rows {
        return None;
    }
    let abs = model.entry_abs_offset(entry)?;
    let bytes = model.primary_bytes();
    if abs + rows * cols > bytes.len() {
        return None;
    }
    let key = (model.uid() as usize, idx);
    let lay = layer_of_name(&model.tensors[idx].name);
    if !weight_resident(c, key) {
        res_who(&model.tensors[idx].name);
        if let Some(v) = host_tier_get(key) {
            return weight_buffer_l(c, key, &v, lay);
        }
        if let Some(v) = pread_range(model, abs, rows * cols) {
            let v = std::sync::Arc::new(v);
            host_tier_put(key, v.clone());
            return weight_buffer_l(c, key, &v, lay);
        }
        if host_tier().is_some() {
            let v = std::sync::Arc::new(bytes[abs..abs + rows * cols].to_vec());
            host_tier_put(key, v.clone());
            return weight_buffer_l(c, key, &v, lay);
        }
    }
    weight_buffer_l(c, key, &bytes[abs..abs + rows * cols], lay)
}

/// `tensor_weight` for tile-packed dtypes whose payload length differs
/// from rows·cols (q4_tiled: 18 B per 32-weight group).
fn tensor_weight_sized(
    c: &Ctx,
    model: &Arc<CmfModel>,
    idx: usize,
    rows: usize,
    payload: usize,
) -> Option<wgpu::Buffer> {
    let entry = &model.tensors[idx];
    if entry.shape.first().copied().unwrap_or(0) < rows {
        return None;
    }
    let abs = model.entry_abs_offset(entry)?;
    let bytes = model.primary_bytes();
    if abs + payload > bytes.len() {
        return None;
    }
    let key = (model.uid() as usize, idx);
    let lay = layer_of_name(&model.tensors[idx].name);
    if !weight_resident(c, key) {
        if let Some(v) = host_tier_get(key) {
            return weight_buffer_l(c, key, &v, lay);
        }
        if let Some(v) = pread_range(model, abs, payload) {
            let v = std::sync::Arc::new(v);
            host_tier_put(key, v.clone());
            return weight_buffer_l(c, key, &v, lay);
        }
        // mmap source: still worth keeping a tier copy — the next miss
        // of this tensor must not fault the pages again.
        if host_tier().is_some() {
            let v = std::sync::Arc::new(bytes[abs..abs + payload].to_vec());
            host_tier_put(key, v.clone());
            return weight_buffer_l(c, key, &v, lay);
        }
    }
    weight_buffer_l(c, key, &bytes[abs..abs + payload], lay)
}

/// Encodes q8-matvec (row0=0) into the given encoder, writes to `y`. The bind
/// group and uniform are ref-counted by the command buffer until submit.
fn encode_matvec(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    rs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
) {
    let p_buf = uniform_u32x4(c, [(cols / 4) as u32, rows as u32, 0, 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &c.layout,
        entries: &[
            bind_buf(0, weight),
            bind_buf(1, xs),
            bind_buf(2, rs),
            bind_buf(3, y),
            bind_buf(4, &p_buf),
        ],
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(&c.matvec);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
}

/// q1 cousin of `encode_matvec`: the q1 pipeline + `layout_q1` (4 bindings,
/// no row-scale — q1 carries its scales inside the tiles). params = the
/// `dispatch_q1` layout `[gpr/2, rows, 0, 0]`. Lets q1 QKV share one encoder.
fn encode_matvec_q1(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
) {
    let gpr = cols / 32;
    let p_buf = uniform_u32x4(c, [(gpr / 2) as u32, rows as u32, 0, 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &c.layout_q1,
        entries: &[
            bind_buf(0, weight),
            bind_buf(1, xs),
            bind_buf(2, y),
            bind_buf(3, &p_buf),
        ],
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(&c.q1);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((rows as u32).div_ceil(c.q1_rows).min(MAX_WG), 1, 1);
}

/// Encode a resident q1 GEMM (batched prefill): Y[k,rows] = X[k,cols] @ Wᵀ, all
/// buffers already on the device. q1_mul_mm omits binding 2 (no row scale).
#[allow(dead_code)] // wired by forward_batch_graph (batched prefill, in progress)
/// Batched q4_tiled / q4tp GEMM into `enc` — the tile GEMMs the imagegen
/// path already used, wired for the graph.
///
/// `ematb` matched kind 0 and sent EVERYTHING else to the q1 kernel, which
/// for a q4tp weight is simply the wrong decoder. Together with `gemmable`
/// admitting only kinds 0 and 1, that shut batched prefill out of every
/// q4t and q4tp file — the same shape of bug as the `prep()` hole, and the
/// reason prefill ran one position at a time at 33 tok/s against 54 on
/// decode.
/// One-shot reason the BATCHED graph declined. Three silent fallbacks in a
/// row this session cost hours; a refusal that says nothing is the most
/// expensive kind of bug in this file.
fn bgraph_refused(why: &'static str) {
    use std::sync::atomic::{AtomicBool, Ordering};
    static SAID: AtomicBool = AtomicBool::new(false);
    if !SAID.swap(true, Ordering::Relaxed) {
        tracing::warn!("batch graph declined: {why}");
    }
}

fn token_graph_outcome(o1_started: bool, ok: bool) -> crate::gpu::TokenGraphOutcome {
    if ok {
        crate::gpu::TokenGraphOutcome::Completed
    } else if o1_started {
        crate::gpu::TokenGraphOutcome::Failed
    } else {
        crate::gpu::TokenGraphOutcome::Declined
    }
}

fn batch_outcome(o1_started: bool, ok: bool) -> crate::gpu::BatchGraphOutcome {
    if ok {
        crate::gpu::BatchGraphOutcome::Completed
    } else if o1_started {
        crate::gpu::BatchGraphOutcome::Failed
    } else {
        crate::gpu::BatchGraphOutcome::Declined
    }
}

// Keep this in sync with `oa_scr` in the WGSL above.  It is the total number
// of exact near entries (sinks + ring window), not a silent enlargement of
// the landmark budget or the portable CPU/Metal state.
const O1_MAX_NEAR: usize = 2052;

fn o1_view_valid(v: &crate::nystrom::O1DeviceView<'_>, nh: usize, nkv: usize, hd: usize) -> bool {
    let hpg = nh.checked_div(nkv).unwrap_or(0);
    let md = v.m_eff.checked_mul(v.d).unwrap_or(usize::MAX);
    let mdv = v.m_eff.checked_mul(v.dv).unwrap_or(usize::MAX);
    let mm = v.m_eff.checked_mul(v.m_eff).unwrap_or(usize::MAX);
    let wd = v.w.checked_mul(v.d).unwrap_or(usize::MAX);
    let wdv = v.w.checked_mul(v.dv).unwrap_or(usize::MAX);
    !v.exact_only
        && nkv > 0
        && nh % nkv == 0
        && v.heads.len() == hpg
        && (4..=32).contains(&v.m_eff)
        && v.w > 0
        && v.sink_len.saturating_add(v.w) <= O1_MAX_NEAR
        && v.d == hd
        && v.dv == hd
        && v.d <= 256
        && v.dv <= 256
        && v.win_len <= v.w
        && v.win_head < v.w
        && v.win_k.len() >= v.win_len.saturating_mul(v.d)
        && v.win_k.len() <= wd
        && v.win_v.len() >= v.win_len.saturating_mul(v.dv)
        && v.win_v.len() <= wdv
        && v.sink_k.len() == v.sink_len.saturating_mul(v.d)
        && v.sink_v.len() == v.sink_len.saturating_mul(v.dv)
        && v.k_tilde.len() == md
        && v.heads.iter().all(|h| {
            h.t_hat.len() == mdv
                && h.z_hat.len() == v.m_eff
                && h.m_max.len() == v.m_eff
                && h.q_tilde.len() == md
                && h.mu.len() == mm
        })
}

fn o1_views_valid(
    views: &[crate::nystrom::O1DeviceView<'_>],
    nh: usize,
    nkv: usize,
    hd: usize,
) -> bool {
    let Some(first) = views.first() else {
        return false;
    };
    o1_view_valid(first, nh, nkv, hd)
        && views.iter().all(|v| {
            o1_view_valid(v, nh, nkv, hd)
                && v.m_eff == first.m_eff
                && v.w == first.w
                && v.sink_len == first.sink_len
                && v.d == first.d
                && v.dv == first.dv
                && v.heads.len() == first.heads.len()
        })
}

/// Device mirror of one layer's sealed o1 state.
struct O1Dev {
    epoch: u64,
    /// Absolute position of the next row that may mutate this state.  None
    /// means the first post-seal graph call has not advanced it yet.
    next_pos: Option<usize>,
    /// Logical device bytes owned by this sealed state.  wgpu does not expose
    /// a portable buffer-size query, so keep the exact upload footprint next
    /// to the handles for diagnostics and reset checks.
    bytes: u64,
    meta: wgpu::Buffer,
    ring_k: wgpu::Buffer,
    ring_v: wgpu::Buffer,
    sink_k: wgpu::Buffer,
    sink_v: wgpu::Buffer,
    k_tilde: wgpu::Buffer,
    qt: wgpu::Buffer,
    mu: wgpu::Buffer,
    mz: wgpu::Buffer,
    that: wgpu::Buffer,
    g: usize,
    h: usize,
    m: usize,
    w: usize,
    ns: usize,
    scale: f32,
}

/// Return the number of O(1) device mirrors currently owned by `kv_id` and
/// their logical upload footprint.  The map is keyed by sequence id and layer
/// so a pooled server can prove that a fresh request did not inherit the prior
/// request's sealed state.
pub fn o1_device_stats(kv_id: u64) -> (usize, u64) {
    let Some(c) = ctx() else { return (0, 0) };
    let map = c.o1m.lock().unwrap();
    map.iter()
        .filter(|((id, _), _)| *id == kv_id)
        .fold((0usize, 0u64), |(layers, bytes), (_, state)| {
            (layers + 1, bytes.saturating_add(state.bytes))
        })
}

/// Upload (or reuse) a layer's o1 state. One upload per seal epoch: the
/// window ring and far skeleton then live and MUTATE on the device, and
/// the CPU copy is stale by design — the same one-way discipline as the
/// KV mirror.
fn o1_ensure(
    c: &Ctx,
    kv_id: u64,
    li: usize,
    views: &[crate::nystrom::O1DeviceView<'_>],
    epoch: u64,
) -> Option<()> {
    {
        let m = c.o1m.lock().unwrap();
        if let Some(d) = m.get(&(kv_id, li)) {
            if d.epoch == epoch {
                return Some(());
            }
        }
    }
    tracing::info!("o1_ensure: UPLOADING layer {li} (epoch {epoch})");
    let g0 = views.first()?;
    let (gcnt, hcnt, m, w, ns) = (views.len(), g0.heads.len(), g0.m_eff, g0.w, g0.sink_len);
    // Landmark threads park at lane 200+ in the attend kernel.
    //
    // A refusal here silently costs the whole token its graph — 6.4
    // against 49.2 tok/s on Qwen3.8 — so every branch says which limit
    // it is and with what numbers. Silent capability gates are how a
    // day was once spent reading one KV ceiling as three model bugs.
    if ns.saturating_add(w) > O1_MAX_NEAR || m > 32 || g0.d > 256 || g0.dv > 256 {
        graph_refused("o1 geometry over kernel limits");
        tracing::warn!(
            "o1_ensure L{li}: sink+window {}+{} (cap {O1_MAX_NEAR}), landmarks {m} (cap 32), \
             d {} dv {} (cap 256)",
            ns,
            w,
            g0.d,
            g0.dv
        );
        return None;
    }
    for v in views {
        if v.m_eff != m || v.w != w || v.sink_len != ns || v.heads.len() != hcnt {
            graph_refused("o1 groups disagree on geometry");
            tracing::warn!(
                "o1_ensure L{li}: group m_eff {} w {} sink {} heads {} vs first {m}/{w}/{ns}/{hcnt}",
                v.m_eff,
                v.w,
                v.sink_len,
                v.heads.len()
            );
            return None;
        }
    }
    let (d, dv) = (g0.d, g0.dv);
    let stor_f = |data: &[f32], label: &str| -> wgpu::Buffer {
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size: ((data.len() * 4).max(4)) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        c.queue.write_buffer(&b, 0, bytemuck::cast_slice(data));
        b
    };
    let mut meta = Vec::with_capacity(gcnt * 4);
    let (mut rk, mut rv, mut sk, mut sv, mut kt) = (vec![], vec![], vec![], vec![], vec![]);
    let (mut qt, mut mu, mut mz, mut th) = (vec![], vec![], vec![], vec![]);
    for v in views {
        meta.extend_from_slice(&[v.win_len as u32, v.win_head as u32, v.far_len as u32, 0]);
        // Ring buffers are cap-sized already (cap = w in skeleton mode).
        rk.extend_from_slice(v.win_k);
        rk.resize(rk.len() + (w * d - v.win_k.len().min(w * d)), 0.0);
        rv.extend_from_slice(v.win_v);
        rv.resize(rv.len() + (w * dv - v.win_v.len().min(w * dv)), 0.0);
        sk.extend_from_slice(v.sink_k);
        sv.extend_from_slice(v.sink_v);
        kt.extend_from_slice(v.k_tilde);
        for hh in &v.heads {
            qt.extend_from_slice(hh.q_tilde);
            mu.extend_from_slice(hh.mu);
            mz.extend_from_slice(hh.m_max);
            mz.extend_from_slice(hh.z_hat);
            th.extend_from_slice(hh.t_hat);
        }
    }
    let meta_b = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("o1-meta"),
        size: (meta.len() * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    c.queue
        .write_buffer(&meta_b, 0, bytemuck::cast_slice(&meta));
    let dev = O1Dev {
        epoch,
        next_pos: None,
        bytes: [
            meta.len(),
            rk.len(),
            rv.len(),
            sk.len(),
            sv.len(),
            kt.len(),
            qt.len(),
            mu.len(),
            mz.len(),
            th.len(),
        ]
        .into_iter()
        .map(|n| n.max(1) as u64 * std::mem::size_of::<f32>() as u64)
        .sum(),
        meta: meta_b,
        ring_k: stor_f(&rk, "o1-rk"),
        ring_v: stor_f(&rv, "o1-rv"),
        sink_k: stor_f(&sk, "o1-sk"),
        sink_v: stor_f(&sv, "o1-sv"),
        k_tilde: stor_f(&kt, "o1-kt"),
        qt: stor_f(&qt, "o1-qt"),
        mu: stor_f(&mu, "o1-mu"),
        mz: stor_f(&mz, "o1-mz"),
        that: stor_f(&th, "o1-th"),
        g: gcnt,
        h: hcnt,
        m,
        w,
        ns,
        scale: g0.scale,
    };
    c.o1m.lock().unwrap().insert((kv_id, li), dev);
    Some(())
}

fn encode_q4_tile_mm(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    pipeline: &wgpu::ComputePipeline,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    k: usize,
) {
    encode_q4_tile_mm_scaled(c, enc, pipeline, weight, xs, y, rows, cols, k, 0.0)
}

/// Encode the scalar q2tp prefill GEMM with its explicit descriptor center.
/// q2tp reuses the four-word MM parameter block, but word 3 is a boolean
/// affine selector (`0 = code-1.5`, `1 = code-1`), not the f32 activation
/// scale used by the q4 cooperative variants.  Keeping this wrapper separate
/// prevents the generic q4 helper from erasing that format bit.
fn encode_q2_tile_mm(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    k: usize,
    affine: bool,
    use_coop: bool,
) {
    let pipeline = if use_coop {
        c.q2tp_mm_coop.as_ref().unwrap_or(&c.q2tp_mm)
    } else {
        &c.q2tp_mm
    };
    let p_buf = uniform_u32x4(c, [
        (cols / 4) as u32,
        rows as u32,
        k as u32,
        u32::from(affine),
    ]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("q2tp-mm"),
        layout: &pipeline.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, weight),
            bind_buf(1, xs),
            bind_buf(2, y),
            bind_buf(3, &p_buf),
        ],
    });
    let active_bytes = cortiq_core::quant::expected_nbytes(
        cortiq_core::TensorDtype::Q2TiledP,
        &[rows, cols],
    )
    .unwrap_or(0) as u64
        + ((k * cols + k * rows) * std::mem::size_of::<f32>()) as u64;
    let tsw = batch_kernel_ts_pair(c, 0, active_bytes);
    let mut pass = begin_pass_with(
        enc,
        Some(if use_coop { "q2tp-mm-coop" } else { "q2tp-mm" }),
        tsw,
    );
    pass.set_pipeline(pipeline);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(
        (rows as u32).div_ceil(64).min(MAX_WG),
        (k as u32).div_ceil(64),
        1,
    );
}

/// The same, carrying the activation scale the cooperative kernel reads
/// out of `pmm.pad`. 0 = no scaling (the scalar arm's contract). A fused
/// chain that skipped this let f16 operands overflow at render batch
/// sizes — correct at 256×160, nearly blank frames at 512×288.
#[allow(clippy::too_many_arguments)]
fn encode_q4_tile_mm_scaled(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    pipeline: &wgpu::ComputePipeline,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    k: usize,
    ascale: f32,
) {
    encode_q4_tile_mm_full(c, enc, pipeline, weight, xs, y, rows, cols, k, ascale, None)
}

/// The full form: `scale_buf` carries a DEVICE-computed activation scale
/// (see `act_absmax`) for a GEMM whose input never reaches the host.
#[allow(clippy::too_many_arguments)]
fn encode_q4_tile_mm_full(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    pipeline: &wgpu::ComputePipeline,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    k: usize,
    ascale: f32,
    scale_buf: Option<&wgpu::Buffer>,
) {
    encode_q4_tile_mm_full_offset(
        c, enc, pipeline, weight, xs, 0, y, rows, cols, k, ascale, scale_buf,
    )
}

/// The same Q4TP dispatch with a byte offset into a larger resident
/// activation panel.  Qwen's joint attention writes `[text, image]` into one
/// token-major buffer; its two output projections bind the corresponding
/// windows directly so the panel never needs a host split or device copy.
#[allow(clippy::too_many_arguments)]
fn encode_q4_tile_mm_full_offset(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    pipeline: &wgpu::ComputePipeline,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    xs_offset: u64,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    k: usize,
    ascale: f32,
    scale_buf: Option<&wgpu::Buffer>,
) {
    // Is this one of the cooperative kernels (both declare binding 4)?
    // wgpu 30 exposes no identity on pipeline handles, so the layouts are
    // compared by their entry count instead: the coop layouts have five
    // bindings, the scalar tile GEMMs four.
    let coop_kernel = c
        .q4tp_mm_coop_s
        .as_ref()
        .map(|p| std::ptr::eq(p, pipeline))
        .unwrap_or(false)
        || c.q4tp_mm_coop_f16
            .as_ref()
            .map(|p| std::ptr::eq(p, pipeline))
            .unwrap_or(false);
    // 0xFFFFFFFF is not a float the host would ever pass; the kernel
    // reads it as "take the scale from the buffer instead".
    let pad = if scale_buf.is_some() {
        0xFFFF_FFFFu32
    } else {
        ascale.to_bits()
    };
    let p_buf = uniform_u32x4(c, [(cols / 4) as u32, rows as u32, k as u32, pad]);
    let layout = pipeline.get_bind_group_layout(0);
    let dummy;
    let sbuf = match scale_buf {
        Some(b) => b,
        None => {
            dummy = c
                .device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytemuck::cast_slice(&[0.0f32]),
                    usage: wgpu::BufferUsages::STORAGE,
                });
            &dummy
        }
    };
    let src_bytes = (k as u64).saturating_mul(cols as u64).saturating_mul(4);
    let src_entry = if xs_offset == 0 {
        bind_buf(1, xs)
    } else {
        bind_buf_off(1, xs, xs_offset, src_bytes)
    };
    let mut entries = vec![bind_buf(0, weight), src_entry, bind_buf(2, y), bind_buf(3, &p_buf)];
    // Only the f16 twin declares binding 4. Deciding that by comparing
    // pipeline ADDRESSES was wrong — the handle passed in is not the
    // same object as the one in the Ctx, so the entry was never added
    // and every dispatch failed layout validation. Ask the pipeline's
    // own layout instead: it is the thing that knows.
    // wgpu 30 exposes no identity on either handle, so the caller says
    // so: `scale_buf` is passed exactly when the f16 twin is the
    // pipeline, and that twin is the only one with binding 4.
    // Both cooperative kernels declare binding 4 now (the in-kernel one
    // grew it for the batched prefill); the scalar tile GEMMs do not.
    // The layout is the pipeline's own — ask it whether slot 4 exists
    // by trying: a caller passing `scale_buf` for a coop kernel, or a
    // coop kernel with a host scale (dummy buffer), both need the entry.
    if scale_buf.is_some() || coop_kernel {
        entries.push(bind_buf(4, sbuf));
    }
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &layout,
        entries: &entries,
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(pipeline);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(
        (rows as u32).div_ceil(64).min(MAX_WG),
        (k as u32).div_ceil(64),
        1,
    );
}

/// The GEMM pipeline to use for this weight layout: the tensor-core one
/// when the device brought it up, the scalar one otherwise.
fn mm_pipeline(c: &Ctx, q4tp: bool, two_bit: bool) -> &wgpu::ComputePipeline {
    if two_bit {
        // No cooperative variant: the tensor-core kernel is written
        // against the 4-bit plane, and a shader compiled against a
        // layout the weights do not have is a silently wrong answer.
        return &c.q2tp_mm;
    }
    if q4tp {
        if let Some(p) = c.q4tp_mm_coop.as_ref() {
            return p;
        }
        return &c.q4tp_mm;
    }
    &c.q4t_mm
}

fn encode_q1_mm(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    k: usize,
) {
    let p_buf = uniform_u32x4(c, [(cols / 4) as u32, rows as u32, k as u32, 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &c.layout_q1mm,
        entries: &[
            bind_buf(0, weight),
            bind_buf(1, xs),
            bind_buf(3, y),
            bind_buf(4, &p_buf),
        ],
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(&c.q1_mm);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(
        (rows as u32).div_ceil(64).min(MAX_WG),
        (k as u32).div_ceil(64),
        1,
    );
}

/// Encode a resident q8 GEMM (int8 weight + per-row f32 scale) into `enc`.
#[allow(dead_code)] // wired by forward_batch_graph (batched prefill, in progress)
fn encode_q8_mm(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    rs: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    k: usize,
) {
    let p_buf = uniform_u32x4(c, [(cols / 4) as u32, rows as u32, k as u32, 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &c.layout_mmm,
        entries: &[
            bind_buf(0, weight),
            bind_buf(1, xs),
            bind_buf(2, rs),
            bind_buf(3, y),
            bind_buf(4, &p_buf),
        ],
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(&c.mul_mm);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(
        (rows as u32).div_ceil(64).min(MAX_WG),
        (k as u32).div_ceil(64),
        1,
    );
}

/// Encode a plain f32 matvec (small unquantized projections) into `enc`.
/// The keyed twin of `encode_f32matvec`: every buffer at the call site is
/// stable for the layer's lifetime, so the group is built once per
/// (tag, sequence, layer) and reused — fresh groups here were a measurable
/// share of the chain's host encoding.
#[allow(clippy::too_many_arguments)]
fn encode_f32matvec_k(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    bkey: (u8, u64, usize),
) {
    let mut pass = begin_pass(enc);
    encode_f32matvec_k_p(&mut pass, c, weight, xs, y, rows, cols, bkey);
}

#[allow(clippy::too_many_arguments)]
/// The chain's f32 matvec: 256 threads a row instead of 64.
#[allow(clippy::too_many_arguments)]
fn encode_f32matvec_w_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    bkey: (u8, u64, usize),
) {
    // Rows are what give this kernel its workgroups. With very few of them,
    // split the shared axis instead: rows·splits workgroups and one cheap
    // merge pass. CMF_DSV4_F32SPLIT=0 reverts.
    // 512, not 4096: the stands' hc·dim is 512, and a threshold that no toy
    // can reach is a path no gate can check — this one reached the release
    // unexercised and failed there on the first dispatch.
    if rows <= 32 && cols >= 512 && f32_split() {
        let nsplit = 8usize;
        let chunk = cols.div_ceil(nsplit);
        // The caller's ROLE tag has to be in the key. Both hyper-connection
        // mixes of a layer arrive here with the same (sequence, layer) and
        // different roles, and sharing one slot means the second one's
        // parameters reach the first one's dispatch — queue writes all land
        // before the submission. CMF_DSV4_SLOT_CHECK=1 named this exactly:
        // "slot (186, 1, 0) written 2 times".
        let key = bkey.2 * 256 + bkey.0 as usize;
        // The partials buffer likewise: two roles in one submission would
        // otherwise write the same scratch.
        let part = frame_buf(c, 115 + (bkey.0 & 1), rows * nsplit * 4, false);
        let p = uni_slot(
            c,
            186,
            bkey.1,
            key,
            [cols as u32, rows as u32, nsplit as u32, chunk as u32],
        );
        let b1 = cached_bind(c, (188, bkey.1, key), || {
            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &c.f32_mv_split.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, weight),
                    bind_buf(1, xs),
                    bind_buf(2, y),
                    bind_buf(3, &p),
                    bind_buf(4, &part),
                ],
            })
        });
        pass.set_pipeline(&c.f32_mv_split);
        pass.set_bind_group(0, &b1, &[]);
        pass.dispatch_workgroups(rows as u32, nsplit as u32, 1);
        let b2 = cached_bind(c, (190, bkey.1, key), || {
            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &c.f32_mv_merge.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, weight),
                    bind_buf(1, xs),
                    bind_buf(2, y),
                    bind_buf(3, &p),
                    bind_buf(4, &part),
                ],
            })
        });
        pass.set_pipeline(&c.f32_mv_merge);
        pass.set_bind_group(0, &b2, &[]);
        pass.dispatch_workgroups(rows as u32, 1, 1);
        return;
    }
    let pipe = if rows < 64 {
        &c.f32_matvec_x
    } else {
        &c.f32_matvec_w
    };
    let bind = cached_bind(c, bkey, || {
        let p = uniform_u32x4(c, [cols as u32, rows as u32, 0, 0]);
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &pipe.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, weight),
                bind_buf(1, xs),
                bind_buf(2, y),
                bind_buf(3, &p),
            ],
        })
    });
    pass.set_pipeline(pipe);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
}

fn encode_f32matvec_k_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    bkey: (u8, u64, usize),
) {
    let bind = cached_bind(c, bkey, || {
        let p_buf = uniform_u32x4(c, [cols as u32, rows as u32, 0, 0]);
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.layout_f32,
            entries: &[
                bind_buf(0, weight),
                bind_buf(1, xs),
                bind_buf(2, y),
                bind_buf(3, &p_buf),
            ],
        })
    });
    pass.set_pipeline(&c.f32_matvec);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
}

fn encode_f32matvec(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
) {
    let p_buf = uniform_u32x4(c, [cols as u32, rows as u32, 0, 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &c.layout_f32,
        entries: &[
            bind_buf(0, weight),
            bind_buf(1, xs),
            bind_buf(2, y),
            bind_buf(3, &p_buf),
        ],
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(&c.f32_matvec);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
}

/// `encode_f32matvec` with byte offsets into `xs` and `y` — the batched MoE
/// router runs the SAME f32 kernel per token (bit-for-bit the logits the
/// parity-proven path produced) but reads its token's row of the batch
/// hidden and writes its token's slice of the logit plane directly. Both
/// offsets land on 256-byte boundaries (t·hidden·4 and t·n_exp·4 with
/// hidden=2048, n_exp≤256), which is all wgpu asks of a buffer binding.
#[allow(clippy::too_many_arguments)]
/// Content-keyed cache for 8-word uniforms — the per-token GDN params of a
/// batched chunk repeat every chunk, and `unif` mints a fresh buffer per
/// call (the OOM lesson of the folded-gate work).
fn uniform_u32x8(c: &Ctx, v: [u32; 8]) -> wgpu::Buffer {
    let mut u = c.uniforms8.lock().unwrap();
    if let Some(b) = u.get(&v) {
        return b.clone();
    }
    let b = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&v),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    u.insert(v, b.clone());
    b
}

fn encode_f32matvec_off(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    xs_off: u64,
    y: &wgpu::Buffer,
    y_off: u64,
    y_len: u64,
    rows: usize,
    cols: usize,
) {
    let p_buf = uniform_u32x4(c, [cols as u32, rows as u32, 0, 0]);
    let entries = [
        wgpu::BindGroupEntry {
            binding: 0,
            resource: weight.as_entire_binding(),
        },
        wgpu::BindGroupEntry {
            binding: 1,
            resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                buffer: xs,
                offset: xs_off,
                size: wgpu::BufferSize::new((cols * 4) as u64),
            }),
        },
        wgpu::BindGroupEntry {
            binding: 2,
            resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                buffer: y,
                offset: y_off,
                size: wgpu::BufferSize::new(y_len),
            }),
        },
        wgpu::BindGroupEntry {
            binding: 3,
            resource: p_buf.as_entire_binding(),
        },
    ];
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &c.layout_f32,
        entries: &entries,
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(&c.f32_matvec);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
}

/// Encode a q4_tiled or q1t matvec into `enc` (same 4-slot layout as q1, but
/// params are [gpr, rows, cols]; q1t reads its sparse overlay from the tail of
/// the same buffer). `pipeline` is c.q4b or c.q1t.
/// Attend kernel flavor by head_dim: stride-129 (16.5 KB of workgroup
/// memory, exists on every device) for hd <= 128, stride-257 for larger
/// heads (desktop-only — see Ctx::hd_cap).
fn attend_pipes(c: &Ctx, hd: usize) -> (&wgpu::ComputePipeline, &wgpu::BindGroupLayout) {
    if hd <= 128 {
        (&c.gqa_attend_s, &c.layout_attend_s)
    } else {
        (&c.gqa_attend, &c.layout_attend)
    }
}

fn attend_part_pipes(c: &Ctx, hd: usize) -> (&wgpu::ComputePipeline, &wgpu::BindGroupLayout) {
    if hd <= 128 {
        (&c.attend_part_s, &c.layout_attend_part_s)
    } else {
        (&c.attend_part, &c.layout_attend_part)
    }
}

/// `q4tp_matvec4`: the q1t-like binding set plus the weight buffer AGAIN at
/// slot 4 as the kernel's vec4 nibble view.
/// Parameters for the q4tp matvec PAIR (`q4tp_matvec16` / `q4tp_matvec4`).
///
/// Word 2 is the batch count for these two kernels and `cols` for everything
/// else bound to `Q1Params`. Building it by hand at each call site put `cols`
/// into the batch slot of the chain's busiest projection encoder — the
/// kernel then ran every row block four thousand times and decoding fell
/// from 27 tok/s to 0.35. One constructor, so the mistake has nowhere to
/// live.
fn q4tp_mv_params(c: &Ctx, gpr: usize, rows: usize, batch: usize) -> wgpu::Buffer {
    q4tp_mv_params_w(c, gpr, rows, batch, 0)
}

/// The same, with word 3 — the grouped projection's row window, which slides
/// the activation with the row. Batch and window are mutually exclusive
/// modes; the kernel reads the batch first.
fn q4tp_mv_params_w(c: &Ctx, gpr: usize, rows: usize, batch: usize, window: usize) -> wgpu::Buffer {
    debug_assert!((1..=64).contains(&batch), "batch {batch} out of range");
    debug_assert!(
        batch == 1 || window == 0,
        "batch {batch} with window {window}: the kernel honours one or the other"
    );
    uniform_u32x4(c, [gpr as u32, rows as u32, batch as u32, window as u32])
}

fn encode_q4tp_mv4(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
) {
    encode_q4tp_mv4_b(c, enc, weight, xs, y, rows, cols, 1);
}

/// The same projection against `batch` activation vectors laid end to end,
/// writing `batch * rows` outputs. The weight is read once for all of them,
/// which is the point: a prompt chunk and a speculative verify both need
/// exactly this and nothing else.
///
/// Always encodes: the row blocking sits inside one batch element by
/// construction, so no shape is refused.
#[allow(clippy::too_many_arguments)]
/// Two batched projections of one input batch in ONE dispatch (the bku
/// x2 kernel): both q4tp, wide (gpr > 64), batch 2..=8 (arm 2's range).
/// False = does not qualify; the caller issues two `encode_q4tp_mv4_b`.
#[allow(clippy::too_many_arguments)]
fn encode_q4tp_mv4_b_x2(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    wa: &wgpu::Buffer,
    wb: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    ya: &wgpu::Buffer,
    yb: &wgpu::Buffer,
    rows_a: usize,
    rows_b: usize,
    cols: usize,
    batch: usize,
) -> bool {
    let gpr = cols / 32;
    if !c.use_mv_x2 || c.use_mv_bk < 2 || !(2..=8).contains(&batch) || gpr <= 64 || cols % 32 != 0 {
        return false;
    }
    let p_buf = uniform_u32x4(c, [gpr as u32, rows_a as u32, batch as u32, rows_b as u32]);
    let layout = c.q4tp_mv4_bku_x2.get_bind_group_layout(0);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("mv-bku-x2"),
        layout: &layout,
        entries: &[
            bind_buf(0, wa),
            bind_buf(2, ya),
            bind_buf(3, &p_buf),
            bind_buf(4, wa),
            bind_buf(5, xs),
            bind_buf(6, wb),
            bind_buf(7, wb),
            bind_buf(8, yb),
        ],
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(&c.q4tp_mv4_bku_x2);
    pass.set_bind_group(0, &bind, &[]);
    let wg = (rows_a as u32).div_ceil(16) + (rows_b as u32).div_ceil(16);
    pass.dispatch_workgroups(mv_grid(wg), 1, 1);
    true
}

fn encode_q4tp_mv4_b(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    batch: usize,
) -> bool {
    encode_q4tp_mv4_b_with(c, enc, weight, xs, y, rows, cols, batch, verify_i8_on())
}

/// `encode_q4tp_mv4_b` with the int8-activation arm chosen explicitly
/// (the parity tests pin the f32 kernels; production reads the switch).
#[allow(clippy::too_many_arguments)]
fn encode_q4tp_mv4_b_with(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    batch: usize,
    i8: bool,
) -> bool {
    let gpr = cols / 32;
    // A small batch over WIDE rows goes to the register-blocked kernel:
    // it reads the weight once for the whole batch, where the arms below
    // read it once per element and leave the reuse to L2. On the dense
    // 5120/17408 FFN planes that reuse never materialized — a k=3 verify
    // cost 2.95x one row's, which is the flat "no reuse" number.
    // Arm 2 chunks the batch internally, so it covers up to 8; arm 1
    // has four accumulator components and stops at 4.
    let bk_max = if c.use_mv_bk >= 2 { 8 } else { 4 };
    if i8 && (2..=8).contains(&batch) && gpr > 64 && cols % 32 == 0 {
        encode_q4tp_mv4_b_i8(c, enc, weight, xs, y, rows, cols, batch);
        return true;
    }
    if c.use_mv_bk > 0 && (2..=bk_max).contains(&batch) && gpr > 64 {
        // A kernel swap that cannot be OBSERVED is a kernel swap that
        // gets credited with someone else's timing. One line, once.
        if std::env::var("CMF_MV_BK_TRACE").is_ok() {
            use std::sync::atomic::{AtomicBool, Ordering};
            static SAID: AtomicBool = AtomicBool::new(false);
            if !SAID.swap(true, Ordering::Relaxed) {
                eprintln!(
                    "mv-bk{}: engaged ({rows}x{cols}, batch {batch}, gpr {gpr})",
                    c.use_mv_bk
                );
            }
        }
        let p_buf = q4tp_mv_params(c, gpr, rows, batch);
        // The layout comes from the pipeline that will RUN. wgpu's auto
        // layouts are built per entry point and are not interchangeable
        // even between kernels declaring the same bindings — borrowing a
        // sibling's cost a day once already.
        let bk = if c.use_mv_bk >= 2 {
            &c.q4tp_mv4_bku
        } else {
            &c.q4tp_mv4_bk
        };
        let layout = bk.get_bind_group_layout(0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layout,
            entries: &[
                bind_buf(0, weight),
                bind_buf(2, y),
                bind_buf(3, &p_buf),
                bind_buf(4, weight),
                bind_buf(5, xs),
            ],
        });
        let mut pass = begin_pass(enc);
        pass.set_pipeline(bk);
        pass.set_bind_group(0, &bind, &[]);
        // The batch lives INSIDE the workgroup now, so the grid is the
        // row blocks alone — the same grid, and the same sixteen rows a
        // workgroup, that the one-vector kernel runs.
        pass.dispatch_workgroups(mv_grid((rows as u32).div_ceil(16)), 1, 1);
        return true;
    }
    // Narrow shapes (one group per lane in the 8-row kernel) go 16-rows.
    // With a batch the divisibility decides first: a block that straddled
    // two tokens would read one token's weights against the other's x.
    let (pipe, per_wg) = if gpr <= 64 {
        (&c.q4tp_mv16, 16u32)
    } else if batch == 1 {
        (&c.q4tp_mv16w, 16u32)
    } else {
        (
            if c.use_mv_nored {
                &c.q4tp_mv4_nored
            } else {
                c.q4tp_mv4_sg.as_ref().unwrap_or(if c.use_mv_u2 {
                    &c.q4tp_mv4_u2
                } else {
                    &c.q4tp_mv4
                })
            },
            8u32,
        )
    };
    let p_buf = q4tp_mv_params(c, gpr, rows, batch);
    let layout = pipe.get_bind_group_layout(0);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &layout,
        entries: &[
            bind_buf(0, weight),
            bind_buf(2, y),
            bind_buf(3, &p_buf),
            bind_buf(4, weight),
            bind_buf(5, xs),
        ],
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(pipe);
    pass.set_bind_group(0, &bind, &[]);
    // One workgroup per (row block, batch element); the batch is the fast
    // axis inside the kernel, so a row block is never split across two.
    pass.dispatch_workgroups(mv_grid((rows as u32).div_ceil(per_wg) * batch as u32), 1, 1);
    true
}

/// The one-row q4tp kernel, which is the one `gpu_q4tp_parity` blesses.
/// `encode_q4tp_mv4` picks a wider variant by shape; when a frame has to
/// agree with the CPU to the last bit, agreement beats throughput.
#[allow(clippy::too_many_arguments)]
fn encode_q4tp_mv1(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    bkey: (u8, u64, usize),
) {
    let mut pass = begin_pass(enc);
    encode_q4tp_mv1_p(&mut pass, c, weight, xs, y, rows, cols, bkey);
}

#[allow(clippy::too_many_arguments)]
fn encode_q4tp_mv1_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    bkey: (u8, u64, usize),
) {
    let bind = cached_bind(c, bkey, || {
        let p_buf = uniform_u32x4(c, [(cols / 32) as u32, rows as u32, cols as u32, 0]);
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.q4tp_mv.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, weight),
                bind_buf(1, xs),
                bind_buf(2, y),
                bind_buf(3, &p_buf),
            ],
        })
    });
    pass.set_pipeline(&c.q4tp_mv);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
}

/// The wide q4tp matvec, as a dispatch inside a caller's pass.
///
/// The chain used the one-row kernel throughout — one workgroup a row — while
/// `encode_q4tp_mv4` picks an 8- or 16-row variant by shape and is the
/// default everywhere else, including the head that measured 0.32 ms against
/// the host's 9.43. The projections inside a chained layer are the same
/// shapes; there is no reason for them to take the narrow kernel.
///
/// It sums the same products in a different lane order, so it is a contract
/// change and carries `CMF_DSV4_MV4=0`.
#[allow(clippy::too_many_arguments)]
fn encode_q4tp_mvw_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    bkey: (u8, u64, usize),
) {
    if !chain_mv4() {
        encode_q4tp_mv1_p(pass, c, weight, xs, y, rows, cols, bkey);
        return;
    }
    let gpr = cols / 32;
    let (pipe, per_wg) = if gpr <= 64 {
        (&c.q4tp_mv16, 16u32)
    } else {
        (
            if c.use_mv_nored {
                &c.q4tp_mv4_nored
            } else {
                c.q4tp_mv4_sg.as_ref().unwrap_or(if c.use_mv_u2 {
                    &c.q4tp_mv4_u2
                } else {
                    &c.q4tp_mv4
                })
            },
            8u32,
        )
    };
    let bind = cached_bind(c, bkey, || {
        let p_buf = q4tp_mv_params(c, gpr, rows, 1);
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &pipe.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, weight),
                bind_buf(2, y),
                bind_buf(3, &p_buf),
                bind_buf(4, weight),
                bind_buf(5, xs),
            ],
        })
    });
    pass.set_pipeline(pipe);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((rows as u32).div_ceil(per_wg).min(MAX_WG), 1, 1);
}

/// The grouped low-rank projection through the EIGHT-ROW q4tp kernel.
///
/// Its weights are ordinary q4tp — the one-row kernel it had was a copy of
/// q4tp_matvec with one added term in the activation index — so the only
/// thing standing between it and the register-blocked kernel was that
/// sliding window, which `_p1` now carries. Measured 3.82 ms against
/// wo_b's 1.24 on comparable weights through the fast one.
///
/// Only when the group width is a multiple of 8: the 8 rows a workgroup
/// owns must share one window, or the pair-blocked x fetch is wrong.
#[allow(clippy::too_many_arguments)]
fn encode_o_lora_mv4_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
    lora: usize,
    bkey: (u8, u64, usize),
) -> bool {
    if lora == 0 || cols % 32 != 0 {
        return false;
    }
    let gpr = cols / 32;
    let (pipe, per_wg) = if gpr <= 64 {
        (&c.q4tp_mv16, 16u32)
    } else {
        (
            if c.use_mv_nored {
                &c.q4tp_mv4_nored
            } else {
                c.q4tp_mv4_sg.as_ref().unwrap_or(if c.use_mv_u2 {
                    &c.q4tp_mv4_u2
                } else {
                    &c.q4tp_mv4
                })
            },
            8u32,
        )
    };
    // The rows a workgroup owns must share one activation window, and a
    // workgroup owns `per_wg` consecutive rows starting at a multiple of
    // `per_wg`. The 16-row kernel therefore needs a width divisible by 16,
    // not by 8 — which is exactly the case the stands hit (their gpr is 2,
    // so they take the 16-row path) and exactly why this read 133.433
    // against the host's 133.396.
    if lora % per_wg as usize != 0 {
        return false;
    }
    let bind = cached_bind(c, bkey, || {
        let p = q4tp_mv_params_w(c, gpr, rows, 1, lora);
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &pipe.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, weight),
                bind_buf(2, y),
                bind_buf(3, &p),
                bind_buf(4, weight),
                bind_buf(5, xs),
            ],
        })
    });
    pass.set_pipeline(pipe);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((rows as u32).div_ceil(per_wg).min(MAX_WG), 1, 1);
    true
}

/// `CMF_DSV4_MOE4=0` puts the q2tp experts back on one row a workgroup.
/// `CMF_DSV4_F32SPLIT=1` splits the few-row f32 matvec over its columns.
///
/// OFF: it measured 27.1 tok/s against the one-dispatch path's 27.2. The
/// commit that claimed to turn it off did not: its patch made two edits and
/// the second failed an assertion, so the file was never written and only
/// the message changed. A review reading the source rather than the log
/// caught it.
fn f32_split() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| std::env::var("CMF_DSV4_F32SPLIT").is_ok_and(|v| v != "0"))
}

/// `CMF_DSV4_OLORA_MV4=0` keeps the grouped projection on its own kernel.
///
/// It is 3.82 ms of a 30.6 ms chain while wo_b — comparable weights through
/// the register-blocked q4tp kernel — is 1.24, and the only difference is
/// the activation window that slides with the row. `_p1` carries it now, in
/// BOTH the 8-row and the 16-row kernel: teaching only the 8-row one is
/// what made this read 133.433 against the host's 133.396, because the
/// stands have gpr = 2 and take the 16-row path. The window was simply
/// missing there, and the rows kept reading the first group's activations.
fn olora_mv4() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| {
        std::env::var("CMF_DSV4_OLORA_MV4")
            .map(|v| v != "0")
            .unwrap_or(true)
    })
}

/// OFF by default: the A/B put four rows a workgroup at 24.9 tok/s against
/// 25.7 without, with the chain's wait identical to within 0.12 ms — so the
/// difference is host noise and the change buys nothing measurable. Kept
/// switchable; the default takes the path that has been measured longer.
fn moe4() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| std::env::var("CMF_DSV4_MOE4").is_ok_and(|v| v != "0"))
}

/// `CMF_DSV4_OLORA=1|2|3` pins the grouped projection to the one-row,
/// 256-thread or four-row kernel, for the A/B that decides which.
fn olora_pick() -> Option<u32> {
    static P: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
    *P.get_or_init(|| {
        std::env::var("CMF_DSV4_OLORA")
            .ok()
            .and_then(|v| v.parse().ok())
    })
}

/// `CMF_DSV4_MV4=0` puts the chain back on the one-row q4tp kernel.
fn chain_mv4() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| {
        std::env::var("CMF_DSV4_MV4")
            .map(|v| v != "0")
            .unwrap_or(true)
    })
}

fn encode_q1t_like(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    pipeline: &wgpu::ComputePipeline,
    weight: &wgpu::Buffer,
    xs: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
) {
    let gpr = cols / 32;
    let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, cols as u32, 0]);
    let layout = pipeline.get_bind_group_layout(0);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &layout,
        entries: &[
            bind_buf(0, weight),
            bind_buf(1, xs),
            bind_buf(2, y),
            bind_buf(3, &p_buf),
        ],
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(pipeline);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
}

/// Fused SiLU(gate)·up → Q4Block down-proj: one dispatch instead of silu + matvec.
fn encode_silu_down(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    weight: &wgpu::Buffer,
    gate: &wgpu::Buffer,
    up: &wgpu::Buffer,
    y: &wgpu::Buffer,
    rows: usize,
    cols: usize,
) {
    let gpr = cols / 32;
    let p_buf = uniform_u32x4(c, [gpr as u32, rows as u32, cols as u32, 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &c.layout_silu_down,
        entries: &[
            bind_buf(0, weight),
            bind_buf(1, gate),
            bind_buf(2, up),
            bind_buf(3, y),
            bind_buf(4, &p_buf),
        ],
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(&c.silu_down);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
}

/// q1 batched matvec: N q1 projections (e.g. QKV) in ONE submit + one
/// readback — the chain-fusion that `matvec_batch` does for q8, now for
/// 1-bit weights. Bails to `false` (→ CPU) on any budget/shape refusal so
/// the caller's fallback stays intact.
fn matvec_batch_q1(model: &Arc<CmfModel>, jobs: &[BatchJob], out: &mut [&mut [f32]]) -> bool {
    let Some(c) = ctx() else { return false };
    let bytes = model.primary_bytes();
    // Resident weight per job (VRAM cache; over-budget/oob → honest CPU).
    let mut weights = Vec::with_capacity(jobs.len());
    for j in jobs {
        let gpr = j.cols / 32;
        if j.rows == 0 || j.cols % 32 != 0 || gpr % 2 != 0 || j.xs.len() < j.cols {
            return false;
        }
        let entry = &model.tensors[j.idx];
        if entry.shape.first().copied().unwrap_or(0) < j.rows {
            return false;
        }
        let Some(abs) = model.entry_abs_offset(entry) else {
            return false;
        };
        let plen = j.rows * gpr * 6;
        if abs + plen > bytes.len() {
            return false;
        }
        let Some(w) = weight_buffer(c, (model.uid() as usize, j.idx), &bytes[abs..abs + plen])
        else {
            return false;
        };
        weights.push(w);
    }
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("q1-batch"),
        });
    let mut y_bufs = Vec::with_capacity(jobs.len());
    for (j, w) in jobs.iter().zip(&weights) {
        let xs_b = storage_bytes(c, bytemuck::cast_slice(&j.xs[..j.cols]));
        let y_b = rw_f32(c, j.rows, true);
        encode_matvec_q1(c, &mut enc, w, &xs_b, &y_b, j.rows, j.cols);
        y_bufs.push(y_b);
    }
    // ONE pooled staging buffer for all outputs, one map (mirror the q8 path).
    let total: u64 = jobs.iter().map(|j| (j.rows * 4) as u64).sum();
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        total,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "q1-batch-stage",
    );
    let mut off = 0u64;
    for (y_b, j) in y_bufs.iter().zip(jobs) {
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(y_b, 0, &stage, off, (j.rows * 4) as u64);
        off += (j.rows * 4) as u64;
    }
    submit(c, finish_enc(enc));
    stage.slice(..total).map_async(wgpu::MapMode::Read, |_| {});
    if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        return false;
    }
    {
        let Ok(data) = stage.slice(..total).get_mapped_range() else {
            return false;
        };
        let mut off = 0usize;
        for (j, o) in jobs.iter().zip(out.iter_mut()) {
            o[..j.rows].copy_from_slice(bytemuck::cast_slice(&data[off..off + j.rows * 4]));
            off += j.rows * 4;
        }
    }
    stage.unmap();
    drop(sc);
    true
}

/// Layer MoE-FFN in a single submission: for each expert gate/up-matvec →
/// silu·mul·col_down → down-matvec → y += w·d. Intermediate buffers are
/// GPU-resident, one sync per layer.
pub fn moe_block(model: &Arc<CmfModel>, jobs: &[MoeJob], out: &mut [f32]) -> bool {
    if jobs.iter().any(|j| j.q1) {
        return false; // q1 WGSL kernel not implemented yet — honest CPU
    }
    if jobs.iter().any(|j| j.gu_q2) {
        // Mixed 2-bit gate/up: this per-op path has no q2tp matvec WGSL —
        // the 2-bit lanes live in the graphs. Honest refusal, CPU or the
        // Metal jobs path own it.
        return false;
    }
    let Some(c) = ctx() else { return false };
    if jobs.is_empty() {
        return false;
    }
    let q4t = jobs[0].q4t;
    let q4tp = jobs[0].q4tp;
    if jobs.iter().any(|j| j.q4t != q4t || j.q4tp != q4tp) {
        return false; // mixed job kinds — honest CPU
    }
    if q4t && q4tp {
        return false; // a trio is one layout or the other
    }
    let inter = jobs[0].gate.1;
    let hidden = jobs[0].down.1;
    if out.len() != hidden {
        return false;
    }
    // Resident weights of all triples — validate first (fail → CPU entirely).
    let fetch = |idx: usize, rows: usize, cols: usize| -> Option<wgpu::Buffer> {
        if q4tp {
            // Three planes, not a flat tile: nibbles, then the per-row
            // (lo, step) pair, then the 5-bit rung codes. Only the layout
            // owner knows the total, so ask it rather than re-deriving.
            let n = cortiq_core::quant::expected_nbytes(
                cortiq_core::TensorDtype::Q4TiledP,
                &[rows, cols],
            )?;
            tensor_weight_sized(c, model, idx, rows, n)
        } else if q4t {
            tensor_weight_sized(c, model, idx, rows, rows * (cols / 32) * 18)
        } else {
            tensor_weight(c, model, idx, rows, cols)
        }
    };
    let mut w3 = Vec::with_capacity(jobs.len());
    for j in jobs {
        let (gi, gr, gc, _) = j.gate;
        let (ui, ur, uc, _) = j.up;
        let (di, dr, dc, _) = j.down;
        let align = if q4t || q4tp { 32 } else { 4 };
        if gc % align != 0 || uc % align != 0 || dc % align != 0 {
            return false;
        }
        let (Some(gw), Some(uw), Some(dw)) =
            (fetch(gi, gr, gc), fetch(ui, ur, uc), fetch(di, dr, dc))
        else {
            return false;
        };
        w3.push((gw, uw, dw));
    }

    let g_buf = rw_f32(c, inter, false);
    let u_buf = rw_f32(c, inter, false);
    let a_buf = rw_f32(c, inter, false);
    let d_buf = rw_f32(c, hidden, false);
    let y_buf = rw_f32(c, hidden, true);

    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("moe") });

    // y = 0
    {
        let np = uniform_u32x4(c, [hidden as u32, 0, 0, 0]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.layout_zero,
            entries: &[bind_buf(0, &y_buf), bind_buf(1, &np)],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.zero);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((hidden as u32).div_ceil(256), 1, 1);
    }

    for (j, (gw, uw, dw)) in jobs.iter().zip(&w3) {
        let (_, gr, gc, grs) = &j.gate;
        let (_, ur, uc, urs) = &j.up;
        let (_, dr, dc, drs) = &j.down;
        // Per-tensor scale/col buffers are stable across tokens — cache
        // them like the matvec row-scales instead of re-uploading.
        let mut rs_map = c.rs_bufs.lock().unwrap();
        let mut cached = |tag: usize, idx: usize, data: &[f32]| -> wgpu::Buffer {
            rs_map
                .entry((model.uid() as usize, (idx, tag)))
                .or_insert_with(|| {
                    crate::gpu::probe_note_cold();
                    storage_bytes(c, bytemuck::cast_slice(data))
                })
                .clone()
        };
        let grs_b = cached(1, j.gate.0, grs);
        let urs_b = cached(2, j.up.0, urs);
        let drs_b = cached(3, j.down.0, drs);
        let has_col = !j.down_col.is_empty();
        let col_b = if has_col {
            cached(4, j.down.0, j.down_col)
        } else {
            cached(5, usize::MAX, &[0f32]) // dummy, gated by f=0
        };
        drop(rs_map);
        let xsg = storage_bytes(c, bytemuck::cast_slice(&j.xs_gate));
        let xsu = storage_bytes(c, bytemuck::cast_slice(&j.xs_up));

        if q4tp {
            encode_q1t_like(c, &mut enc, &c.q4tp_mv, gw, &xsg, &g_buf, *gr, *gc);
            encode_q1t_like(c, &mut enc, &c.q4tp_mv, uw, &xsu, &u_buf, *ur, *uc);
        } else if q4t {
            encode_q1t_like(c, &mut enc, &c.q4t_mv, gw, &xsg, &g_buf, *gr, *gc);
            encode_q1t_like(c, &mut enc, &c.q4t_mv, uw, &xsu, &u_buf, *ur, *uc);
        } else {
            encode_matvec(c, &mut enc, gw, &xsg, &grs_b, &g_buf, *gr, *gc);
            encode_matvec(c, &mut enc, uw, &xsu, &urs_b, &u_buf, *ur, *uc);
        }
        // act = silu(g)·u·col_down
        {
            let np = uniform_u32x4(
                c,
                [inter as u32, has_col as u32, j.swiglu_limit.to_bits(), 0],
            );
            let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &c.layout_silu,
                entries: &[
                    bind_buf(0, &g_buf),
                    bind_buf(1, &u_buf),
                    bind_buf(2, &col_b),
                    bind_buf(3, &a_buf),
                    bind_buf(4, &np),
                ],
            });
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(&c.silu);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups_flat((inter as u32).div_ceil(256));
        }
        if q4tp {
            encode_q1t_like(c, &mut enc, &c.q4tp_mv, dw, &a_buf, &d_buf, *dr, *dc);
        } else if q4t {
            encode_q1t_like(c, &mut enc, &c.q4t_mv, dw, &a_buf, &d_buf, *dr, *dc);
        } else {
            encode_matvec(c, &mut enc, dw, &a_buf, &drs_b, &d_buf, *dr, *dc);
        }
        // y += w·d
        {
            let wp = uniform_u32x4(c, [j.w.to_bits(), hidden as u32, 0, 0]);
            let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &c.layout_axpy,
                entries: &[bind_buf(0, &d_buf), bind_buf(1, &y_buf), bind_buf(2, &wp)],
            });
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(&c.axpy);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups((hidden as u32).div_ceil(256), 1, 1);
        }
    }
    // Hold the scratch lock across the readback: with concurrent server
    // slots two ops must not share the staging buffer mid-flight.
    let mut sc = c.scratch.lock().unwrap();
    let stage_buf = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        (hidden * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "moe-stage",
    );
    let ok = readback(c, enc, &y_buf, &stage_buf, (hidden * 4) as u64, out);
    drop(sc);
    ok
}

/// N independent q8-matvec (GDN projections of one input) in a single submission.
pub fn matvec_batch(model: &Arc<CmfModel>, jobs: &[BatchJob], out: &mut [&mut [f32]]) -> bool {
    let Some(c) = ctx() else { return false };
    if jobs.is_empty() || jobs.len() != out.len() {
        return false;
    }
    // q1 jobs carry tile-embedded scales (empty row_scale) and need the q1
    // pipeline — route the whole batch to the q1 encoder. Mixed batches
    // (shouldn't happen: QKV share a dtype) fall to the CPU path.
    // wgpu has a q1 batched kernel and no q4t/q4tp twin, so those layouts
    // keep the CPU path here rather than being fed to the wrong kernel.
    if jobs.iter().any(|j| {
        matches!(
            j.layout,
            crate::gpu::BatchLayout::Q4t | crate::gpu::BatchLayout::Q4tp
        )
    }) {
        return false;
    }
    let n_q1 = jobs
        .iter()
        .filter(|j| j.layout == crate::gpu::BatchLayout::Q1)
        .count();
    if n_q1 == jobs.len() {
        return matvec_batch_q1(model, jobs, out);
    }
    if n_q1 != 0 {
        return false;
    }
    let mut weights = Vec::with_capacity(jobs.len());
    for j in jobs {
        if j.cols % 4 != 0 {
            return false;
        }
        let Some(w) = tensor_weight(c, model, j.idx, j.rows, j.cols) else {
            return false;
        };
        weights.push(w);
    }
    let mut y_bufs = Vec::with_capacity(jobs.len());
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("batch"),
        });
    for (j, w) in jobs.iter().zip(&weights) {
        let rs_b = storage_bytes(c, bytemuck::cast_slice(j.row_scale));
        let xs_b = storage_bytes(c, bytemuck::cast_slice(&j.xs));
        let y_b = rw_f32(c, j.rows, true);
        encode_matvec(c, &mut enc, w, &xs_b, &rs_b, &y_b, j.rows, j.cols);
        y_bufs.push(y_b);
    }
    // ONE pooled staging buffer for all outputs (per-job offsets),
    // one map — instead of N fresh MAP_READ allocations per call.
    let total: u64 = jobs.iter().map(|j| (j.rows * 4) as u64).sum();
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        total,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "batch-stage",
    );
    let mut off = 0u64;
    for (y_b, j) in y_bufs.iter().zip(jobs) {
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(y_b, 0, &stage, off, (j.rows * 4) as u64);
        off += (j.rows * 4) as u64;
    }
    submit(c, finish_enc(enc));
    stage.slice(..total).map_async(wgpu::MapMode::Read, |_| {});
    if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        return false;
    }
    {
        let Ok(data) = stage.slice(..total).get_mapped_range() else {
            return false;
        };
        let mut off = 0usize;
        for (j, o) in jobs.iter().zip(out.iter_mut()) {
            o[..j.rows].copy_from_slice(bytemuck::cast_slice(&data[off..off + j.rows * 4]));
            off += j.rows * 4;
        }
    }
    stage.unmap();
    drop(sc);
    true
}

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

    /// wgpu defers destruction of device objects until a device poll.  The
    /// q2tp tests deliberately create and map several short-lived buffers;
    /// polling after those locals drop keeps driver deferred frees inside the
    /// test process instead of racing harness exit.
    ///
    /// `Ctx` is process-global and `ctx()` publishes references into boxes
    /// owned by `CTXS`.  A test teardown must therefore never call the public
    /// process-final `shutdown`: another test can still be using the same
    /// context while the Rust test harness runs tests in parallel.  Retaining
    /// only a cloned device here gives the teardown a safe poll target without
    /// draining the shared context map.
    struct TestGpuDrain {
        device: wgpu::Device,
    }

    impl TestGpuDrain {
        fn new(c: &Ctx) -> Self {
            Self {
                device: c.device.clone(),
            }
        }
    }

    impl Drop for TestGpuDrain {
        fn drop(&mut self) {
            let _ = self.device.poll(wgpu::PollType::wait_indefinitely());
        }
    }

    /// Test teardown may poll a shared device, but it must not invalidate a
    /// context acquired by the surrounding test (or by a parallel test).
    #[test]
    fn gpu_test_drain_keeps_shared_context_live() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping shared-context teardown test");
            return;
        };
        {
            let drain = TestGpuDrain::new(c);
            let probe = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("gpu-test-drain-probe"),
                size: 4,
                usage: wgpu::BufferUsages::STORAGE,
                mapped_at_creation: false,
            });
            drop(probe);
            drop(drain);
        }
        let probe_after_drain = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("gpu-test-drain-after"),
            size: 4,
            usage: wgpu::BufferUsages::STORAGE,
            mapped_at_creation: false,
        });
        drop(probe_after_drain);
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
    }

    #[test]
    fn o1_admission_accepts_bounded_extended_window() {
        let (d, dv, m, w, sink, hpg, t) =
            (8usize, 8usize, 4usize, 2048usize, 4usize, 2usize, 2064usize);
        let row = |salt: usize, width: usize| {
            (0..width)
                .map(|i| ((i * 17 + salt * 13) % 97) as f32 / 97.0 - 0.5)
                .collect::<Vec<_>>()
        };
        let qs: Vec<Vec<f32>> = (0..hpg)
            .map(|h| (0..t).flat_map(|i| row(i + h, d)).collect())
            .collect();
        let ks: Vec<f32> = (0..t).flat_map(|i| row(i + 31, d)).collect();
        let vs: Vec<f32> = (0..t).flat_map(|i| row(i + 71, dv)).collect();
        let qrefs: Vec<&[f32]> = qs.iter().map(Vec::as_slice).collect();
        let mut st = crate::nystrom::NystromState::new_group(m, w, sink, hpg);
        st.prefill_group(&qrefs, &ks, &vs, t, d, dv);
        let view = st.device_view();
        assert!(o1_view_valid(&view, hpg * 2, 2, d));

        // The shader scratch is a total near-entry cap, so one additional
        // row must remain an explicit decline rather than indexing past it.
        let mut over = crate::nystrom::NystromState::new_group(m, w + 1, sink, hpg);
        over.prefill_group(&qrefs, &ks, &vs, t, d, dv);
        let over_view = over.device_view();
        assert!(!o1_view_valid(&over_view, hpg * 2, 2, d));
    }

    #[test]
    fn kv_capacity_grows_without_reserving_the_advertised_context() {
        assert_eq!(kv_capacity(131_072, 1), 512);
        assert_eq!(kv_capacity(131_072, 512), 512);
        assert_eq!(kv_capacity(131_072, 513), 1024);
        assert_eq!(kv_capacity(131_072, 65_537), 131_072);
        assert_eq!(kv_capacity(1000, 999), 1000);
    }

    #[test]
    fn qwen_f32_const_keeps_recorded_bias_snapshot() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping qwen constant lifetime test");
            return;
        };
        let _gate = c.mm_gate.lock().unwrap();
        let mut bias = vec![1.0f32, 2.0, 3.0, 4.0];
        let zeros = [0.0f32; 4];
        let make_values = |label: &'static str| {
            let b = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some(label),
                size: (zeros.len() * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_DST
                    | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });
            c.queue.write_buffer(&b, 0, bytemuck::cast_slice(&zeros));
            b
        };
        let out_a = make_values("qwen-const-test-a");
        let out_b = make_values("qwen-const-test-b");
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("qwen-const-test"),
            });

        let old_bias = qwen_f32_const(c, &bias, "qwen-const-test-bias");
        encode_qwen_gelu_bias(c, &mut enc, &out_a, &old_bias, 4, 4, false);

        // Keep the Vec allocation and pointer unchanged so this exercises
        // the fingerprint-change branch rather than a new cache key.
        bias.copy_from_slice(&[5.0, 6.0, 7.0, 8.0]);
        let new_bias = qwen_f32_const(c, &bias, "qwen-const-test-bias");
        encode_qwen_gelu_bias(c, &mut enc, &out_b, &new_bias, 4, 4, false);

        let mut got_a = [0.0f32; 4];
        let mut got_b = [0.0f32; 4];
        assert!(readback2(
            c,
            enc,
            (&out_a, &mut got_a),
            (&out_b, &mut got_b),
        ));
        assert_eq!(got_a, [1.0, 2.0, 3.0, 4.0]);
        assert_eq!(got_b, [5.0, 6.0, 7.0, 8.0]);
    }

    #[test]
    fn qwen_layernorm_mod_matches_cpu_across_warps() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping qwen layernorm parity test");
            return;
        };
        let _gate = c.mm_gate.lock().unwrap();
        let eps = 1.0e-6f32;
        let batch = 4usize;
        for width in [64usize, 257, 3072] {
            let row_means = [3.25f32, -7.5, 19.0, -31.0];
            let x: Vec<f32> = (0..batch * width)
                .map(|j| {
                    let row = j / width;
                    let col = j % width;
                    let centered = ((col * 37 + row * 19 + width) % 101) as f32 / 50.0 - 1.0;
                    row_means[row] + centered + (col % 7) as f32 * 0.013
                })
                .collect();
            let shift: Vec<f32> = (0..width)
                .map(|i| ((i * 17 + width) % 101) as f32 * 0.002 - 0.1)
                .collect();
            let scale: Vec<f32> = (0..width)
                .map(|i| ((i * 19 + 7) % 89) as f32 * 0.003 - 0.132)
                .collect();
            let mut modulation = shift.clone();
            modulation.extend_from_slice(&scale);

            let mut want = vec![0.0f32; x.len()];
            for row in 0..batch {
                let xr = &x[row * width..(row + 1) * width];
                let mean = xr.iter().map(|&v| v as f64).sum::<f64>() / width as f64;
                let var = xr
                    .iter()
                    .map(|&v| {
                        let d = v as f64 - mean;
                        d * d
                    })
                    .sum::<f64>()
                    / width as f64;
                let inv = 1.0 / (var + eps as f64).sqrt();
                for i in 0..width {
                    want[row * width + i] =
                        (((xr[i] as f64 - mean) * inv) as f32) * (1.0 + scale[i]) + shift[i];
                }
            }

            let src = storage_bytes(c, bytemuck::cast_slice(&x));
            let dst = rw_f32(c, x.len(), true);
            assert!(qwen_layernorm_mod_keep(
                c,
                &src,
                &modulation,
                &dst,
                batch,
                width,
                eps,
            ));
            let mut got = vec![0.0f32; x.len()];
            let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("qwen-layernorm-test-stage"),
                size: (x.len() * 4) as u64,
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            let enc = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("qwen-layernorm-test-readback"),
                });
            assert!(readback(
                c,
                enc,
                &dst,
                &stage,
                (x.len() * 4) as u64,
                &mut got
            ));
            assert!(
                got.iter().all(|value| value.is_finite()),
                "qwen layernorm width={width} produced non-finite output"
            );
            let max_d = want
                .iter()
                .zip(&got)
                .map(|(a, b)| (a - b).abs())
                .fold(0.0f32, f32::max);
            assert!(
                max_d < 2.0e-3,
                "qwen layernorm width={width} ≠ CPU: max|Δ| = {max_d}"
            );
        }
    }

    #[test]
    fn dit_coop_matmul_preserves_odd_reduction_tail() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping DiT cooperative odd-tail test");
            return;
        };
        let _gate = c.mm_gate.lock().unwrap();
        let Some(pipe) = c.dit_gemm_coop.as_ref() else {
            eprintln!("cooperative matrix pipeline unavailable — skipping odd-tail test");
            return;
        };

        // Both the reduction and the two tiled output dimensions are odd;
        // nonzero input offsets also exercise the per-head packed layout.
        for (case, &(k, m, n)) in [(65usize, 65usize, 67usize), (129, 129, 131)]
            .iter()
            .enumerate()
        {
            let a_off = 5 + case * 3;
            let b_off = 7 + case * 5;
            let mut a = vec![0.0f32; a_off + m * k + 1];
            let mut b = vec![0.0f32; b_off + n * k + 1];
            for row in 0..m {
                for col in 0..k {
                    a[a_off + row * k + col] =
                        (((row * 13 + col * 7 + case * 5) % 17) as f32 - 8.0) * 0.125;
                }
            }
            for row in 0..n {
                for col in 0..k {
                    b[b_off + row * k + col] =
                        (((row * 11 + col * 3 + case * 9) % 19) as f32 - 9.0) * 0.125;
                }
            }
            let mut want = vec![0.0f32; m * n];
            for row in 0..m {
                for col in 0..n {
                    let mut sum = 0.0f32;
                    for i in 0..k {
                        sum += a[a_off + row * k + i] * b[b_off + col * k + i];
                    }
                    want[row * n + col] = sum;
                }
            }

            let wb = storage_bytes(c, bytemuck::cast_slice(&b));
            let xb = storage_bytes(c, bytemuck::cast_slice(&a));
            let yb = rw_f32(c, m * n, true);
            let p = uniform_u32x8(
                c,
                [
                    k.div_ceil(4) as u32,
                    n as u32,
                    m as u32,
                    1.0f32.to_bits(),
                    a_off as u32,
                    b_off as u32,
                    0,
                    k as u32,
                ],
            );
            let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("dit-coop-odd-tail-test-bg"),
                layout: &pipe.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, &wb),
                    bind_buf(1, &xb),
                    bind_buf(2, &yb),
                    bind_buf(3, &p),
                ],
            });
            let mut enc = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("dit-coop-odd-tail-test"),
                });
            {
                let mut pass = begin_pass(&mut enc);
                pass.set_pipeline(pipe);
                pass.set_bind_group(0, &bind, &[]);
                pass.dispatch_workgroups((n as u32).div_ceil(64), (m as u32).div_ceil(64), 1);
            }
            let mut got = vec![0.0f32; m * n];
            let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("dit-coop-odd-tail-test-stage"),
                size: (got.len() * 4) as u64,
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            assert!(readback(c, enc, &yb, &stage, (got.len() * 4) as u64, &mut got));
            assert!(
                got.iter().all(|value| value.is_finite()),
                "DiT cooperative odd-tail K={k}, M={m}, N={n} produced non-finite output"
            );
            let max_d = want
                .iter()
                .zip(&got)
                .map(|(a, b)| (a - b).abs())
                .fold(0.0f32, f32::max);
            assert!(
                max_d < 1.0e-5,
                "DiT cooperative odd-tail K={k}, M={m}, N={n} ≠ CPU: max|Δ| = {max_d}"
            );
        }
    }

    #[test]
    fn wgpu_q8_matvec_matches_cpu_reference() {
        // Force the wgpu path on (Metal-via-wgpu locally; Vulkan on the server).
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping parity test");
            return;
        };
        let (rows, cols) = (256usize, 64usize); // cols % 4 == 0
                                                // Synthetic int8 weights + row scales + pre-scaled activations.
        let mut q = vec![0i8; rows * cols];
        for (i, v) in q.iter_mut().enumerate() {
            *v = (((i * 37 + 11) % 255) as i32 - 127) as i8;
        }
        let rs: Vec<f32> = (0..rows).map(|r| 0.01 + (r % 7) as f32 * 0.003).collect();
        let xs: Vec<f32> = (0..cols).map(|i| ((i % 13) as f32 - 6.0) * 0.1).collect();

        // CPU reference: y[o] = rs[o] * Σ q[o,i]·xs[i].
        let mut want = vec![0f32; rows];
        for o in 0..rows {
            let mut acc = 0f32;
            for i in 0..cols {
                acc += q[o * cols + i] as f32 * xs[i];
            }
            want[o] = acc * rs[o];
        }

        let qbytes: &[u8] = bytemuck::cast_slice(&q);
        let mut got = vec![0f32; rows];
        assert!(dispatch_matvec(
            c, None, qbytes, 0, &rs, &xs, rows, cols, &mut got
        ));

        let max_d = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(max_d < 1e-3, "wgpu q8_matvec ≠ CPU: max|Δ| = {max_d}");

        // Also check the row0 offset: the range [rows/2, rows) of the full
        // tensor must match the tail of the reference.
        let r0 = rows / 2;
        let mut got2 = vec![0f32; rows - r0];
        assert!(dispatch_matvec(
            c,
            None,
            qbytes,
            r0,
            &rs[r0..],
            &xs,
            rows - r0,
            cols,
            &mut got2
        ));
        let max_d2 = want[r0..]
            .iter()
            .zip(&got2)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(max_d2 < 1e-3, "wgpu row0 offset ≠ CPU: max|Δ| = {max_d2}");
    }

    #[test]
    fn wgpu_q8_2f_odd_rows_matches_cpu_reference() {
        use wgpu::util::DeviceExt;

        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping q8_2f parity test");
            return;
        };
        // Three rows put the column-scale plane in the high half of the
        // final row-scale word. This is the packed layout that the old
        // word-aligned address calculation decoded one half-word late.
        let (rows, cols) = (3usize, 8usize);
        let q: Vec<i8> = (0..rows * cols)
            .map(|i| ((i * 17 + 5) % 29) as i8 - 14)
            .collect();
        let rs_src = [0.25f32, 0.5, 0.75];
        let cs_src = [0.2f32, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9];
        let xs = [0.7f32, -0.4, 0.2, 1.1, -0.8, 0.3, 0.6, -0.5];
        let rs16: Vec<u16> = rs_src
            .iter()
            .map(|&v| cortiq_core::quant::f32_to_f16(v))
            .collect();
        let cs16: Vec<u16> = cs_src
            .iter()
            .map(|&v| cortiq_core::quant::f32_to_f16(v))
            .collect();
        let mut payload = bytemuck::cast_slice::<i8, u8>(&q).to_vec();
        payload.extend_from_slice(bytemuck::cast_slice(&rs16));
        payload.extend_from_slice(bytemuck::cast_slice(&cs16));
        payload.resize(payload.len().next_multiple_of(4), 0);

        let weights = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("q8-2f-odd-w"),
                contents: &payload,
                usage: wgpu::BufferUsages::STORAGE,
            });
        let xb = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("q8-2f-odd-x"),
                contents: bytemuck::cast_slice(&xs),
                usage: wgpu::BufferUsages::STORAGE,
            });
        let yb = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("q8-2f-odd-y"),
            size: (rows * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let params = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("q8-2f-odd-p"),
                contents: bytemuck::cast_slice(&[(cols / 4) as u32, rows as u32, cols as u32, 0]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let layout = c.q8_2f_mv.get_bind_group_layout(0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("q8-2f-odd-bg"),
            layout: &layout,
            entries: &[
                bind_buf(0, &weights),
                bind_buf(1, &xb),
                bind_buf(2, &yb),
                bind_buf(3, &params),
            ],
        });
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("q8-2f-odd-stage"),
            size: (rows * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(&c.q8_2f_mv);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(rows as u32, 1, 1);
        }
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&yb, 0, &stage, 0, (rows * 4) as u64);
        submit(c, finish_enc(enc));
        let (tx, rx) = std::sync::mpsc::channel();
        stage.map_async(wgpu::MapMode::Read, .., move |r| tx.send(r).unwrap());
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        rx.recv().unwrap().unwrap();
        let got: Vec<f32> = bytemuck::cast_slice(&stage.get_mapped_range(..).unwrap()).to_vec();

        let rs: Vec<f32> = rs16
            .iter()
            .map(|&v| cortiq_core::quant::f16_to_f32(v))
            .collect();
        let cs: Vec<f32> = cs16
            .iter()
            .map(|&v| cortiq_core::quant::f16_to_f32(v))
            .collect();
        let want: Vec<f32> = (0..rows)
            .map(|r| {
                rs[r]
                    * (0..cols)
                        .map(|i| q[r * cols + i] as f32 * xs[i] * cs[i])
                        .sum::<f32>()
            })
            .collect();
        let max_d = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(max_d < 2e-2, "wgpu q8_2f odd rows ≠ CPU: max|Δ| = {max_d}");
    }

    /// Real-model regression for vocabulary heads wider than Vulkan's 65,535
    /// workgroup limit.  Set CMF_Q8_2F_MODEL to a CMF whose lm_head is q8_2f.
    /// The small synthetic parity case above cannot catch dispatch/scale-plane
    /// mistakes that only appear at a 100k-row vocabulary.
    #[test]
    #[ignore]
    fn wgpu_q8_2f_large_vocab_head_matches_cpu_reference() {
        let Ok(path) = std::env::var("CMF_Q8_2F_MODEL") else {
            eprintln!("CMF_Q8_2F_MODEL is not set — skipping");
            return;
        };
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping q8_2f model parity test");
            return;
        };
        let model = Arc::new(CmfModel::open(path).expect("open CMF_Q8_2F_MODEL"));
        let (idx, entry) = model
            .tensors
            .iter()
            .enumerate()
            .find(|(_, e)| e.name == "lm_head.weight")
            .expect("lm_head.weight");
        assert_eq!(entry.dtype, cortiq_core::TensorDtype::Q8_2f);
        let (rows, cols) = (entry.shape[0], entry.shape[1]);
        let abs = model.entry_abs_offset(entry).expect("primary tensor");
        let payload = &model.primary_bytes()[abs..abs + entry.nbytes as usize];
        let xs: Vec<f32> = (0..cols)
            .map(|i| ((i * 37 % 101) as f32 - 50.0) * 0.002)
            .collect();

        // Reproduce the production order: a batched/per-op q8 GEMM first
        // primes residency with only the int8 body, then the whole-token
        // q8_2f graph asks for the same tensor including both scale planes.
        let key = (model.uid() as usize, idx);
        let body = &payload[..rows * cols];
        let body_buf = weight_buffer_l(c, key, body, layer_of_name(&entry.name))
            .expect("resident body-only q8 head");
        assert_eq!(body_buf.size(), body.len().next_multiple_of(4) as u64);
        assert!(!q8_resident_or_upload(&model, idx, true));
        let weights = c
            .weight_bufs
            .lock()
            .unwrap()
            .get(&key)
            .expect("resident q8_2f head")
            .buf
            .clone();
        assert_eq!(weights.size(), entry.nbytes);
        let xb = storage_bytes(c, bytemuck::cast_slice(&xs));
        let yb = rw_f32(c, rows, true);
        let params = uniform_u32x4(c, [(cols / 4) as u32, rows as u32, cols as u32, 0]);
        let layout = c.q8_2f_mv.get_bind_group_layout(0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("q8-2f-model-head-bg"),
            layout: &layout,
            entries: &[
                bind_buf(0, &weights),
                bind_buf(1, &xb),
                bind_buf(2, &yb),
                bind_buf(3, &params),
            ],
        });
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("q8-2f-model-head-stage"),
            size: (rows * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(&c.q8_2f_mv);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
        }
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&yb, 0, &stage, 0, (rows * 4) as u64);
        submit(c, finish_enc(enc));
        stage.slice(..).map_async(wgpu::MapMode::Read, |_| {});
        c.device
            .poll(wgpu::PollType::wait_indefinitely())
            .expect("poll");
        let got: Vec<f32> = bytemuck::cast_slice(&stage.get_mapped_range(..).unwrap()).to_vec();

        let qn = rows * cols;
        let rs = &payload[qn..qn + rows * 2];
        let cs = &payload[qn + rows * 2..qn + rows * 2 + cols * 2];
        let f16_at = |p: &[u8], i: usize| {
            cortiq_core::quant::f16_to_f32(u16::from_le_bytes([p[i * 2], p[i * 2 + 1]]))
        };
        let mut max_d = 0.0f32;
        let mut max_want = 0.0f32;
        for r in (0..rows).step_by((rows / 31).max(1)).take(32) {
            let mut acc = 0.0f32;
            for i in 0..cols {
                acc += payload[r * cols + i] as i8 as f32 * xs[i] * f16_at(cs, i);
            }
            let want = acc * f16_at(rs, r);
            max_want = max_want.max(want.abs());
            max_d = max_d.max((got[r] - want).abs());
        }
        eprintln!(
            "q8_2f real head idx={idx} {rows}x{cols}: max|want|={max_want:.6} max|Δ|={max_d:.6} gpu_nonzero={}",
            got.iter().filter(|v| **v != 0.0).count()
        );
        assert!(max_want > 0.0 && max_d < 2e-2, "large q8_2f head mismatch");
    }

    /// Quantifies the whole-token-graph ceiling on THIS device: K chained
    /// matvecs run as K separate submit+readback ops (today's per-op path)
    /// vs the same K dispatches in ONE command buffer with a single readback
    /// (intermediates stay on the GPU — what the graph does). The ratio is how
    /// much the submit/PCIe-readback wall is costing per token.
    /// Run: `CMF_GPU=wgpu cargo test -p cortiq-engine --release --features gpu
    ///       --test-threads 1 wgpu_chain_probe -- --ignored --nocapture`
    #[test]
    #[ignore]
    fn wgpu_chain_probe() {
        use std::time::Instant;
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let n: usize = std::env::var("CMF_CHAIN_N")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(896);
        let k: usize = std::env::var("CMF_CHAIN_K")
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(100);
        assert!(n % 4 == 0);
        // Resident n×n q8 weights + row scales (values irrelevant — timing only).
        let q = vec![1i8; n * n];
        let w = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("probe-w"),
                contents: bytemuck::cast_slice(&q),
                usage: wgpu::BufferUsages::STORAGE,
            });
        let rs = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("probe-rs"),
                contents: bytemuck::cast_slice(&vec![1f32; n]),
                usage: wgpu::BufferUsages::STORAGE,
            });
        let p = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("probe-p"),
                contents: bytemuck::cast_slice(&[(n / 4) as u32, n as u32, 0u32, 0u32]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let mkbuf = |lbl| {
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some(lbl),
                size: (n * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_SRC
                    | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            })
        };
        let a = mkbuf("probe-a");
        let b = mkbuf("probe-b");
        c.queue
            .write_buffer(&a, 0, bytemuck::cast_slice(&vec![0.01f32; n]));
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("probe-stage"),
            size: (n * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let bg = |xs: &wgpu::Buffer, y: &wgpu::Buffer| {
            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("probe-bg"),
                layout: &c.layout,
                entries: &[
                    bind_buf(0, &w),
                    bind_buf(1, xs),
                    bind_buf(2, &rs),
                    bind_buf(3, y),
                    bind_buf(4, &p),
                ],
            })
        };
        let bg_ab = bg(&a, &b);
        let bg_ba = bg(&b, &a);
        let wg = (n as u32).min(MAX_WG);
        let readback = |buf: &wgpu::Buffer, enc: wgpu::CommandEncoder| {
            let mut enc = enc;
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(buf, 0, &stage, 0, (n * 4) as u64);
            submit(c, finish_enc(enc));
            stage.slice(..).map_async(wgpu::MapMode::Read, |_| {});
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            let _ = stage.slice(..).get_mapped_range();
            stage.unmap();
        };
        let dispatch = |enc: &mut wgpu::CommandEncoder, even: bool| {
            let mut pass = begin_pass(enc);
            pass.set_pipeline(&c.matvec);
            pass.set_bind_group(0, if even { &bg_ab } else { &bg_ba }, &[]);
            pass.dispatch_workgroups(wg, 1, 1);
        };
        // Warm.
        for _ in 0..3 {
            let mut e = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
            dispatch(&mut e, true);
            readback(&b, e);
        }
        // Per-op: K submits + K readbacks.
        let t = Instant::now();
        for i in 0..k {
            let mut e = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
            dispatch(&mut e, i % 2 == 0);
            readback(if i % 2 == 0 { &b } else { &a }, e);
        }
        let per_op = t.elapsed().as_secs_f64();
        // Fused: K dispatches, ONE submit + ONE readback.
        let t = Instant::now();
        let mut e = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        for i in 0..k {
            dispatch(&mut e, i % 2 == 0);
        }
        readback(if (k - 1) % 2 == 0 { &b } else { &a }, e);
        let fused = t.elapsed().as_secs_f64();
        eprintln!(
            "CHAIN PROBE n={n} k={k}: per-op {:.2} ms ({:.3} ms/op) | fused {:.2} ms | speedup {:.2}× | submit+readback wall ≈ {:.3} ms/op",
            per_op * 1e3,
            per_op * 1e3 / k as f64,
            fused * 1e3,
            per_op / fused,
            (per_op - fused) * 1e3 / (k - 1) as f64,
        );
    }

    #[test]
    fn wgpu_q1_matvec_matches_cpu_reference() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping q1 parity test");
            return;
        };
        let (rows, cols) = (33usize, 256usize); // gpr = 8 (even), odd rows
        let gpr = cols / 32;
        let mut payload = Vec::new();
        for t in 0..rows * gpr {
            let sc = 0.005 + (t % 9) as f32 * 0.004;
            payload.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
            for j in 0..4 {
                payload.push(((t * 41 + j * 71 + 13) % 253) as u8);
            }
        }
        let xs: Vec<f32> = (0..cols)
            .map(|i| ((i * 7 + 3) % 29) as f32 / 29.0 - 0.5)
            .collect();
        let mut w = vec![0f32; rows * cols];
        cortiq_core::quant::dequant_q1(&payload, &mut w);
        let mut want = vec![0f32; rows];
        for o in 0..rows {
            want[o] = (0..cols).map(|i| w[o * cols + i] * xs[i]).sum();
        }
        let mut got = vec![0f32; rows];
        assert!(dispatch_q1(c, None, &payload, &xs, rows, cols, &mut got));
        let max_d = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(max_d < 1e-3, "wgpu q1_matvec ≠ CPU: max|Δ| = {max_d}");
    }

    #[test]
    fn wgpu_rmsnorm_matches_cpu() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        if ctx().is_none() {
            eprintln!("no wgpu adapter — skipping rmsnorm parity test");
            return;
        }
        let n = 896usize;
        let eps = 1e-6f32;
        let x: Vec<f32> = (0..n)
            .map(|i| ((i * 13 + 7) % 101) as f32 / 101.0 - 0.5)
            .collect();
        let w: Vec<f32> = (0..n)
            .map(|i| 0.5 + ((i * 5 + 1) % 17) as f32 / 17.0)
            .collect();
        let ss: f32 = x.iter().map(|v| v * v).sum();
        let inv = 1.0 / (ss / n as f32 + eps).sqrt();
        // plain RMSNorm
        let want: Vec<f32> = (0..n).map(|i| x[i] * inv * w[i]).collect();
        let mut got = vec![0f32; n];
        assert!(rmsnorm_row(&x, &w, &mut got, false, eps));
        let md = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(md < 1e-4, "wgpu rmsnorm ≠ CPU: max|Δ| = {md}");
        // gemma variant: w' = 1 + w
        let wantg: Vec<f32> = (0..n).map(|i| x[i] * inv * (1.0 + w[i])).collect();
        let mut gotg = vec![0f32; n];
        assert!(rmsnorm_row(&x, &w, &mut gotg, true, eps));
        let mdg = wantg
            .iter()
            .zip(&gotg)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(mdg < 1e-4, "wgpu rmsnorm(gemma) ≠ CPU: max|Δ| = {mdg}");
    }

    #[test]
    fn wgpu_add_rmsnorm_matches_cpu() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let n = 896usize;
        let eps = 1e-6f32;
        let h: Vec<f32> = (0..n)
            .map(|i| ((i * 13 + 7) % 101) as f32 / 101.0 - 0.5)
            .collect();
        let d: Vec<f32> = (0..n)
            .map(|i| ((i * 7 + 3) % 61) as f32 / 61.0 - 0.5)
            .collect();
        let w: Vec<f32> = (0..n)
            .map(|i| 0.5 + ((i * 5 + 1) % 17) as f32 / 17.0)
            .collect();
        // CPU reference: h += d, then rmsnorm(h, w)
        let hd: Vec<f32> = (0..n).map(|i| h[i] + d[i]).collect();
        let ss: f32 = hd.iter().map(|x| x * x).sum();
        let inv = 1.0 / (ss / n as f32 + eps).sqrt();
        let want: Vec<f32> = (0..n).map(|i| hd[i] * inv * w[i]).collect();
        // GPU add_rmsnorm
        let hb = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&h),
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            });
        let db = storage_bytes(c, bytemuck::cast_slice(&d));
        let wb = storage_bytes(c, bytemuck::cast_slice(&w));
        let ob = rw_f32(c, n, true);
        let pb = uniform_u32x4(c, [n as u32, 0, eps.to_bits(), 0]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.layout_add_rmsnorm,
            entries: &[
                bind_buf(0, &hb),
                bind_buf(1, &db),
                bind_buf(2, &wb),
                bind_buf(3, &ob),
                bind_buf(4, &pb),
            ],
        });
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        {
            let mut p = begin_pass(&mut enc);
            p.set_pipeline(&c.add_rmsnorm);
            p.set_bind_group(0, &bind, &[]);
            p.dispatch_workgroups(1, 1, 1);
        }
        let mut got = vec![0f32; n];
        let sz = (n * 4) as u64;
        let mut sc = c.scratch.lock().unwrap();
        let stage = Scratch::ensure(
            &c.device,
            &mut sc.stage,
            sz,
            wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            "arn-stage",
        );
        assert!(readback(c, enc, &ob, &stage, sz, &mut got));
        drop(sc);
        let md = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(md < 1e-4, "wgpu add_rmsnorm ≠ CPU: max|Δ| = {md}");
    }

    #[test]
    fn wgpu_attn_rope_qkn_matches_cpu() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        if ctx().is_none() {
            eprintln!("no wgpu adapter — skipping attn_rope parity test");
            return;
        }
        // head_dim 256 with partial RoPE (rd=64) — the Qwen3.5 geometry: nt=8
        // (>4-slot xv) and hlf=32 exercise the paths that broke the graph.
        let (nh, nkv, hd, rd, pos) = (4usize, 2usize, 256usize, 64usize, 5usize);
        let eps = 1e-6f32;
        let flags = 1u32 | 2u32 | 4u32; // gate + qnorm + knorm, non-gemma
        let jitter = |a: usize, b: usize| ((a * 31 + b * 17 + 7) % 97) as f32 / 97.0 - 0.5;
        // qraw: nh heads × 2·hd (q part || gate part); k: nkv × hd
        let qraw: Vec<f32> = (0..nh * 2 * hd).map(|i| jitter(i, 1)).collect();
        let k_in: Vec<f32> = (0..nkv * hd).map(|i| jitter(i, 2)).collect();
        let qnw: Vec<f32> = (0..hd).map(|d| 0.7 + jitter(d, 3)).collect();
        let knw: Vec<f32> = (0..hd).map(|d| 0.7 + jitter(d, 4)).collect();
        let invf: Vec<f32> = (0..rd / 2)
            .map(|i| 1.0 / (10000f32).powf(2.0 * i as f32 / rd as f32))
            .collect();
        // CPU reference: qk-norm then half-split partial RoPE.
        let norm_rope = |v: &mut [f32], w: &[f32]| {
            let ss: f32 = v.iter().map(|x| x * x).sum();
            let inv = 1.0 / (ss / hd as f32 + eps).sqrt();
            for d in 0..hd {
                v[d] = v[d] * inv * w[d];
            }
            let hlf = rd / 2;
            for i in 0..hlf {
                let ang = pos as f32 * invf[i];
                let (c, s) = (ang.cos(), ang.sin());
                let (x0, x1) = (v[i], v[i + hlf]);
                v[i] = x0 * c - x1 * s;
                v[i + hlf] = x0 * s + x1 * c;
            }
        };
        let mut want_q = vec![0f32; nh * hd];
        let mut want_g = vec![0f32; nh * hd];
        for h in 0..nh {
            let mut q: Vec<f32> = qraw[h * 2 * hd..h * 2 * hd + hd].to_vec();
            norm_rope(&mut q, &qnw);
            want_q[h * hd..(h + 1) * hd].copy_from_slice(&q);
            want_g[h * hd..(h + 1) * hd]
                .copy_from_slice(&qraw[h * 2 * hd + hd..h * 2 * hd + 2 * hd]);
        }
        let mut want_k = k_in.clone();
        for kh in 0..nkv {
            let mut kk = want_k[kh * hd..(kh + 1) * hd].to_vec();
            norm_rope(&mut kk, &knw);
            want_k[kh * hd..(kh + 1) * hd].copy_from_slice(&kk);
        }
        let mut got_q = vec![0f32; nh * hd];
        let mut got_k = vec![0f32; nkv * hd];
        let mut got_g = vec![0f32; nh * hd];
        assert!(attn_rope_qkn_gpu(
            &qraw, &k_in, &qnw, &knw, &invf, nh, nkv, hd, rd, pos, flags, eps, &mut got_q,
            &mut got_k, &mut got_g,
        ));
        let md = |a: &[f32], b: &[f32]| {
            a.iter()
                .zip(b)
                .map(|(x, y)| (x - y).abs())
                .fold(0.0f32, f32::max)
        };
        assert!(
            md(&want_q, &got_q) < 1e-4,
            "q mismatch: {}",
            md(&want_q, &got_q)
        );
        assert!(
            md(&want_k, &got_k) < 1e-4,
            "k mismatch: {}",
            md(&want_k, &got_k)
        );
        assert!(
            md(&want_g, &got_g) < 1e-4,
            "gate mismatch: {}",
            md(&want_g, &got_g)
        );
    }

    #[test]
    fn wgpu_gqa_attend_matches_cpu() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping gqa_attend parity test");
            return;
        };
        // hd=128 exercises the stride-129 kernel (exists everywhere);
        // hd=256 exercises stride-257 where the device's workgroup
        // storage allows it (32 KB devices — Adreno/Mali/wgpu-Metal —
        // honestly refuse: hd_cap gates them to the small kernel).
        attend_case(128);
        if c.hd_cap >= 256 {
            attend_case(256);
        } else {
            eprintln!("hd_cap {} — skipping hd=256 attend case", c.hd_cap);
        }
    }

    fn attend_case(hd: usize) {
        let (nh, hpk, cap, n) = (4usize, 2usize, 16usize, 5usize);
        let nkv = nh / hpk;
        let jit = |a: usize, b: usize| ((a * 29 + b * 13 + 5) % 89) as f32 / 89.0 - 0.5;
        let q: Vec<f32> = (0..nh * hd).map(|i| jit(i, 1)).collect();
        // caches laid out [nkv, cap, hd]; only first n rows are valid.
        let mut kc = vec![0f32; nkv * cap * hd];
        let mut vc = vec![0f32; nkv * cap * hd];
        for kh in 0..nkv {
            for p in 0..n {
                for d in 0..hd {
                    kc[(kh * cap + p) * hd + d] = jit(kh * 1000 + p * 10 + d, 2);
                    vc[(kh * cap + p) * hd + d] = jit(kh * 1000 + p * 10 + d, 3);
                }
            }
        }
        // CPU reference: scaled softmax attention per head.
        let scale = 1.0 / (hd as f32).sqrt();
        let mut want = vec![0f32; nh * hd];
        for h in 0..nh {
            let kh = h / hpk;
            let mut sc: Vec<f32> = (0..n)
                .map(|p| {
                    (0..hd)
                        .map(|d| q[h * hd + d] * kc[(kh * cap + p) * hd + d])
                        .sum::<f32>()
                        * scale
                })
                .collect();
            let mx = sc.iter().cloned().fold(f32::MIN, f32::max);
            let mut den = 0.0;
            for s in sc.iter_mut() {
                *s = (*s - mx).exp();
                den += *s;
            }
            for d in 0..hd {
                want[h * hd + d] = (0..n)
                    .map(|p| sc[p] * vc[(kh * cap + p) * hd + d])
                    .sum::<f32>()
                    / den;
            }
        }
        let mut got = vec![0f32; nh * hd];
        assert!(gqa_attend_gpu(&q, &kc, &vc, nh, hpk, hd, cap, n, &mut got));
        let md = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(md < 1e-4, "wgpu gqa_attend hd={hd} ≠ CPU: max|Δ| = {md}");
    }

    /// The GQA-shared split-K attend against the per-head split kernel
    /// and the CPU reference, at Qwen3.8's geometry (24 heads on 4 kv
    /// heads, hd 256) across several 256-position chunks with a partial
    /// last one — the regime the per-head kernel streamed six times over.
    #[test]
    fn wgpu_gqa_attend_gpart_matches_split_and_cpu() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        if c.attend_gpart.is_none() {
            eprintln!("no gqa gpart kernel on this device (small workgroup storage) — skipping");
            return;
        }
        for (nh, hpk, hd, cap, n) in [
            (24usize, 6usize, 256usize, 1024usize, 700usize),
            (8, 2, 128, 512, 300),
            (6, 6, 64, 300, 257),
        ] {
            let nkv = nh / hpk;
            let jit = |a: usize, b: usize| ((a * 29 + b * 13 + 5) % 89) as f32 / 89.0 - 0.5;
            let q: Vec<f32> = (0..nh * hd).map(|i| jit(i, 1)).collect();
            let mut kc = vec![0f32; nkv * cap * hd];
            let mut vc = vec![0f32; nkv * cap * hd];
            for kh in 0..nkv {
                for p in 0..n {
                    for d in 0..hd {
                        kc[(kh * cap + p) * hd + d] = jit(kh * 1000 + p * 10 + d, 2) * 3.0;
                        vc[(kh * cap + p) * hd + d] = jit(kh * 1000 + p * 10 + d, 3);
                    }
                }
            }
            let scale = 1.0 / (hd as f32).sqrt();
            let mut want = vec![0f32; nh * hd];
            for h in 0..nh {
                let kh = h / hpk;
                let mut sc: Vec<f32> = (0..n)
                    .map(|p| {
                        (0..hd)
                            .map(|d| q[h * hd + d] as f64 * kc[(kh * cap + p) * hd + d] as f64)
                            .sum::<f64>() as f32
                            * scale
                    })
                    .collect();
                let mx = sc.iter().cloned().fold(f32::MIN, f32::max);
                let mut den = 0.0f64;
                for s in sc.iter_mut() {
                    *s = (*s - mx).exp();
                    den += *s as f64;
                }
                for d in 0..hd {
                    want[h * hd + d] = ((0..n)
                        .map(|p| sc[p] as f64 * vc[(kh * cap + p) * hd + d] as f64)
                        .sum::<f64>()
                        / den) as f32;
                }
            }
            let mut per_head = vec![0f32; nh * hd];
            let mut shared = vec![0f32; nh * hd];
            assert!(gqa_attend_split_gpu(
                &q,
                &kc,
                &vc,
                nh,
                hpk,
                hd,
                cap,
                n,
                false,
                &mut per_head
            ));
            assert!(gqa_attend_split_gpu(
                &q,
                &kc,
                &vc,
                nh,
                hpk,
                hd,
                cap,
                n,
                true,
                &mut shared
            ));
            let md = |a: &[f32], b: &[f32]| {
                a.iter()
                    .zip(b)
                    .map(|(x, y)| (x - y).abs())
                    .fold(0.0f32, f32::max)
            };
            let (d_ph, d_sh, d_xx) = (
                md(&want, &per_head),
                md(&want, &shared),
                md(&per_head, &shared),
            );
            eprintln!(
                "gpart nh={nh} hpk={hpk} hd={hd} n={n}: |per-head−cpu| {d_ph:.2e} |shared−cpu| {d_sh:.2e} |per-head−shared| {d_xx:.2e}"
            );
            assert!(d_ph < 1e-4, "per-head split drifted from CPU: {d_ph}");
            assert!(d_sh < 1e-4, "GQA-shared split drifted from CPU: {d_sh}");
        }
    }

    #[test]
    fn wgpu_o1_step_matches_cpu() {
        // End-to-end: the REAL NystromState is the reference — prefill a
        // group, clone it, advance the clone one token on the CPU, and
        // require the device mirror (upload + o1_far/o1_push/o1_attend)
        // to produce the same attention output from the same state.
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping o1 test");
            return;
        };
        // Production geometry (Qwen3.6): the small-dim version passed while
        // the real model garbled, so the test runs BOTH.
        for (d, dv, m, w, sink, hpg, t) in [
            (8usize, 8usize, 4usize, 8usize, 2usize, 2usize, 40usize),
            (256, 256, 32, 128, 4, 8, 430),
            // Exercise the extended near grid with a real >256 score count
            // while keeping the fixture small enough for a unit test.  The
            // production-sized d=256 case above still guards the normal
            // upload/layout geometry; this case specifically catches stale
            // one-lane indexing and the landmark-lane collision at w2048.
            (8, 8, 4, 2048, 4, 2, 2064),
        ] {
            let jit = |a: usize, b: usize| ((a * 37 + b * 13 + 3) % 83) as f32 / 83.0 - 0.5;
            let ks: Vec<f32> = (0..t * d).map(|i| jit(i, 1)).collect();
            let vs: Vec<f32> = (0..t * dv).map(|i| jit(i, 2)).collect();
            let qs_own: Vec<Vec<f32>> = (0..hpg)
                .map(|h| (0..t * d).map(|i| jit(i, 3 + h)).collect())
                .collect();
            let qs_refs: Vec<&[f32]> = qs_own.iter().map(|v| v.as_slice()).collect();
            let mut st = crate::nystrom::NystromState::new_group(m, w, sink, hpg);
            st.prefill_group(&qs_refs, &ks, &vs, t, d, dv);
            // CPU ground truth for the next token (built below per group).
            // TWO groups — the production model has nkv=2, and the group
            // concatenation in the upload plus every g-offset in the kernels
            // is exactly what a single-group test cannot catch.
            let mut st2 = crate::nystrom::NystromState::new_group(m, w, sink, hpg);
            let qs2_own: Vec<Vec<f32>> = (0..hpg)
                .map(|h| (0..t * d).map(|i| jit(i, 23 + h)).collect())
                .collect();
            let qs2_refs: Vec<&[f32]> = qs2_own.iter().map(|v| v.as_slice()).collect();
            let ks2: Vec<f32> = (0..t * d).map(|i| jit(i, 21)).collect();
            let vs2: Vec<f32> = (0..t * dv).map(|i| jit(i, 22)).collect();
            st2.prefill_group(&qs2_refs, &ks2, &vs2, t, d, dv);
            let gcnt = 2usize;
            let q_new: Vec<f32> = (0..gcnt * hpg * d).map(|i| jit(i, 5)).collect();
            let k_new: Vec<f32> = (0..gcnt * d).map(|i| jit(i, 6)).collect();
            let v_new: Vec<f32> = (0..gcnt * dv).map(|i| jit(i, 7)).collect();
            let mut want = vec![0f32; gcnt * hpg * dv];
            let mut cpu1 = st.clone();
            let mut cpu2 = st2.clone();
            cpu1.step_group(
                &q_new[..hpg * d],
                &k_new[..d],
                &v_new[..dv],
                &mut want[..hpg * dv],
            );
            cpu2.step_group(
                &q_new[hpg * d..],
                &k_new[d..],
                &v_new[dv..],
                &mut want[hpg * dv..],
            );
            // Device: upload the PRE-step states, run the three kernels.
            let views = vec![st.device_view(), st2.device_view()];
            assert!(!views[0].exact_only, "t must exceed w+8 for this test");
            let mv = views[0].m_eff;
            o1_ensure(c, u64::MAX, usize::MAX, &views, 1).expect("o1 upload");
            let dev_bufs = {
                let map = c.o1m.lock().unwrap();
                let dref = map.get(&(u64::MAX, usize::MAX)).unwrap();
                (
                    dref.meta.clone(),
                    dref.ring_k.clone(),
                    dref.ring_v.clone(),
                    dref.sink_k.clone(),
                    dref.sink_v.clone(),
                    dref.k_tilde.clone(),
                    dref.qt.clone(),
                    dref.mu.clone(),
                    dref.mz.clone(),
                    dref.that.clone(),
                    dref.scale,
                )
            };
            let (dmeta, drk, drv, dsk, dsv, dkt, dqt, dmu, dmz, dth, sc) = dev_bufs;
            let rect_fm = views[0].heads[0].rect_fm;
            let stor = |data: &[f32]| {
                use wgpu::util::DeviceExt;
                c.device
                    .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: None,
                        contents: bytemuck::cast_slice(data),
                        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                    })
            };
            let qb = stor(&q_new);
            let kb = stor(&k_new);
            let vb = stor(&v_new);
            let ob = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (gcnt * hpg * dv * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });
            let o1_u = uniform_u32x8(
                c,
                [
                    hpg as u32,
                    mv as u32,
                    w as u32,
                    (sink as u32) | (u32::from(rect_fm) << 8),
                    d as u32,
                    dv as u32,
                    sc.to_bits(),
                    0,
                ],
            );
            let bgf = |layout: &wgpu::BindGroupLayout, bufs: &[&wgpu::Buffer]| {
                let entries: Vec<_> = bufs
                    .iter()
                    .enumerate()
                    .map(|(i, b)| bind_buf(i as u32, b))
                    .collect();
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout,
                    entries: &entries,
                })
            };
            let bg_far = bgf(
                &c.layout_o1_far,
                &[&dmeta, &drk, &drv, &dqt, &dmz, &dth, &o1_u],
            );
            let bg_push = bgf(&c.layout_o1_push, &[&dmeta, &kb, &vb, &drk, &drv, &o1_u]);
            let bg_att = bgf(
                &c.layout_o1_attend,
                &[
                    &dmeta, &qb, &drk, &drv, &dsk, &dsv, &dkt, &dmu, &dmz, &dth, &ob, &o1_u,
                ],
            );
            let mut enc = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
            {
                let mut pass = begin_pass(&mut enc);
                pass.set_pipeline(&c.o1_far);
                pass.set_bind_group(0, &bg_far, &[]);
                pass.dispatch_workgroups((gcnt * hpg * mv) as u32, 1, 1);
                pass.set_pipeline(&c.o1_push);
                pass.set_bind_group(0, &bg_push, &[]);
                pass.dispatch_workgroups(gcnt as u32, 1, 1);
                pass.set_pipeline(&c.o1_attend);
                pass.set_bind_group(0, &bg_att, &[]);
                pass.dispatch_workgroups((gcnt * hpg) as u32, 1, 1);
            }
            let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (gcnt * hpg * dv * 4) as u64,
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(&ob, 0, &stage, 0, (gcnt * hpg * dv * 4) as u64);
            submit(c, finish_enc(enc));
            let (tx, rx) = std::sync::mpsc::channel();
            stage.map_async(wgpu::MapMode::Read, .., move |r| tx.send(r).unwrap());
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            rx.recv().unwrap().unwrap();
            let got: Vec<f32> = bytemuck::cast_slice(&stage.get_mapped_range(..).unwrap()).to_vec();
            stage.unmap();
            // Exercise the far-empty branch on the same real GPU dispatch.
            // The production admission normally supplies farl > 0, but a
            // retained/handed-off state can reach the attend shader with an
            // empty far field.  Reuse the post-push ring contents and compare
            // against exact softmax over sinks + ring; this also covers m <
            // 32 without manufacturing a second shader or changing state.
            if m < 32 {
                let mut empty_meta = Vec::with_capacity(gcnt * 4);
                let mut near_k = Vec::with_capacity(gcnt);
                let mut near_v = Vec::with_capacity(gcnt);
                for (g, view) in views.iter().enumerate() {
                    let post_len = if view.win_len == w {
                        w
                    } else {
                        view.win_len + 1
                    };
                    let post_head = if view.win_len == w {
                        (view.win_head + 1) % w
                    } else {
                        view.win_head
                    };
                    empty_meta.extend_from_slice(&[post_len as u32, post_head as u32, 0, 0]);
                    let slot = if view.win_len == w {
                        view.win_head
                    } else {
                        view.win_len
                    };
                    let mut rk = view.win_k.to_vec();
                    let mut rv = view.win_v.to_vec();
                    rk[slot * d..(slot + 1) * d].copy_from_slice(&k_new[g * d..(g + 1) * d]);
                    rv[slot * dv..(slot + 1) * dv].copy_from_slice(&v_new[g * dv..(g + 1) * dv]);
                    near_k.push((rk, post_len));
                    near_v.push(rv);
                }
                let mut near_want = vec![0.0f32; gcnt * hpg * dv];
                let scale = views[0].scale;
                for g in 0..gcnt {
                    let view = &views[g];
                    let (rk, post_len) = &near_k[g];
                    let rv = &near_v[g];
                    for h in 0..hpg {
                        let q = &q_new[(g * hpg + h) * d..(g * hpg + h + 1) * d];
                        let mut scores = Vec::with_capacity(view.sink_len + *post_len);
                        for s in 0..view.sink_len {
                            let mut acc = 0.0f32;
                            for j in 0..d {
                                acc += q[j] * view.sink_k[s * d + j];
                            }
                            scores.push(acc * scale);
                        }
                        for s in 0..*post_len {
                            let mut acc = 0.0f32;
                            for j in 0..d {
                                acc += q[j] * rk[s * d + j];
                            }
                            scores.push(acc * scale);
                        }
                        let c = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);
                        let weights: Vec<f32> = scores.iter().map(|&x| (x - c).exp()).collect();
                        let den = weights.iter().sum::<f32>();
                        for j in 0..dv {
                            let mut acc = 0.0f32;
                            for s in 0..view.sink_len {
                                acc += weights[s] * view.sink_v[s * dv + j];
                            }
                            for s in 0..*post_len {
                                acc += weights[view.sink_len + s] * rv[s * dv + j];
                            }
                            near_want[(g * hpg + h) * dv + j] = acc / den;
                        }
                    }
                }
                c.queue
                    .write_buffer(&dmeta, 0, bytemuck::cast_slice(&empty_meta));
                let mut empty_enc = c
                    .device
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
                {
                    let mut pass = begin_pass(&mut empty_enc);
                    pass.set_pipeline(&c.o1_attend);
                    pass.set_bind_group(0, &bg_att, &[]);
                    pass.dispatch_workgroups((gcnt * hpg) as u32, 1, 1);
                }
                flush_pass(&mut empty_enc);
                empty_enc.copy_buffer_to_buffer(&ob, 0, &stage, 0, (gcnt * hpg * dv * 4) as u64);
                submit(c, finish_enc(empty_enc));
                let (tx, rx) = std::sync::mpsc::channel();
                stage.map_async(wgpu::MapMode::Read, .., move |r| tx.send(r).unwrap());
                let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
                rx.recv().unwrap().unwrap();
                let empty_got: Vec<f32> =
                    bytemuck::cast_slice(&stage.get_mapped_range(..).unwrap()).to_vec();
                stage.unmap();
                assert!(
                    near_want.iter().all(|x| x.is_finite()),
                    "far-empty CPU reference produced non-finite output"
                );
                assert!(
                    empty_got.iter().all(|x| x.is_finite()),
                    "far-empty GPU output produced non-finite output"
                );
                let md_empty = near_want
                    .iter()
                    .zip(&empty_got)
                    .map(|(a, b)| (a - b).abs())
                    .fold(0.0f32, f32::max);
                assert!(
                    md_empty < 1e-3,
                    "wgpu o1 far-empty ≠ exact near (d={d} m={m} w={w} hpg={hpg}): max|Δ| = {md_empty}"
                );
            }
            c.o1m.lock().unwrap().remove(&(u64::MAX, usize::MAX));
            assert!(
                want.iter().all(|x| x.is_finite()),
                "O1 CPU reference produced non-finite output"
            );
            assert!(
                got.iter().all(|x| x.is_finite()),
                "O1 GPU output produced non-finite output"
            );
            let md = want
                .iter()
                .zip(&got)
                .map(|(a, b)| (a - b).abs())
                .fold(0.0f32, f32::max);
            assert!(
                md < 1e-3,
                "wgpu o1 step ≠ CPU (d={d} m={m} w={w} hpg={hpg}): max|Δ| = {md}"
            );
        }
    }

    #[test]
    #[ignore = "requires CMF_O1_REPLAY_TAPE retained external tape"]
    fn wgpu_o1_replay_real_tape_corrected() {
        // Actual first-Full-layer operator gate.  This is deliberately
        // separate from the old whole-tape/duplicate replay: the tape
        // contains positions 0..1523, the state seals 0..255, and rows
        // 256..1523 are each inserted once as a real post-seal step.
        let path = std::env::var("CMF_O1_REPLAY_TAPE")
            .expect("CMF_O1_REPLAY_TAPE is required for corrected replay");
        unsafe {
            std::env::set_var("CMF_GPU", "wgpu");
        }
        assert_eq!(
            std::env::var("WGPU_BACKEND").as_deref(),
            Ok("vulkan"),
            "corrected replay must run with WGPU_BACKEND=vulkan"
        );
        let data = std::fs::read(&path).expect("read retained real activation tape");
        let mut off = 0usize;
        fn u32le(data: &[u8], off: &mut usize) -> u32 {
            let end = *off + 4;
            assert!(end <= data.len(), "real tape u32 truncated");
            let x = u32::from_le_bytes(data[*off..end].try_into().unwrap());
            *off = end;
            x
        }
        assert!(data.len() >= 24, "real tape header truncated");
        let magic = u32le(&data, &mut off);
        let version = u32le(&data, &mut off);
        let layer = u32le(&data, &mut off) as usize;
        let nh = u32le(&data, &mut off) as usize;
        let nkv = u32le(&data, &mut off) as usize;
        let hd = u32le(&data, &mut off) as usize;
        assert_eq!(magic, 0x434d4652, "real tape magic");
        assert_eq!(version, 1, "real tape version");
        assert_eq!(layer, 3, "earliest Full layer");
        assert_eq!((nh, nkv, hd), (24, 4, 256));
        assert_eq!(nh % nkv, 0);
        let hpg = nh / nkv;
        let mut q = Vec::<f32>::new();
        let mut k = Vec::<f32>::new();
        let mut v = Vec::<f32>::new();
        let mut positions = Vec::<usize>::new();
        while off < data.len() {
            assert!(off + 16 <= data.len(), "real tape record header truncated");
            let pos = u32le(&data, &mut off) as usize;
            let qn = u32le(&data, &mut off) as usize;
            let kn = u32le(&data, &mut off) as usize;
            let vn = u32le(&data, &mut off) as usize;
            assert_eq!((qn, kn, vn), (nh * hd, nkv * hd, nkv * hd));
            let read_f32 = |off: &mut usize, n: usize| {
                let end = *off + n * 4;
                assert!(end <= data.len(), "real tape payload truncated");
                let out = data[*off..end]
                    .chunks_exact(4)
                    .map(|x| f32::from_le_bytes(x.try_into().unwrap()))
                    .collect::<Vec<_>>();
                *off = end;
                out
            };
            q.extend_from_slice(&read_f32(&mut off, qn));
            k.extend_from_slice(&read_f32(&mut off, kn));
            v.extend_from_slice(&read_f32(&mut off, vn));
            positions.push(pos);
        }
        let t = positions.len();
        assert_eq!(t, 1524, "retained replay tape length");
        assert_eq!(positions, (0..t).collect::<Vec<_>>(), "absolute positions");
        assert_eq!(off, data.len(), "real tape trailing bytes");
        assert!(q.iter().chain(&k).chain(&v).all(|x| x.is_finite()));

        const PREFIX: usize = 256;
        const M: usize = 32;
        const W: usize = 128;
        const SINK: usize = 4;
        const CHECKS: [usize; 4] = [256, 511, 1023, 1523];
        let scale = 1.0f64 / (hd as f64).sqrt();

        // Use the same f64 landmark and ridge routines as runtime seal, then
        // keep an independent f64 Aggregate stream for the four checkpoints.
        fn seg64(xs: &[f32], rows: usize, d: usize, m: usize) -> Vec<f64> {
            let mut out = vec![0.0; m * d];
            for i in 0..m {
                let lo = i * rows / m;
                let hi = (i + 1) * rows / m;
                for p in lo..hi {
                    for c in 0..d {
                        out[i * d + c] += xs[p * d + c] as f64;
                    }
                }
                let inv = 1.0 / (hi - lo) as f64;
                for c in 0..d {
                    out[i * d + c] *= inv;
                }
            }
            out
        }
        struct F64Head {
            q_tilde: Vec<f64>,
            k_tilde: Vec<f64>,
            mu: Vec<f64>,
            t_hat: Vec<f64>,
            z_hat: Vec<f64>,
            m_max: Vec<f64>,
        }
        let insert_f64 = |h: &mut F64Head, kr: &[f32], vr: &[f32]| {
            for a in 0..M {
                let mut l = 0.0;
                for d0 in 0..hd {
                    l += h.q_tilde[a * hd + d0] * kr[d0] as f64;
                }
                l *= scale;
                if l > h.m_max[a] {
                    let r = (h.m_max[a] - l).exp();
                    h.z_hat[a] *= r;
                    for x in &mut h.t_hat[a * hd..(a + 1) * hd] {
                        *x *= r;
                    }
                    h.m_max[a] = l;
                }
                let e = (l - h.m_max[a]).exp();
                h.z_hat[a] += e;
                for d0 in 0..hd {
                    h.t_hat[a * hd + d0] += e * vr[d0] as f64;
                }
            }
        };

        let mut groups = Vec::with_capacity(nkv);
        let mut f64_heads = Vec::with_capacity(nkv * hpg);
        for g in 0..nkv {
            let mut qh = vec![0.0f32; hpg * PREFIX * hd];
            let mut kg = vec![0.0f32; PREFIX * hd];
            let mut vg = vec![0.0f32; PREFIX * hd];
            for p in 0..PREFIX {
                kg[p * hd..(p + 1) * hd]
                    .copy_from_slice(&k[(p * nkv + g) * hd..(p * nkv + g + 1) * hd]);
                vg[p * hd..(p + 1) * hd]
                    .copy_from_slice(&v[(p * nkv + g) * hd..(p * nkv + g + 1) * hd]);
                for hh in 0..hpg {
                    let head = g * hpg + hh;
                    qh[(hh * PREFIX + p) * hd..(hh * PREFIX + p + 1) * hd]
                        .copy_from_slice(&q[(p * nh + head) * hd..(p * nh + head + 1) * hd]);
                }
            }
            let qrefs: Vec<&[f32]> = (0..hpg)
                .map(|hh| &qh[hh * PREFIX * hd..(hh + 1) * PREFIX * hd])
                .collect();
            let mut st = crate::nystrom::NystromState::new_group(M, W, SINK, hpg)
                .with_rect(crate::nystrom::O1Rect::Aggregate);
            st.prefill_group(&qrefs, &kg, &vg, PREFIX, hd, hd);
            let view = st.device_view();
            assert!(!view.exact_only, "p256 corrected replay must seal");
            assert_eq!(
                (view.m_eff, view.w, view.sink_len, view.far_len),
                (M, W, SINK, 124)
            );

            let k_tilde = seg64(&kg, PREFIX, hd, M);
            for hh in 0..hpg {
                let q_tilde = seg64(&qh[hh * PREFIX * hd..(hh + 1) * PREFIX * hd], PREFIX, hd, M);
                let mut au = vec![0.0; M * M];
                for a in 0..M {
                    for b in 0..M {
                        let mut dot = 0.0;
                        for d0 in 0..hd {
                            dot += q_tilde[a * hd + d0] * k_tilde[b * hd + d0];
                        }
                        au[a * M + b] = (dot * scale).exp();
                    }
                }
                let mu = crate::nystrom::ridge_pinv(&au, M);
                let mut fh = F64Head {
                    q_tilde,
                    k_tilde: k_tilde.clone(),
                    mu,
                    t_hat: vec![0.0; M * hd],
                    z_hat: vec![0.0; M],
                    m_max: vec![f64::NEG_INFINITY; M],
                };
                // At seal, prompt rows 4..127 are already in the far field.
                for p in SINK..PREFIX - W {
                    insert_f64(
                        &mut fh,
                        &k[(p * nkv + g) * hd..(p * nkv + g + 1) * hd],
                        &v[(p * nkv + g) * hd..(p * nkv + g + 1) * hd],
                    );
                }
                f64_heads.push(fh);
            }
            groups.push(st);
        }

        let c = ctx().expect("actual Vulkan adapter required for corrected replay");
        let uid = u64::MAX - 17;
        let views: Vec<_> = groups.iter().map(|st| st.device_view()).collect();
        o1_ensure(c, uid, layer, &views, 1).expect("upload corrected replay state");
        let dev_bufs = {
            let map = c.o1m.lock().unwrap();
            let dref = map.get(&(uid, layer)).expect("uploaded O1 state");
            (
                dref.meta.clone(),
                dref.ring_k.clone(),
                dref.ring_v.clone(),
                dref.sink_k.clone(),
                dref.sink_v.clone(),
                dref.k_tilde.clone(),
                dref.qt.clone(),
                dref.mu.clone(),
                dref.mz.clone(),
                dref.that.clone(),
                dref.scale,
            )
        };
        let (dmeta, drk, drv, dsk, dsv, dkt, dqt, dmu, dmz, dth, scale_f32) = dev_bufs;
        use wgpu::util::DeviceExt;
        let input = |data: &[f32], label: &str| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some(label),
                    contents: bytemuck::cast_slice(data),
                    usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                })
        };
        let q_all = input(&q, "o1-corrected-replay-q");
        let k_all = input(&k, "o1-corrected-replay-k");
        let v_all = input(&v, "o1-corrected-replay-v");
        let work = |bytes: u64, label: &str| {
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some(label),
                size: bytes,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            })
        };
        let q_work = work((nh * hd * 4) as u64, "o1-corrected-replay-q-work");
        let k_work = work((nkv * hd * 4) as u64, "o1-corrected-replay-k-work");
        let v_work = work((nkv * hd * 4) as u64, "o1-corrected-replay-v-work");
        let out_bytes = (nh * hd * 4) as u64;
        let out = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("o1-corrected-replay-out"),
            size: out_bytes,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let stage_bytes = out_bytes * CHECKS.len() as u64;
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("o1-corrected-replay-stage"),
            size: stage_bytes,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let uniform = uniform_u32x8(
            c,
            [
                hpg as u32,
                M as u32,
                W as u32,
                SINK as u32,
                hd as u32,
                hd as u32,
                scale_f32.to_bits(),
                0,
            ],
        );
        let bgf = |layout: &wgpu::BindGroupLayout, bufs: &[&wgpu::Buffer]| {
            let entries: Vec<_> = bufs
                .iter()
                .enumerate()
                .map(|(i, b)| bind_buf(i as u32, b))
                .collect();
            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("o1-corrected-replay-bind"),
                layout,
                entries: &entries,
            })
        };
        let bg_far = bgf(
            &c.layout_o1_far,
            &[&dmeta, &drk, &drv, &dqt, &dmz, &dth, &uniform],
        );
        let bg_push = bgf(
            &c.layout_o1_push,
            &[&dmeta, &k_work, &v_work, &drk, &drv, &uniform],
        );
        let bg_att = bgf(
            &c.layout_o1_attend,
            &[
                &dmeta, &q_work, &drk, &drv, &dsk, &dsv, &dkt, &dmu, &dmz, &dth, &out, &uniform,
            ],
        );
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("o1-corrected-replay-encoder"),
            });
        let mut cpu_checkpoints = Vec::<Vec<f32>>::new();
        let mut f64_errors = Vec::<f64>::new();
        for row in PREFIX..t {
            let mut cpu_row = vec![0.0f32; nh * hd];
            for g in 0..nkv {
                groups[g].step_group(
                    &q[row * nh * hd + g * hpg * hd..row * nh * hd + (g + 1) * hpg * hd],
                    &k[(row * nkv + g) * hd..(row * nkv + g + 1) * hd],
                    &v[(row * nkv + g) * hd..(row * nkv + g + 1) * hd],
                    &mut cpu_row[g * hpg * hd..(g + 1) * hpg * hd],
                );
            }
            let evicted = row - W;
            for g in 0..nkv {
                for hh in 0..hpg {
                    insert_f64(
                        &mut f64_heads[g * hpg + hh],
                        &k[(evicted * nkv + g) * hd..(evicted * nkv + g + 1) * hd],
                        &v[(evicted * nkv + g) * hd..(evicted * nkv + g + 1) * hd],
                    );
                }
            }

            // The CPU-side f64 checkpoint uses the same Aggregate partition:
            // sinks 0..3, far 4..row-128, and near row-127..row.
            if let Some(ci) = CHECKS.iter().position(|&x| x == row) {
                let mut max_err = 0.0f64;
                for g in 0..nkv {
                    for hh in 0..hpg {
                        let head = g * hpg + hh;
                        let qrow = &q[(row * nh + head) * hd..(row * nh + head + 1) * hd];
                        let fh = &f64_heads[head];
                        let mut near = Vec::<(f64, usize)>::with_capacity(SINK + W);
                        for p in 0..SINK {
                            let kr = &k[(p * nkv + g) * hd..(p * nkv + g + 1) * hd];
                            let mut s = 0.0;
                            for d0 in 0..hd {
                                s += qrow[d0] as f64 * kr[d0] as f64;
                            }
                            near.push((s * scale, p));
                        }
                        for p in row + 1 - W..=row {
                            let kr = &k[(p * nkv + g) * hd..(p * nkv + g + 1) * hd];
                            let mut s = 0.0;
                            for d0 in 0..hd {
                                s += qrow[d0] as f64 * kr[d0] as f64;
                            }
                            near.push((s * scale, p));
                        }
                        let c_near = near
                            .iter()
                            .map(|&(s, _)| s)
                            .fold(f64::NEG_INFINITY, f64::max);
                        let mut landmark = vec![0.0f64; M];
                        let mut f = f64::NEG_INFINITY;
                        for a in 0..M {
                            let mut s = 0.0;
                            for d0 in 0..hd {
                                s += qrow[d0] as f64 * fh.k_tilde[a * hd + d0];
                            }
                            landmark[a] = s * scale;
                            f = f.max(landmark[a]);
                        }
                        let mut c_all = c_near;
                        for a in 0..M {
                            c_all = c_all.max(f + fh.m_max[a]);
                        }
                        let mut far_den = 0.0;
                        let mut out64 = vec![0.0f64; hd];
                        for b in 0..M {
                            let mut u = 0.0;
                            for a in 0..M {
                                u += (landmark[a] - f).exp() * fh.mu[a * M + b];
                            }
                            let gain = u * (f + fh.m_max[b] - c_all).exp();
                            far_den += gain * fh.z_hat[b];
                            for d0 in 0..hd {
                                out64[d0] += gain * fh.t_hat[b * hd + d0];
                            }
                        }
                        if far_den < 0.0 {
                            // Aggregate's signed far readout drops an
                            // unusable negative denominator and numerator
                            // before adding the exact near field.
                            far_den = 0.0;
                            out64.fill(0.0);
                        }
                        let mut den = far_den;
                        for &(s, p) in &near {
                            let wt = (s - c_all).exp();
                            den += wt;
                            let vr = &v[(p * nkv + g) * hd..(p * nkv + g + 1) * hd];
                            for d0 in 0..hd {
                                out64[d0] += wt * vr[d0] as f64;
                            }
                        }
                        assert!(den.is_finite() && den >= 0.0);
                        for x in &mut out64 {
                            *x /= den.max(1e-30);
                        }
                        let mut logits = vec![0.0f64; row + 1];
                        for p in 0..=row {
                            let kr = &k[(p * nkv + g) * hd..(p * nkv + g + 1) * hd];
                            let mut s = 0.0;
                            for d0 in 0..hd {
                                s += qrow[d0] as f64 * kr[d0] as f64;
                            }
                            logits[p] = s * scale;
                        }
                        let mx = logits.iter().copied().fold(f64::NEG_INFINITY, f64::max);
                        let mut den_exact = 0.0;
                        for x in &mut logits {
                            *x = (*x - mx).exp();
                            den_exact += *x;
                        }
                        let mut exact = vec![0.0f64; hd];
                        for (p, wt) in logits.iter().enumerate() {
                            let vr = &v[(p * nkv + g) * hd..(p * nkv + g + 1) * hd];
                            for d0 in 0..hd {
                                exact[d0] += *wt * vr[d0] as f64;
                            }
                        }
                        let mut local = 0.0f64;
                        let mut cpu_local = 0.0f32;
                        for d0 in 0..hd {
                            exact[d0] /= den_exact;
                            local = local.max((out64[d0] - exact[d0]).abs());
                            cpu_local =
                                cpu_local.max((cpu_row[head * hd + d0] - out64[d0] as f32).abs());
                        }
                        assert!(out64.iter().all(|x| x.is_finite()));
                        assert!(cpu_local.is_finite());
                        max_err = max_err.max(local);
                    }
                }
                f64_errors.push(max_err);
                cpu_checkpoints.push(cpu_row);
                eprintln!("o1-corrected-replay: row={row} f64_dense_max={max_err:.6e}");
                let src_q = (row * nh * hd * 4) as u64;
                let src_k = (row * nkv * hd * 4) as u64;
                let src_v = (row * nkv * hd * 4) as u64;
                flush_pass(&enc);
                enc.copy_buffer_to_buffer(&q_all, src_q, &q_work, 0, (nh * hd * 4) as u64);
                enc.copy_buffer_to_buffer(&k_all, src_k, &k_work, 0, (nkv * hd * 4) as u64);
                enc.copy_buffer_to_buffer(&v_all, src_v, &v_work, 0, (nkv * hd * 4) as u64);
                {
                    let mut pass = begin_pass(&mut enc);
                    pass.set_pipeline(&c.o1_far);
                    pass.set_bind_group(0, &bg_far, &[]);
                    pass.dispatch_workgroups((nkv * hpg * M) as u32, 1, 1);
                    pass.set_pipeline(&c.o1_push);
                    pass.set_bind_group(0, &bg_push, &[]);
                    pass.dispatch_workgroups(nkv as u32, 1, 1);
                    pass.set_pipeline(&c.o1_attend);
                    pass.set_bind_group(0, &bg_att, &[]);
                    pass.dispatch_workgroups((nkv * hpg) as u32, 1, 1);
                }
                flush_pass(&enc);
                enc.copy_buffer_to_buffer(&out, 0, &stage, ci as u64 * out_bytes, out_bytes);
            } else {
                let src_q = (row * nh * hd * 4) as u64;
                let src_k = (row * nkv * hd * 4) as u64;
                let src_v = (row * nkv * hd * 4) as u64;
                flush_pass(&enc);
                enc.copy_buffer_to_buffer(&q_all, src_q, &q_work, 0, (nh * hd * 4) as u64);
                enc.copy_buffer_to_buffer(&k_all, src_k, &k_work, 0, (nkv * hd * 4) as u64);
                enc.copy_buffer_to_buffer(&v_all, src_v, &v_work, 0, (nkv * hd * 4) as u64);
                {
                    let mut pass = begin_pass(&mut enc);
                    pass.set_pipeline(&c.o1_far);
                    pass.set_bind_group(0, &bg_far, &[]);
                    pass.dispatch_workgroups((nkv * hpg * M) as u32, 1, 1);
                    pass.set_pipeline(&c.o1_push);
                    pass.set_bind_group(0, &bg_push, &[]);
                    pass.dispatch_workgroups(nkv as u32, 1, 1);
                    pass.set_pipeline(&c.o1_attend);
                    pass.set_bind_group(0, &bg_att, &[]);
                    pass.dispatch_workgroups((nkv * hpg) as u32, 1, 1);
                }
            }
        }
        flush_pass(&enc);
        submit(c, finish_enc(enc));
        let (tx, rx) = std::sync::mpsc::channel();
        stage.map_async(wgpu::MapMode::Read, ..stage_bytes, move |r| {
            tx.send(r).unwrap()
        });
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        rx.recv().unwrap().unwrap();
        let mapped = stage.get_mapped_range(..stage_bytes).unwrap();
        let gpu: Vec<f32> = bytemuck::cast_slice(&mapped).to_vec();
        drop(mapped);
        stage.unmap();
        for (ci, &row) in CHECKS.iter().enumerate() {
            let got = &gpu[ci * nh * hd..(ci + 1) * nh * hd];
            let want = &cpu_checkpoints[ci];
            let md = want
                .iter()
                .zip(got)
                .map(|(a, b)| (a - b).abs())
                .fold(0.0f32, f32::max);
            let expected = [3.61143, 1.16262, 1.10457, 1.03681][ci];
            eprintln!(
                "o1-corrected-replay: row={row} cpu_gpu_maxabs={md:.6e} f64_dense_max={:.6e}",
                f64_errors[ci]
            );
            assert!(got.iter().all(|x| x.is_finite()));
            assert!(
                md < 1e-3,
                "corrected Vulkan/CPU mismatch at row {row}: {md}"
            );
            assert!(
                (f64_errors[ci] - expected).abs() < 0.05 * expected,
                "f64 checkpoint changed at row {row}: {} vs {expected}",
                f64_errors[ci]
            );
        }
        c.o1m.lock().unwrap().remove(&(uid, layer));
    }

    #[test]
    fn wgpu_gdn_step_matches_cpu() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        if ctx().is_none() {
            eprintln!("no wgpu adapter — skipping gdn_step test");
            return;
        }
        let (nv, nk, dk, dv) = (4usize, 2usize, 8usize, 8usize);
        let kd = nk * dk;
        let rep = nv / nk;
        let cdim = 2 * kd + nv * dv;
        let eps = 1e-6f32;
        let jit = |a: usize, b: usize| ((a * 23 + b * 11 + 5) % 71) as f32 / 71.0 - 0.5;
        let cq: Vec<f32> = (0..cdim).map(|i| jit(i, 1)).collect();
        let z: Vec<f32> = (0..nv * dv).map(|i| jit(i, 2)).collect();
        let a: Vec<f32> = (0..nv).map(|i| jit(i, 3)).collect();
        let b: Vec<f32> = (0..nv).map(|i| jit(i, 4)).collect();
        let alog: Vec<f32> = (0..nv).map(|i| jit(i, 5) - 0.5).collect();
        let dtb: Vec<f32> = (0..nv).map(|i| jit(i, 6)).collect();
        let norm: Vec<f32> = (0..dv).map(|i| 0.8 + jit(i, 7)).collect();
        let s0: Vec<f32> = (0..nv * dk * dv).map(|i| jit(i, 8) * 0.3).collect();
        // CPU reference (mirrors linear_core::gdn_step).
        let sp = |x: f32| if x > 20.0 { x } else { (1.0 + x.exp()).ln() };
        let sig = |x: f32| 1.0 / (1.0 + (-x).exp());
        let silu = |x: f32| x / (1.0 + (-x).exp());
        let mut sc = s0.clone();
        let mut want = vec![0f32; nv * dv];
        for h in 0..nv {
            let ko = h / rep;
            let (qs, ks) = (ko * dk, kd + ko * dk);
            let nq: f32 = (0..dk).map(|d| cq[qs + d] * cq[qs + d]).sum();
            let nkn: f32 = (0..dk).map(|d| cq[ks + d] * cq[ks + d]).sum();
            let invq = 1.0 / ((nq + 1e-6).sqrt() * (dk as f32).sqrt());
            let invk = 1.0 / (nkn + 1e-6).sqrt();
            let qf: Vec<f32> = (0..dk).map(|d| cq[qs + d] * invq).collect();
            let kf: Vec<f32> = (0..dk).map(|d| cq[ks + d] * invk).collect();
            let g = (-(alog[h].exp()) * sp(a[h] + dtb[h])).exp();
            let beta = sig(b[h]);
            let sbase = h * dk * dv;
            let mut o = vec![0f32; dv];
            for dj in 0..dv {
                let vt = cq[2 * kd + h * dv + dj];
                let mut kv = 0.0;
                for di in 0..dk {
                    kv += sc[sbase + di * dv + dj] * kf[di];
                }
                let delta = (vt - g * kv) * beta;
                for di in 0..dk {
                    let idx = sbase + di * dv + dj;
                    let cell = g * sc[idx] + kf[di] * delta;
                    sc[idx] = cell;
                    o[dj] += qf[di] * cell;
                }
            }
            let ss: f32 = o.iter().map(|v| v * v).sum();
            let inv = 1.0 / (ss / dv as f32 + eps).sqrt();
            for dj in 0..dv {
                want[h * dv + dj] = o[dj] * inv * norm[dj] * silu(z[h * dv + dj]);
            }
        }
        // GPU
        let mut sg = s0.clone();
        let mut got = vec![0f32; nv * dv];
        assert!(gdn_step_gpu(
            &cq, &z, &a, &b, &alog, &dtb, &norm, &mut sg, nv, dk, dv, kd, rep, cdim, eps, &mut got
        ));
        let mo = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        let msd = sc
            .iter()
            .zip(&sg)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(mo < 2e-3, "wgpu gdn_step o ≠ CPU: max|Δ| = {mo}");
        assert!(msd < 2e-3, "wgpu gdn_step S ≠ CPU: max|Δ| = {msd}");
    }

    #[test]
    fn wgpu_gdn_conv_matches_cpu() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        if ctx().is_none() {
            eprintln!("no wgpu adapter — skipping gdn_conv test");
            return;
        }
        let (cdim, kk) = (48usize, 4usize);
        let jit = |a: usize, b: usize| ((a * 19 + b * 7 + 3) % 61) as f32 / 61.0 - 0.5;
        let qkv: Vec<f32> = (0..cdim).map(|i| jit(i, 1)).collect();
        let taps: Vec<f32> = (0..cdim * kk).map(|i| jit(i, 2)).collect();
        let ring0: Vec<f32> = (0..(kk - 1) * cdim).map(|i| jit(i, 3)).collect();
        let silu = |x: f32| x / (1.0 + (-x).exp());
        // CPU reference
        let mut rc = ring0.clone();
        let mut want_cq = vec![0f32; cdim];
        for c in 0..cdim {
            let t = &taps[c * kk..(c + 1) * kk];
            let mut acc = qkv[c] * t[kk - 1];
            for j in 0..kk - 1 {
                acc += rc[j * cdim + c] * t[j];
            }
            want_cq[c] = silu(acc);
        }
        rc.copy_within(cdim.., 0);
        let tail = (kk - 2) * cdim;
        rc[tail..tail + cdim].copy_from_slice(&qkv[..cdim]);
        // GPU
        let mut rg = ring0.clone();
        let mut got_cq = vec![0f32; cdim];
        assert!(gdn_conv_gpu(&qkv, &taps, &mut rg, cdim, kk, &mut got_cq));
        let mc = want_cq
            .iter()
            .zip(&got_cq)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        let mr = rc
            .iter()
            .zip(&rg)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(mc < 1e-5, "wgpu gdn_conv cq ≠ CPU: {mc}");
        assert!(mr < 1e-6, "wgpu gdn_conv ring ≠ CPU: {mr}");
    }

    // Build a deterministic q1 payload for a [rows, cols] weight + its dequant.
    #[cfg(test)]
    fn mk_q1(rows: usize, cols: usize, seed: usize) -> (Vec<u8>, Vec<f32>) {
        let gpr = cols / 32;
        let mut payload = Vec::new();
        for t in 0..rows * gpr {
            let sc = 0.004 + ((t + seed) % 9) as f32 * 0.003;
            payload.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
            for j in 0..4 {
                payload.push(((t * 37 + j * 53 + seed * 7 + 11) % 251) as u8);
            }
        }
        let mut w = vec![0f32; rows * cols];
        cortiq_core::quant::dequant_q1(&payload, &mut w);
        (payload, w)
    }

    #[test]
    fn wgpu_attn_block_matches_cpu() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping attn_block test");
            return;
        };
        let (nh, nkv, hd, rd, hidden, cap, stored) =
            (4usize, 2usize, 64usize, 64usize, 128usize, 8usize, 2usize);
        let hpk = nh / nkv;
        let eps = 1e-6f32;
        let flags = 2u32 | 4u32; // qnorm + knorm, no gate
        let jit = |a: usize, b: usize| ((a * 31 + b * 17 + 3) % 83) as f32 / 83.0 - 0.5;
        let h_in: Vec<f32> = (0..hidden).map(|i| jit(i, 1)).collect();
        let norm_w: Vec<f32> = (0..hidden).map(|i| 0.8 + jit(i, 2)).collect();
        let (wq_p, wq) = mk_q1(nh * hd, hidden, 1);
        let (wk_p, wk) = mk_q1(nkv * hd, hidden, 2);
        let (wv_p, wv) = mk_q1(nkv * hd, hidden, 3);
        let (wo_p, wo) = mk_q1(hidden, nh * hd, 4);
        let qnw: Vec<f32> = (0..hd).map(|d| 0.7 + jit(d, 5)).collect();
        let knw: Vec<f32> = (0..hd).map(|d| 0.7 + jit(d, 6)).collect();
        let invf: Vec<f32> = (0..rd / 2)
            .map(|i| 1.0 / (10000f32).powf(2.0 * i as f32 / rd as f32))
            .collect();
        // Pre-filled device K/V caches [nkv, cap, hd] (first `stored` rows valid).
        let mut kc = vec![0f32; nkv * cap * hd];
        let mut vc = vec![0f32; nkv * cap * hd];
        for kh in 0..nkv {
            for p in 0..stored {
                for d in 0..hd {
                    kc[(kh * cap + p) * hd + d] = jit(kh * 900 + p * 30 + d, 7);
                    vc[(kh * cap + p) * hd + d] = jit(kh * 900 + p * 30 + d, 8);
                }
            }
        }
        let mkcache = |data: &[f32]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some("cache"),
                    contents: bytemuck::cast_slice(data),
                    usage: wgpu::BufferUsages::STORAGE
                        | wgpu::BufferUsages::COPY_SRC
                        | wgpu::BufferUsages::COPY_DST,
                })
        };
        let kbuf = mkcache(&kc);
        let vbuf = mkcache(&vc);
        // ---- CPU reference ----
        let ss: f32 = h_in.iter().map(|x| x * x).sum();
        let rinv = 1.0 / (ss / hidden as f32 + eps).sqrt();
        let normed: Vec<f32> = (0..hidden).map(|i| h_in[i] * rinv * norm_w[i]).collect();
        let matvec = |w: &[f32], rows: usize, cols: usize, x: &[f32]| -> Vec<f32> {
            (0..rows)
                .map(|o| (0..cols).map(|i| w[o * cols + i] * x[i]).sum())
                .collect()
        };
        let qraw = matvec(&wq, nh * hd, hidden, &normed);
        let kv_k = matvec(&wk, nkv * hd, hidden, &normed);
        let kv_v = matvec(&wv, nkv * hd, hidden, &normed);
        let norm_rope = |v: &mut [f32], w: &[f32]| {
            let s: f32 = v.iter().map(|x| x * x).sum();
            let inv = 1.0 / (s / hd as f32 + eps).sqrt();
            for d in 0..hd {
                v[d] = v[d] * inv * w[d];
            }
            for i in 0..rd / 2 {
                let ang = stored as f32 * invf[i];
                let (co, si) = (ang.cos(), ang.sin());
                let (x0, x1) = (v[i], v[i + rd / 2]);
                v[i] = x0 * co - x1 * si;
                v[i + rd / 2] = x0 * si + x1 * co;
            }
        };
        let mut qout = vec![0f32; nh * hd];
        for h in 0..nh {
            let mut q = qraw[h * hd..(h + 1) * hd].to_vec();
            norm_rope(&mut q, &qnw);
            qout[h * hd..(h + 1) * hd].copy_from_slice(&q);
        }
        for kh in 0..nkv {
            let mut kk = kv_k[kh * hd..(kh + 1) * hd].to_vec();
            norm_rope(&mut kk, &knw);
            kc[(kh * cap + stored) * hd..(kh * cap + stored) * hd + hd].copy_from_slice(&kk);
            vc[(kh * cap + stored) * hd..(kh * cap + stored) * hd + hd]
                .copy_from_slice(&kv_v[kh * hd..(kh + 1) * hd]);
        }
        let n = stored + 1;
        let scale = 1.0 / (hd as f32).sqrt();
        let mut attn = vec![0f32; nh * hd];
        for h in 0..nh {
            let kh = h / hpk;
            let mut sc: Vec<f32> = (0..n)
                .map(|p| {
                    (0..hd)
                        .map(|d| qout[h * hd + d] * kc[(kh * cap + p) * hd + d])
                        .sum::<f32>()
                        * scale
                })
                .collect();
            let mx = sc.iter().cloned().fold(f32::MIN, f32::max);
            let mut den = 0.0;
            for s in sc.iter_mut() {
                *s = (*s - mx).exp();
                den += *s;
            }
            for d in 0..hd {
                attn[h * hd + d] = (0..n)
                    .map(|p| sc[p] * vc[(kh * cap + p) * hd + d])
                    .sum::<f32>()
                    / den;
            }
        }
        let o = matvec(&wo, hidden, nh * hd, &attn);
        let want: Vec<f32> = (0..hidden).map(|i| h_in[i] + o[i]).collect();
        // ---- GPU block ----
        let mut got = vec![0f32; hidden];
        assert!(attn_block_gpu(
            &h_in, &norm_w, &wq_p, &wk_p, &wv_p, &wo_p, &qnw, &knw, &invf, &kbuf, &vbuf, nh, nkv,
            hd, rd, hidden, cap, stored, flags, eps, &mut got,
        ));
        let md = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(md < 2e-3, "wgpu attn_block ≠ CPU: max|Δ| = {md}");
    }

    // Payoff microbench: the resident attention block (ONE submit) vs the same
    // steps as separate submit+readback ops (today's per-op decode). Run with
    //   cargo test -p cortiq-engine --release --features gpu attn_block_timing -- --ignored --nocapture
    fn env_usize(k: &str, d: usize) -> usize {
        std::env::var(k)
            .ok()
            .and_then(|v| v.parse().ok())
            .unwrap_or(d)
    }

    /// CPU model of the q2tp subgroup reduction.  Keeping this independent
    /// from the WGSL tree/subgroup implementation catches a lane-grouping
    /// mistake before an optional GPU A/B is allowed to touch model output.
    fn q2tp_subgroup_cpu(acc: &[[f32; 4]], subgroup_width: usize) -> Option<[[f32; 4]; 4]> {
        if acc.len() != 256 || !matches!(subgroup_width, 32 | 64) {
            return None;
        }
        let subgroups_per_row_group = 64 / subgroup_width;
        let mut out = [[0.0f32; 4]; 4];
        for row_group in 0..4 {
            let first_subgroup = row_group * subgroups_per_row_group;
            for subgroup in 0..subgroups_per_row_group {
                let lo = (first_subgroup + subgroup) * subgroup_width;
                let hi = lo + subgroup_width;
                for lane in lo..hi {
                    for component in 0..4 {
                        out[row_group][component] += acc[lane][component];
                    }
                }
            }
        }
        Some(out)
    }

    #[test]
    fn q2tp_subgroup_cpu_reduction_matches_64_lane_rows() {
        let acc: Vec<[f32; 4]> = (0..256)
            .map(|lane| {
                [
                    (lane as f32 + 1.0) * 0.03125,
                    ((lane * 7 % 53) as f32 - 26.0) * 0.125,
                    if lane & 1 == 0 { 1.0 } else { -1.0 },
                    (lane * lane % 97) as f32 * 0.0078125,
                ]
            })
            .collect();
        for width in [32usize, 64] {
            let got = q2tp_subgroup_cpu(&acc, width).expect("supported subgroup width");
            for row_group in 0..4 {
                let lo = row_group * 64;
                let hi = lo + 64;
                for component in 0..4 {
                    let want: f32 = acc[lo..hi].iter().map(|v| v[component]).sum();
                    assert_eq!(got[row_group][component].to_bits(), want.to_bits());
                }
            }
        }
        assert!(q2tp_subgroup_cpu(&acc, 16).is_none());
        assert!(q2tp_subgroup_cpu(&acc, 128).is_none());
    }

    /// Direct f64 Walsh-sign reference for the device FWHT tests.  This is
    /// deliberately not `cortiq_core::hadamard::fwht_f32`, so a shared
    /// butterfly bug cannot make the GPU test pass on both sides.
    fn fwht_reference_f64(values: &[f32], signs: &[f32], block: usize) -> Vec<f32> {
        assert_eq!(values.len(), signs.len());
        assert!(block.is_power_of_two() && values.len() % block == 0);
        let inv = 1.0f64 / (block as f64).sqrt();
        let mut out = vec![0.0f32; values.len()];
        for base in (0..values.len()).step_by(block) {
            for row in 0..block {
                let mut sum = 0.0f64;
                for col in 0..block {
                    let h = if (row & col).count_ones() & 1 == 0 { 1.0 } else { -1.0 };
                    sum += h * values[base + col] as f64 * signs[base + col] as f64;
                }
                out[base + row] = (sum * inv) as f32;
            }
        }
        out
    }

    /// The cooperative GEMM on the smallest shape that still exercises it:
    /// one K-step, one tile, sixteen tokens. Everything the full kernel
    /// does, with few enough numbers to read.
    ///
    /// `cargo test -p cortiq-engine --release --features gpu
    ///  wgpu_coop_small_gemm -- --ignored --nocapture`
    #[test]
    #[ignore]
    fn wgpu_coop_small_gemm() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        unsafe { std::env::set_var("CMF_COOP", "1") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu device — skipping");
            return;
        };
        let Some(pipe) = c.q4tp_mm_coop.as_ref() else {
            eprintln!("no cooperative pipeline — skipping");
            return;
        };
        let (rows, cols, nb) = (64usize, 32usize, 16usize);
        let total =
            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
                .unwrap();
        let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
        let mut wb: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
        let lo = cortiq_core::quant::f32_to_f16(-4.0);
        let step = cortiq_core::quant::f32_to_f16(0.1);
        for r in 0..rows {
            let o = params_off + r * 4;
            wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
            wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
        }
        let mut w = vec![0f32; rows * cols];
        cortiq_core::quant::dequant_q4tp(&wb, rows, cols, &mut w);
        let xs: Vec<f32> = (0..nb * cols)
            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
            .collect();

        use wgpu::util::DeviceExt;
        let mk = |b: &[u8], u: wgpu::BufferUsages| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: b,
                    usage: u,
                })
        };
        let wbuf = mk(&wb, wgpu::BufferUsages::STORAGE);
        let xbuf = mk(bytemuck::cast_slice(&xs), wgpu::BufferUsages::STORAGE);
        let ybytes = (nb * rows * 4) as u64;
        let ybuf = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: ybytes,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: ybytes,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let pbuf = mk(
            bytemuck::cast_slice(&[(cols / 4) as u32, rows as u32, nb as u32, 0u32]),
            wgpu::BufferUsages::UNIFORM,
        );
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &pipe.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wbuf),
                bind_buf(1, &xbuf),
                bind_buf(2, &ybuf),
                bind_buf(3, &pbuf),
            ],
        });
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(pipe);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups((rows as u32).div_ceil(64), (nb as u32).div_ceil(64), 1);
        }
        let mut got = vec![0f32; nb * rows];
        assert!(
            readback(c, enc, &ybuf, &stage, ybytes, &mut got),
            "readback"
        );
        let want: Vec<f32> = (0..nb)
            .flat_map(|t| {
                let x = &xs[t * cols..(t + 1) * cols];
                (0..rows)
                    .map(|r| (0..cols).map(|k| w[r * cols + k] * x[k]).sum::<f32>())
                    .collect::<Vec<_>>()
            })
            .collect();
        println!("  w[0][0..4] {:?}", &w[..4]);
        println!("  x[0][0..4] {:?}", &xs[..4]);
        for t in [0usize, 1, 15] {
            println!("token {t}: got {:?}", &got[t * rows..t * rows + 4]);
            println!("         want {:?}", &want[t * rows..t * rows + 4]);
        }
        let worst = got
            .iter()
            .zip(&want)
            .map(|(g, w)| (g - w).abs() / w.abs().max(1e-3))
            .fold(0f32, f32::max);
        println!("worst relative {worst:.3e}");
    }

    /// What do `coopLoad`, `coopLoadT` and `coopStore` actually mean?
    /// A 16x16 product of two matrices whose answer is known by hand, so
    /// the layout question is settled by arithmetic instead of by trying
    /// variants on a GEMM with a thousand other moving parts.
    ///
    /// `cargo test -p cortiq-engine --release --features gpu
    ///  wgpu_coop_layout_probe -- --ignored --nocapture`
    #[test]
    #[ignore]
    fn wgpu_coop_layout_probe() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        unsafe { std::env::set_var("CMF_COOP", "1") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu device — skipping");
            return;
        };
        // a[i][k] = i + k/16, b[k][j] = (k == j) ? 1 : 0  → c == a.
        // With b the identity, any mix-up in B's layout still returns a,
        // so b is made asymmetric: b[k][j] = k*16 + j, and the expected
        // product is computed on the host below.
        let src = r#"
enable wgpu_cooperative_matrix;
enable f16;
@group(0) @binding(0) var<storage, read> ain: array<f32>;
@group(0) @binding(1) var<storage, read> bin: array<f32>;
@group(0) @binding(2) var<storage, read_write> out: array<f32>;
var<workgroup> sa: array<f16, 32 * 32>;
var<workgroup> sb: array<f16, 32 * 32>;
var<workgroup> sc: array<f32, 32 * 32>;
@compute @workgroup_size(128)
fn main(@builtin(local_invocation_index) tid: u32,
        @builtin(subgroup_id) sg: u32) {
    for (var i = tid; i < 1024u; i = i + 128u) {
        sa[i] = f16(0.0); sb[i] = f16(0.0); sc[i] = 0.0;
    }
    workgroupBarrier();
    for (var i = tid; i < 256u; i = i + 128u) {
        let r = i / 16u;
        let cc = i % 16u;
        sa[16u + r * 32u + cc] = f16(ain[i]);
        sb[16u + r * 32u + cc] = f16(bin[i]);
    }
    workgroupBarrier();
    // The same product twice, accumulated across a loop iteration — which
    // is what a GEMM does and what the isolated probe never did.
    var acc: coop_mat16x16<f32, C>;
    for (var it = 0u; it < 2u; it = it + 1u) {
        let a = coopLoad<coop_mat16x16<f16, A>>(&sa[16u], 32u);
        let b = coopLoad<coop_mat16x16<f16, B>>(&sb[16u], 32u);
        acc = coopMultiplyAdd(a, b, acc);
    }
    // Each subgroup writes its OWN row of the output. If `subgroup_id`
    // does not actually distinguish them, three of the four blocks stay
    // zero — which the probe with a single shared store could not see.
    if (sg == 0u) { coopStore(acc, &sc[16u], 32u); }
    workgroupBarrier();
    for (var i = tid; i < 256u; i = i + 128u) {
        let r = i / 16u;
        let cc = i % 16u;
        out[i] = sc[16u + r * 32u + cc];
    }
    out[256u + tid] = f32(sg);
}
"#;
        let scope = c.device.push_error_scope(wgpu::ErrorFilter::Validation);
        let m = c.device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("coop-layout"),
            source: wgpu::ShaderSource::Wgsl(src.into()),
        });
        if let Some(e) = pollster::block_on(scope.pop()) {
            println!("module rejected: {e}");
            return;
        }
        let pipe = c
            .device
            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("coop-layout"),
                layout: None,
                module: &m,
                entry_point: Some("main"),
                compilation_options: Default::default(),
                cache: c.pipeline_cache.as_ref(),
            });
        // What the driver actually gives a 128-thread workgroup: the GEMM
        // derives its subgroup index as `tid / 32`, which is only right if
        // this says 32 and the assignment is linear.
        {
            let probe = r#"
@group(0) @binding(0) var<storage, read_write> o: array<f32>;
@compute @workgroup_size(128)
fn main(@builtin(local_invocation_index) tid: u32,
        @builtin(subgroup_size) ssz: u32,
        @builtin(subgroup_invocation_id) sid: u32) {
    o[tid] = f32(ssz) * 1000.0 + f32(sid);
}
"#;
            let sc2 = c.device.push_error_scope(wgpu::ErrorFilter::Validation);
            let pm = c.device.create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("sg-probe"),
                source: wgpu::ShaderSource::Wgsl(probe.into()),
            });
            if let Some(e) = pollster::block_on(sc2.pop()) {
                println!(
                    "subgroup probe rejected: {}",
                    format!("{e}")
                        .lines()
                        .take(2)
                        .collect::<Vec<_>>()
                        .join(" | ")
                );
            } else {
                let pp = c
                    .device
                    .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                        label: Some("sg-probe"),
                        layout: None,
                        module: &pm,
                        entry_point: Some("main"),
                        compilation_options: Default::default(),
                        cache: c.pipeline_cache.as_ref(),
                    });
                let ob2 = c.device.create_buffer(&wgpu::BufferDescriptor {
                    label: None,
                    size: 512,
                    usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                    mapped_at_creation: false,
                });
                let st2 = c.device.create_buffer(&wgpu::BufferDescriptor {
                    label: None,
                    size: 512,
                    usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                });
                let bg2 = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &pp.get_bind_group_layout(0),
                    entries: &[bind_buf(0, &ob2)],
                });
                let mut e2 = c
                    .device
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
                {
                    let mut ps = begin_pass(&mut e2);
                    ps.set_pipeline(&pp);
                    ps.set_bind_group(0, &bg2, &[]);
                    ps.dispatch_workgroups(1, 1, 1);
                }
                let mut sg = vec![0f32; 128];
                if readback(c, e2, &ob2, &st2, 512, &mut sg) {
                    let size = (sg[0] / 1000.0) as u32;
                    let linear = (0..128).all(|i| (sg[i] as u32 % 1000) == (i as u32 % size));
                    println!("subgroup size {size}, tid/{size} is the subgroup index: {linear}");
                }
            }
        }
        let a: Vec<f32> = (0..256)
            .map(|i| (i / 16) as f32 + (i % 16) as f32 / 16.0)
            .collect();
        let b: Vec<f32> = (0..256)
            .map(|i| ((i / 16) * 16 + (i % 16)) as f32 / 64.0)
            .collect();
        let mk = |d: &[f32], usage: wgpu::BufferUsages| {
            use wgpu::util::DeviceExt;
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytemuck::cast_slice(d),
                    usage,
                })
        };
        let ab = mk(&a, wgpu::BufferUsages::STORAGE);
        let bb = mk(&b, wgpu::BufferUsages::STORAGE);
        let ob = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: 1536,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: 1536,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &pipe.get_bind_group_layout(0),
            entries: &[bind_buf(0, &ab), bind_buf(1, &bb), bind_buf(2, &ob)],
        });
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(&pipe);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(1, 1, 1);
        }
        let mut all = vec![0f32; 384];
        assert!(readback(c, enc, &ob, &stage, 1536, &mut all), "readback");
        let ids: Vec<u32> = all[256..384].iter().map(|v| *v as u32).collect();
        let distinct: std::collections::BTreeSet<u32> = ids.iter().copied().collect();
        println!(
            "subgroup_id over 128 lanes: {distinct:?}, lane0..3 {:?}, lane32 {}",
            &ids[..4],
            ids[32]
        );
        let got: Vec<f32> = all[..256].to_vec();
        // Row-major on both sides: c[i][j] = sum_k a[i][k] * b[k][j].
        let mut want = vec![0f32; 256];
        for i in 0..16 {
            for j in 0..16 {
                want[i * 16 + j] = (0..16).map(|k| a[i * 16 + k] * b[k * 16 + j]).sum();
            }
        }
        let worst = got
            .iter()
            .zip(&want)
            .map(|(g, w)| (g - w).abs() / w.abs().max(1.0))
            .fold(0f32, f32::max);
        // The kernel now accumulates the product TWICE, so the answer is
        // 2x the single product; anything else means the accumulator did
        // not survive the loop.
        println!(
            "loop-carried accumulator: got[0]={} want 2x={}  ratio {:.4}",
            got[0],
            2.0 * want[0],
            got[0] / (2.0 * want[0])
        );
        println!(
            "row-major both: worst {worst:.3e}   got[0]={} want[0]={}",
            got[0], want[0]
        );
        // And what it would be if B were read column-major.
        let mut wantt = vec![0f32; 256];
        for i in 0..16 {
            for j in 0..16 {
                wantt[i * 16 + j] = (0..16).map(|k| a[i * 16 + k] * b[j * 16 + k]).sum();
            }
        }
        let worstt = got
            .iter()
            .zip(&wantt)
            .map(|(g, w)| (g - w).abs() / w.abs().max(1.0))
            .fold(0f32, f32::max);
        println!("B column-major:  worst {worstt:.3e}  wantT[0]={}", wantt[0]);
    }

    /// Does wgpu accept the shape the hardware actually implements —
    /// f16 16x16 with an f32 accumulator? Its docs say 8x8 f32 only, but
    /// naga's type generator knows `coop_mat16x16`, and a doc sentence is
    /// not a test. If this compiles, the native Vulkan lane is unnecessary.
    #[test]
    #[ignore]
    fn wgpu_coop_f16_16x16_probe() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        unsafe { std::env::set_var("CMF_COOP", "1") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu device — skipping");
            return;
        };
        for (label, src) in [
            (
                "16x16 f16 in, f32 out",
                r#"
enable wgpu_cooperative_matrix;
enable f16;
@group(0) @binding(0) var<storage, read_write> y: array<f32>;
var<workgroup> sa: array<f16, 256>;
var<workgroup> sb: array<f16, 256>;
var<workgroup> sc: array<f32, 256>;
@compute @workgroup_size(32)
fn main() {
    let a = coopLoad<coop_mat16x16<f16, A>>(&sa[0], 16u);
    let b = coopLoad<coop_mat16x16<f16, B>>(&sb[0], 16u);
    var acc: coop_mat16x16<f32, C>;
    acc = coopMultiplyAdd(a, b, acc);
    coopStore(acc, &sc[0], 16u);
    y[0] = sc[0];
}
"#,
            ),
            (
                "8x8 f32, what the docs promise",
                r#"
enable wgpu_cooperative_matrix;
@group(0) @binding(0) var<storage, read_write> y: array<f32>;
var<workgroup> sa: array<f32, 64>;
var<workgroup> sb: array<f32, 64>;
var<workgroup> sc: array<f32, 64>;
@compute @workgroup_size(32)
fn main() {
    let a = coopLoad<coop_mat8x8<f32, A>>(&sa[0], 8u);
    let b = coopLoad<coop_mat8x8<f32, B>>(&sb[0], 8u);
    var acc: coop_mat8x8<f32, C>;
    acc = coopMultiplyAdd(a, b, acc);
    coopStore(acc, &sc[0], 8u);
    y[0] = sc[0];
}
"#,
            ),
        ] {
            let scope = c.device.push_error_scope(wgpu::ErrorFilter::Validation);
            let _m = c.device.create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some(label),
                source: wgpu::ShaderSource::Wgsl(src.into()),
            });
            match pollster::block_on(scope.pop()) {
                None => println!("{label}: ACCEPTED"),
                Some(e) => {
                    let t = format!("{e}");
                    println!(
                        "{label}: rejected — {}",
                        t.lines().take(8).collect::<Vec<_>>().join(" | ")
                    );
                }
            }
        }
    }

    /// What the selected adapter offers for matrix math: cooperative
    /// (tensor-core) matrices and f16. Run:
    /// `cargo test -p cortiq-engine --release --features gpu
    ///  wgpu_matrix_features -- --ignored --nocapture`
    #[test]
    #[ignore]
    fn wgpu_matrix_features() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let inst = wgpu::Instance::new(wgpu::InstanceDescriptor {
            backends: wgpu::Backends::all(),
            flags: wgpu::InstanceFlags::default(),
            memory_budget_thresholds: Default::default(),
            backend_options: Default::default(),
            display: None,
        });
        for a in pollster::block_on(inst.enumerate_adapters(wgpu::Backends::all())) {
            let info = a.get_info();
            let f = a.features();
            let l = a.limits();
            println!(
                "{:?} {} | coop_matrix {} | f16 {} | subgroup {} | wg storage {} B",
                info.backend,
                info.name,
                f.contains(wgpu::Features::EXPERIMENTAL_COOPERATIVE_MATRIX),
                f.contains(wgpu::Features::SHADER_F16),
                f.contains(wgpu::Features::SUBGROUP),
                l.max_compute_workgroup_storage_size,
            );
        }
    }

    /// The DiT's own GEMM shape, on the device, with nothing else in the
    /// submission: 9216x2304 weights against 2085 tokens, the Lumina FFN
    /// at 512x512. Iterating on the kernel through a full render measures
    /// twenty other things; this measures one.
    ///
    /// `cargo test -p cortiq-engine --release --features gpu
    ///  wgpu_q4tp_mm_throughput -- --ignored --nocapture`
    #[test]
    #[ignore]
    fn wgpu_q4tp_mm_throughput() {
        use std::time::Instant;
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let rows: usize = env_usize("CMF_MM_ROWS", 9216);
        let cols: usize = env_usize("CMF_MM_COLS", 2304);
        let n: usize = env_usize("CMF_MM_N", 2085);
        let gpr = cols / 32;
        let total =
            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
                .unwrap();
        let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
        let mut wb: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
        let lo = cortiq_core::quant::f32_to_f16(-4.0);
        let step = cortiq_core::quant::f32_to_f16(0.1);
        for r in 0..rows {
            let o = params_off + r * 4;
            wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
            wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
        }
        let xs: Vec<f32> = (0..n * cols)
            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
            .collect();
        let mk = |bytes: &[u8]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytes,
                    usage: wgpu::BufferUsages::STORAGE,
                })
        };
        let wbuf = mk(&wb);
        let xbuf = mk(bytemuck::cast_slice(&xs));
        let ybuf = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (n * rows * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE,
            mapped_at_creation: false,
        });
        let _ = gpr;
        let reps = env_usize("CMF_MM_REPS", 20);
        let run = || {
            let mut enc = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
            for _ in 0..reps {
                encode_q4_tile_mm(
                    c,
                    &mut enc,
                    mm_pipeline(c, true, false),
                    &wbuf,
                    &xbuf,
                    &ybuf,
                    rows,
                    cols,
                    n,
                );
            }
            submit(c, finish_enc(enc));
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        };
        run();
        let mut best = f64::MAX;
        for _ in 0..3 {
            let t = Instant::now();
            run();
            best = best.min(t.elapsed().as_secs_f64());
        }
        let flops = 2.0 * rows as f64 * cols as f64 * n as f64 * reps as f64;
        println!(
            "wgpu q4tp mm {rows}x{cols} n={n}: {:.2} ms/call  {:.0} GFLOP/s",
            best * 1e3 / reps as f64,
            flops / best / 1e9
        );
    }

    /// The two-weight 16w kernel must equal two single 16w dispatches BIT
    /// FOR BIT — the fusion moves a dispatch boundary, not a rounding.
    /// Uneven row counts on purpose (a partial last block on each side).
    #[test]
    fn wgpu_q4tp_matvec16w_x2_matches_two_singles() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let cols = 4096usize; // gpr 128 > 64: the 16w regime
        let mk_w = |rows: usize, seed: usize| -> Vec<u8> {
            let total = cortiq_core::quant::expected_nbytes(
                cortiq_core::TensorDtype::Q4TiledP,
                &[rows, cols],
            )
            .unwrap();
            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
            let mut wb: Vec<u8> = (0..total)
                .map(|i| ((i * 37 + seed * 11) % 251) as u8)
                .collect();
            let lo = cortiq_core::quant::f32_to_f16(-4.0);
            let step = cortiq_core::quant::f32_to_f16(0.1);
            for r in 0..rows {
                let o = params_off + r * 4;
                wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
                wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
            }
            wb
        };
        let (rows_a, rows_b) = (1000usize, 296usize);
        let (wa, wb) = (mk_w(rows_a, 1), mk_w(rows_b, 2));
        let xs: Vec<f32> = (0..cols).map(|i| ((i % 97) as f32 - 48.0) / 48.0).collect();
        let mk = |bytes: &[u8]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytes,
                    usage: wgpu::BufferUsages::STORAGE,
                })
        };
        let (wab, wbb, xb) = (mk(&wa), mk(&wb), mk(bytemuck::cast_slice(&xs)));
        let out = |n: usize| {
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (n * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            })
        };
        let (ya1, yb1, ya2, yb2) = (out(rows_a), out(rows_b), out(rows_a), out(rows_b));
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        // two singles (the 16w path: batch 1, gpr > 64)
        encode_q4tp_mv4(c, &mut enc, &wab, &xb, &ya1, rows_a, cols);
        encode_q4tp_mv4(c, &mut enc, &wbb, &xb, &yb1, rows_b, cols);
        // the pair
        let (bind, wg) = mv_x2_bind(c, &wab, &wbb, &xb, &ya2, &yb2, rows_a, rows_b, cols);
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(&c.q4tp_mv16w_x2);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(wg, 1, 1);
        }
        let n = rows_a + rows_b;
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (2 * n * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&ya1, 0, &stage, 0, (rows_a * 4) as u64);
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&yb1, 0, &stage, (rows_a * 4) as u64, (rows_b * 4) as u64);
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&ya2, 0, &stage, (n * 4) as u64, (rows_a * 4) as u64);
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(
            &yb2,
            0,
            &stage,
            ((n + rows_a) * 4) as u64,
            (rows_b * 4) as u64,
        );
        submit(c, finish_enc(enc));
        let slice = stage.slice(..);
        slice.map_async(wgpu::MapMode::Read, |_| {});
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        let data = slice.get_mapped_range().expect("map");
        let all: &[f32] = bytemuck::cast_slice(&data);
        let (single, pair) = all.split_at(n);
        let mism = single
            .iter()
            .zip(pair)
            .filter(|(a, b)| a.to_bits() != b.to_bits())
            .count();
        let nz = single.iter().filter(|v| **v != 0.0).count();
        drop(data);
        stage.unmap();
        assert!(
            nz > n / 2,
            "the singles produced mostly zeros — the harness is wrong"
        );
        assert_eq!(
            mism, 0,
            "{mism} of {n} outputs differ between the x2 kernel and two singles"
        );
    }

    /// The 2-bit decode matvec against the CPU dequant of the same
    /// payload (random chunks, per-row ladders, random 5-bit rungs
    /// including rung 0 = exact zero).
    #[test]
    fn wgpu_q2tp_matvec16w_matches_cpu_dequant() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let sg_requested = q2tp_sg_env_admitted();
        eprintln!(
            "q2tp matvec pipeline: subgroup_requested={sg_requested} admitted={} diagnostics={}",
            c.q2tp_mv16w_sg.is_some(),
            q2tp_sg_diag()
        );
        if sg_requested {
            assert!(
                c.q2tp_mv16w_sg.is_some(),
                "requested q2tp subgroup optimization was not admitted: {}",
                q2tp_sg_diag()
            );
        }
        let _drain = TestGpuDrain::new(c);
        let (rows, cols) = (300usize, 4096usize);
        let gpr = cols / 32;
        let total =
            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q2TiledP, &[rows, cols])
                .unwrap();
        let (params_off, codes_off, stride) = cortiq_core::quant::q2tp_sections(rows, cols);
        let mut wb: Vec<u8> = (0..total).map(|i| ((i * 37 + 11) % 251) as u8).collect();
        let lo = cortiq_core::quant::f32_to_f16(-4.0);
        let step = cortiq_core::quant::f32_to_f16(0.1);
        for r in 0..rows {
            let o = params_off + r * 4;
            wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
            wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
            let crow = &mut wb[codes_off + r * stride..codes_off + (r + 1) * stride];
            crow.fill(0);
            for g in 0..gpr {
                cortiq_core::quant::q4tp_put_code(crow, g, (g * 7 + r) % 32);
            }
        }
        let mut w = vec![0f32; rows * cols];
        cortiq_core::quant::dequant_q2tp(&wb, rows, cols, &mut w);
        let xs: Vec<f32> = (0..cols).map(|i| ((i % 97) as f32 - 48.0) / 48.0).collect();
        let want: Vec<f32> = (0..rows)
            .map(|r| {
                (0..cols)
                    .map(|i| w[r * cols + i] as f64 * xs[i] as f64)
                    .sum::<f64>() as f32
            })
            .collect();
        let mk = |bytes: &[u8]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytes,
                    usage: wgpu::BufferUsages::STORAGE,
                })
        };
        let wbuf = mk(&wb);
        let xbuf = mk(bytemuck::cast_slice(&xs));
        let ybuf = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (rows * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (rows * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        encode_q2tp_mv16w(c, &mut enc, &wbuf, &xbuf, &ybuf, rows, cols, false);
        let mut got = vec![0f32; rows];
        assert!(readback(c, enc, &ybuf, &stage, (rows * 4) as u64, &mut got));
        let (mut num, mut den) = (0f64, 0f64);
        for (a, b) in got.iter().zip(&want) {
            num += ((a - b) as f64).powi(2);
            den += (*b as f64).powi(2);
        }
        let rel = (num / den.max(1e-30)).sqrt();
        eprintln!("q2tp matvec vs cpu dequant: rel rms {rel:.2e}");
        assert!(
            rel < 1e-5,
            "q2tp kernel drifted from the CPU dequant: {rel:.2e}"
        );

        // The same physical payload must also match the explicit affine
        // center `(c-1)·s`; this catches a shader that silently applies the
        // ordinary 1.5 center to a q2tp_affine descriptor.
        let params = &wb[params_off..params_off + rows * 4];
        let want_aff: Vec<f32> = (0..rows)
            .map(|r| {
                let tab = cortiq_core::quant::q2tp_ladder(params, r);
                let codes = &wb[codes_off + r * stride..codes_off + (r + 1) * stride];
                let mut acc = 0f64;
                for g in 0..gpr {
                    let s = tab[cortiq_core::quant::q4tp_code(codes, g)] as f64;
                    let chunk = &wb[(r * gpr + g) * 8..(r * gpr + g + 1) * 8];
                    for (k, &byte) in chunk.iter().enumerate() {
                        for j in 0..4 {
                            let c = ((byte >> (2 * j)) & 3) as f64;
                            let i = g * 32 + k * 4 + j;
                            acc += (c - 1.0) * s * xs[i] as f64;
                        }
                    }
                }
                acc as f32
            })
            .collect();
        let ya = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (rows * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let stage_a = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (rows * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let mut enc_a = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        encode_q2tp_mv16w(c, &mut enc_a, &wbuf, &xbuf, &ya, rows, cols, true);
        let mut got_aff = vec![0f32; rows];
        assert!(readback(
            c,
            enc_a,
            &ya,
            &stage_a,
            (rows * 4) as u64,
            &mut got_aff,
        ));
        let (mut num_aff, mut den_aff) = (0f64, 0f64);
        for (a, b) in got_aff.iter().zip(&want_aff) {
            num_aff += ((a - b) as f64).powi(2);
            den_aff += (*b as f64).powi(2);
        }
        let rel_aff = (num_aff / den_aff.max(1e-30)).sqrt();
        eprintln!("q2tp affine matvec vs cpu dequant: rel rms {rel_aff:.2e}");
        assert!(
            rel_aff < 1e-5,
            "q2tp affine kernel drifted from the CPU affine decode: {rel_aff:.2e}"
        );
    }

    /// The original-engine transfer's first decode gate: the dedicated NB=1
    /// Q8/DP4A affine kernel must agree with an independent CPU Q8 oracle. The
    /// scalar affine q2tp result is also reported, but is intentionally not
    /// used as the Q8 acceptance target because activation quantization is the
    /// explicitly measured approximation in this experiment.
    #[test]
    fn wgpu_q2tp_affine_dp4a_matches_independent_q8_oracle() {
        unsafe {
            std::env::set_var("CMF_GPU", "wgpu");
            std::env::set_var("CMF_Q2_DP4A", "1");
        }
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let _drain = TestGpuDrain::new(c);
        let (rows, cols) = (129usize, 4096usize);
        let gpr = cols / cortiq_core::quant::GROUP_SIZE;
        let total = cortiq_core::quant::expected_nbytes(
            cortiq_core::TensorDtype::Q2TiledP,
            &[rows, cols],
        )
        .unwrap();
        let (params_off, codes_off, stride) = cortiq_core::quant::q2tp_sections(rows, cols);
        let mut wb = vec![0u8; total];
        let lo = cortiq_core::quant::f32_to_f16(-3.0);
        let step = cortiq_core::quant::f32_to_f16(0.125);
        for r in 0..rows {
            let p = params_off + r * 4;
            wb[p..p + 2].copy_from_slice(&lo.to_le_bytes());
            wb[p + 2..p + 4].copy_from_slice(&step.to_le_bytes());
            let codes = &mut wb[codes_off + r * stride..codes_off + (r + 1) * stride];
            codes.fill(0);
            for g in 0..gpr {
                cortiq_core::quant::q4tp_put_code(codes, g, 1 + ((r * 11 + g * 7) % 30));
            }
            for g in 0..gpr {
                let chunk = &mut wb[(r * gpr + g) * 8..(r * gpr + g + 1) * 8];
                for (k, byte) in chunk.iter_mut().enumerate() {
                    let base = r.wrapping_mul(13).wrapping_add(g * 5).wrapping_add(k * 3);
                    *byte = (0..4)
                        .map(|j| ((base + j * 2) % 3) as u8)
                        .enumerate()
                        .fold(0u8, |acc, (j, v)| acc | (v << (2 * j)));
                }
            }
        }
        let xs: Vec<f32> = (0..cols)
            .map(|i| ((i.wrapping_mul(7919) % 4093) as f32 - 2046.0) / 513.0)
            .collect();
        let params = &wb[params_off..params_off + rows * 4];
        let mut want_q8 = vec![0f32; rows];
        let mut want_f32 = vec![0f32; rows];
        for r in 0..rows {
            let tab = cortiq_core::quant::q2tp_ladder(params, r);
            let codes = &wb[codes_off + r * stride..codes_off + (r + 1) * stride];
            // The test's x is one activation vector, so each 32-group uses
            // its own scale, exactly as x_quant_i8 does on the device.
            let mut q8 = 0f64;
            let mut f32v = 0f64;
            for g in 0..gpr {
                let base = g * 32;
                let mut gmax = 0f32;
                for &v in &xs[base..base + 32] {
                    gmax = gmax.max(v.abs());
                }
                let sx = if gmax == 0.0 { 1.0 } else { gmax / 127.0 };
                let code = cortiq_core::quant::q4tp_code(codes, g);
                let scale = tab[code];
                let chunk = &wb[(r * gpr + g) * 8..(r * gpr + g + 1) * 8];
                for (k, &byte) in chunk.iter().enumerate() {
                    for j in 0..4 {
                        let i = base + k * 4 + j;
                        let sym = ((byte >> (2 * j)) & 3) as f32 - 1.0;
                        let q = (xs[i] / sx).round().clamp(-127.0, 127.0);
                        q8 += q as f64 * sx as f64 * sym as f64 * scale as f64;
                        f32v += xs[i] as f64 * sym as f64 * scale as f64;
                    }
                }
            }
            want_q8[r] = q8 as f32;
            want_f32[r] = f32v as f32;
        }
        let mk = |bytes: &[u8]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some("q2-dp4a-component"),
                    contents: bytes,
                    usage: wgpu::BufferUsages::STORAGE,
                })
        };
        let wbuf = mk(&wb);
        let xbuf = mk(bytemuck::cast_slice(&xs));
        let y_dp = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("q2-dp4a-y"),
            size: (rows * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("q2-dp4a-stage"),
            size: (rows * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("q2-dp4a-component"),
            });
        encode_q2tp_mv1_i8(c, &mut enc, &wbuf, &xbuf, &y_dp, rows, cols);
        let mut got = vec![0f32; rows];
        assert!(readback(
            c,
            enc,
            &y_dp,
            &stage,
            (rows * 4) as u64,
            &mut got,
        ));
        let mut num = 0f64;
        let mut den = 0f64;
        let mut num_f32 = 0f64;
        let mut den_f32 = 0f64;
        for ((&a, &b), &f) in got.iter().zip(&want_q8).zip(&want_f32) {
            assert!(a.is_finite(), "q2 dp4a produced non-finite output");
            num += (a as f64 - b as f64).powi(2);
            den += (b as f64).powi(2);
            num_f32 += (a as f64 - f as f64).powi(2);
            den_f32 += (f as f64).powi(2);
        }
        let rel_q8 = (num / den.max(1e-30)).sqrt();
        let rel_f32 = (num_f32 / den_f32.max(1e-30)).sqrt();
        eprintln!(
            "q2 affine dp4a NB1: q8_oracle_rel={rel_q8:.3e} f32_route_rel={rel_f32:.3e}"
        );
        assert!(rel_q8 <= 1e-3, "q2 DP4A drifted from independent Q8 oracle: {rel_q8:.3e}");
    }

    /// The optional Q2TP affine cooperative GEMM is a declared F16 operand
    /// boundary, not a retagged Q4 path.  Run with
    /// `CMF_Q2_COOP=1 CMF_COOP=1` and this ignored test to require admission,
    /// compare every output against the scalar F32 decoder, and check a small
    /// prefix against an independent F16-rounded-weight reference.  The
    /// default matrix exercises real Prism FFN/down widths plus a row tail;
    /// `CMF_Q2_COOP_HEAD=1` adds the 248320-row vocabulary head.
    #[test]
    #[ignore]
    fn wgpu_q2tp_mm_coop_affine_component_matrix() {
        unsafe {
            std::env::set_var("CMF_GPU", "wgpu");
            std::env::set_var("CMF_COOP", "1");
            std::env::set_var("CMF_Q2_COOP", "1");
        }
        let Some(c) = ctx() else {
            panic!("Q2 cooperative component requires a wgpu adapter");
        };
        let _drain = TestGpuDrain::new(c);
        assert!(
            c.q2tp_mm_coop.is_some(),
            "CMF_Q2_COOP=1 did not admit the isolated q2tp cooperative pipeline"
        );
        let mut shapes = vec![
            (17_408usize, 5_120usize, 16usize, "gate-up-k16"),
            (17_408, 5_120, 32, "gate-up-k32"),
            (5_120, 17_408, 32, "down-k32"),
            (513, 5_120, 32, "row-tail-k32"),
            (5_120, 6_144, 32, "out-k32"),
        ];
        if std::env::var("CMF_Q2_COOP_HEAD").as_deref() == Ok("1") {
            shapes.push((248_320, 5_120, 16, "lm-head-k16"));
        }
        for (rows, cols, batch, label) in shapes {
            let total = cortiq_core::quant::expected_nbytes(
                cortiq_core::TensorDtype::Q2TiledP,
                &[rows, cols],
            )
            .unwrap();
            let (params_off, codes_off, stride) = cortiq_core::quant::q2tp_sections(rows, cols);
            let gpr = cols / cortiq_core::quant::GROUP_SIZE;
            let mut wb_host = vec![0u8; total];
            let lo = cortiq_core::quant::f32_to_f16(-4.0);
            let step = cortiq_core::quant::f32_to_f16(0.1);
            for r in 0..rows {
                let p = params_off + r * 4;
                wb_host[p..p + 2].copy_from_slice(&lo.to_le_bytes());
                wb_host[p + 2..p + 4].copy_from_slice(&step.to_le_bytes());
                let codes = &mut wb_host[codes_off + r * stride..codes_off + (r + 1) * stride];
                codes.fill(0);
                // Affine Prism symbols are ternary 0/1/2; keep the reserved
                // two-bit symbol 3 out of this candidate's payload.
                for g in 0..gpr {
                    cortiq_core::quant::q4tp_put_code(codes, g, (r * 13 + g * 7) % 31);
                }
                let chunk = &mut wb_host[(r * gpr) * 8..(r * gpr + gpr) * 8];
                for (i, byte) in chunk.iter_mut().enumerate() {
                    let base = r.wrapping_mul(17).wrapping_add(i * 5);
                    *byte = (0..4)
                        .map(|j| ((base + j * 3) % 3) as u8)
                        .enumerate()
                        .fold(0u8, |acc, (j, v)| acc | (v << (2 * j)));
                }
            }
            let xs: Vec<f32> = (0..batch * cols)
                .map(|i| ((i.wrapping_mul(7919) % 4093) as f32 - 2046.0) / 257.0)
                .collect();
            let mk = |bytes: &[u8]| {
                c.device
                    .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: Some("q2tp-coop-component"),
                        contents: bytes,
                        usage: wgpu::BufferUsages::STORAGE,
                    })
            };
            let wb = mk(&wb_host);
            let xb = mk(bytemuck::cast_slice(&xs));
            let out_bytes = (batch * rows * 4) as u64;
            let y_scalar = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("q2tp-coop-scalar"),
                size: out_bytes,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });
            let y_coop = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("q2tp-coop-output"),
                size: out_bytes,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });
            let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("q2tp-coop-stage"),
                size: out_bytes * 2,
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            let mut enc = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("q2tp-coop-component"),
                });
            encode_q2_tile_mm(
                c, &mut enc, &wb, &xb, &y_scalar, rows, cols, batch, true, false,
            );
            encode_q2_tile_mm(
                c, &mut enc, &wb, &xb, &y_coop, rows, cols, batch, true, true,
            );
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(&y_scalar, 0, &stage, 0, out_bytes);
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(&y_coop, 0, &stage, out_bytes, out_bytes);
            submit(c, finish_enc(enc));
            let slice = stage.slice(..);
            slice.map_async(wgpu::MapMode::Read, |_| {});
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            let data = slice.get_mapped_range().expect("q2 coop map");
            let all: &[f32] = bytemuck::cast_slice(&data);
            let (scalar, coop) = all.split_at(batch * rows);
            let mut num = 0f64;
            let mut den = 0f64;
            let mut max_abs = 0f32;
            let mut max_ref = 0f32;
            for (&a, &b) in scalar.iter().zip(coop) {
                assert!(a.is_finite() && b.is_finite(), "{label}: non-finite output");
                let d = a as f64 - b as f64;
                num += d * d;
                den += (a as f64) * (a as f64);
                max_abs = max_abs.max(d.abs() as f32);
                max_ref = max_ref.max(a.abs());
            }
            let rel = (num / den.max(1e-30)).sqrt();
            let nlinf = max_abs / max_ref.max(1e-30);
            eprintln!(
                "q2tp coop {label} rows={rows} cols={cols} k={batch}: rel_l2={rel:.3e} normalized_linf={nlinf:.3e}"
            );
            assert!(rel <= 1e-3 && nlinf <= 1e-3, "{label}: scalar/coop drift exceeds 1e-3");
            // A bounded independent oracle: enough full columns to exercise
            // every scale/code pattern while keeping CPU work finite for the
            // vocabulary-head option.
            let params = &wb_host[..];
            let (params_off, codes_off, stride) =
                cortiq_core::quant::q2tp_sections(rows, cols);
            let sample_rows = rows.min(8);
            let mut sample_num = 0f64;
            let mut sample_den = 0f64;
            for r in 0..sample_rows {
                let tab = cortiq_core::quant::q2tp_ladder(&params[params_off..], r);
                let codes = &params[codes_off + r * stride..codes_off + (r + 1) * stride];
                for t in 0..batch {
                    let mut want = 0f64;
                    for g in 0..gpr {
                        let scale = tab[cortiq_core::quant::q4tp_code(codes, g)];
                        let chunk = &params[(r * gpr + g) * 8..(r * gpr + g + 1) * 8];
                        for (i, &byte) in chunk.iter().enumerate() {
                            for j in 0..4 {
                                let sym = ((byte >> (2 * j)) & 3) as f32;
                                let w = (sym - 1.0) * scale;
                                let wf = cortiq_core::quant::f16_to_f32(
                                    cortiq_core::hadamard::prism_f32_to_f16_rne(w),
                                );
                                let xf = cortiq_core::quant::f16_to_f32(
                                    cortiq_core::hadamard::prism_f32_to_f16_rne(xs[t * cols + g * 32 + i * 4 + j]),
                                );
                                want += wf as f64 * xf as f64;
                            }
                        }
                    }
                    let got = coop[t * rows + r] as f64;
                    let d = got - want;
                    sample_num += d * d;
                    sample_den += want * want;
                }
            }
            let sample_rel = (sample_num / sample_den.max(1e-30)).sqrt();
            eprintln!("q2tp coop {label}: independent_f16_sample_rel={sample_rel:.3e}");
            assert!(sample_rel <= 1e-3, "{label}: independent F16 sample drift {sample_rel:.3e}");
            drop(data);
            stage.unmap();
        }
    }

    /// The resident graph's activation boundary must be the same signed,
    /// normalized FWHT as the CPU Prism oracle, including the source f16
    /// round.  Keep this as a device test: a visually plausible generation
    /// can survive a wrong butterfly/sign placement for several tokens.
    #[test]
    fn wgpu_prism_fwht_matches_cpu_oracle() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let _drain = TestGpuDrain::new(c);
        let Some(pipe) = c.fwht.as_ref() else {
            eprintln!("shader-f16 unavailable — skipping Prism FWHT test");
            return;
        };
        let width = 1024usize;
        let signs: Vec<f32> = (0..width)
            .map(|i| if (i * 17 + i / 7) & 1 == 0 { 1.0 } else { -1.0 })
            .collect();
        let xs: Vec<f32> = (0..width)
            .map(|i| ((i * 7919 % 1000) as f32 - 500.0) / 137.0)
            .collect();
        let mut want = xs.clone();
        cortiq_core::hadamard::signed_fwht_forward(&mut want, &signs, width).unwrap();
        for v in &mut want {
            *v = cortiq_core::quant::f16_to_f32(cortiq_core::hadamard::prism_f32_to_f16_rne(*v));
        }
        let mk = |bytes: &[u8], usage: wgpu::BufferUsages| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some("prism-fwht-test"),
                    contents: bytes,
                    usage,
                })
        };
        let xb = mk(
            bytemuck::cast_slice(&xs),
            wgpu::BufferUsages::STORAGE,
        );
        let sb = mk(
            bytemuck::cast_slice(&signs),
            wgpu::BufferUsages::STORAGE,
        );
        let yb = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("prism-fwht-test-y"),
            size: (width * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("prism-fwht-test-stage"),
            size: (width * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let p = uniform_u32x8(c, [width as u32, 1024, 0, 0, 1, 0, 0, 0]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("prism-fwht-test-bind"),
            layout: &pipe.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &xb),
                bind_buf(1, &yb),
                bind_buf(2, &sb),
                bind_buf(3, &p),
            ],
        });
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("prism-fwht-test-enc"),
            });
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(pipe);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(1, 1, 1);
        }
        let mut got = vec![0f32; width];
        assert!(readback(
            c,
            enc,
            &yb,
            &stage,
            (width * 4) as u64,
            &mut got,
        ));
        let mut max_abs = 0.0f32;
        let mut rms = 0.0f64;
        for (a, b) in got.iter().zip(&want) {
            max_abs = max_abs.max((a - b).abs());
            rms += (*a as f64 - *b as f64).powi(2);
        }
        let rms = (rms / width as f64).sqrt();
        eprintln!("Prism FWHT GPU/CPU f16 oracle: rms={rms:.3e} max={max_abs:.3e}");
        assert!(rms < 2e-3 && max_abs < 1e-2, "FWHT drift rms={rms:.3e} max={max_abs:.3e}");
    }

    /// Independent full-width FWHT coverage.  The existing smoke above is
    /// retained for the historical gate; this test separates raw f32 error
    /// from the declared f16 boundary and exercises every Prism activation
    /// width used by the model.
    #[test]
    fn wgpu_prism_fwht_matches_independent_f32_and_f16_widths() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let _drain = TestGpuDrain::new(c);
        let Some(pipe) = c.fwht.as_ref() else {
            eprintln!("shader-f16 unavailable — skipping Prism FWHT width test");
            return;
        };
        let run = |width: usize, round16: bool| -> (Vec<f32>, Vec<f32>) {
            let signs: Vec<f32> = (0..width)
                .map(|i| if (i * 17 + i / 7 + width) & 1 == 0 { 1.0 } else { -1.0 })
                .collect();
            let xs: Vec<f32> = (0..width)
                .map(|i| {
                    if width == 1024 && i == 127 {
                        f32::from_bits(1)
                    } else if width == 1024 && i == 255 {
                        f32::from_bits(0x3380_0000)
                    } else if width == 1024 && i == 767 {
                        65504.0
                    } else {
                        ((i * 7919 % 4093) as f32 - 2046.0) / 257.0
                    }
                })
                .collect();
            let mut want = fwht_reference_f64(&xs, &signs, 1024);
            if round16 {
                for v in &mut want {
                    *v = cortiq_core::quant::f16_to_f32(cortiq_core::hadamard::prism_f32_to_f16_rne(*v));
                }
            }
            let mk = |bytes: &[u8], usage: wgpu::BufferUsages| {
                c.device
                    .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: Some("prism-fwht-independent"),
                        contents: bytes,
                        usage,
                    })
            };
            let xb = mk(bytemuck::cast_slice(&xs), wgpu::BufferUsages::STORAGE);
            let sb = mk(bytemuck::cast_slice(&signs), wgpu::BufferUsages::STORAGE);
            let yb = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("prism-fwht-independent-y"),
                size: (width * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });
            let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("prism-fwht-independent-stage"),
                size: (width * 4) as u64,
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            let p = uniform_u32x8(
                c,
                [width as u32, 1024, 0, 0, u32::from(round16), 0, 0, 0],
            );
            let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("prism-fwht-independent-bind"),
                layout: &pipe.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, &xb),
                    bind_buf(1, &yb),
                    bind_buf(2, &sb),
                    bind_buf(3, &p),
                ],
            });
            let mut enc = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("prism-fwht-independent-enc"),
                });
            {
                let mut pass = begin_pass(&mut enc);
                pass.set_pipeline(pipe);
                pass.set_bind_group(0, &bind, &[]);
                pass.dispatch_workgroups((width / 1024) as u32, 1, 1);
            }
            let mut got = vec![0f32; width];
            assert!(readback(
                c,
                enc,
                &yb,
                &stage,
                (width * 4) as u64,
                &mut got,
            ));
            (got, want)
        };
        for &width in &[1024usize, 5120, 6144, 17408] {
            // First establish the raw f32 device result against an
            // independent f64 CPU oracle.  Then use that same device result
            // as the input to the local RNE oracle for the second dispatch:
            // this isolates the declared f16 boundary from valid f32
            // butterfly-association/driver differences, while still checking
            // every output half bit rather than only an aggregate RMS.
            let (got, want) = run(width, false);
            let (mut num, mut den, mut max_abs) = (0.0f64, 0.0f64, 0.0f32);
            for (&a, &b) in got.iter().zip(&want) {
                let d = a as f64 - b as f64;
                num += d * d;
                den += (b as f64) * (b as f64);
                max_abs = max_abs.max(d.abs() as f32);
            }
            let rel = (num / den.max(1e-30)).sqrt();
            eprintln!(
                "Prism FWHT width={width} f32: rel_rms={rel:.3e} max={max_abs:.3e}"
            );
            assert!(rel <= 1e-5 && max_abs <= 2e-3);

            let (got16, source_want16) = run(width, true);
            let want16: Vec<f32> = got
                .iter()
                .map(|&v| {
                    cortiq_core::quant::f16_to_f32(
                        cortiq_core::hadamard::prism_f32_to_f16_rne(v),
                    )
                })
                .collect();
            let mismatches: Vec<(usize, u32, u32)> = got16
                .iter()
                .zip(&want16)
                .enumerate()
                .filter_map(|(i, (&a, &b))| {
                    (a.to_bits() != b.to_bits()).then_some((i, a.to_bits(), b.to_bits()))
                })
                .take(8)
                .collect();
            assert!(
                mismatches.is_empty(),
                "Prism FWHT f16 per-element mismatch width={width}: {mismatches:?}"
            );

            let (mut num, mut den, mut max_abs, mut max_scaled) =
                (0.0f64, 0.0f64, 0.0f32, 0.0f32);
            for (&a, &b) in got16.iter().zip(&source_want16) {
                let d = a as f64 - b as f64;
                num += d * d;
                den += (b as f64) * (b as f64);
                max_abs = max_abs.max(d.abs() as f32);
                // A final RNE f16 cast has an absolute error proportional
                // to the exponent (the 65504 impulse intentionally
                // exercises that boundary), so a fixed absolute bound
                // would reject a correct half-round at large outputs.
                max_scaled = max_scaled.max((d.abs() as f32) / b.abs().max(1.0));
            }
            let rel = (num / den.max(1e-30)).sqrt();
            eprintln!(
                "Prism FWHT width={width} f16: rel_rms={rel:.3e} max={max_abs:.3e} scaled={max_scaled:.3e}"
            );
            assert!(rel <= 3e-3 && max_scaled <= 2e-3);
        }
    }

    /// The q2tp prefill GEMM must carry the same explicit center bit as the
    /// single-token kernel. This is a direct shader parity check; no model
    /// descriptor is needed because the caller supplies the validated bit.
    #[test]
    fn wgpu_q2tp_mm_affine_matches_cpu_decode() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let _drain = TestGpuDrain::new(c);
        let (rows, cols, batch) = (64usize, 512usize, 3usize);
        let gpr = cols / 32;
        let total =
            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q2TiledP, &[rows, cols])
                .unwrap();
        let (params_off, codes_off, stride) = cortiq_core::quant::q2tp_sections(rows, cols);
        let mut wb: Vec<u8> = (0..total).map(|i| ((i * 19 + 7) % 251) as u8).collect();
        let lo = cortiq_core::quant::f32_to_f16(-4.0);
        let step = cortiq_core::quant::f32_to_f16(0.1);
        for r in 0..rows {
            let o = params_off + r * 4;
            wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
            wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
            let crow = &mut wb[codes_off + r * stride..codes_off + (r + 1) * stride];
            crow.fill(0);
            for g in 0..gpr {
                cortiq_core::quant::q4tp_put_code(crow, g, (g * 11 + r * 3) % 32);
            }
        }
        let xs: Vec<f32> = (0..batch * cols)
            .map(|i| ((i * 29 % 127) as f32 - 63.0) / 63.0)
            .collect();
        let mk = |bytes: &[u8]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytes,
                    usage: wgpu::BufferUsages::STORAGE,
                })
        };
        let wbuf = mk(&wb);
        let xbuf = mk(bytemuck::cast_slice(&xs));
        let params = &wb[params_off..params_off + rows * 4];
        let want = |affine: bool| -> Vec<f32> {
            let mut out = Vec::with_capacity(batch * rows);
            for bi in 0..batch {
                for r in 0..rows {
                    let tab = cortiq_core::quant::q2tp_ladder(params, r);
                    let codes = &wb[codes_off + r * stride..codes_off + (r + 1) * stride];
                    let mut acc = 0.0f32;
                    for i in 0..cols {
                        let g = i / 32;
                        let p = i % 32;
                        let byte = wb[(r * gpr + g) * 8 + p / 4];
                        let code = cortiq_core::quant::q4tp_code(codes, g);
                        let s = tab[code];
                        let center = if affine { 1.0 } else { 1.5 };
                        acc += (f32::from((byte >> (2 * (p % 4))) & 3) - center)
                            * s
                            * xs[bi * cols + i];
                    }
                    out.push(acc);
                }
            }
            out
        };
        let run = |affine: bool| -> Vec<f32> {
            let y = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (batch * rows * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });
            let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (batch * rows * 4) as u64,
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            let p = uniform_u32x4(
                c,
                [(cols / 4) as u32, rows as u32, batch as u32, affine as u32],
            );
            let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("q2tp-mm-test"),
                layout: &c.q2tp_mm.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, &wbuf),
                    bind_buf(1, &xbuf),
                    bind_buf(2, &y),
                    bind_buf(3, &p),
                ],
            });
            let mut enc = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
            {
                let mut pass = begin_pass(&mut enc);
                pass.set_pipeline(&c.q2tp_mm);
                pass.set_bind_group(0, &bind, &[]);
                pass.dispatch_workgroups(
                    (rows as u32).div_ceil(64),
                    (batch as u32).div_ceil(64),
                    1,
                );
            }
            let mut out = vec![0f32; batch * rows];
            assert!(readback(
                c,
                enc,
                &y,
                &stage,
                (batch * rows * 4) as u64,
                &mut out,
            ));
            out
        };
        for affine in [false, true] {
            let got = run(affine);
            let expected = want(affine);
            let (mut num, mut den) = (0f64, 0f64);
            for (a, b) in got.iter().zip(&expected) {
                num += ((*a - *b) as f64).powi(2);
                den += (*b as f64).powi(2);
            }
            let rel = (num / den.max(1e-30)).sqrt();
            eprintln!(
                "q2tp {} mm vs cpu dequant: rel rms {rel:.2e}",
                if affine { "affine" } else { "raw" }
            );
            assert!(rel < 1e-5, "q2tp mm drifted from CPU decode: {rel:.2e}");
        }
    }

    /// The int8-activation batched kernel (dp4a) against the f32 singles:
    /// not bit-exact by design — the activations carry a per-32 int8
    /// rounding — but the drift must stay in the Q8_1 band (rel rms
    /// well under 1e-2 on these random rows; llama.cpp's activation grid).
    #[test]
    fn wgpu_q4tp_matvec4_bk8_tracks_the_f32_singles() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let (rows, cols, batch) = (1000usize, 4096usize, 5usize);
        let total =
            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
                .unwrap();
        let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
        let mut wb: Vec<u8> = (0..total).map(|i| ((i * 37 + 5) % 251) as u8).collect();
        let lo = cortiq_core::quant::f32_to_f16(-4.0);
        let step = cortiq_core::quant::f32_to_f16(0.1);
        for r in 0..rows {
            let o = params_off + r * 4;
            wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
            wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
        }
        // activations with a spread of magnitudes (an int8 grid is per 32-group)
        let xs: Vec<f32> = (0..batch * cols)
            .map(|i| {
                let t = ((i * 7919) % 1000) as f32 / 1000.0 - 0.5;
                t * (1.0 + ((i / 32) % 5) as f32) + (i / cols) as f32 * 0.01
            })
            .collect();
        let mk = |bytes: &[u8]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytes,
                    usage: wgpu::BufferUsages::STORAGE,
                })
        };
        let wbuf = mk(&wb);
        let xall = mk(bytemuck::cast_slice(&xs));
        let out = |n: usize| {
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (n * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            })
        };
        let yb = out(batch * rows);
        let mut singles = Vec::new();
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        encode_q4tp_mv4_b_i8(c, &mut enc, &wbuf, &xall, &yb, rows, cols, batch);
        for e in 0..batch {
            let xe = mk(bytemuck::cast_slice(&xs[e * cols..(e + 1) * cols]));
            let ye = out(rows);
            encode_q4tp_mv4(c, &mut enc, &wbuf, &xe, &ye, rows, cols);
            singles.push((xe, ye));
        }
        let n = batch * rows;
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (2 * n * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&yb, 0, &stage, 0, (n * 4) as u64);
        for (e, (_, ye)) in singles.iter().enumerate() {
            enc.copy_buffer_to_buffer(
                ye,
                0,
                &stage,
                ((n + e * rows) * 4) as u64,
                (rows * 4) as u64,
            );
        }
        submit(c, finish_enc(enc));
        let slice = stage.slice(..);
        slice.map_async(wgpu::MapMode::Read, |_| {});
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        let data = slice.get_mapped_range().expect("map");
        let all: &[f32] = bytemuck::cast_slice(&data);
        let (b, sgl) = all.split_at(n);
        let nz = sgl.iter().filter(|v| **v != 0.0).count();
        let (mut num, mut den, mut worst) = (0f64, 0f64, 0f64);
        for (x, y) in b.iter().zip(sgl) {
            num += ((x - y) as f64).powi(2);
            den += (*y as f64).powi(2);
            worst = worst.max((x - y).abs() as f64);
        }
        let amax = sgl.iter().fold(0f32, |m, v| m.max(v.abs())) as f64;
        drop(data);
        stage.unmap();
        let rel = (num / den.max(1e-30)).sqrt();
        eprintln!(
            "bk8 (int8 x, dp4a) vs f32 singles: rel rms {rel:.2e}, worst |Δ| {worst:.2e} of {amax:.2e}"
        );
        assert!(nz > n / 2, "singles mostly zero — harness wrong");
        assert!(
            rel < 5e-3,
            "int8-activation kernel drifted from f32: {rel:.2e}"
        );
    }

    /// The batched kernel's rows against the one-vector kernel, element
    /// by element — the verify must land on the plain token's bits.
    #[test]
    fn wgpu_q4tp_matvec4_bku_matches_single_bits() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        if c.use_mv_bk < 2 {
            eprintln!("bku arm off — skipping");
            return;
        }
        let (rows, cols, batch) = (1000usize, 4096usize, 3usize);
        let total =
            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
                .unwrap();
        let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
        let mut wb: Vec<u8> = (0..total).map(|i| ((i * 37 + 5) % 251) as u8).collect();
        let lo = cortiq_core::quant::f32_to_f16(-4.0);
        let step = cortiq_core::quant::f32_to_f16(0.1);
        for r in 0..rows {
            let o = params_off + r * 4;
            wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
            wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
        }
        let xs: Vec<f32> = (0..batch * cols)
            .map(|i| ((i % 97) as f32 - 48.0) / 48.0 + (i / cols) as f32 * 0.01)
            .collect();
        let mk = |bytes: &[u8]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytes,
                    usage: wgpu::BufferUsages::STORAGE,
                })
        };
        let wbuf = mk(&wb);
        let xall = mk(bytemuck::cast_slice(&xs));
        let out = |n: usize| {
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (n * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            })
        };
        let yb = out(batch * rows);
        let mut singles = Vec::new();
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        assert!(encode_q4tp_mv4_b_with(
            c, &mut enc, &wbuf, &xall, &yb, rows, cols, batch, false
        ));
        for e in 0..batch {
            let xe = mk(bytemuck::cast_slice(&xs[e * cols..(e + 1) * cols]));
            let ye = out(rows);
            encode_q4tp_mv4(c, &mut enc, &wbuf, &xe, &ye, rows, cols);
            singles.push((xe, ye));
        }
        let n = batch * rows;
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (2 * n * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        enc.copy_buffer_to_buffer(&yb, 0, &stage, 0, (n * 4) as u64);
        for (e, (_, ye)) in singles.iter().enumerate() {
            enc.copy_buffer_to_buffer(
                ye,
                0,
                &stage,
                ((n + e * rows) * 4) as u64,
                (rows * 4) as u64,
            );
        }
        submit(c, finish_enc(enc));
        let slice = stage.slice(..);
        slice.map_async(wgpu::MapMode::Read, |_| {});
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        let data = slice.get_mapped_range().expect("map");
        let all: &[f32] = bytemuck::cast_slice(&data);
        let (b, sgl) = all.split_at(n);
        let mism = b
            .iter()
            .zip(sgl)
            .filter(|(x, y)| x.to_bits() != y.to_bits())
            .count();
        let nz = sgl.iter().filter(|v| **v != 0.0).count();
        let (mut num, mut den) = (0f64, 0f64);
        for (x, y) in b.iter().zip(sgl) {
            num += ((x - y) as f64).powi(2);
            den += (*y as f64).powi(2);
        }
        drop(data);
        stage.unmap();
        assert!(nz > n / 2, "singles mostly zero — harness wrong");
        // Metal compiles WGSL with fast math (contraction AND
        // reassociation), so the two kernels' bits are not comparable
        // there; on Vulkan they must agree exactly.
        if cfg!(target_os = "macos") {
            let rel = (num / den.max(1e-30)).sqrt();
            eprintln!(
                "bku vs single on Metal: {mism} of {n} bits differ, rel rms {rel:.2e} (fast math)"
            );
            assert!(
                rel < 1e-5,
                "batched drifted from the one-vector kernel: rel rms {rel:.2e}"
            );
        } else {
            assert_eq!(
                mism, 0,
                "{mism} of {n} batched outputs differ from the one-vector kernel"
            );
        }
    }

    /// The batched two-weight kernel against two batched singles (bku
    /// arm), batch 3, uneven rows — bit for bit.
    #[test]
    fn wgpu_q4tp_matvec4_bku_x2_matches_two_singles() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        if c.use_mv_bk < 2 || !c.use_mv_x2 {
            eprintln!("bku arm or x2 off in this environment — skipping");
            return;
        }
        let (cols, batch) = (4096usize, 3usize);
        let mk_w = |rows: usize, seed: usize| -> Vec<u8> {
            let total = cortiq_core::quant::expected_nbytes(
                cortiq_core::TensorDtype::Q4TiledP,
                &[rows, cols],
            )
            .unwrap();
            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
            let mut wb: Vec<u8> = (0..total)
                .map(|i| ((i * 37 + seed * 11) % 251) as u8)
                .collect();
            let lo = cortiq_core::quant::f32_to_f16(-4.0);
            let step = cortiq_core::quant::f32_to_f16(0.1);
            for r in 0..rows {
                let o = params_off + r * 4;
                wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
                wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
            }
            wb
        };
        let (rows_a, rows_b) = (1000usize, 296usize);
        let (wa, wb) = (mk_w(rows_a, 1), mk_w(rows_b, 2));
        let xs: Vec<f32> = (0..batch * cols)
            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
            .collect();
        let mk = |bytes: &[u8]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytes,
                    usage: wgpu::BufferUsages::STORAGE,
                })
        };
        let (wab, wbb, xb) = (mk(&wa), mk(&wb), mk(bytemuck::cast_slice(&xs)));
        let out = |n: usize| {
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (n * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            })
        };
        let (na, nb) = (batch * rows_a, batch * rows_b);
        let (ya1, yb1, ya2, yb2) = (out(na), out(nb), out(na), out(nb));
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        assert!(encode_q4tp_mv4_b_with(
            c, &mut enc, &wab, &xb, &ya1, rows_a, cols, batch, false
        ));
        assert!(encode_q4tp_mv4_b_with(
            c, &mut enc, &wbb, &xb, &yb1, rows_b, cols, batch, false
        ));
        assert!(encode_q4tp_mv4_b_x2(
            c, &mut enc, &wab, &wbb, &xb, &ya2, &yb2, rows_a, rows_b, cols, batch
        ));
        let n = na + nb;
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (2 * n * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&ya1, 0, &stage, 0, (na * 4) as u64);
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&yb1, 0, &stage, (na * 4) as u64, (nb * 4) as u64);
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&ya2, 0, &stage, (n * 4) as u64, (na * 4) as u64);
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&yb2, 0, &stage, ((n + na) * 4) as u64, (nb * 4) as u64);
        submit(c, finish_enc(enc));
        let slice = stage.slice(..);
        slice.map_async(wgpu::MapMode::Read, |_| {});
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        let data = slice.get_mapped_range().expect("map");
        let all: &[f32] = bytemuck::cast_slice(&data);
        let (single, pair) = all.split_at(n);
        let mism = single
            .iter()
            .zip(pair)
            .filter(|(a, b)| a.to_bits() != b.to_bits())
            .count();
        let nz = single.iter().filter(|v| **v != 0.0).count();
        drop(data);
        stage.unmap();
        assert!(nz > n / 2, "singles produced mostly zeros — harness wrong");
        assert_eq!(
            mism, 0,
            "{mism} of {n} batched outputs differ between bku_x2 and two singles"
        );
    }

    /// The fused gate+up+SiLU kernel against gate matvec + up matvec +
    /// `silu_mul_pre` on the device — bit for bit.
    #[test]
    fn wgpu_q4tp_matvec16w_gu_matches_three_dispatches() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let (inter, cols) = (1000usize, 4096usize);
        let mk_w = |seed: usize| -> Vec<u8> {
            let total = cortiq_core::quant::expected_nbytes(
                cortiq_core::TensorDtype::Q4TiledP,
                &[inter, cols],
            )
            .unwrap();
            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(inter, cols);
            let mut wb: Vec<u8> = (0..total)
                .map(|i| ((i * 37 + seed * 11) % 251) as u8)
                .collect();
            let lo = cortiq_core::quant::f32_to_f16(-4.0);
            let step = cortiq_core::quant::f32_to_f16(0.1);
            for r in 0..inter {
                let o = params_off + r * 4;
                wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
                wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
            }
            wb
        };
        let (wg, wu) = (mk_w(1), mk_w(2));
        let xs: Vec<f32> = (0..cols)
            .map(|i| ((i % 97) as f32 - 48.0) / 480.0)
            .collect();
        let mk = |bytes: &[u8]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytes,
                    usage: wgpu::BufferUsages::STORAGE,
                })
        };
        let (gb, ub, xb) = (mk(&wg), mk(&wu), mk(bytemuck::cast_slice(&xs)));
        let out = |n: usize| {
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (n * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            })
        };
        let (gy, uy, act_ref, act) = (out(inter), out(inter), out(inter), out(inter));
        let dummy = out(4);
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        encode_q4tp_mv4(c, &mut enc, &gb, &xb, &gy, inter, cols);
        encode_q4tp_mv4(c, &mut enc, &ub, &xb, &uy, inter, cols);
        let silu_u = uniform_u32x4(c, [inter as u32, 0, 0, 0]);
        let bg_silu = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.layout_silu,
            entries: &[
                bind_buf(0, &gy),
                bind_buf(1, &uy),
                bind_buf(2, &dummy),
                bind_buf(3, &act_ref),
                bind_buf(4, &silu_u),
            ],
        });
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(&c.silu);
            pass.set_bind_group(0, &bg_silu, &[]);
            pass.dispatch_workgroups((inter as u32).div_ceil(256), 1, 1);
        }
        let (bind, wgc) = mv_gu_bind(c, &gb, &ub, &xb, &act, inter, cols);
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(&c.q4tp_mv16w_gu);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(wgc, 1, 1);
        }
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (2 * inter * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&act_ref, 0, &stage, 0, (inter * 4) as u64);
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&act, 0, &stage, (inter * 4) as u64, (inter * 4) as u64);
        submit(c, finish_enc(enc));
        let slice = stage.slice(..);
        slice.map_async(wgpu::MapMode::Read, |_| {});
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        let data = slice.get_mapped_range().expect("map");
        let all: &[f32] = bytemuck::cast_slice(&data);
        let (a, b) = all.split_at(inter);
        let mism = a
            .iter()
            .zip(b)
            .filter(|(x, y)| x.to_bits() != y.to_bits())
            .count();
        let nz = a.iter().filter(|v| **v != 0.0).count();
        drop(data);
        stage.unmap();
        assert!(
            nz > inter / 2,
            "reference activations mostly zero — harness wrong"
        );
        assert_eq!(
            mism, 0,
            "{mism} of {inter} activations differ between the fused kernel and gate/up/silu"
        );
    }

    /// Decode-matvec bandwidth by SHAPE, in one submit: the wide FFN
    /// gate/up shape (17408 x 5120) against the narrow down/out shape
    /// (5120 x 17408, the same bytes) through `encode_q4tp_mv4` — the
    /// exact dispatch the token graph issues. If the narrow shape streams
    /// visibly slower, the card is short of workgroups there (320 blocks
    /// of 16 rows on 170 SMs) and split-K is the medicine; if both sit at
    /// the same fraction of the bus, the kernel is. `cargo test --release
    /// --features gpu -- --ignored wgpu_q4tp_matvec_shape_bandwidth --nocapture`.
    #[test]
    #[ignore]
    fn wgpu_q4tp_matvec_shape_bandwidth() {
        use std::time::Instant;
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let reps = env_usize("CMF_MM_REPS", 40);
        // Per-dispatch floor: a tiny matvec (16 rows) issued 700 times in
        // ONE submit — the launch + drain cost a serial chain pays per
        // dispatch on this device, times the ~700 dispatches of a token.
        {
            let (rows, cols) = (16usize, 128usize);
            let total = cortiq_core::quant::expected_nbytes(
                cortiq_core::TensorDtype::Q4TiledP,
                &[rows, cols],
            )
            .unwrap();
            let wb: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
            let xs: Vec<f32> = (0..cols).map(|i| i as f32 * 0.01).collect();
            let mk = |bytes: &[u8]| {
                c.device
                    .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: None,
                        contents: bytes,
                        usage: wgpu::BufferUsages::STORAGE,
                    })
            };
            let wbuf = mk(&wb);
            let xbuf = mk(bytemuck::cast_slice(&xs));
            let ybuf = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (rows * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE,
                mapped_at_creation: false,
            });
            let n = 700usize;
            let run = || {
                let mut enc = c
                    .device
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
                for _ in 0..n {
                    encode_q4tp_mv4(c, &mut enc, &wbuf, &xbuf, &ybuf, rows, cols);
                }
                let t = Instant::now();
                submit(c, finish_enc(enc));
                let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
                t.elapsed().as_secs_f64()
            };
            run();
            let mut best = f64::MAX;
            for _ in 0..3 {
                best = best.min(run());
            }
            println!(
                "null-dispatch chain: {n} tiny matvecs in one submit = {:.2} ms  ({:.1} us/dispatch)",
                best * 1e3,
                best * 1e6 / n as f64
            );
            // The same 700 dispatches inside ONE compute pass: what a
            // pass boundary costs is the difference.
            let (bind1, _) = {
                let gpr = cols / 32;
                let p_buf = q4tp_mv_params(c, gpr, rows, 1);
                let layout = c.q4tp_mv16.get_bind_group_layout(0);
                (
                    c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &layout,
                        entries: &[
                            bind_buf(0, &wbuf),
                            bind_buf(2, &ybuf),
                            bind_buf(3, &p_buf),
                            bind_buf(4, &wbuf),
                            bind_buf(5, &xbuf),
                        ],
                    }),
                    0,
                )
            };
            let run1 = || {
                let mut enc = c
                    .device
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
                {
                    let mut pass = begin_pass(&mut enc);
                    pass.set_pipeline(&c.q4tp_mv16);
                    pass.set_bind_group(0, &bind1, &[]);
                    for _ in 0..n {
                        pass.dispatch_workgroups(1, 1, 1);
                    }
                }
                let t = Instant::now();
                submit(c, finish_enc(enc));
                let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
                t.elapsed().as_secs_f64()
            };
            run1();
            let mut best1 = f64::MAX;
            for _ in 0..3 {
                best1 = best1.min(run1());
            }
            println!(
                "same {n} dispatches in ONE pass = {:.2} ms  ({:.1} us/dispatch) — pass boundary ≈ {:.1} us",
                best1 * 1e3,
                best1 * 1e6 / n as f64,
                (best - best1) * 1e6 / n as f64
            );
        }
        for (rows, cols, label) in [
            (17408usize, 5120usize, "gate/up 17408x5120"),
            (5120, 17408, "down 5120x17408"),
            (16384, 5120, "gdn qkv 16384x5120"),
            (5120, 6144, "o/out 5120x6144"),
            (248320, 5120, "lm_head 248320x5120"),
        ] {
            let total = cortiq_core::quant::expected_nbytes(
                cortiq_core::TensorDtype::Q4TiledP,
                &[rows, cols],
            )
            .unwrap();
            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
            let mut wb: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
            let lo = cortiq_core::quant::f32_to_f16(-4.0);
            let step = cortiq_core::quant::f32_to_f16(0.1);
            for r in 0..rows {
                let o = params_off + r * 4;
                wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
                wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
            }
            let xs: Vec<f32> = (0..cols).map(|i| ((i % 97) as f32 - 48.0) / 48.0).collect();
            let mk = |bytes: &[u8]| {
                c.device
                    .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: None,
                        contents: bytes,
                        usage: wgpu::BufferUsages::STORAGE,
                    })
            };
            let wbuf = mk(&wb);
            let xbuf = mk(bytemuck::cast_slice(&xs));
            let ybuf = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (rows * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE,
                mapped_at_creation: false,
            });
            let run = || {
                let mut enc = c
                    .device
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
                for _ in 0..reps {
                    encode_q4tp_mv4(c, &mut enc, &wbuf, &xbuf, &ybuf, rows, cols);
                }
                submit(c, finish_enc(enc));
                let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            };
            run();
            let mut best = f64::MAX;
            for _ in 0..3 {
                let t = Instant::now();
                run();
                best = best.min(t.elapsed().as_secs_f64());
            }
            let per = best / reps as f64;
            println!(
                "{label}: {:.1} MB  {:.1} us/dispatch  {:.0} GB/s",
                total as f64 / 1e6,
                per * 1e6,
                total as f64 / per / 1e9
            );
        }
        // The 2-bit plane on the gate/up shape (its only decode shape).
        {
            let (rows, cols) = (17408usize, 5120usize);
            let total = cortiq_core::quant::expected_nbytes(
                cortiq_core::TensorDtype::Q2TiledP,
                &[rows, cols],
            )
            .unwrap();
            let (params_off, _, _) = cortiq_core::quant::q2tp_sections(rows, cols);
            let mut wb: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
            let lo = cortiq_core::quant::f32_to_f16(-4.0);
            let step = cortiq_core::quant::f32_to_f16(0.1);
            for r in 0..rows {
                let o = params_off + r * 4;
                wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
                wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
            }
            let xs: Vec<f32> = (0..cols).map(|i| ((i % 97) as f32 - 48.0) / 48.0).collect();
            let mk = |bytes: &[u8]| {
                c.device
                    .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: None,
                        contents: bytes,
                        usage: wgpu::BufferUsages::STORAGE,
                    })
            };
            let wbuf = mk(&wb);
            let xbuf = mk(bytemuck::cast_slice(&xs));
            let ybuf = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (rows * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE,
                mapped_at_creation: false,
            });
            {
                let run = || {
                    let mut enc = c
                        .device
                        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
                    for _ in 0..reps {
                        let (bind, wg) = q2tp_mv_bind(c, &wbuf, &xbuf, &ybuf, rows, cols, false);
                        let mut pass = begin_pass(&mut enc);
                        pass.set_pipeline(&c.q2tp_mv16w);
                        pass.set_bind_group(0, &bind, &[]);
                        pass.dispatch_workgroups(wg, 1, 1);
                    }
                    submit(c, finish_enc(enc));
                    let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
                };
                run();
                let mut best = f64::MAX;
                for _ in 0..3 {
                    let t = Instant::now();
                    run();
                    best = best.min(t.elapsed().as_secs_f64());
                }
                let per = best / reps as f64;
                println!(
                    "q2tp gate/up 17408x5120: {:.1} MB  {:.1} us/dispatch  {:.0} GB/s",
                    total as f64 / 1e6,
                    per * 1e6,
                    total as f64 / per / 1e9
                );
            }
        }
    }

    /// The f16-PLANE tensor-core GEMM with the DEVICE-computed activation
    /// scale (the fused-FFN / DiT-block path) against the scalar GEMM,
    /// with activations past the 1000 threshold so the scale is not 1.0.
    /// The staging must shrink the operand by the same factor the store
    /// grows the result by — a kernel that reads the scale from the
    /// buffer at one end and from the (sentinel) uniform at the other is
    /// off by exactly max|x|/1000, and only on the inputs large enough to
    /// have needed the scale in the first place.
    #[test]
    fn wgpu_q4tp_mm_coop_f16_device_scale_matches_scalar() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        let (Some(pipe), Some(_)) = (c.q4tp_mm_coop_f16.as_ref(), c.q4tp_dq_f16.as_ref()) else {
            eprintln!("no f16 plane GEMM here — skipping");
            return;
        };
        if c.act_amax_part.is_none() && c.act_absmax.is_none() {
            eprintln!("no device absmax here — skipping");
            return;
        }
        let (rows, cols, n) = (256usize, 128usize, 96usize);
        let total =
            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
                .unwrap();
        let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
        let mut wb: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
        let lo = cortiq_core::quant::f32_to_f16(-4.0);
        let step = cortiq_core::quant::f32_to_f16(0.1);
        for r in 0..rows {
            let o = params_off + r * 4;
            wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
            wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
        }
        // |x| up to ~3000 in a few lanes: max|x| > 1000, so the device
        // scale is 1000/max and the two ends of the kernel must agree.
        let xs: Vec<f32> = (0..n * cols)
            .map(|i| {
                let b = ((i % 97) as f32 - 48.0) / 48.0;
                if i % 11 == 0 {
                    b * 3000.0
                } else {
                    b
                }
            })
            .collect();
        let mk = |bytes: &[u8]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytes,
                    usage: wgpu::BufferUsages::STORAGE,
                })
        };
        let wbuf = mk(&wb);
        let xbuf = mk(bytemuck::cast_slice(&xs));
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        // scalar reference
        let y_ref = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (n * rows * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        encode_q4_tile_mm(c, &mut enc, &c.q4tp_mm, &wbuf, &xbuf, &y_ref, rows, cols, n);
        // plane + device scale + f16 GEMM
        let (plane, fresh) = plane_cached(c, (usize::MAX - 7, 0), &wbuf, rows, cols, 8192).unwrap();
        if let Some(bind_dq) = fresh {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(c.q4tp_dq_f16.as_ref().unwrap());
            pass.set_bind_group(0, &bind_dq, &[]);
            let wgs = ((rows * cols / 2) as u32).div_ceil(256);
            pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
        }
        let asc = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("test-ascale"),
            size: 4,
            usage: wgpu::BufferUsages::STORAGE,
            mapped_at_creation: false,
        });
        assert!(encode_act_absmax(c, &mut enc, &xbuf, n * cols, &asc));
        let y_f16 = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (n * rows * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });
        encode_q4_tile_mm_full(
            c,
            &mut enc,
            pipe,
            &plane,
            &xbuf,
            &y_f16,
            rows,
            cols,
            n,
            0.0,
            Some(&asc),
        );
        // read both back
        let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: (2 * n * rows * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&y_ref, 0, &stage, 0, (n * rows * 4) as u64);
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(
            &y_f16,
            0,
            &stage,
            (n * rows * 4) as u64,
            (n * rows * 4) as u64,
        );
        submit(c, finish_enc(enc));
        let slice = stage.slice(..);
        slice.map_async(wgpu::MapMode::Read, |_| {});
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        let data = slice.get_mapped_range().expect("map");
        let all: &[f32] = bytemuck::cast_slice(&data);
        let (a, b) = all.split_at(n * rows);
        let (mut num, mut den) = (0f64, 0f64);
        for (x, y) in b.iter().zip(a) {
            num += ((x - y) as f64).powi(2);
            den += (*y as f64).powi(2);
        }
        let r = (num / den.max(1e-30)).sqrt();
        println!("coop f16 (device scale) vs scalar: relative rms {r:.3e}");
        drop(data);
        stage.unmap();
        assert!(
            r < 1e-2,
            "f16 plane GEMM with the device scale drifted from the scalar arm: {r:.3e}"
        );
    }

    /// The tensor-core GEMM must agree with the scalar one it replaced.
    /// This is the gate the dequantize-once path needed and did not
    /// have: the neighbouring `wgpu_coop_small_gemm` only PRINTS its
    /// worst relative error, so the f16 operands, the activation scale
    /// and the unpacked plane could all drift unwatched. Runs on any
    /// machine with a device; skips cleanly without one.
    #[test]
    fn wgpu_q4tp_mm_coop_matches_scalar() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        if c.q4tp_mm_coop_f16.is_none() {
            // A device WITH cooperative matrices that has no f16 pipeline
            // means the shader was rejected — which is how a wrong
            // binding hid for several commits while the path silently
            // fell back. Skipping is only honest where the hardware
            // cannot do it at all.
            assert!(
                c.q4tp_mm_coop.is_none(),
                "cooperative matrices are up but the f16 pipeline is missing — \
                 the shader was rejected (check the init warning)"
            );
            eprintln!("no cooperative-matrix hardware here — skipping");
            return;
        }
        // Small enough to run anywhere, wide enough to cross several
        // 64-row tiles and more than one k-chunk.
        let (rows, cols, n) = (256usize, 128usize, 96usize);
        let total =
            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
                .unwrap();
        let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
        let mut wb: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
        let lo = cortiq_core::quant::f32_to_f16(-4.0);
        let step = cortiq_core::quant::f32_to_f16(0.1);
        for r in 0..rows {
            let o = params_off + r * 4;
            wb[o..o + 2].copy_from_slice(&lo.to_le_bytes());
            wb[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
        }
        // Activations spanning the range that used to overflow f16.
        let xs: Vec<f32> = (0..n * cols)
            .map(|i| {
                let b = ((i % 97) as f32 - 48.0) / 48.0;
                if i % 11 == 0 {
                    b * 3000.0
                } else {
                    b
                }
            })
            .collect();
        let mk = |bytes: &[u8]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytes,
                    usage: wgpu::BufferUsages::STORAGE,
                })
        };
        let wbuf = mk(&wb);
        let xbuf = mk(bytemuck::cast_slice(&xs));
        let run = |pipe: &wgpu::ComputePipeline| -> Vec<f32> {
            let ybuf = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (n * rows * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });
            let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: None,
                size: (n * rows * 4) as u64,
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            let mut enc = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
            encode_q4_tile_mm(c, &mut enc, pipe, &wbuf, &xbuf, &ybuf, rows, cols, n);
            let mut out = vec![0f32; n * rows];
            assert!(readback(
                c,
                enc,
                &ybuf,
                &stage,
                (n * rows * 4) as u64,
                &mut out
            ));
            out
        };
        let scalar = run(&c.q4tp_mm);
        let coop = run(c.q4tp_mm_coop.as_ref().unwrap());
        let rel = |a: &[f32], b: &[f32]| -> f64 {
            let (mut num, mut den) = (0f64, 0f64);
            for (x, y) in a.iter().zip(b) {
                num += ((x - y) as f64).powi(2);
                den += (*y as f64).powi(2);
            }
            (num / den.max(1e-30)).sqrt()
        };
        let r = rel(&coop, &scalar);
        println!("coop vs scalar: relative rms {r:.3e}");
        // f16 operands carry ~11 bits of mantissa; anything past 1e-2 is
        // a broken layout or a lost activation scale, not rounding.
        assert!(r < 1e-2, "coop GEMM drifted from the scalar arm: {r:.3e}");
    }

    #[test]
    #[ignore]
    fn wgpu_attn_block_timing() {
        use std::time::Instant;
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping");
            return;
        };
        // 1.7B-ish attention geometry.
        let (nh, nkv, hd, rd, hidden, cap, stored) = (
            16usize, 8usize, 128usize, 128usize, 2048usize, 256usize, 128usize,
        );
        let hpk = nh / nkv;
        let eps = 1e-6f32;
        let flags = 2u32 | 4u32;
        let h_in = vec![0.01f32; hidden];
        let norm_w = vec![1.0f32; hidden];
        let (wq_p, _) = mk_q1(nh * hd, hidden, 1);
        let (wk_p, _) = mk_q1(nkv * hd, hidden, 2);
        let (wv_p, _) = mk_q1(nkv * hd, hidden, 3);
        let (wo_p, _) = mk_q1(hidden, nh * hd, 4);
        let qnw = vec![1.0f32; hd];
        let knw = vec![1.0f32; hd];
        let invf: Vec<f32> = (0..rd / 2)
            .map(|i| 1.0 / (10000f32).powf(2.0 * i as f32 / rd as f32))
            .collect();
        let kc = vec![0.01f32; nkv * cap * hd];
        let vc = vec![0.01f32; nkv * cap * hd];
        let mkc = |d: &[f32]| {
            c.device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: None,
                    contents: bytemuck::cast_slice(d),
                    usage: wgpu::BufferUsages::STORAGE
                        | wgpu::BufferUsages::COPY_SRC
                        | wgpu::BufferUsages::COPY_DST,
                })
        };
        let (kbuf, vbuf) = (mkc(&kc), mkc(&vc));
        let iters = 200;
        let mut hout = vec![0f32; hidden];
        // FUSED: the resident block, one submit + one readback per call.
        for _ in 0..20 {
            attn_block_gpu(
                &h_in, &norm_w, &wq_p, &wk_p, &wv_p, &wo_p, &qnw, &knw, &invf, &kbuf, &vbuf, nh,
                nkv, hd, rd, hidden, cap, stored, flags, eps, &mut hout,
            );
        }
        let t0 = Instant::now();
        for _ in 0..iters {
            attn_block_gpu(
                &h_in, &norm_w, &wq_p, &wk_p, &wv_p, &wo_p, &qnw, &knw, &invf, &kbuf, &vbuf, nh,
                nkv, hd, rd, hidden, cap, stored, flags, eps, &mut hout,
            );
        }
        let fused = t0.elapsed().as_secs_f64() * 1000.0 / iters as f64;
        // UNFUSED: each step its own submit+readback (rmsnorm, QKV×3, rope, attend, O).
        let mut normed = vec![0f32; hidden];
        let mut qraw = vec![0f32; nh * hd];
        let mut kk = vec![0f32; nkv * hd];
        let mut vv = vec![0f32; nkv * hd];
        let mut qout = vec![0f32; nh * hd];
        let mut kout = vec![0f32; nkv * hd];
        let mut gout = vec![0f32; nh * hd];
        let mut attn = vec![0f32; nh * hd];
        let mut oout = vec![0f32; hidden];
        let unfused_once = |normed: &mut [f32],
                            qraw: &mut [f32],
                            kk: &mut [f32],
                            vv: &mut [f32],
                            qout: &mut [f32],
                            kout: &mut [f32],
                            gout: &mut [f32],
                            attn: &mut [f32],
                            oout: &mut [f32]| {
            rmsnorm_row(&h_in, &norm_w, normed, false, eps);
            dispatch_q1(c, None, &wq_p, normed, nh * hd, hidden, qraw);
            dispatch_q1(c, None, &wk_p, normed, nkv * hd, hidden, kk);
            dispatch_q1(c, None, &wv_p, normed, nkv * hd, hidden, vv);
            attn_rope_qkn_gpu(
                qraw, kk, &qnw, &knw, &invf, nh, nkv, hd, rd, stored, flags, eps, qout, kout, gout,
            );
            gqa_attend_gpu(qout, &kc, &vc, nh, hpk, hd, cap, stored + 1, attn);
            dispatch_q1(c, None, &wo_p, attn, hidden, nh * hd, oout);
        };
        for _ in 0..20 {
            unfused_once(
                &mut normed,
                &mut qraw,
                &mut kk,
                &mut vv,
                &mut qout,
                &mut kout,
                &mut gout,
                &mut attn,
                &mut oout,
            );
        }
        let t1 = Instant::now();
        for _ in 0..iters {
            unfused_once(
                &mut normed,
                &mut qraw,
                &mut kk,
                &mut vv,
                &mut qout,
                &mut kout,
                &mut gout,
                &mut attn,
                &mut oout,
            );
        }
        let unfused = t1.elapsed().as_secs_f64() * 1000.0 / iters as f64;
        eprintln!(
            "ATTN BLOCK 1.7B-dims: fused(1 submit) {fused:.3} ms/layer | unfused(per-op) {unfused:.3} ms/layer | speedup {:.2}×",
            unfused / fused
        );
    }

    #[test]
    fn wgpu_q1t_matvec_matches_cpu_reference() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping q1t parity test");
            return;
        };
        use cortiq_core::quant::{f32_to_f16, q1t_pack, GROUP_SIZE};
        let (rows, cols) = (33usize, 256usize);
        let gpr = cols / GROUP_SIZE;
        let outliers: [(usize, f32); 3] = [(5, 3.0), (300, -2.0), (600, 1.5)]; // sorted
        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i == flat);
        let mut payload = Vec::new();
        for r in 0..rows {
            for g in 0..gpr {
                let s = 0.02 + ((r + g) % 7) as f32 * 0.01;
                payload.extend_from_slice(&f32_to_f16(s).to_le_bytes());
                let mut cc = [0u8; 7];
                for k in 0..GROUP_SIZE {
                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
                        0
                    } else {
                        ((k * 7 + r + g) % 3) as u8
                    };
                    q1t_pack(&mut cc, k, code);
                }
                payload.extend_from_slice(&cc);
            }
        }
        let mut row_ptr = vec![0u32; rows + 1];
        for &(idx, _) in &outliers {
            row_ptr[idx / cols + 1] += 1;
        }
        for r in 0..rows {
            row_ptr[r + 1] += row_ptr[r];
        }
        for &p in &row_ptr {
            payload.extend_from_slice(&p.to_le_bytes());
        }
        for &(idx, v) in &outliers {
            payload.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
            payload.extend_from_slice(&f32_to_f16(v).to_le_bytes());
        }
        let xs: Vec<f32> = (0..cols)
            .map(|i| ((i * 7 + 3) % 29) as f32 / 29.0 - 0.5)
            .collect();
        let mut w = vec![0f32; rows * cols];
        cortiq_core::quant::dequant_q1t(&payload, rows, cols, &mut w);
        let mut want = vec![0f32; rows];
        for o in 0..rows {
            want[o] = (0..cols).map(|i| w[o * cols + i] * xs[i]).sum();
        }
        let mut got = vec![0f32; rows];
        assert!(dispatch_q1t(
            c, &c.q1t, None, &payload, &xs, rows, cols, &mut got
        ));
        let max_d = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0f32, f32::max);
        assert!(max_d < 1e-2, "wgpu q1t_matvec ≠ CPU: max|Δ| = {max_d}");
    }

    #[test]
    fn wgpu_q4b_matvec_matches_cpu_reference() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping q4b parity test");
            return;
        };
        use cortiq_core::quant::{f32_to_f16, GROUP_SIZE};
        let (rows, cols) = (33usize, 256usize);
        let n_groups = rows * (cols / GROUP_SIZE);
        let mut payload = vec![0u8; n_groups * 16]; // packed nibbles
        for g in 0..n_groups {
            for k in 0..16 {
                let lo = ((g * 3 + k) % 16) as u8;
                let hi = ((g * 5 + k * 2) % 16) as u8;
                payload[g * 16 + k] = lo | (hi << 4);
            }
        }
        for g in 0..n_groups {
            let s = 0.02 + (g % 7) as f32 * 0.01;
            payload.extend_from_slice(&f32_to_f16(s).to_le_bytes());
        }
        let xs: Vec<f32> = (0..cols)
            .map(|i| ((i * 7 + 3) % 29) as f32 / 29.0 - 0.5)
            .collect();
        let mut w = vec![0f32; rows * cols];
        cortiq_core::quant::dequant_q4_block(&payload, &mut w);
        let mut want = vec![0f32; rows];
        for o in 0..rows {
            want[o] = (0..cols).map(|i| w[o * cols + i] * xs[i]).sum();
        }
        let mut got = vec![0f32; rows];
        assert!(dispatch_q1t(
            c, &c.q4b, None, &payload, &xs, rows, cols, &mut got
        ));
        let max_d = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0f32, f32::max);
        assert!(max_d < 1e-2, "wgpu q4b_matvec ≠ CPU: max|Δ| = {max_d}");
    }

    #[test]
    fn wgpu_q1t_matmat_matches_cpu_reference() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping q1t GEMM parity test");
            return;
        };
        use cortiq_core::quant::{f32_to_f16, q1t_pack, GROUP_SIZE};
        let (b, rows, cols) = (40usize, 64usize, 256usize);
        let gpr = cols / GROUP_SIZE;
        let outliers: [(usize, f32); 4] = [(5, 3.0), (300, -2.0), (600, 1.5), (2000, -1.0)];
        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i == flat);
        let mut payload = Vec::new();
        for r in 0..rows {
            for g in 0..gpr {
                let s = 0.02 + ((r + g) % 7) as f32 * 0.01;
                payload.extend_from_slice(&f32_to_f16(s).to_le_bytes());
                let mut cc = [0u8; 7];
                for k in 0..GROUP_SIZE {
                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
                        0
                    } else {
                        ((k * 7 + r + g) % 3) as u8
                    };
                    q1t_pack(&mut cc, k, code);
                }
                payload.extend_from_slice(&cc);
            }
        }
        let mut row_ptr = vec![0u32; rows + 1];
        for &(idx, _) in &outliers {
            row_ptr[idx / cols + 1] += 1;
        }
        for r in 0..rows {
            row_ptr[r + 1] += row_ptr[r];
        }
        for &p in &row_ptr {
            payload.extend_from_slice(&p.to_le_bytes());
        }
        for &(idx, v) in &outliers {
            payload.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
            payload.extend_from_slice(&f32_to_f16(v).to_le_bytes());
        }
        let xs: Vec<f32> = (0..b * cols)
            .map(|i| ((i * 13 + 7) % 31) as f32 / 31.0 - 0.5)
            .collect();
        let mut w = vec![0f32; rows * cols];
        cortiq_core::quant::dequant_q1t(&payload, rows, cols, &mut w);
        let mut want = vec![0f32; b * rows];
        for bi in 0..b {
            for o in 0..rows {
                want[bi * rows + o] = (0..cols).map(|i| w[o * cols + i] * xs[bi * cols + i]).sum();
            }
        }
        let mut got = vec![0f32; b * rows];
        assert!(dispatch_q1t_mm(
            c, None, &payload, &xs, b, rows, cols, &mut got
        ));
        let max_d = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0f32, f32::max);
        assert!(max_d < 2e-2, "wgpu q1t_mul_mm ≠ CPU: max|Δ| = {max_d}");
    }

    #[test]
    fn wgpu_q8_matmat_matches_cpu_reference() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping matmat test");
            return;
        };
        let (rows, cols, b) = (128usize, 64usize, 5usize);
        let mut q = vec![0i8; rows * cols];
        for (i, v) in q.iter_mut().enumerate() {
            *v = (((i * 53 + 3) % 255) as i32 - 127) as i8;
        }
        let rs: Vec<f32> = (0..rows).map(|r| 0.01 + (r % 5) as f32 * 0.004).collect();
        let pre: Vec<f32> = (0..b * cols)
            .map(|i| ((i % 17) as f32 - 8.0) * 0.05)
            .collect();
        // CPU ref: out[bi, o] = rs[o]·Σ q[o,i]·pre[bi,i].
        let mut want = vec![0f32; b * rows];
        for bi in 0..b {
            for o in 0..rows {
                let mut acc = 0f32;
                for i in 0..cols {
                    acc += q[o * cols + i] as f32 * pre[bi * cols + i];
                }
                want[bi * rows + o] = acc * rs[o];
            }
        }
        let qbytes: &[u8] = bytemuck::cast_slice(&q);
        let mut got = vec![0f32; b * rows];
        assert!(dispatch_matmat(
            c, None, qbytes, &rs, &pre, b, rows, cols, &mut got
        ));
        let max_d = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(max_d < 1e-3, "wgpu q8_matmat ≠ CPU: max|Δ| = {max_d}");
    }

    /// The tiled kernel (b ≥ 32) on deliberately awkward shapes: rows
    /// not a multiple of the 64-tile, cols not a multiple of the K-step
    /// — every edge guard fires.
    #[test]
    fn wgpu_q8_mul_mm_matches_cpu_reference() {
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping mul_mm test");
            return;
        };
        let (rows, cols, b) = (100usize, 52usize, 70usize);
        let mut q = vec![0i8; rows * cols];
        for (i, v) in q.iter_mut().enumerate() {
            *v = (((i * 31 + 7) % 255) as i32 - 127) as i8;
        }
        let rs: Vec<f32> = (0..rows).map(|r| 0.01 + (r % 7) as f32 * 0.003).collect();
        let pre: Vec<f32> = (0..b * cols)
            .map(|i| ((i % 19) as f32 - 9.0) * 0.04)
            .collect();
        let mut want = vec![0f32; b * rows];
        for bi in 0..b {
            for o in 0..rows {
                let mut acc = 0f32;
                for i in 0..cols {
                    acc += q[o * cols + i] as f32 * pre[bi * cols + i];
                }
                want[bi * rows + o] = acc * rs[o];
            }
        }
        let qbytes: &[u8] = bytemuck::cast_slice(&q);
        let mut got = vec![0f32; b * rows];
        assert!(dispatch_matmat(
            c, None, qbytes, &rs, &pre, b, rows, cols, &mut got
        ));
        let max_d = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(max_d < 1e-3, "wgpu q8_mul_mm ≠ CPU: max|Δ| = {max_d}");
    }

    // Tiled q1 GEMM on an awkward shape (rows/batch not 64-multiples, cols a
    // 64-multiple as the format requires): the prefill / speculative-batch path.
    #[test]
    fn wgpu_q1_mul_mm_matches_cpu_reference() {
        use cortiq_core::quant::{f16_to_f32, f32_to_f16};
        unsafe { std::env::set_var("CMF_GPU", "wgpu") };
        let Some(c) = ctx() else {
            eprintln!("no wgpu adapter — skipping q1_mul_mm test");
            return;
        };
        let (rows, cols, b) = (100usize, 128usize, 70usize); // cols % 64 == 0
        let np = cols / 64;
        let jit = |a: usize| ((a * 2654435761usize) >> 13) as u32; // cheap hash → bits
                                                                   // Build the q1 weight blob + a decoded f32 reference weight in lock-step.
        let mut q1w = vec![0u32; rows * np * 3];
        let mut wref = vec![0f32; rows * cols];
        for o in 0..rows {
            for pi in 0..np {
                let s0 = 0.02 + ((o * 7 + pi) % 11) as f32 * 0.005;
                let s1 = 0.03 + ((o * 3 + pi * 5) % 9) as f32 * 0.004;
                let (h0, h1) = (f32_to_f16(s0), f32_to_f16(s1));
                let (sf0, sf1) = (f16_to_f32(h0), f16_to_f32(h1));
                let bits0 = jit(o * 131 + pi * 17 + 1);
                let bits1 = jit(o * 131 + pi * 17 + 2);
                let base = o * np * 3 + pi * 3;
                q1w[base] = (h0 as u32) | ((bits0 & 0xFFFF) << 16);
                q1w[base + 1] = (bits0 >> 16) | ((h1 as u32) << 16);
                q1w[base + 2] = bits1;
                for j in 0..32usize {
                    let sgn0 = if (bits0 >> j) & 1 != 0 { sf0 } else { -sf0 };
                    let sgn1 = if (bits1 >> j) & 1 != 0 { sf1 } else { -sf1 };
                    wref[o * cols + pi * 64 + j] = sgn0;
                    wref[o * cols + pi * 64 + 32 + j] = sgn1;
                }
            }
        }
        let x: Vec<f32> = (0..b * cols)
            .map(|i| ((i % 23) as f32 - 11.0) * 0.03)
            .collect();
        let mut want = vec![0f32; b * rows];
        for bi in 0..b {
            for o in 0..rows {
                let mut acc = 0f32;
                for i in 0..cols {
                    acc += wref[o * cols + i] * x[bi * cols + i];
                }
                want[bi * rows + o] = acc;
            }
        }
        // GPU dispatch (inline — q1_mm is not yet wired into a public entry).
        let qbuf = storage_bytes(c, bytemuck::cast_slice(&q1w));
        let xbuf = storage_bytes(c, bytemuck::cast_slice(&x));
        let ybuf = rw_f32(c, b * rows, true);
        let pbuf = uniform_u32x4(c, [(cols / 4) as u32, rows as u32, b as u32, 0]);
        // q1_mul_mm never reads rs → its auto layout omits binding 2.
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.layout_q1mm,
            entries: &[
                bind_buf(0, &qbuf),
                bind_buf(1, &xbuf),
                bind_buf(3, &ybuf),
                bind_buf(4, &pbuf),
            ],
        });
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(&c.q1_mm);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups((rows as u32).div_ceil(64), (b as u32).div_ceil(64), 1);
        }
        let mut sc = c.scratch.lock().unwrap();
        let stage = Scratch::ensure(
            &c.device,
            &mut sc.stage,
            (b * rows * 4) as u64,
            wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            "q1mm-stage",
        );
        let mut got = vec![0f32; b * rows];
        assert!(readback(
            c,
            enc,
            &ybuf,
            &stage,
            (b * rows * 4) as u64,
            &mut got
        ));
        drop(sc);
        let max_d = want
            .iter()
            .zip(&got)
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(max_d < 1e-3, "wgpu q1_mul_mm ≠ CPU: max|Δ| = {max_d}");
    }
}

/// Was the wgpu backend ASKED for, and did it come up? The two halves must
/// be told apart: a machine nobody pointed at wgpu is a legitimate skip, a
/// machine that was pointed at it and produced no context is a failure. A
/// reserved word in one shader once took the whole context down and every
/// GPU test reported success by skipping.
/// A GPU test that found no device: FAIL when the environment asked for
/// one, skip loudly when it did not.
///
/// A bare `return` here is why a column of `ok` could mean "every GPU
/// test skipped" — cargo's summary has no way to say "passed without
/// running anything", so the skip has to be made visible by the test
/// itself, and a skip that happens after someone explicitly set
/// `CMF_GPU` is not a skip, it is a failure to honour the request.
pub fn skip_or_fail(what: &str) {
    // Only the value that NAMES this backend can make a skip a failure.
    // `CMF_GPU=1` selects Metal on macOS, where wgpu being absent is
    // correct, not a broken request — a wider test here would turn every
    // Mac run red for the right reason on the wrong platform.
    let asked = std::env::var("CMF_GPU")
        .map(|v| v == "wgpu")
        .unwrap_or(false);
    assert!(
        !asked,
        "{what}: CMF_GPU=wgpu was set but no wgpu adapter came up — \
         refusing to report a skip as a pass"
    );
    eprintln!("SKIPPED ({what}): no wgpu device — set CMF_GPU=wgpu to run it");
}

pub fn selected_and_up() -> Option<bool> {
    // "Asked" means an EXPLICIT request. The wgpu path also self-selects
    // by default on Linux/Windows, but a default selection on a box with
    // no device is the designed CPU fallback, not a failure — parity
    // tests skip there (what headless CI looks like) instead of dying.
    if std::env::var("CMF_GPU").is_err() || !selected() {
        return None; // nobody asked
    }
    Some(ctx().is_some())
}

/// Cheap device probe for `gpu::backend_available`: can wgpu bring an
/// adapter up here at all? One instance, no device/queue, no caching —
/// the caller caches.
/// Every adapter wgpu can see, and which one would be chosen. Three times in
/// one night the question "is the GPU actually visible?" was answered by
/// inference from a missing log line; this answers it directly.
/// Per-head RMS and the rope tail on the device — `rms` for the queries'
/// second normalisation, `inverse` for attention's output.
#[allow(clippy::too_many_arguments)]
pub fn rope_heads_for_test(
    x: &mut [f32],
    inv_freq: &[f32],
    nh: usize,
    hd: usize,
    rd: usize,
    pos: usize,
    eps: f32,
    rms: bool,
    inverse: bool,
) -> bool {
    let Some(c) = ctx() else { return false };
    if x.len() != nh * hd || rd > hd || rd % 2 != 0 || inv_freq.len() * 2 < rd {
        return false;
    }
    // In place: seeded from the host and read back afterwards, so the buffer
    // needs both directions. rw_f32 gives STORAGE|COPY_SRC and refuses the
    // write; storage_bytes gives STORAGE and refuses the read.
    let xb = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("rope-x"),
            contents: bytemuck::cast_slice(x),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        });
    let fb = storage_bytes(c, bytemuck::cast_slice(&inv_freq[..rd / 2]));
    let pb = storage_bytes(c, bytemuck::cast_slice(&[pos as f32, eps]));
    let flags = (rms as u32) | ((inverse as u32) << 1);
    let p = uniform_u32x4(c, [nh as u32, hd as u32, rd as u32, flags]);
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("rope"),
        });
    {
        let layout = c.rope_heads.get_bind_group_layout(0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layout,
            entries: &[
                bind_buf(0, &xb),
                bind_buf(1, &fb),
                bind_buf(2, &p),
                bind_buf(3, &pb),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.rope_heads);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(nh as u32, 1, 1);
    }
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        (nh * hd * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "rope-stage",
    );
    let ok = readback(c, enc, &xb, &stage, (nh * hd * 4) as u64, x);
    drop(sc);
    ok
}

/// Stage A of the grouped low-rank output projection on the device.
///
/// `lora` rows share each group's slice of `attn`; `rows` is `groups * lora`.
/// q4tp only — the release stores wo_a that way in both published variants,
/// and guessing at a layout is how a kernel returns plausible nonsense.
pub fn o_lora_a_for_test(
    model: &Arc<CmfModel>,
    idx: usize,
    attn: &[f32],
    rows: usize,
    lora: usize,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if rows == 0 || lora == 0 || rows % lora != 0 || out.len() < rows {
        return false;
    }
    let entry = &model.tensors[idx];
    if entry.dtype != cortiq_core::TensorDtype::Q4TiledP || entry.shape.len() != 2 {
        return false;
    }
    let (trows, cols) = (entry.shape[0], entry.shape[1]);
    let groups = rows / lora;
    if trows < rows || cols % 32 != 0 || attn.len() < groups * cols {
        return false;
    }
    let Some(abs) = model.entry_abs_offset(entry) else {
        return false;
    };
    let bytes = model.primary_bytes();
    let plen = entry.nbytes as usize;
    if abs + plen > bytes.len() {
        return false;
    }
    let Some(w) = weight_buffer_l(
        c,
        (model.uid() as usize, idx),
        &bytes[abs..abs + plen],
        layer_of_name(&model.tensors[idx].name),
    ) else {
        return false; // over budget → the caller keeps it on the CPU
    };
    let xb = storage_bytes(c, bytemuck::cast_slice(&attn[..groups * cols]));
    let yb = rw_f32(c, rows, true);
    let p = uniform_u32x4(c, [(cols / 32) as u32, rows as u32, lora as u32, 0]);
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("o-lora-a"),
        });
    {
        let layout = c.o_lora_a.get_bind_group_layout(0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layout,
            entries: &[
                bind_buf(0, &w),
                bind_buf(1, &xb),
                bind_buf(2, &yb),
                bind_buf(3, &p),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.o_lora_a);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
    }
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        (rows * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "o-lora-stage",
    );
    let ok = readback(c, enc, &yb, &stage, (rows * 4) as u64, &mut out[..rows]);
    drop(sc);
    ok
}

/// The compressor's pooling step on the device — `overlap` picks the folding
/// the release uses at ratio 4, `ape` the positional bias the plain one adds.
#[allow(clippy::too_many_arguments)]
pub fn kv_pool_for_test(
    prev_kv: &[f32],
    prev_score: &[f32],
    cur_kv: &[f32],
    cur_score: &[f32],
    ape: Option<&[f32]>,
    ratio: usize,
    width: usize,
    overlap: bool,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if width == 0 || ratio == 0 || out.len() < width {
        return false;
    }
    let slots = if overlap { 2 * ratio } else { ratio };
    let stride = if overlap { 2 * width } else { width };
    if cur_kv.len() < ratio * stride || cur_score.len() < ratio * stride {
        return false;
    }
    let have_prev =
        overlap && prev_kv.len() >= ratio * stride && prev_score.len() >= ratio * stride;
    // Unused bindings still have to point somewhere; the current window is as
    // good a placeholder as an empty buffer and costs no allocation.
    let ckv = storage_bytes(c, bytemuck::cast_slice(cur_kv));
    let csc = storage_bytes(c, bytemuck::cast_slice(cur_score));
    let pkv = if have_prev {
        storage_bytes(c, bytemuck::cast_slice(prev_kv))
    } else {
        ckv.clone()
    };
    let psc = if have_prev {
        storage_bytes(c, bytemuck::cast_slice(prev_score))
    } else {
        csc.clone()
    };
    let use_ape = ape.is_some_and(|a| a.len() >= ratio * width) && !overlap;
    let apb = if use_ape {
        storage_bytes(c, bytemuck::cast_slice(ape.unwrap()))
    } else {
        csc.clone()
    };
    let yb = rw_f32(c, width, true);
    let flags = (overlap as u32) | ((have_prev as u32) << 1) | ((use_ape as u32) << 2);
    let p = uniform_u32x4(c, [slots as u32, width as u32, ratio as u32, flags]);
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("kv-pool"),
        });
    {
        let layout = c.kv_pool.get_bind_group_layout(0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layout,
            entries: &[
                bind_buf(0, &pkv),
                bind_buf(1, &psc),
                bind_buf(2, &ckv),
                bind_buf(3, &csc),
                bind_buf(4, &apb),
                bind_buf(5, &yb),
                bind_buf(6, &p),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.kv_pool);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((width as u32).div_ceil(256), 1, 1);
    }
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        (width * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "kv-pool-stage",
    );
    let ok = readback(c, enc, &yb, &stage, (width * 4) as u64, &mut out[..width]);
    drop(sc);
    ok
}

/// The indexer's scoring pass on the device.
#[allow(clippy::too_many_arguments)]
pub fn index_scores_for_test(
    q: &[f32],
    kv: &[f32],
    hw: &[f32],
    nh: usize,
    hd: usize,
    n_pos: usize,
    limit: usize,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if q.len() < nh * hd || kv.len() < n_pos * hd || hw.len() < nh || out.len() < n_pos {
        return false;
    }
    if n_pos == 0 {
        return true;
    }
    let qb = storage_bytes(c, bytemuck::cast_slice(&q[..nh * hd]));
    let kb = storage_bytes(c, bytemuck::cast_slice(&kv[..n_pos * hd]));
    let wb = storage_bytes(c, bytemuck::cast_slice(&hw[..nh]));
    let yb = rw_f32(c, n_pos, true);
    let p = uniform_u32x4(c, [nh as u32, hd as u32, n_pos as u32, limit as u32]);
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("ix") });
    {
        let layout = c.index_scores.get_bind_group_layout(0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layout,
            entries: &[
                bind_buf(0, &qb),
                bind_buf(1, &kb),
                bind_buf(2, &wb),
                bind_buf(3, &yb),
                bind_buf(4, &p),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.index_scores);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((n_pos as u32).min(MAX_WG), 1, 1);
    }
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        (n_pos * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "ix-stage",
    );
    let ok = readback(c, enc, &yb, &stage, (n_pos * 4) as u64, &mut out[..n_pos]);
    drop(sc);
    ok
}

/// Top-k positions on the device, in index order — the list `sparse_attend`
/// consumes. Bounded by the kernel's workgroup array; beyond it the caller
/// keeps the CPU's version rather than getting a truncated answer.
/// `y += w·x` (or `y = w·x` when `set`) on the device, for the parity test.
/// This kernel had no test at all, and its uniform was written in a
/// different field order than the shader reads — so it silently did nothing.
pub fn axpy_for_test(x: &[f32], y: &mut [f32], w: f32, set: bool, soff: usize) -> bool {
    let Some(c) = ctx() else { return false };
    let n = y.len();
    if n == 0 || x.len() < soff + n {
        return false;
    }
    let xb = storage_bytes(c, bytemuck::cast_slice(x));
    // COPY_DST as well: this one is written from the host before the
    // dispatch, and `rw_f32` only asks for STORAGE | COPY_SRC.
    let yb = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("axpy-y"),
        size: (n * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE
            | wgpu::BufferUsages::COPY_SRC
            | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    c.queue.write_buffer(&yb, 0, bytemuck::cast_slice(y));
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("axpy"),
        });
    {
        let mut pass = begin_pass(&mut enc);
        encode_axpy_full_p(&mut pass, c, &xb, &yb, w, n, set, soff, None);
    }
    let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: None,
        size: (n * 4) as u64,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    readback(c, enc, &yb, &stage, (n * 4) as u64, y)
}

pub fn top_k_for_test(scores: &[f32], k: usize, out: &mut Vec<u32>) -> bool {
    let Some(c) = ctx() else { return false };
    let n = scores.len();
    if n == 0 || n > 4096 || k == 0 {
        out.clear();
        return n == 0;
    }
    let kk = k.min(n);
    let sb = storage_bytes(c, bytemuck::cast_slice(scores));
    let ib = rw_f32(c, kk, true);
    let cb = rw_f32(c, 1, true);
    let p = uniform_u32x4(c, [n as u32, k as u32, 0, 0]);
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("topk"),
        });
    {
        let layout = c.top_k_index.get_bind_group_layout(0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layout,
            entries: &[
                bind_buf(0, &sb),
                bind_buf(1, &ib),
                bind_buf(2, &cb),
                bind_buf(3, &p),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.top_k_index);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(1, 1, 1);
    }
    let bytes = ((kk + 1) * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "topk-stage",
    );
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&ib, 0, &stage, 0, (kk * 4) as u64);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&cb, 0, &stage, (kk * 4) as u64, 4);
    submit(c, finish_enc(enc));
    let slice = stage.slice(..bytes);
    slice.map_async(wgpu::MapMode::Read, |_| {});
    if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        return false;
    }
    let mut ok = false;
    if let Ok(data) = slice.get_mapped_range() {
        let words: &[u32] = bytemuck::cast_slice(&data[..bytes as usize]);
        let cnt = (words[kk] as usize).min(kk);
        out.clear();
        out.extend_from_slice(&words[..cnt]);
        ok = true;
    }
    stage.unmap();
    drop(sc);
    ok
}

#[allow(clippy::too_many_arguments)]
fn encode_hc_fold_k(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    state: &wgpu::Buffer,
    mixes: &wgpu::Buffer,
    sc: &wgpu::Buffer,
    base: &wgpu::Buffer,
    fold: &wgpu::Buffer,
    post: &wgpu::Buffer,
    comb: &wgpu::Buffer,
    p: &wgpu::Buffer,
    bkey: (u8, u64, usize),
) {
    let mut pass = begin_pass(enc);
    encode_hc_fold_k_p(
        &mut pass, c, state, mixes, sc, base, fold, post, comb, p, None, bkey,
    );
}

#[allow(clippy::too_many_arguments)]
fn encode_hc_fold_k_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    state: &wgpu::Buffer,
    mixes: &wgpu::Buffer,
    sc: &wgpu::Buffer,
    base: &wgpu::Buffer,
    fold: &wgpu::Buffer,
    post: &wgpu::Buffer,
    comb: &wgpu::Buffer,
    p: &wgpu::Buffer,
    // The norm that always follows: same workgroup, same reduction machinery,
    // one dispatch instead of two.
    nrm: Option<(&wgpu::Buffer, &wgpu::Buffer)>,
    bkey: (u8, u64, usize),
) {
    // The norm's weight and output must be BOUND either way — a bind group
    // has to satisfy the layout — so the absent case binds the fold buffer
    // twice and the kernel skips the phase.
    let (nw, out) = nrm.unwrap_or((fold, fold));
    let bind = cached_bind(c, bkey, || {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.hc_pre_fold.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, state),
                bind_buf(1, mixes),
                bind_buf(2, sc),
                bind_buf(3, base),
                bind_buf(4, fold),
                bind_buf(5, post),
                bind_buf(6, comb),
                bind_buf(7, p),
                bind_buf(8, nw),
                bind_buf(9, out),
            ],
        })
    });
    pass.set_pipeline(&c.hc_pre_fold);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(1, 1, 1);
}

fn encode_hc_fold(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    state: &wgpu::Buffer,
    mixes: &wgpu::Buffer,
    sc: &wgpu::Buffer,
    base: &wgpu::Buffer,
    fold: &wgpu::Buffer,
    post: &wgpu::Buffer,
    comb: &wgpu::Buffer,
    p: &wgpu::Buffer,
) {
    let norm_out = frame_buf(c, 123, fold.size() as usize, false);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &c.hc_pre_fold.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, state),
            bind_buf(1, mixes),
            bind_buf(2, sc),
            bind_buf(3, base),
            bind_buf(4, fold),
            bind_buf(5, post),
            bind_buf(6, comb),
            bind_buf(7, p),
            // The fused-normalisation phase is disabled in this legacy
            // wrapper (`HcP.nrm == 0`) and runs as the next dispatch, but the
            // pipeline layout still requires every declared binding.
            bind_buf(8, base),
            bind_buf(9, &norm_out),
        ],
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(&c.hc_pre_fold);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(1, 1, 1);
}

#[allow(clippy::too_many_arguments)]
fn encode_hc_expand_k(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    x: &wgpu::Buffer,
    res: &wgpu::Buffer,
    post: &wgpu::Buffer,
    comb: &wgpu::Buffer,
    out: &wgpu::Buffer,
    p: &wgpu::Buffer,
    hc: usize,
    dim: usize,
    bkey: (u8, u64, usize),
) {
    let mut pass = begin_pass(enc);
    encode_hc_expand_k_p(&mut pass, c, x, res, post, comb, out, p, hc, dim, bkey);
}

#[allow(clippy::too_many_arguments)]
fn encode_hc_expand_k_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    x: &wgpu::Buffer,
    res: &wgpu::Buffer,
    post: &wgpu::Buffer,
    comb: &wgpu::Buffer,
    out: &wgpu::Buffer,
    p: &wgpu::Buffer,
    hc: usize,
    dim: usize,
    bkey: (u8, u64, usize),
) {
    let bind = cached_bind(c, bkey, || {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.hc_post_expand.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, x),
                bind_buf(1, res),
                bind_buf(2, post),
                bind_buf(3, comb),
                bind_buf(4, out),
                bind_buf(5, p),
            ],
        })
    });
    pass.set_pipeline(&c.hc_post_expand);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(((hc * dim) as u32).div_ceil(256), 1, 1);
}

fn encode_hc_expand(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    x: &wgpu::Buffer,
    res: &wgpu::Buffer,
    post: &wgpu::Buffer,
    comb: &wgpu::Buffer,
    out: &wgpu::Buffer,
    p: &wgpu::Buffer,
    hc: usize,
    dim: usize,
) {
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &c.hc_post_expand.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, x),
            bind_buf(1, res),
            bind_buf(2, post),
            bind_buf(3, comb),
            bind_buf(4, out),
            bind_buf(5, p),
        ],
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(&c.hc_post_expand);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(((hc * dim) as u32).div_ceil(256), 1, 1);
}

/// The attention chain from an already-normed LoRA vector to the block output.
#[allow(clippy::too_many_arguments)]
fn encode_attn_chain(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    wb: &[wgpu::Buffer],
    qn: &wgpu::Buffer,
    q: &wgpu::Buffer,
    attn: &wgpu::Buffer,
    mid: &wgpu::Buffer,
    out: &wgpu::Buffer,
    cache: &wgpu::Buffer,
    ixb: &wgpu::Buffer,
    sink: &wgpu::Buffer,
    freq: &wgpu::Buffer,
    posb: &wgpu::Buffer,
    g: Dsv4AttnGeom,
    kv_id: u64,
    li: usize,
    m: usize,
    q_ready: bool,
) {
    let rows = g.o_groups * g.o_lora;
    let cols = g.nh * g.hd / g.o_groups;
    // By SHAPE, not by flag. The 256-thread twin measured 21.3 tok/s against
    // the 64-thread kernel's 21.8: this projection's columns are
    // nh·hd/o_groups, so `gpr` is 128 on the release and half of 256 threads
    // would stride past the end. Wide only where there is width to use.
    // By shape: 256 threads only where a row has 256 groups to give them;
    // otherwise four rows at once, which buys overlap instead of width.
    // CMF_DSV4_OLORA picks explicitly for the A/B.
    let o_pipe = match olora_pick() {
        Some(1) => &c.o_lora_a,
        Some(2) => &c.o_lora_a_w,
        Some(3) => &c.o_lora_a_m,
        _ if !chain_mv4() => &c.o_lora_a,
        _ if cols / 32 >= 256 => &c.o_lora_a_w,
        _ => &c.o_lora_a_m,
    };
    let o_rows_per_wg = if std::ptr::eq(o_pipe, &c.o_lora_a_m) {
        4u32
    } else {
        1u32
    };
    let o_bind = cached_bind(c, (63, kv_id, li), || {
        let p = uniform_u32x4(c, [(cols / 32) as u32, rows as u32, g.o_lora as u32, 0]);
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &o_pipe.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wb[2]),
                bind_buf(1, attn),
                bind_buf(2, mid),
                bind_buf(3, &p),
            ],
        })
    });
    // Six dependent dispatches, no copy between them: one pass. Split back
    // into six only when the timestamp query set exists, because a per-stage
    // timing needs a pass boundary to be written at — and that is the whole
    // reason attention had six in the first place.
    if c.ts_query.is_none() && !sa_split() {
        let sa_bind = sa_bind_single(
            c,
            q,
            cache,
            ixb,
            sink,
            attn,
            g.nh,
            g.hd,
            m,
            g.scale,
            Some((kv_id, li)),
        );
        let mut pass = begin_pass(enc);
        if !q_ready && !dsv4_skip("qproj") {
            encode_q4tp_mvw_p(
                &mut pass,
                c,
                &wb[1],
                qn,
                q,
                g.nh * g.hd,
                g.q_lora,
                (60, kv_id, li),
            );
        }
        encode_rope_heads_p(
            &mut pass,
            c,
            q,
            freq,
            posb,
            g.nh,
            g.hd,
            g.rd,
            true,
            false,
            (61, kv_id, li),
        );
        if !dsv4_skip("sa") {
            if sa_split_k() {
                encode_sa_split_p(
                    &mut pass, c, q, cache, ixb, sink, attn, g.nh, g.hd, m, g.scale, kv_id, li,
                );
            } else {
                pass.set_pipeline(&c.sparse_attend);
                pass.set_bind_group(0, &sa_bind, &[]);
                pass.dispatch_workgroups(g.nh as u32, 1, 1);
            }
        }
        encode_rope_heads_p(
            &mut pass,
            c,
            attn,
            freq,
            posb,
            g.nh,
            g.hd,
            g.rd,
            false,
            true,
            (62, kv_id, li),
        );
        // Split apart: the pair measured 5.0 ms of a 30.6 ms chain against a
        // bandwidth floor near 0.7, and they are different kernels reading
        // differently-shaped weights.
        if !dsv4_skip("olora")
            && !dsv4_skip("oproj")
            && !(olora_mv4()
                && encode_o_lora_mv4_p(
                    &mut pass,
                    c,
                    &wb[2],
                    attn,
                    mid,
                    rows,
                    cols,
                    g.o_lora,
                    (65, kv_id, li),
                ))
        {
            pass.set_pipeline(o_pipe);
            pass.set_bind_group(0, &o_bind, &[]);
            pass.dispatch_workgroups((rows as u32).div_ceil(o_rows_per_wg).min(MAX_WG), 1, 1);
        }
        if !dsv4_skip("wob") && !dsv4_skip("oproj") {
            encode_q4tp_mvw_p(
                &mut pass,
                c,
                &wb[3],
                mid,
                out,
                g.dim,
                g.o_groups * g.o_lora,
                (64, kv_id, li),
            );
        }
        return;
    }
    if !q_ready {
        encode_q4tp_mv1(
            c,
            enc,
            &wb[1],
            qn,
            q,
            g.nh * g.hd,
            g.q_lora,
            (60, kv_id, li),
        );
    }
    encode_rope_heads(
        c,
        enc,
        q,
        freq,
        posb,
        g.nh,
        g.hd,
        g.rd,
        true,
        false,
        (61, kv_id, li),
    );
    encode_sparse_attend2(
        c,
        enc,
        q,
        cache,
        ixb,
        sink,
        attn,
        g.nh,
        g.hd,
        m,
        g.scale,
        Some((kv_id, li)),
    );
    encode_rope_heads(
        c,
        enc,
        attn,
        freq,
        posb,
        g.nh,
        g.hd,
        g.rd,
        false,
        true,
        (62, kv_id, li),
    );
    {
        let mut pass = begin_pass(enc);
        pass.set_pipeline(o_pipe);
        pass.set_bind_group(0, &o_bind, &[]);
        pass.dispatch_workgroups((rows as u32).div_ceil(o_rows_per_wg).min(MAX_WG), 1, 1);
    }
    encode_q4tp_mv1(
        c,
        enc,
        &wb[3],
        mid,
        out,
        g.dim,
        g.o_groups * g.o_lora,
        (64, kv_id, li),
    );
}

/// Route, then the chosen experts and the shared one.
#[allow(clippy::too_many_arguments)]
fn encode_moe_chain(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    logits: &wgpu::Buffer,
    x: &wgpu::Buffer,
    msel: &wgpu::Buffer,
    mwt: &wgpu::Buffer,
    mcnt: &wgpu::Buffer,
    mact: &wgpu::Buffer,
    out: &wgpu::Buffer,
    gate_all: &wgpu::Buffer,
    up_all: &wgpu::Buffer,
    down_all: &wgpu::Buffer,
    w: &Dsv4LayerW,
    g: Dsv4MoeGeom,
    n_pack: usize,
    n_route: usize,
    slots: usize,
    bkey: (u64, usize),
) {
    let mut pass = begin_pass(enc);
    encode_moe_chain_p(
        &mut pass, c, logits, x, msel, mwt, mcnt, mact, out, gate_all, up_all, down_all, w, g,
        n_pack, n_route, slots, bkey,
    );
}

#[allow(clippy::too_many_arguments)]
fn encode_moe_chain_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    logits: &wgpu::Buffer,
    x: &wgpu::Buffer,
    msel: &wgpu::Buffer,
    mwt: &wgpu::Buffer,
    mcnt: &wgpu::Buffer,
    mact: &wgpu::Buffer,
    out: &wgpu::Buffer,
    gate_all: &wgpu::Buffer,
    up_all: &wgpu::Buffer,
    down_all: &wgpu::Buffer,
    w: &Dsv4LayerW,
    g: Dsv4MoeGeom,
    n_pack: usize,
    n_route: usize,
    slots: usize,
    bkey: (u64, usize),
) {
    // `remap` also matters for a full hot-first permutation: every expert is
    // resident, but global router id N need not live in slot N.
    let subset = w.moe.remap.is_some();
    // The bias now lives in the PACK, whose address is stable for the life
    // of the process — so the const cache is sound for it, and each layer
    // gets its own device buffer. The per-call pool here was the many-layer
    // clobber: every queue write lands before the run's single submit.
    let bs = match w.moe.bias {
        Some(b) if b.len() >= n_route => const_buf(c, bytemuck::cast_slice(&b[..n_route])),
        _ => logits.clone(),
    };
    let mk = match w.moe.mask {
        Some(m) if m.len() >= n_route => const_buf(c, bytemuck::cast_slice(&m[..n_route])),
        _ => frame_buf(c, 17, n_route.max(1) * 4, true),
    };
    // Per-layer, not pooled: a run holding two hash layers wrote both lists
    // into one buffer before the single submit, and both routed with the
    // second one's experts.
    let fc = match w.moe.forced {
        Some(f) if f.len() >= g.top_k => {
            let v: Vec<u32> = f[..g.top_k].iter().map(|&i| i as u32).collect();
            store_slot(c, 18, bkey.0, bkey.1, bytemuck::cast_slice(&v))
        }
        _ => store_slot(c, 18, bkey.0, bkey.1, &vec![0u8; g.top_k * 4]),
    };
    let rflags = (w.moe.bias.is_some_and(|b| b.len() >= n_route) as u32)
        | ((w.moe.mask.is_some_and(|m| m.len() >= n_route) as u32) << 1)
        | ((w.moe.forced.is_some_and(|f| f.len() >= g.top_k) as u32) << 2)
        | 8
        | ((subset as u32) << 4)
        | (w.moe.global.map_or(n_pack as u32, |gl| gl.shared_slot) << 8);
    let rp = uniform_mixed(c, [n_route as u32, g.top_k as u32, rflags], g.route_scale);
    let stride16 = |rows: usize, cols: usize, q2: bool| -> u32 {
        let dt = if q2 {
            cortiq_core::TensorDtype::Q2TiledP
        } else {
            cortiq_core::TensorDtype::Q4TiledP
        };
        (cortiq_core::quant::expected_nbytes(dt, &[rows, cols]).unwrap_or(0) / 2) as u32
    };
    let gu_u = uniform_u32x8(
        c,
        [
            (g.hidden / 32) as u32,
            g.inter as u32,
            slots as u32,
            stride16(g.inter, g.hidden, g.gu_q2),
            g.swiglu_limit.to_bits(),
            0,
            0,
            0,
        ],
    );
    let dn_u = uniform_u32x4(
        c,
        [
            (g.inter / 32) as u32,
            g.hidden as u32,
            slots as u32,
            stride16(g.hidden, g.inter, false),
        ],
    );
    // Four rows to a workgroup where the layout allows it: the columns give
    // gpr = 128, so a row cannot use more than 64 lanes, and the only width
    // left is overlap between rows. CMF_DSV4_MOE4=0 reverts.
    let gu_r4 = g.inter % 4 == 0 && bt_gu_r4_on();
    let (p_gu, p_dn, _l_gu, _l_dn) = if g.gu_q2 {
        (
            if gu_r4 {
                // The token-axis kernel is also the one-token kernel at z=0:
                // four q2tp rows share each activation load instead of four
                // 64-lane subgroups rereading it independently.
                &c.bt_moe_gate_up_q2tp_r4
            } else if moe4() {
                &c.moe_gu_q2tp_m
            } else {
                &c.moe_gate_up_q2tp
            },
            if moe4() {
                &c.moe_dn_q4tp_m
            } else {
                &c.moe_down_q4tp
            },
            &c.layout_moe_gu_q2tp,
            &c.layout_moe_dn_q4tp,
        )
    } else {
        (
            if gu_r4 {
                &c.moe_gate_up_q4tp_b_r4
            } else {
                &c.moe_gate_up_q4tp_b
            },
            &c.moe_down_q4tp_b,
            &c.layout_moe_gu_b,
            &c.layout_moe_dn_b,
        )
    };
    // The live remap rides a per-(kv, layer) store slot: stable identity
    // for the cached bind group, fresh CONTENT every call — a layer stays
    // subset (or full) for the life of the process, so the identity never
    // flips under the cache.
    let rm = if subset {
        let r = w.moe.remap.unwrap();
        store_slot(c, 26, bkey.0, bkey.1, bytemuck::cast_slice(&r[..n_route]))
    } else {
        frame_buf(c, 26, n_pack.max(1) * 4, true)
    };
    // The cold list must survive a MERGED run: a pooled slot is written by
    // every layer of one submission and read once at the end, so each
    // (kv, layer) owns its own. The u32::MAX fill is the "no cold" marker
    // the reader checks per pair.
    let cold_slot = if subset {
        store_slot(
            c,
            27,
            bkey.0,
            bkey.1,
            bytemuck::cast_slice(&vec![u32::MAX; 4 * g.top_k]),
        )
    } else {
        frame_buf(c, 27, 4 * g.top_k * 4, false)
    };
    let bind_r = cached_bind(c, (127, bkey.0, bkey.1), || {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.moe_route.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, logits),
                bind_buf(1, &bs),
                bind_buf(2, &mk),
                bind_buf(3, &fc),
                bind_buf(4, msel),
                bind_buf(5, mwt),
                bind_buf(6, mcnt),
                bind_buf(7, &rp),
                bind_buf(8, &rm),
                bind_buf(9, &cold_slot),
            ],
        })
    });
    // The layer-frame/chain-of-one twin of `dsv4_moe_frame`'s global arm.
    // Routing, descriptor-indexed gate/up and descriptor-indexed down stay
    // inside THIS compute pass, so the common pool does not give back the
    // two barriers per layer that chain-of-one removed.
    if let Some(gl) = w.moe.global {
        let Some(gb) = c.dsv4_global_moe.lock().unwrap().get(&gl.pool_uid).cloned() else {
            return;
        };
        if gb.gu_q2 != g.gu_q2 {
            return;
        }
        let Some((p_gu, p_dn)) = dsv4_global_moe_pipelines(c, g.gu_q2, gb.segments) else {
            return;
        };
        let gu_gp = uniform_u32x8(
            c,
            [
                (g.hidden / 32) as u32,
                g.inter as u32,
                slots as u32,
                stride16(g.inter, g.hidden, g.gu_q2),
                g.swiglu_limit.to_bits(),
                gl.segment_slots,
                0,
                0,
            ],
        );
        let dn_gp = uniform_u32x8(
            c,
            [
                (g.inter / 32) as u32,
                g.hidden as u32,
                slots as u32,
                stride16(g.hidden, g.inter, false),
                gl.segment_slots,
                0,
                0,
                0,
            ],
        );
        let gate_bindings: Vec<_> = gb
            .gate
            .iter()
            .map(wgpu::Buffer::as_entire_buffer_binding)
            .collect();
        let up_bindings: Vec<_> = gb
            .up
            .iter()
            .map(wgpu::Buffer::as_entire_buffer_binding)
            .collect();
        let down_bindings: Vec<_> = gb
            .down
            .iter()
            .map(wgpu::Buffer::as_entire_buffer_binding)
            .collect();
        let bg_gu = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("dsv4-global-chain-gu"),
            layout: &p_gu.get_bind_group_layout(0),
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::BufferArray(&gate_bindings),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::BufferArray(&up_bindings),
                },
                bind_buf(2, x),
                bind_buf(3, msel),
                bind_buf(4, mact),
                bind_buf(5, mwt),
            ],
        });
        let bg_gu_p = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("dsv4-global-chain-gu-p"),
            layout: &p_gu.get_bind_group_layout(1),
            entries: &[bind_buf(0, &gu_gp)],
        });
        let bg_dn = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("dsv4-global-chain-dn"),
            layout: &p_dn.get_bind_group_layout(0),
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::BufferArray(&down_bindings),
                },
                bind_buf(1, mact),
                bind_buf(2, msel),
                bind_buf(3, mwt),
                bind_buf(4, out),
            ],
        });
        let bg_dn_p = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("dsv4-global-chain-dn-p"),
            layout: &p_dn.get_bind_group_layout(1),
            entries: &[bind_buf(0, &dn_gp)],
        });
        if !dsv4_skip("route") {
            pass.set_pipeline(&c.moe_route);
            pass.set_bind_group(0, &bind_r, &[]);
            pass.dispatch_workgroups(1, 1, 1);
        }
        if !dsv4_skip("gu") {
            pass.set_pipeline(p_gu);
            pass.set_bind_group(0, &bg_gu, &[]);
            pass.set_bind_group(1, &bg_gu_p, &[]);
            pass.dispatch_workgroups(g.inter as u32, slots as u32, 1);
        }
        if !dsv4_skip("dn") {
            pass.set_pipeline(p_dn);
            pass.set_bind_group(0, &bg_dn, &[]);
            pass.set_bind_group(1, &bg_dn_p, &[]);
            pass.dispatch_workgroups(g.hidden as u32, 1, 1);
        }
        return;
    }
    // The layout comes from the PIPELINE, not the cached one: wgpu treats an
    // auto-derived layout as exclusive to the pipeline that produced it, so a
    // group built against a twin's layout is rejected at dispatch — with a
    // message about "exclusive pipelines", not about layouts.
    let bg_gu = cached_bind(c, (128, bkey.0, bkey.1), || {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &p_gu.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, gate_all),
                bind_buf(1, up_all),
                bind_buf(2, x),
                bind_buf(3, msel),
                bind_buf(4, mact),
                bind_buf(5, &gu_u),
            ],
        })
    });
    let bg_dn = cached_bind(c, (129, bkey.0, bkey.1), || {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &p_dn.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, down_all),
                bind_buf(1, mact),
                bind_buf(2, msel),
                bind_buf(3, mwt),
                bind_buf(4, out),
                bind_buf(5, &dn_u),
            ],
        })
    });
    // Three separately-skippable parts: the router ranks 256 experts in one
    // workgroup, the gate/up pair streams the experts' weights, and the down
    // projection streams them again. The MoE measured 5.5 ms of a 30.6 ms
    // chain against a weight-bandwidth floor near 2.5, and which of the
    // three owns that gap is not something to reason about.
    if !dsv4_skip("route") {
        pass.set_pipeline(&c.moe_route);
        pass.set_bind_group(0, &bind_r, &[]);
        pass.dispatch_workgroups(1, 1, 1);
    }
    if !dsv4_skip("gu") {
        pass.set_pipeline(p_gu);
        pass.set_bind_group(0, &bg_gu, &[]);
        let gu_per_wg = if gu_r4 || (g.gu_q2 && moe4()) {
            4u32
        } else {
            1u32
        };
        pass.dispatch_workgroups((g.inter as u32).div_ceil(gu_per_wg), slots as u32, 1);
    }
    if !dsv4_skip("dn") {
        pass.set_pipeline(p_dn);
        pass.set_bind_group(0, &bg_dn, &[]);
        pass.dispatch_workgroups(g.hidden as u32, 1, 1);
    }
}

/// Everything one DeepSeek-V4 layer does, in ONE submission.
///
/// The two frames before this cost two barriers a layer — 30 ms of a 76 ms
/// token, spent waiting rather than computing. Fusing them means the
/// hyper-connection glue between the halves has to run on the device too,
/// which is what `hc_pre_fold` (mixes, Sinkhorn and the fold in one kernel)
/// and `hc_post_expand` were built for.
///
/// The frame is shifted by one on purpose: it ENDS by folding and norming the
/// state for the NEXT layer's attention half and projecting its LoRA vector,
/// then reads back that normed hidden. The host needs exactly that one vector
/// — for the kv projection, the compressor and the indexer, which still live
/// there — and nothing else. Layer zero's opening fold is done on the host
/// once, which costs nothing at all.
#[derive(Clone)]
pub struct Dsv4LayerW<'a> {
    pub attn: Dsv4AttnW<'a>,
    pub moe: Dsv4MoeW<'a>,
    /// Hyper-connection projection of the FFN half: `[mix_hc, hc*dim]` f32.
    pub hc_ffn_fn: &'a [f32],
    pub hc_ffn_scale: &'a [f32; 3],
    pub hc_ffn_base: &'a [f32],
    /// The same for the NEXT layer's attention half — absent on the last.
    pub hc_next_fn: Option<&'a [f32]>,
    pub hc_next_scale: &'a [f32; 3],
    pub hc_next_base: &'a [f32],
    pub ffn_norm: &'a [f32],
    /// The next layer's input norm and its q_norm, for the tail that
    /// prepares the following frame.
    pub next_norm: &'a [f32],
    pub next_q_norm: &'a [f32],
    /// The next layer's wq_a, by directory index.
    pub next_wq_a: Option<usize>,
    /// Router logits weight, f32 `[n_exp, dim]`.
    pub router: &'a [f32],
}

#[derive(Clone, Copy)]
pub struct Dsv4LayerGeom {
    pub attn: Dsv4AttnGeom,
    pub moe: Dsv4MoeGeom,
    pub hc: usize,
    pub hc_eps: f32,
    pub sinkhorn_iters: usize,
}

/// Returns the next layer's normed hidden in `folded_next` (or this layer's
/// state contribution when there is no next layer).
#[allow(clippy::too_many_arguments)]
/// Encode one layer's whole preparation — window, compressor, the indexer's
/// compressor, the indexer — into the caller's encoder, and return the
/// attended list's length.
///
/// The order matters and matches the host's: the compressors see the token
/// BEFORE the window append changes what `filled` means, and the indexer
/// scores against the cache its own compressor has just extended.
#[allow(clippy::too_many_arguments)]
pub fn dsv4_encode_prep(
    model: &Arc<CmfModel>,
    p: &Dsv4Prep,
    kv_id: u64,
    li: usize,
    hidden: &wgpu::Buffer,
    qn: &wgpu::Buffer,
    cache: &wgpu::Buffer,
    idx_out: &wgpu::Buffer,
    hd: usize,
    window: usize,
    dim: usize,
    rope_dim: usize,
    eps: f32,
    pos: usize,
    inv_freq: &[f32],
    enc: &mut wgpu::CommandEncoder,
) -> Option<usize> {
    let c = ctx()?;
    // The compressors first: both read this token's hidden state, and the
    // attention one appends into the same cache buffer the window lives in,
    // past the window's capacity.
    let mut n_comp = p.n_comp;
    if let Some((cw, cg)) = &p.comp {
        if !dsv4_skip("comp")
            && dsv4_compressor_frame(
                model,
                cw,
                *cg,
                0,
                kv_id,
                li,
                hidden,
                pos,
                inv_freq,
                cache,
                p.comp_dst_off,
                enc,
            )
            .is_some()
        {
            n_comp += 1;
        }
    }
    let mut n_ix = p.n_ix;
    let mut m = None;
    if let Some((iw, ig, ixw, ixg)) = &p.ix {
        let ixkv = dsv4_index_cache(kv_id, li, ixg.idim * (n_ix + 1))?;
        if !dsv4_skip("comp")
            && dsv4_compressor_frame(
                model,
                iw,
                *ig,
                1,
                kv_id,
                li,
                hidden,
                pos,
                inv_freq,
                &ixkv,
                p.ix_dst_off,
                enc,
            )
            .is_some()
        {
            n_ix += 1;
        }
        // Nothing compressed yet means nothing to score: attention sees the
        // window alone, which is what the host does at the start of a
        // sequence too.
        if n_comp > 0 && n_ix > 0 {
            m = if dsv4_skip("ix") {
                None
            } else {
                dsv4_indexer_frame(
                    model,
                    ixw,
                    *ixg,
                    kv_id,
                    li,
                    hidden,
                    qn,
                    &ixkv,
                    n_ix,
                    n_comp,
                    // The window AFTER this token's append, which is what the
                    // host reads off its own vector: clamp the incremented
                    // count, not the count. `filled.min(window) + 1` is right
                    // until the window fills and then hands attention one
                    // position more than the window holds, for the rest of the
                    // sequence.
                    (p.filled + 1).min(window),
                    pos,
                    inv_freq,
                    idx_out,
                    enc,
                )
            };
        }
    }
    if dsv4_skip("win") {
        return Some(m.unwrap_or(p.filled.min(window)));
    }
    let filled = dsv4_window_append(
        model, p.wkv, p.kv_norm, hidden, cache, hd, window, p.filled, dim, rope_dim, eps, pos,
        inv_freq, kv_id, li, enc,
    )?;
    match m {
        Some(m) => Some(m),
        // No indexer, or nothing to score: the list is the window, and every
        // compressed position after it when there is no indexer to choose.
        None => {
            let k = if p.ix.is_none() { n_comp } else { 0 };
            let pick: Vec<u32> = (0..k as u32).collect();
            let pb = if pick.is_empty() {
                frame_buf(c, 111, 4, true)
            } else {
                frame_up(c, 111, bytemuck::cast_slice(&pick))
            };
            encode_idx_build(c, enc, &pb, idx_out, filled, window, k, None);
            Some(filled + k)
        }
    }
}

/// The indexer's own compressed cache, on the device and grown as the
/// sequence does — the attention cache's twin, kept apart because the two
/// have different widths.
fn dsv4_index_cache(kv_id: u64, li: usize, floats: usize) -> Option<wgpu::Buffer> {
    let c = ctx()?;
    let cap = floats.next_power_of_two().max(1024);
    let mut m = c.dsv4_ixkv.lock().unwrap();
    // Growing CARRIES the contents. This cache accumulates one compressed
    // entry every `ratio` tokens for the whole sequence, and dropping the old
    // buffer meant the indexer scored every earlier entry against zeros from
    // the next power of two onward — silently, and only on contexts long
    // enough to cross one. The KV cache beside it has always copied; this one
    // did not.
    let carry = match m.get(&(kv_id, li)) {
        Some((old, have)) if *have < cap => Some((old.clone(), *have)),
        _ => None,
    };
    if let Some((old, have)) = carry {
        let bigger = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("dsv4-index-kv"),
            size: (cap * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_SRC
                | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("ixkv-grow"),
            });
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&old, 0, &bigger, 0, (have * 4) as u64);
        submit(c, finish_enc(enc));
        m.insert((kv_id, li), (bigger, cap));
        // The same epoch the KV cache bumps: a cached bind group outlives
        // the buffer it points at otherwise.
        GREW.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }
    let (b, _) = m.entry((kv_id, li)).or_insert_with(|| {
        (
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("dsv4-index-kv"),
                size: (cap * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_SRC
                    | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }),
            cap,
        )
    });
    Some(b.clone())
}

/// Everything a layer needs to build its OWN attention inputs on the card:
/// the window append, the compressor (and the indexer's, and the indexer) —
/// the three producers that used to run on the host between submissions.
///
/// With this present a layer frame needs nothing from the host but the
/// position, so its result never has to come back: `folded_next` stays on
/// the device and the NEXT layer reads it there. That is the whole of step
/// four — the readback at the end of this function was there for the host's
/// prep and for nothing else.
#[derive(Clone)]
pub struct Dsv4Prep<'a> {
    /// The KV projection that feeds the sliding window, and its norm.
    pub wkv: usize,
    pub kv_norm: &'a [f32],
    /// The attention compressor; `None` on the pure sliding-window layers.
    pub comp: Option<(Dsv4CompW<'a>, Dsv4CompGeom)>,
    /// The indexer: its own compressor, then the scoring half.
    pub ix: Option<(Dsv4CompW<'a>, Dsv4CompGeom, Dsv4IxW, Dsv4IxGeom)>,
    /// Cache bookkeeping the host keeps (all of it is derivable from the
    /// position, so none of it costs a readback): how much of the window is
    /// filled BEFORE this token, how many compressed entries each cache
    /// holds, and where the next one goes.
    pub filled: usize,
    /// The window's CAPACITY in slots — where the compressed region starts.
    pub window: usize,
    pub n_comp: usize,
    pub n_ix: usize,
    pub comp_dst_off: usize,
    pub ix_dst_off: usize,
    /// The most positions attention can be asked to read: window capacity
    /// plus the indexer's budget (or every compressed entry, without one).
    /// The list buffer is sized by THIS — sizing it by the host list, which
    /// is empty when the device builds its own, cut the release's 640-entry
    /// list to 64: the writes past that are silently clamped and the reads
    /// return zero, so attention quietly stares at position zero.
    pub idx_cap: usize,
}

#[allow(clippy::too_many_arguments)]
/// One layer, encoded into the CALLER'S encoder and never submitted here.
///
/// `prep` present means the layer builds its own attention inputs on the
/// card and `idxs` is ignored; absent means the host supplied the list, as
/// before. Either way nothing comes back: the folded state for the next
/// layer is left in the frame's own buffer, which is where the next call
/// looks for it. That is what lets a whole token be one submission.
#[allow(clippy::too_many_arguments)]
/// `CMF_DSV4_SKIP=attn|moe|prep|hc` — TIMING ONLY, the answer is garbage.
/// Drops a stage's dispatches while leaving every buffer, pass and shape in
/// place, so the delta in tok/s is that stage's real share of the token.
/// Neither dispatch counting nor pass counting predicted the frame
/// correctly on the qwen graph; there is no reason to trust them here.
/// `CMF_DSV4_HCFUSE=1` runs the fused hyper-connection join — four steps and
/// a Sinkhorn in ONE workgroup instead of four dispatches.
///
/// OFF by default, because it was MEASURED and it loses. The arithmetic
/// argument was right and the conclusion was wrong: 336 launches a token is
/// real, but so is the mix projection — 24 outputs over hc·dim = 16 384
/// inputs, 1.5 MB of weights — and putting that through one workgroup uses
/// one SM of two hundred. On the release the chain's wait went 51.6 → 84.5
/// ms and decode 16.5 → 10.6 tok/s. Kept, and kept switchable, because the
/// same kernel wins wherever launches dominate and the mix is small.
fn hc_fuse() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| std::env::var("CMF_DSV4_HCFUSE").is_ok_and(|v| v != "0"))
}

fn dsv4_skip(what: &str) -> bool {
    static S: std::sync::OnceLock<String> = std::sync::OnceLock::new();
    S.get_or_init(|| std::env::var("CMF_DSV4_SKIP").unwrap_or_default())
        .contains(what)
}

/// How many tokens of one batch a layer's keys can hold apart. Generous:
/// the keys are only cache discriminators, and collisions between layers
/// would be far more expensive to debug than a sparse map is to keep.
const FRAME_TOK_STRIDE: usize = 64;

thread_local! {
    // A batch records several calls before one submit. Every pooled upload,
    // mutable uniform and cached bind group inside a frame must therefore
    // have a token identity; otherwise the final queue.write_buffer wins for
    // all of them. Persistent model/cache buffers do NOT use this salt.
    static DSV4_FRAME_SALT: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
}

struct Dsv4FrameSalt(usize);

impl Dsv4FrameSalt {
    fn enter(salt: usize) -> Self {
        let old = DSV4_FRAME_SALT.with(|s| s.replace(salt));
        Self(old)
    }
}

impl Drop for Dsv4FrameSalt {
    fn drop(&mut self) {
        DSV4_FRAME_SALT.with(|s| s.set(self.0));
    }
}

#[inline]
fn dsv4_frame_salt() -> usize {
    DSV4_FRAME_SALT.with(std::cell::Cell::get)
}

#[inline]
fn dsv4_salted_li(li: usize) -> usize {
    let salt = dsv4_frame_salt();
    if salt == 0 {
        li
    } else {
        li + salt * 1_000_000
    }
}

fn dsv4_layer_frame_enc(
    model: &Arc<CmfModel>,
    w: &Dsv4LayerW,
    g: Dsv4LayerGeom,
    kv_id: u64,
    li: usize,
    // Which token of the batch this frame encodes. The KV cache is the
    // layer's, so it keys on `li`; everything else — uniform slots, cached
    // bind groups, the scratch that carries this token's state between
    // dispatches — has to be the TOKEN's, or two tokens encoded into one
    // submission share a slot and the last write decides for both. That is
    // silent and it is wrong: their rope positions differ.
    tok: usize,
    batch_frame: bool,
    q_ready: bool,
    defer_next_q: bool,
    qn: Option<&[f32]>,
    idxs: &[u32],
    prep: Option<&Dsv4Prep>,
    inv_freq: &[f32],
    pos: usize,
    enc: &mut wgpu::CommandEncoder,
) -> Option<wgpu::Buffer> {
    // Salt 0 preserves every single-token cache key. Batch token zero must
    // still differ from it, hence the one-based value.
    let _frame_salt = Dsv4FrameSalt::enter(if batch_frame { tok + 1 } else { 0 });
    macro_rules! no {
        ($($t:tt)*) => {{
            if std::env::var("CMF_DSV4_FRAME_DEBUG").is_ok() {
                eprintln!("кадр слоя отклонён: {}", format_args!($($t)*));
            }
            return None;
        }};
    }
    let Some(c) = ctx() else {
        no!("нет контекста wgpu")
    };
    // Keys are per (layer, token); the cache lookup below stays per layer.
    let lk = li * FRAME_TOK_STRIDE + tok;
    let a = g.attn;
    let m = g.moe;
    let (hc, dim) = (g.hc, a.dim);
    let mix_hc = (2 + hc) * hc;
    if qn.is_some_and(|v| v.len() < a.q_lora)
        || (prep.is_none() && (idxs.is_empty() || idxs.len() > 1024))
    {
        no!("формы: idx {}", idxs.len());
    }
    if w.hc_ffn_fn.len() < mix_hc * hc * dim || w.router.len() < m.hidden {
        no!("гипер-связи или роутер не той формы");
    }

    // ── weights ──
    let bytes = model.primary_bytes();
    let mut wb = Vec::with_capacity(5);
    for &idx in &[
        w.attn.wq_a,
        w.attn.wq_b,
        w.attn.wo_a,
        w.attn.wo_b,
        w.next_wq_a.unwrap_or(w.attn.wq_a),
    ] {
        let Some(e) = model.tensors.get(idx) else {
            no!("тензора {idx} нет");
        };
        if e.dtype != cortiq_core::TensorDtype::Q4TiledP {
            no!("{} не q4tp", e.name);
        }
        let (Some(abs), plen) = (model.entry_abs_offset(e), e.nbytes as usize) else {
            no!("{} без смещения", e.name);
        };
        let Some(b) = weight_buffer_l(
            c,
            (model.uid() as usize, idx),
            &bytes[abs..abs + plen],
            layer_of_name(&model.tensors[idx].name),
        ) else {
            no!("{} не влез в VRAM", e.name);
        };
        wb.push(b);
    }
    let local_moe = if w.moe.global.is_none() {
        let Some(v) = moe_expert_bufs(
            c,
            model,
            w.moe.experts,
            m.inter,
            m.hidden,
            true,
            m.gu_q2,
            false,
        ) else {
            no!("эксперты не влезли в VRAM");
        };
        Some(v)
    } else {
        None
    };
    // The global branch in `encode_moe_chain_p` ignores these arguments;
    // a real four-byte buffer keeps the common call shape without allocating
    // a second expert bank.
    let moe_dummy = frame_buf(c, 90, 4, true);
    let (gate_all, up_all, down_all) =
        local_moe.unwrap_or_else(|| (moe_dummy.clone(), moe_dummy.clone(), moe_dummy.clone()));
    let cache = {
        let map = c.dsv4_kv.lock().unwrap();
        match map.get(&(kv_id, li)) {
            Some((b, _)) => b.clone(),
            None => no!("кеш ({kv_id}, {li}) не заведён"),
        }
    };
    // The hyper-connection state lives on the card for the whole token; the
    // host seeds it once at layer zero.
    let state = frame_buf_t(c, 40, tok, hc * dim * 4, true);

    // ── constants (model-owned, address keying is sound) ──
    let _qnw = const_buf(c, bytemuck::cast_slice(&w.attn.q_norm[..a.q_lora]));
    let sink = const_buf(c, bytemuck::cast_slice(&w.attn.sink[..a.nh]));
    let freq = const_buf(c, bytemuck::cast_slice(&inv_freq[..a.rd / 2]));
    let ffn_fn = const_buf(c, bytemuck::cast_slice(w.hc_ffn_fn));
    let ffn_sc = const_buf(c, bytemuck::cast_slice(w.hc_ffn_scale));
    let ffn_bs = const_buf(c, bytemuck::cast_slice(&w.hc_ffn_base[..mix_hc]));
    let ffn_nw = const_buf(c, bytemuck::cast_slice(&w.ffn_norm[..dim]));
    let n_exp = w.moe.experts.len().saturating_sub(1);
    // Routing width. A subset pack (live remap present) ranks over EVERY
    // expert and the remap turns winners into slots or cold picks — the
    // same contract the two-frame path runs; a full pack keeps the packed
    // width and the remap stays a dummy.
    let n_route = w.moe.remap.map_or(n_exp, |r| r.len().max(n_exp));
    if w.router.len() < m.hidden * n_route {
        no!("роутер короче {} × {}", n_route, m.hidden);
    }
    let router = const_buf(c, bytemuck::cast_slice(&w.router[..m.hidden * n_route]));
    let next_nw = const_buf(c, bytemuck::cast_slice(&w.next_norm[..dim]));
    let next_qn = const_buf(c, bytemuck::cast_slice(&w.next_q_norm[..a.q_lora]));

    // ── per-call uploads ──
    let posb = frame_up_pos_t(c, 1, tok, pos, a.eps);
    let ixb = {
        // Flat 1024 — the attention list's hard ceiling (the m > 1024
        // guard). 4 KB buys a PERMANENT identity: frame_buf keys on
        // (tag, len), so a growing capacity would mint a new buffer while
        // every cached group kept the old one, with no GREW bump to save it.
        let b = frame_buf_t(c, 2, tok, 1024 * 4, true);
        if !idxs.is_empty() {
            c.queue.write_buffer(&b, 0, bytemuck::cast_slice(idxs));
        }
        b
    };
    let qnb = match qn {
        Some(v) => frame_up(c, 4, bytemuck::cast_slice(&v[..a.q_lora])),
        None => frame_buf_t(c, 4, tok, a.q_lora * 4, true),
    };

    // ── working buffers ──
    let n_pack = w
        .moe
        .global
        .map_or(n_exp, |_| w.moe.remap.map_or(0, <[u32]>::len));
    let slots = m.top_k + 1;
    let q = frame_buf(c, 5, a.nh * a.hd * 4, false);
    let attn = frame_buf(c, 6, a.nh * a.hd * 4, false);
    let mid = frame_buf(c, 7, a.o_groups * a.o_lora * 4, false);
    let ao = frame_buf(c, 8, dim * 4, false);
    let mixes = frame_buf(c, 41, mix_hc * 4, false);
    // `true`: the tail copies the next layer's input INTO this one, so it is
    // a copy destination as well as a kernel output. Without the flag the
    // layer-frame path fails validation at the first token — which is how
    // the release run caught it, and a unit test could not have.
    let folded = frame_buf_t(c, 42, tok, dim * 4, true);
    let hpost = frame_buf_t(c, 43, tok, hc * 4, true);
    let hcomb = frame_buf_t(c, 44, tok, hc * hc * 4, true);
    let x2 = frame_buf_t(c, 45, tok, dim * 4, true);
    let state2 = frame_buf(c, 46, hc * dim * 4, false);
    let logit_b = frame_buf(c, 47, n_route * 4, false);
    let msel = frame_buf(c, 19, slots * 4, false);
    let mwt = frame_buf(c, 20, slots * 4, false);
    let mcnt = frame_buf(c, 21, 4, false);
    let mact = frame_buf(c, 22, slots * m.inter * 4, false);
    let mo = frame_buf(c, 24, dim * 4, false);
    let qr2 = frame_buf(c, 48, a.q_lora * 4, false);
    let qn2 = frame_buf(c, 49, a.q_lora * 4, false);

    // Eight words now: the fifth says whether the fold also norms.
    let hcp = uniform_u32x8(
        c,
        [
            hc as u32,
            dim as u32,
            g.sinkhorn_iters as u32,
            g.hc_eps.to_bits(),
            0,
            0,
            0,
            0,
        ],
    );
    let hcp_n = uniform_u32x8(
        c,
        [
            hc as u32,
            dim as u32,
            g.sinkhorn_iters as u32,
            g.hc_eps.to_bits(),
            1,
            0,
            0,
            0,
        ],
    );

    // ── the layer's own preparation, when it owns it ──
    // `folded` is the previous frame's output and this token's hidden state;
    // for layer zero the caller seeded it.
    let m_attend = match prep {
        None => idxs.len(),
        Some(p) => {
            if dsv4_skip("prep") {
                return None;
            }
            let Some(n) = dsv4_encode_prep(
                // `x2` and NOT `folded`: the previous frame's tail leaves
                // the next layer's NORMED input there, which is exactly what
                // the host used to hand in. `folded` is a mid-frame scratch
                // and holds nothing yet at this point.
                model, p, kv_id, li, &x2, &qnb, &cache, &ixb, a.hd, p.window, dim, a.rd, a.eps, pos,
                inv_freq, enc,
            ) else {
                no!("подготовка слоя не собралась");
            };
            if n == 0 || n > 1024 {
                no!("список позиций длиной {n}");
            }
            n
        }
    };

    // ── attention half (the fold for it was prepared by the previous frame) ──
    if !dsv4_skip("attn") {
        encode_attn_chain(
            c, enc, &wb, &qnb, &q, &attn, &mid, &ao, &cache, &ixb, &sink, &freq, &posb, a, kv_id,
            li, m_attend, q_ready,
        );
    }

    // ── glue: expand, then the FFN half's fold and norm ──
    // Everything from the attention output to the next layer's input is one
    // chain of dependent dispatches with no copy in the middle, so it is ONE
    // pass. Twelve passes' worth of driver bookkeeping a layer went here,
    // and on a small layer that bookkeeping IS the token.
    let mut copy_qn = false;
    {
        let mut pass = begin_pass(enc);
        // Four dispatches — expand, mix, fold, norm — in one. `hc_fuse()`
        // reverts to the four for a bisect.
        if hc_fuse() {
            encode_hc_block_p(
                &mut pass,
                c,
                &ao,
                &state,
                &hpost,
                &hcomb,
                &ffn_fn,
                &ffn_sc,
                &ffn_bs,
                &ffn_nw,
                &state2,
                &folded,
                &x2,
                hc,
                dim,
                mix_hc,
                g.sinkhorn_iters,
                a.eps,
                (172, kv_id, lk),
            );
        } else {
            if !dsv4_skip("hc") {
                if !dsv4_skip("hcexp") {
                    encode_hc_expand_k_p(
                        &mut pass,
                        c,
                        &ao,
                        &state,
                        &hpost,
                        &hcomb,
                        &state2,
                        &hcp,
                        hc,
                        dim,
                        (120, kv_id, lk),
                    );
                }
                if !dsv4_skip("hcmix") {
                    encode_f32matvec_w_p(
                        &mut pass,
                        c,
                        &ffn_fn,
                        &state2,
                        &mixes,
                        mix_hc,
                        hc * dim,
                        (121, kv_id, lk),
                    );
                }
                // The fold norms too: one dispatch, not two.
                if !dsv4_skip("hcfold") {
                    encode_hc_fold_k_p(
                        &mut pass,
                        c,
                        &state2,
                        &mixes,
                        &ffn_sc,
                        &ffn_bs,
                        &folded,
                        &hpost,
                        &hcomb,
                        &hcp_n,
                        Some((&ffn_nw, &x2)),
                        (122, kv_id, lk),
                    );
                }
            }
        }

        // ── MoE half ──
        encode_f32matvec_w_p(
            &mut pass,
            c,
            &router,
            &x2,
            &logit_b,
            n_route,
            m.hidden,
            (123, kv_id, lk),
        );
        if !dsv4_skip("moe") {
            encode_moe_chain_p(
                &mut pass,
                c,
                &logit_b,
                &x2,
                &msel,
                &mwt,
                &mcnt,
                &mact,
                &mo,
                &gate_all,
                &up_all,
                &down_all,
                w,
                m,
                n_pack,
                n_route,
                slots,
                (kv_id, li),
            );
        }

        // ── expand, then prepare the NEXT layer ──
        // A subset pack may return cold winners whose exact correction
        // (`state[j] += post[j] * cold`) needs THIS layer's post — and the
        // next-layer fold below rewrites the canonical slot. A blit the
        // width of the hyper-connection count preserves it; the correction
        // variant (`dsv4_state_add_cold_preserved`) reads the spare.
        if w.moe.remap.is_some() {
            let spare = store_slot(c, 48, kv_id, li, &vec![0u8; hc * 4]);
            encode_blit_p(&mut pass, c, &hpost, &spare, hc, 0, 0, None);
        }
        let fused_next = hc_fuse() && w.hc_next_fn.is_some();
        if !fused_next {
            encode_hc_expand_k_p(
                &mut pass,
                c,
                &mo,
                &state2,
                &hpost,
                &hcomb,
                &state,
                &hcp,
                hc,
                dim,
                (124, kv_id, lk),
            );
        }
        if let Some(nf) = w.hc_next_fn {
            let nfn = const_buf(c, bytemuck::cast_slice(nf));
            let nsc = const_buf(c, bytemuck::cast_slice(w.hc_next_scale));
            let nbs = const_buf(c, bytemuck::cast_slice(&w.hc_next_base[..mix_hc]));
            if fused_next {
                // The same four steps as the FFN half, for the next layer's
                // attention: expand the MoE output into the state, mix, fold,
                // norm. The expand above is folded in here, which is why it
                // is skipped when this branch runs.
                encode_hc_block_p(
                    &mut pass,
                    c,
                    &mo,
                    &state2,
                    &hpost,
                    &hcomb,
                    &nfn,
                    &nsc,
                    &nbs,
                    &next_nw,
                    &state,
                    &folded,
                    &x2,
                    hc,
                    dim,
                    mix_hc,
                    g.sinkhorn_iters,
                    a.eps,
                    (174, kv_id, lk),
                );
            } else {
                encode_f32matvec_w_p(
                    &mut pass,
                    c,
                    &nfn,
                    &state,
                    &mixes,
                    mix_hc,
                    hc * dim,
                    (125, kv_id, lk),
                );
                encode_hc_fold_k_p(
                    &mut pass,
                    c,
                    &state,
                    &mixes,
                    &nsc,
                    &nbs,
                    &folded,
                    &hpost,
                    &hcomb,
                    &hcp_n,
                    Some((&next_nw, &x2)),
                    (126, kv_id, lk),
                );
            }
            // The next layer's LoRA vector, but only when the host is not
            // going to hand it over anyway — the indexer needs `qr` there, so
            // today it projects it regardless and computing it twice is waste.
            if qn.is_none() && !defer_next_q && !dsv4_skip("nextq") {
                encode_q4tp_mvw_p(
                    &mut pass,
                    c,
                    &wb[4],
                    &x2,
                    &qr2,
                    a.q_lora,
                    dim,
                    (52, kv_id, lk),
                );
                encode_rmsnorm_p(
                    &mut pass,
                    c,
                    &qr2,
                    &next_qn,
                    &qn2,
                    a.q_lora,
                    a.eps,
                    (53, kv_id, lk),
                );
                copy_qn = true;
            }
        }
    }
    // Outside the pass: a buffer-to-buffer copy cannot be recorded inside one.
    if copy_qn {
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&qn2, 0, &qnb, 0, (a.q_lora * 4) as u64);
    }

    // The next layer's input. It stays here: `folded` is where the next
    // frame reads its hidden state from, so a chain of layers needs one
    // copy and no round trip. The wrapper below reads it when a caller
    // still wants the value on the host.
    let src = if w.hc_next_fn.is_some() { &x2 } else { &ao };
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(src, 0, &folded, 0, (dim * 4) as u64);
    Some(folded)
}

/// Make sure a layer's cache exists and is at least `cap` floats.
///
/// The host path got this for free: it rewrote the whole cache every token
/// through `dsv4_cache_write`, which creates and grows. A chained layer
/// never calls that — it appends on the card — so without this its buffer is
/// whatever some earlier token happened to size it to, and the compressed
/// region walks off the end as the sequence lengthens.
pub fn dsv4_cache_ensure(kv_id: u64, li: usize, cap: usize) -> bool {
    let Some(c) = ctx() else { return false };
    if (cap * 4) as u64 > c.device.limits().max_storage_buffer_binding_size {
        return false;
    }
    let mut map = c.dsv4_kv.lock().unwrap();
    // Growing must CARRY the contents. The host path rewrites the whole
    // cache every token, so drop-and-recreate cost it nothing; the chain
    // owns the contents on the device, and dropping the buffer there wipes
    // the window and every compressed entry the sequence has accumulated —
    // a drift that grows with length and never fails loudly.
    if let Some((old, have)) = map.get(&(kv_id, li)).map(|(b, h)| (b.clone(), *h)) {
        if have < cap {
            let bigger = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("dsv4-kv"),
                size: (cap * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_DST
                    | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });
            let mut enc = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("dsv4-kv-grow"),
                });
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(&old, 0, &bigger, 0, (have * 4) as u64);
            submit(c, finish_enc(enc));
            map.insert((kv_id, li), (bigger, cap));
            GREW.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        }
        return true;
    }
    map.insert(
        (kv_id, li),
        (
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("dsv4-kv"),
                size: (cap * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_DST
                    | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            }),
            cap,
        ),
    );
    true
}

/// Seed the buffer a chain's FIRST layer reads its hidden state from.
///
/// Every later layer finds it there because the previous frame's tail wrote
/// it. Layer zero has no previous frame — the same hole that once put
/// garbage into `post`/`comb` and cost a perplexity of 1470. Seed it, or the
/// chain starts on whatever the last token left behind.
/// Seed only the fold: the device's qn is already this layer's.
pub fn dsv4_chain_seed_fold(x: &[f32]) -> bool {
    let Some(c) = ctx() else { return false };
    let b = frame_buf(c, 45, x.len() * 4, true);
    c.queue.write_buffer(&b, 0, bytemuck::cast_slice(x));
    true
}

/// Seed one token of a batch: its opening fold and its LoRA vector.
pub fn dsv4_chain_seed_t(x: &[f32], qn: &[f32], tok: usize) -> bool {
    let Some(c) = ctx() else { return false };
    let b = frame_buf_t(c, 45, tok, x.len() * 4, true);
    c.queue.write_buffer(&b, 0, bytemuck::cast_slice(x));
    let q = frame_buf_t(c, 4, tok, qn.len() * 4, true);
    c.queue.write_buffer(&q, 0, bytemuck::cast_slice(qn));
    true
}

pub fn dsv4_chain_seed(x: &[f32], qn: &[f32]) -> bool {
    let Some(c) = ctx() else { return false };
    let b = frame_buf(c, 45, x.len() * 4, true);
    c.queue.write_buffer(&b, 0, bytemuck::cast_slice(x));
    // The LoRA vector too: every frame leaves the NEXT layer's there, and
    // layer zero has no frame before it. Passing it as `qn` instead would
    // stop that frame computing the one after — the tail is guarded on
    // `qn.is_none()`.
    let q = frame_buf(c, 4, qn.len() * 4, true);
    c.queue.write_buffer(&q, 0, bytemuck::cast_slice(qn));
    true
}

/// Encode a run of consecutive layers into ONE encoder and submit it once,
/// returning the last layer's folded output on the host.
///
/// This is the point of the whole exercise: 43 layers used to cost 86
/// submissions and 43 round trips, because the host had to see each layer's
/// output to prepare the next one's attention inputs. It does not any more.
#[allow(clippy::too_many_arguments)]
/// A whole BATCH of known tokens through a run of layers, in one submission.
///
/// The layer is the outer loop and the token the inner one, which is the
/// only ordering that works: token t's attention has to see the window and
/// the compressed entries that tokens before it in the same batch just
/// wrote, and a compute pass orders its dispatches, so encoding them in that
/// order is enough — no fence, no readback between them.
///
/// Nothing here is coupled through CONTENT. A token's prep is the layer's
/// static weights plus counts, and the counts follow from the position: the
/// window fills by one a token to its capacity, and a compressed entry
/// appears every `ratio` tokens. So the whole batch's preps are known before
/// the first dispatch is encoded — which is what lets this be one
/// submission rather than B of them.
///
/// `folded_out` takes `batch * dim`: each token's fold for the head, in
/// order. The head is the caller's business — it is one q4tp matvec with a
/// batch, which the pair already does.
#[allow(clippy::too_many_arguments)]
pub fn dsv4_chain_batch(
    model: &Arc<CmfModel>,
    layers: &[(Dsv4LayerW<'_>, Dsv4LayerGeom, Dsv4Prep<'_>)],
    kv_id: u64,
    first_li: usize,
    inv_freq: &[&[f32]],
    pos: usize,
    batch: usize,
    // Hash layers force their expert list from the TOKEN's id, so the layer
    // description itself varies across a batch — not just the prep. One row
    // per token, each as long as `layers`; None where the layer does not hash.
    forced: Option<&[Vec<Option<Vec<usize>>>]>,
    folded_out: &mut [f32],
    // Optional hyper-connection state for every token, laid out
    // [batch, hc, dim]. A device prefix followed by a host layer needs it;
    // bringing it home beside the folds still costs one fence.
    state_out: Option<&mut [f32]>,
) -> bool {
    let Some(c) = ctx() else { return false };
    let Some((_, g0, _)) = layers.first() else {
        return false;
    };
    let dim = g0.attn.dim;
    let state_len = batch * g0.hc * dim;
    if batch == 0
        || batch > FRAME_TOK_STRIDE
        || folded_out.len() < batch * dim
        || state_out.as_ref().is_some_and(|s| s.len() < state_len)
        || inv_freq.len() != layers.len()
    {
        return false;
    }
    if bt_frame_on() {
        return dsv4_chain_batch_bt(
            model, layers, kv_id, first_li, inv_freq, pos, batch, forced, folded_out, state_out,
        );
    }
    // Grow every index cache to its FINAL batch size before recording the
    // encoder. Growing half way through would copy only the pre-batch data:
    // earlier tokens' appends still live in the unsubmitted encoder and the
    // later tokens would bind a fresh buffer that cannot contain them.
    for (i, (_, _, p)) in layers.iter().enumerate() {
        if let Some((_, cg, _, ixg)) = p.ix.as_ref() {
            let extra = batch.div_ceil(cg.ratio.max(1));
            if dsv4_index_cache(kv_id, first_li + i, ixg.idim * (p.n_ix + extra + 1)).is_none() {
                return false;
            }
        }
    }
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-chain-batch"),
        });
    let gather = frame_buf(c, 117, batch * dim * 4, false);
    let gather_state = state_out
        .as_ref()
        .map(|_| frame_buf(c, 118, state_len * 4, false));
    let t_enc = std::time::Instant::now();
    for (i, (w, g, p)) in layers.iter().enumerate() {
        // The query projection is independent across the known tokens. Pack
        // their normalized LoRA vectors, run the actual B-axis q4tp kernel,
        // then scatter the heads back to the per-token frames. The previous
        // implementation merely recorded B matvecs in one encoder and was
        // slower than the walk despite calling itself batched.
        let bytes = model.primary_bytes();
        let Some(qe) = model.tensors.get(w.attn.wq_b) else {
            return false;
        };
        let (Some(qabs), qlen) = (model.entry_abs_offset(qe), qe.nbytes as usize) else {
            return false;
        };
        let Some(qw) = weight_buffer(
            c,
            (model.uid() as usize, w.attn.wq_b),
            &bytes[qabs..qabs + qlen],
        ) else {
            return false;
        };
        let qn_pack = frame_buf(c, 119, batch * g.attn.q_lora * 4, false);
        let q_pack = frame_buf(c, 120, batch * g.attn.nh * g.attn.hd * 4, false);
        {
            let mut pass = begin_pass(&mut enc);
            for t in 0..batch {
                let qn = frame_buf_t(c, 4, t, g.attn.q_lora * 4, true);
                encode_blit_p(
                    &mut pass,
                    c,
                    &qn,
                    &qn_pack,
                    g.attn.q_lora,
                    0,
                    t * g.attn.q_lora,
                    None,
                );
            }
        }
        if !encode_q4tp_mv4_b(
            c,
            &mut enc,
            &qw,
            &qn_pack,
            &q_pack,
            g.attn.nh * g.attn.hd,
            g.attn.q_lora,
            batch,
        ) {
            return false;
        }
        {
            let mut pass = begin_pass(&mut enc);
            for t in 0..batch {
                let q = frame_buf_t(
                    c,
                    5,
                    FRAME_TOK_STRIDE + t + 1,
                    g.attn.nh * g.attn.hd * 4,
                    false,
                );
                encode_blit_p(
                    &mut pass,
                    c,
                    &q_pack,
                    &q,
                    g.attn.nh * g.attn.hd,
                    t * g.attn.nh * g.attn.hd,
                    0,
                    None,
                );
            }
        }
        for t in 0..batch {
            // The counts as of this token: everything the prep carries that
            // is not a weight.
            let mut pt = p.clone();
            pt.filled = (p.filled + t).min(p.window);
            let advanced = |ratio: usize| -> usize {
                if ratio == 0 {
                    return 0;
                }
                (0..t).filter(|k| (pos + k + 1) % ratio == 0).count()
            };
            if let Some((_, cg)) = p.comp.as_ref() {
                let ew = if cg.overlap { cg.width / 2 } else { cg.width };
                pt.n_comp = p.n_comp + advanced(cg.ratio);
                pt.comp_dst_off = p.comp_dst_off + (pt.n_comp - p.n_comp) * ew;
            }
            if let Some((_, cg, _, _)) = p.ix.as_ref() {
                let ew = if cg.overlap { cg.width / 2 } else { cg.width };
                pt.n_ix = p.n_ix + advanced(cg.ratio);
                pt.ix_dst_off = p.ix_dst_off + (pt.n_ix - p.n_ix) * ew;
            }
            // The layer as this token sees it: identical but for the row a
            // hash layer forces from the token's id.
            let mut wt;
            let w = match forced.and_then(|f| f.get(t)).and_then(|r| r.get(i)) {
                Some(row) => {
                    wt = w.clone();
                    wt.moe.forced = row.as_deref();
                    &wt
                }
                None => w,
            };
            let Some(b) = dsv4_layer_frame_enc(
                model,
                w,
                *g,
                kv_id,
                first_li + i,
                t,
                true,
                true,
                true,
                None,
                &[],
                Some(&pt),
                inv_freq[i],
                pos + t,
                &mut enc,
            ) else {
                return false;
            };
            // The last layer's fold is what the head reads; gather the batch
            // into one buffer so a single readback brings all of it home.
            if i + 1 == layers.len() {
                let mut pass = begin_pass(&mut enc);
                encode_blit_p(&mut pass, c, &b, &gather, dim, 0, t * dim, None);
                if let Some(gs) = gather_state.as_ref() {
                    let sb = frame_buf_t(c, 40, t, g0.hc * dim * 4, true);
                    encode_blit_p(&mut pass, c, &sb, gs, g0.hc * dim, 0, t * g0.hc * dim, None);
                }
            }
        }
        // Likewise, prepare the next device layer's shared LoRA vector once
        // for the batch. A host tail recomputes its own projection, so the
        // final device layer deliberately does not pay for an unused next-q.
        if i + 1 < layers.len() && !dsv4_skip("nextq") {
            let Some(next_idx) = w.next_wq_a else {
                return false;
            };
            let Some(ne) = model.tensors.get(next_idx) else {
                return false;
            };
            let (Some(nabs), nlen) = (model.entry_abs_offset(ne), ne.nbytes as usize) else {
                return false;
            };
            let Some(nw) = weight_buffer(
                c,
                (model.uid() as usize, next_idx),
                &bytes[nabs..nabs + nlen],
            ) else {
                return false;
            };
            let x_pack = frame_buf(c, 121, batch * dim * 4, false);
            let qr_pack = frame_buf(c, 122, batch * g.attn.q_lora * 4, false);
            {
                let mut pass = begin_pass(&mut enc);
                for t in 0..batch {
                    let x = frame_buf_t(c, 45, t, dim * 4, true);
                    encode_blit_p(&mut pass, c, &x, &x_pack, dim, 0, t * dim, None);
                }
            }
            if !encode_q4tp_mv4_b(
                c,
                &mut enc,
                &nw,
                &x_pack,
                &qr_pack,
                g.attn.q_lora,
                dim,
                batch,
            ) {
                return false;
            }
            let norm = const_buf(c, bytemuck::cast_slice(&w.next_q_norm[..g.attn.q_lora]));
            let mut pass = begin_pass(&mut enc);
            for t in 0..batch {
                let qr = frame_buf_t(c, 48, FRAME_TOK_STRIDE + t + 1, g.attn.q_lora * 4, false);
                let qn = frame_buf_t(c, 4, t, g.attn.q_lora * 4, true);
                encode_blit_p(
                    &mut pass,
                    c,
                    &qr_pack,
                    &qr,
                    g.attn.q_lora,
                    t * g.attn.q_lora,
                    0,
                    None,
                );
                encode_rmsnorm_p(
                    &mut pass,
                    c,
                    &qr,
                    &norm,
                    &qn,
                    g.attn.q_lora,
                    g.attn.eps,
                    (
                        53,
                        kv_id,
                        10_000_000 + (first_li + i) * FRAME_TOK_STRIDE + t,
                    ),
                );
            }
        }
    }
    CHAIN_ENC_NS.fetch_add(
        t_enc.elapsed().as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    match (gather_state.as_ref(), state_out) {
        (Some(gs), Some(states)) => readback2(
            c,
            enc,
            (&gather, &mut folded_out[..batch * dim]),
            (gs, &mut states[..state_len]),
        ),
        _ => {
            let bytes = (batch * dim * 4) as u64;
            let mut sc = c.scratch.lock().unwrap();
            let stage = Scratch::ensure(
                &c.device,
                &mut sc.stage,
                bytes,
                wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                "dsv4-batch-stage",
            );
            let ok = readback(
                c,
                enc,
                &gather,
                &stage,
                bytes,
                &mut folded_out[..batch * dim],
            );
            drop(sc);
            ok
        }
    }
}

/// The token-axis batch: one encoder, one frame per LAYER, every dispatch
/// covering all B tokens. The historical path encoded one frame per token
/// and was launch-bound — 204 ms for a 5-wide pass against 37.8 for one
/// token; the token axis exists to put the batch back at one token's
/// dispatch count.
#[allow(clippy::too_many_arguments)]
fn dsv4_chain_batch_bt(
    model: &Arc<CmfModel>,
    layers: &[(Dsv4LayerW<'_>, Dsv4LayerGeom, Dsv4Prep<'_>)],
    kv_id: u64,
    first_li: usize,
    inv_freq: &[&[f32]],
    pos: usize,
    batch: usize,
    forced: Option<&[Vec<Option<Vec<usize>>>]>,
    folded_out: &mut [f32],
    state_out: Option<&mut [f32]>,
) -> bool {
    let Some(c) = ctx() else { return false };
    let Some((_, g0, _)) = layers.first() else {
        return false;
    };
    let dim = g0.attn.dim;
    let state_len = batch * g0.hc * dim;
    // Index caches grow to their final batch size before recording — growing
    // mid-encode would strand the unsubmitted appends (see the per-token
    // path's comment).
    for (i, (_, _, p)) in layers.iter().enumerate() {
        if let Some((_, cg, _, ixg)) = p.ix.as_ref() {
            let extra = batch.div_ceil(cg.ratio.max(1));
            if dsv4_index_cache(kv_id, first_li + i, ixg.idim * (p.n_ix + extra + 1)).is_none() {
                return false;
            }
        }
    }
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-chain-batch-bt"),
        });
    // Fresh slots for this chain's sampled passes: the counter rotates
    // globally, and a wrap inside one chain would pair unrelated stamps.
    if c.ts_query.is_some() && !bt_ts_lis().is_empty() {
        TS_SLOT.store(0, std::sync::atomic::Ordering::Relaxed);
        TS_PAIRS.lock().unwrap().clear();
    }
    let t_enc = std::time::Instant::now();
    // `CMF_DSV4_CHAIN_SPLIT=N`: submit the chain in N pieces so
    // the card starts the first layers while the host still encodes the
    // rest. Same queue, same order — the split changes when work is handed
    // over, never what it computes. N=1 restores the single submission.
    let split_n = {
        static N: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
        *N.get_or_init(|| {
            std::env::var("CMF_DSV4_CHAIN_SPLIT")
                .ok()
                .and_then(|v| v.parse().ok())
                .filter(|&n| n >= 1)
                // Ampere/Ada-sized budgets benefit from shorter recordings;
                // the 96 GB class pays more for the extra submissions. The
                // configured budget is also the right signal for an emulated
                // smaller card and avoids a device-name table.
                .unwrap_or(if c.vram_budget <= 64 * 1024 * 1024 * 1024 {
                    8
                } else {
                    4
                })
        })
    };
    let chunk = layers.len().div_ceil(split_n).max(1);
    let mut last: Option<(wgpu::Buffer, wgpu::Buffer)> = None;
    for (i, (w, g, p)) in layers.iter().enumerate() {
        if i > 0 && i % chunk == 0 {
            flush_pass(&enc);
            let full = std::mem::replace(
                &mut enc,
                c.device
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                        label: Some("dsv4-chain-batch-bt"),
                    }),
            );
            submit(c, finish_enc(full));
        }
        // Per-token counts follow from the position alone (window fills by
        // one, a compressed entry appears every `ratio`), so the whole
        // batch's preps are known before anything is encoded.
        let mut preps = Vec::with_capacity(batch);
        for t in 0..batch {
            let mut pt = p.clone();
            pt.filled = (p.filled + t).min(p.window);
            let advanced = |ratio: usize| -> usize {
                if ratio == 0 {
                    return 0;
                }
                (0..t).filter(|k| (pos + k + 1) % ratio == 0).count()
            };
            if let Some((_, cg)) = p.comp.as_ref() {
                let ew = if cg.overlap { cg.width / 2 } else { cg.width };
                pt.n_comp = p.n_comp + advanced(cg.ratio);
                pt.comp_dst_off = p.comp_dst_off + (pt.n_comp - p.n_comp) * ew;
            }
            if let Some((_, cg, _, _)) = p.ix.as_ref() {
                let ew = if cg.overlap { cg.width / 2 } else { cg.width };
                pt.n_ix = p.n_ix + advanced(cg.ratio);
                pt.ix_dst_off = p.ix_dst_off + (pt.n_ix - p.n_ix) * ew;
            }
            preps.push(pt);
        }
        let rows: Option<Vec<Option<Vec<usize>>>> = forced.map(|f| {
            (0..batch)
                .map(|t| f.get(t).and_then(|r| r.get(i)).cloned().flatten())
                .collect()
        });
        // Speculative verify: photograph every layer's per-token hidden
        // input so a partial acceptance can replay the accepted tokens'
        // state appends without recomputing the pass.
        let (retain_n, caps) = SPEC_RETAIN.with(|v| v.borrow().clone());
        if retain_n > 0 && i < retain_n {
            let x2_bt = frame_buf_t(c, BT_X2, 0, batch * dim * 4, true);
            let retain = frame_buf_t(c, BT_RETAIN, 0, retain_n * batch * dim * 4, false);
            let mut pass = begin_pass(&mut enc);
            encode_blit_p(
                &mut pass,
                c,
                &x2_bt,
                &retain,
                batch * dim,
                0,
                i * batch * dim,
                None,
            );
        }
        let Some(out) = dsv4_layer_frame_bt_enc(
            model,
            w,
            *g,
            kv_id,
            first_li + i,
            batch,
            &preps,
            rows.as_deref(),
            inv_freq[i],
            pos,
            &mut enc,
        ) else {
            return false;
        };
        // The draft's capture: photograph the post-layer state of the armed
        // target layers, every token, for the host to read after acceptance.
        if let Some(slot) = caps.iter().position(|&t| t == first_li + i) {
            let hcd = g0.hc * dim;
            // Partial layers may replace this photograph after exact host
            // cold-expert correction, so the shared capture buffer must also
            // accept queue writes.
            let cap = frame_buf_t(c, BT_CAP, 0, caps.len() * batch * hcd * 4, true);
            let mut pass = begin_pass(&mut enc);
            encode_blit_p(
                &mut pass,
                c,
                &out.1,
                &cap,
                batch * hcd,
                0,
                slot * batch * hcd,
                None,
            );
        }
        last = Some(out);
    }
    let Some((folds, states)) = last else {
        return false;
    };
    // Blit, not copy_buffer_to_buffer: the pooled gather buffers carry no
    // COPY_DST (the per-token path fills them with the same kernel).
    let gather = frame_buf(c, 117, batch * dim * 4, false);
    let gather_state = state_out
        .as_ref()
        .map(|_| frame_buf(c, 118, state_len * 4, false));
    {
        let mut pass = begin_pass(&mut enc);
        encode_blit_p(&mut pass, c, &folds, &gather, batch * dim, 0, 0, None);
        if let Some(gs) = gather_state.as_ref() {
            encode_blit_p(&mut pass, c, &states, gs, state_len, 0, 0, None);
        }
    }
    CHAIN_ENC_NS.fetch_add(
        t_enc.elapsed().as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    // Resolve the sampled layers' pass stamps into the staging buffer; the
    // frame's own fence below pays for the round trip.
    let ts_pairs: Vec<(usize, u32)> = if !bt_ts_lis().is_empty() {
        std::mem::take(&mut TS_PAIRS.lock().unwrap())
    } else {
        Vec::new()
    };
    if let (false, Some((qs, resolve, tstage))) = (ts_pairs.is_empty(), c.ts_query.as_ref()) {
        for (n, (_, slot)) in ts_pairs.iter().enumerate() {
            if (n as u64 + 1) * 16 > 256 * 8 {
                break;
            }
            flush_pass(&enc);
            enc.resolve_query_set(qs, *slot..*slot + 2, resolve, 0);
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(resolve, 0, tstage, (n as u64) * 16, 16);
        }
    }
    let ok = match (gather_state.as_ref(), state_out) {
        (Some(gs), Some(states_out)) => readback2(
            c,
            enc,
            (&gather, &mut folded_out[..batch * dim]),
            (gs, &mut states_out[..state_len]),
        ),
        _ => {
            let bytes = (batch * dim * 4) as u64;
            let mut sc = c.scratch.lock().unwrap();
            let stage = Scratch::ensure(
                &c.device,
                &mut sc.stage,
                bytes,
                wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                "dsv4-batch-stage",
            );
            let ok = readback(
                c,
                enc,
                &gather,
                &stage,
                bytes,
                &mut folded_out[..batch * dim],
            );
            drop(sc);
            ok
        }
    };
    if ok && !ts_pairs.is_empty() {
        if let Some((_, _, tstage)) = &c.ts_query {
            let n_read = ts_pairs.len().min(128);
            let bytes = (n_read as u64) * 16;
            let (tx, rx) = std::sync::mpsc::channel();
            tstage.map_async(wgpu::MapMode::Read, ..bytes, move |r| {
                let _ = tx.send(r);
            });
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            if rx.recv().map(|r| r.is_ok()).unwrap_or(false) {
                if let Ok(raw) = tstage.get_mapped_range(..bytes) {
                    let t: &[u64] = bytemuck::cast_slice(&raw);
                    let mut agg = [(0f64, 0u32); BT_TS_NAMES.len()];
                    for (n, (which, _)) in ts_pairs.iter().enumerate().take(n_read) {
                        let d = t[2 * n + 1].saturating_sub(t[2 * n]);
                        let ms = d as f64 * c.ts_period as f64 / 1e6;
                        if let Some(e) = agg.get_mut(*which) {
                            e.0 += ms;
                            e.1 += 1;
                        }
                    }
                    drop(raw);
                    let n_lis = bt_ts_lis().len().max(1) as f64;
                    let line: Vec<String> = agg
                        .iter()
                        .enumerate()
                        .filter(|(_, e)| e.1 > 0)
                        .map(|(i, e)| format!("{} {:.2}({})", BT_TS_NAMES[i], e.0 / n_lis, e.1))
                        .collect();
                    eprintln!(
                        "[bt-ts] на слой ({} слоёв): {}",
                        bt_ts_lis().len(),
                        line.join(" | ")
                    );
                }
            }
            tstage.unmap();
        }
    }
    ok
}

/// `CMF_DSV4_BT_HCSPLIT=0`: the batch frame joins the halves with the
/// one-workgroup fused block instead of the three-dispatch
/// expand→mix→fold sequence. The fuse saves two links but serializes the
/// mix's 1.5 MB of weights through one SM per token — measured on the
/// release at B=5: fused 2.98 ms of the layer, split 1.15; the chain went
/// 187 → 110 ms. Split is the default; the fuse stays for a bisect.
/// `CMF_DSV4_OLORA_A4=0`: the grouped projection back on the shared-ladder
/// kernel. The four-row register twin computes the same sums in the same
/// order without the ladder's barriers.
fn ol_a4() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| {
        std::env::var("CMF_DSV4_OLORA_A4")
            .map(|v| v != "0")
            .unwrap_or(true)
    })
}

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

/// `CMF_DSV4_DN=b|b2|split|b4` picks the batch frame's down-projection
/// kernel. b4 (default) gives each workgroup four rows so the x span
/// loads once for four weight tiles — the one-row kernel is L2-bound on
/// exactly that traffic. split keeps per-slot partials + an ascending
/// sum (round-off class); b/b2 are the one-row originals for a bisect.
fn bt_dn_mode() -> u8 {
    static M: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
    *M.get_or_init(|| match std::env::var("CMF_DSV4_DN").as_deref() {
        Ok("b") => 0,
        Ok("b2") => 1,
        Ok("split") => 2,
        _ => 3,
    })
}

/// Four Q4TP/Q2TP gate+up rows share one activation span. Enabled by
/// default; `CMF_DSV4_GU_R4=0` keeps the one-row kernel for parity and
/// performance A/B on new adapters.
fn bt_gu_r4_on() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| {
        std::env::var("CMF_DSV4_GU_R4")
            .map(|v| v != "0")
            .unwrap_or(true)
    })
}

/// `CMF_DSV4_COMPFOLD=1`: the fold token as ONE fused dispatch instead of
/// the seven-dispatch per-token step. Bit-exact on every toy and on
/// chunked-prefill perplexity, −4 ms off the verify chain — and OFF by
/// default, because on the release bench the DRAFT's proposals diverge
/// two rounds after an indexer fold inside a verify pass (trace: fed
/// differs at spec@97 while both argmaxes agree) and acceptance pays 5
/// points, more than the chain saves. The verify-only difference is not
/// found yet; the trace and the suspects live in the campaign notes.
fn bt_comp_fold_on() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| std::env::var("CMF_DSV4_COMPFOLD").is_ok_and(|v| v != "0"))
}

/// Whether the batch uses the token-axis frame (default) or the historical
/// one-frame-per-token encoding (`CMF_DSV4_BT=0`, kept for a bisect).
fn bt_frame_on() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| {
        std::env::var("CMF_DSV4_BT")
            .map(|v| v != "0")
            .unwrap_or(true)
    })
}

// ── the batched frame's buffer tags (frame_buf, tok 0) ──
// Working buffers are sized batch*len, so a different batch width lands in a
// different pool entry by length; bind groups additionally salt their key
// with the batch. Registry: frame_buf 180–202, cached_bind 180–199.
const BT_STATE: u8 = 180;
const BT_X2: u8 = 181;
const BT_FOLD: u8 = 182;
const BT_QN: u8 = 183;
const BT_Q: u8 = 184;
const BT_ATTN: u8 = 185;
const BT_MID: u8 = 186;
const BT_AO: u8 = 187;
const BT_MIXES: u8 = 188;
const BT_HPOST: u8 = 189;
const BT_HCOMB: u8 = 190;
const BT_STATE2: u8 = 191;
const BT_LOGIT: u8 = 192;
const BT_MSEL: u8 = 193;
const BT_MWT: u8 = 194;
const BT_MCNT: u8 = 195;
const BT_MACT: u8 = 196;
const BT_MO: u8 = 197;
const BT_QR: u8 = 198;
const BT_ROPE_META: u8 = 199;
const BT_SA_M: u8 = 200;
const BT_IDX: u8 = 201;
const BT_FORCED: u8 = 202;

/// Seed one token of the batched frame: its hyper-connection state, the
/// attention half's post/comb (computed on the host for layer zero), its
/// opening fold and its LoRA vector. The strided twin of
/// `dsv4_state_write_t` + `dsv4_hc_write_t` + `dsv4_chain_seed_t`.
pub fn dsv4_chain_seed_bt(
    tok: usize,
    batch: usize,
    state: &[f32],
    post: &[f32],
    comb: &[f32],
    fold: &[f32],
    qn: &[f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    let up = |tag: u8, len: usize, off: usize, data: &[f32]| {
        let b = frame_buf_t(c, tag, 0, batch * len * 4, true);
        c.queue
            .write_buffer(&b, (off * len * 4) as u64, bytemuck::cast_slice(data));
    };
    up(BT_STATE, state.len(), tok, state);
    up(BT_HPOST, post.len(), tok, post);
    up(BT_HCOMB, comb.len(), tok, comb);
    up(BT_X2, fold.len(), tok, fold);
    up(BT_QN, qn.len(), tok, qn);
    true
}

const BT_RETAIN: u8 = 204;
const BT_SPEC_SCRATCH: u8 = 205;
const BT_DSPARK_IDX: u8 = 207;
const BT_DSPARK_META: u8 = 208;
const BT_DSPARK_KV: u8 = 209;

/// One DSpark stage as the graph consumes it: directory indices for the
/// quantized weights, host slices for the small f32 pieces (all living in
/// the pack or the layer — address-stable, so `const_buf` keying is sound).
pub struct DsparkStageW<'a> {
    pub wq_a: usize,
    pub wq_b: usize,
    pub wo_a: usize,
    pub wo_b: usize,
    pub wkv: usize,
    pub q_norm: &'a [f32],
    pub kv_norm: &'a [f32],
    pub attn_norm: &'a [f32],
    pub ffn_norm: &'a [f32],
    pub sink: &'a [f32],
    pub hc_attn_fn: &'a [f32],
    pub hc_attn_scale: &'a [f32],
    pub hc_attn_base: &'a [f32],
    pub hc_ffn_fn: &'a [f32],
    pub hc_ffn_scale: &'a [f32],
    pub hc_ffn_base: &'a [f32],
    pub router: &'a [f32],
    pub bias: Option<&'a [f32]>,
    pub experts: &'a [(usize, usize, usize)],
    /// Resident bitmap over ALL experts (the router still ranks them all).
    pub mask_u32: &'a [u32],
    /// Global id → pack slot, u32, 0xFFFFFFFF where cold (never chosen —
    /// the mask forbids it).
    pub map_u32: &'a [u32],
    /// Protected draft residents inside the trunk's unified physical pool.
    pub global: Option<Dsv4GlobalMoe>,
}

#[derive(Clone, Copy)]
pub struct DsparkGeom {
    pub dim: usize,
    pub hc: usize,
    pub nh: usize,
    pub hd: usize,
    pub rd: usize,
    pub q_lora: usize,
    pub o_lora: usize,
    pub o_groups: usize,
    pub inter: usize,
    pub n_experts: usize,
    pub top_k: usize,
    pub window: usize,
    pub eps: f32,
    pub hc_eps: f32,
    pub sinkhorn_iters: usize,
    pub route_scale: f32,
    pub swiglu_limit: f32,
    pub scale: f32,
    pub gu_q2: bool,
    pub dn_q2: bool,
}

/// The five-position DSpark block, whole, in ONE submission.
///
/// Every dispatch carries the block on a grid axis; the stage rings live on
/// the device under `(kv_id, 1000 + stage)`; the block's own keys occupy
/// rows `[window, window+block)` of the same buffer and are rewritten
/// every call. The draft does not owe the walk bit-exactness — a worse draft
/// costs acceptance, never correctness — so every kernel here is the fast
/// batched twin. Returns the last stage's states, `block * hc * dim`.
#[allow(clippy::too_many_arguments)]
pub fn dspark_graph(
    model: &Arc<CmfModel>,
    stages: &[DsparkStageW<'_>],
    g: DsparkGeom,
    kv_id: u64,
    main_proj: usize,
    main_norm: &[f32],
    mh: &[f32],
    states0: &[f32],
    pos: usize,
    filled: usize,
    inv_freq: &[f32],
    block: usize,
    states_out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    let (hc, dim) = (g.hc, g.dim);
    let mix_hc = (2 + hc) * hc;
    if states0.len() < block * hc * dim || states_out.len() < block * hc * dim {
        return false;
    }
    let bytes = model.primary_bytes();
    let wbuf = |idx: usize| -> Option<wgpu::Buffer> {
        let e = model.tensors.get(idx)?;
        let abs = model.entry_abs_offset(e)?;
        weight_buffer(
            c,
            (model.uid() as usize, idx),
            &bytes[abs..abs + e.nbytes as usize],
        )
    };
    let dk = |si: usize| 100_000 + si; // the graph's li-namespace
    let fb = |tag: u8, len: usize, upload: bool| frame_buf_t(c, tag, 0, block * len * 4, upload);
    let state_bt = fb(BT_STATE, hc * dim, true);
    let x2_bt = fb(BT_X2, dim, true);
    let fold_bt = fb(BT_FOLD, dim, false);
    let mixes_bt = fb(BT_MIXES, mix_hc, false);
    let hpost_bt = fb(BT_HPOST, hc, true);
    let hcomb_bt = fb(BT_HCOMB, hc * hc, true);
    let state2_bt = fb(BT_STATE2, hc * dim, false);
    let q_bt = fb(BT_Q, g.nh * g.hd, false);
    let attn_bt = fb(BT_ATTN, g.nh * g.hd, false);
    let mid_bt = fb(BT_MID, g.o_groups * g.o_lora, false);
    let ao_bt = fb(BT_AO, dim, false);
    let qr_bt = fb(BT_QR, g.q_lora, false);
    let logit_bt = fb(BT_LOGIT, g.n_experts, false);
    let slots = g.top_k + 1;
    let msel_bt = fb(BT_MSEL, slots, false);
    let mwt_bt = fb(BT_MWT, slots, false);
    let mcnt_bt = fb(BT_MCNT, 1, false);
    let mact_bt = fb(BT_MACT, slots * g.inter, false);
    let mo_bt = fb(BT_MO, dim, false);
    let cold_bt = fb(203, 4 * g.top_k, false);
    // The stage rings must exist and hold window + block rows.
    for si in 0..stages.len() {
        if !dsv4_cache_ensure(kv_id, dk(si), (g.window + block) * g.hd) {
            return false;
        }
    }
    // Whole-call uploads: initial states, captures, the attended list and
    // the per-position rope table. One submission per draft call means one
    // write each — no write-before-submit aliasing.
    c.queue.write_buffer(
        &state_bt,
        0,
        bytemuck::cast_slice(&states0[..block * hc * dim]),
    );
    let mh_buf = frame_buf_t(c, BT_DSPARK_KV, 0, mh.len() * 4, true);
    c.queue.write_buffer(&mh_buf, 0, bytemuck::cast_slice(mh));
    let m_attend = filled + block;
    if m_attend > 1024 {
        return false;
    }
    let idx: Vec<u32> = (0..filled as u32)
        .chain(g.window as u32..(g.window + block) as u32)
        .collect();
    let idx_b = frame_buf_t(c, BT_DSPARK_IDX, 0, 1024 * 4, true);
    c.queue.write_buffer(&idx_b, 0, bytemuck::cast_slice(&idx));
    let ms = vec![m_attend as u32; block];
    let sa_m = frame_buf_t(c, BT_DSPARK_META, 0, block * 4, true);
    c.queue.write_buffer(&sa_m, 0, bytemuck::cast_slice(&ms));
    let metas: Vec<f32> = (0..block)
        .flat_map(|i| [(pos + 1 + i) as f32, g.eps])
        .collect();
    let rope_meta = fb(BT_ROPE_META, 2, true);
    c.queue
        .write_buffer(&rope_meta, 0, bytemuck::cast_slice(&metas));
    let tip_meta = frame_up_pos_t(c, 1, FRAME_TOK_STRIDE - 1, pos, g.eps);

    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dspark-graph"),
        });
    let freq = const_buf(c, bytemuck::cast_slice(&inv_freq[..g.rd / 2]));
    let hcp = uniform_u32x8(
        c,
        [
            hc as u32,
            dim as u32,
            g.sinkhorn_iters as u32,
            g.hc_eps.to_bits(),
            0,
            0,
            0,
            0,
        ],
    );
    let hcp_n = uniform_u32x8(
        c,
        [
            hc as u32,
            dim as u32,
            g.sinkhorn_iters as u32,
            g.hc_eps.to_bits(),
            1,
            mix_hc as u32,
            0,
            0,
        ],
    );

    // ── the tip's ring entry: main_x = main_norm(main_proj(captures)),
    //    then each stage's kv row at slot pos % window. ──
    let Some(mp_w) = wbuf(main_proj) else {
        return false;
    };
    let main_raw = frame_buf_t(c, BT_DSPARK_KV, 9, dim * 4, false);
    let main_x = frame_buf_t(c, BT_DSPARK_KV, 10, dim * 4, false);
    {
        let mnw = const_buf(c, bytemuck::cast_slice(main_norm));
        let mut pass = begin_pass(&mut enc);
        encode_q4tp_mvw_p(
            &mut pass,
            c,
            &mp_w,
            &mh_buf,
            &main_raw,
            dim,
            mh.len(),
            (231, kv_id, 0),
        );
        encode_rmsnorm_p(
            &mut pass,
            c,
            &main_raw,
            &mnw,
            &main_x,
            dim,
            g.eps,
            (232, kv_id, 0),
        );
    }
    for (si, s) in stages.iter().enumerate() {
        let Some(wkv_w) = wbuf(s.wkv) else {
            return false;
        };
        let cache = {
            let map = c.dsv4_kv.lock().unwrap();
            match map.get(&(kv_id, dk(si))) {
                Some((b, _)) => b.clone(),
                None => return false,
            }
        };
        let kvw = {
            let Some(e) = model.tensors.get(s.wkv) else {
                return false;
            };
            e.shape[0]
        };
        let kv_raw = frame_buf_t(c, BT_DSPARK_KV, 11, kvw * 4, false);
        let kv_row = frame_buf_t(c, BT_DSPARK_KV, 12 + si, kvw * 4, false);
        let knw = const_buf(c, bytemuck::cast_slice(s.kv_norm));
        let mut pass = begin_pass(&mut enc);
        encode_q4tp_mvw_p(
            &mut pass,
            c,
            &wkv_w,
            &main_x,
            &kv_raw,
            kvw,
            dim,
            (233, kv_id, dk(si)),
        );
        encode_rmsnorm_p(
            &mut pass,
            c,
            &kv_raw,
            &knw,
            &kv_row,
            kvw,
            g.eps,
            (234, kv_id, dk(si)),
        );
        encode_rope_heads_p(
            &mut pass,
            c,
            &kv_row,
            &freq,
            &tip_meta,
            1,
            kvw,
            g.rd,
            false,
            false,
            (235, kv_id, dk(si)),
        );
        encode_blit_p(
            &mut pass,
            c,
            &kv_row,
            &cache,
            g.hd,
            kvw - g.hd,
            (pos % g.window) * g.hd,
            None,
        );
    }

    // ── the block through the three stages ──
    for (si, s) in stages.iter().enumerate() {
        let li = dk(si);
        let bk = |tag: u8| (tag, kv_id, li);
        let (Some(wq_a), Some(wq_b), Some(wo_a), Some(wo_b), Some(wkv_w)) = (
            wbuf(s.wq_a),
            wbuf(s.wq_b),
            wbuf(s.wo_a),
            wbuf(s.wo_b),
            wbuf(s.wkv),
        ) else {
            return false;
        };
        let global_bufs = s
            .global
            .and_then(|gl| c.dsv4_global_moe.lock().unwrap().get(&gl.pool_uid).cloned());
        if s.global.is_some() && global_bufs.is_none() {
            return false;
        }
        if global_bufs.as_ref().is_some_and(|b| b.gu_q2 != g.gu_q2) {
            return false;
        }
        let local_bufs = if global_bufs.is_none() {
            let Some(v) = (if g.gu_q2 {
                moe_expert_bufs_requant_gu(model, s.experts, g.inter, dim)
            } else {
                moe_expert_bufs(c, model, s.experts, g.inter, dim, true, false, false)
            }) else {
                return false;
            };
            Some(v)
        } else {
            None
        };
        let cache = {
            let map = c.dsv4_kv.lock().unwrap();
            match map.get(&(kv_id, li)) {
                Some((b, _)) => b.clone(),
                None => return false,
            }
        };
        let kvw = {
            let Some(e) = model.tensors.get(s.wkv) else {
                return false;
            };
            e.shape[0]
        };
        let a_fn = const_buf(c, bytemuck::cast_slice(s.hc_attn_fn));
        let a_sc = const_buf(c, bytemuck::cast_slice(s.hc_attn_scale));
        let a_bs = const_buf(c, bytemuck::cast_slice(s.hc_attn_base));
        let a_nw = const_buf(c, bytemuck::cast_slice(s.attn_norm));
        let f_fn = const_buf(c, bytemuck::cast_slice(s.hc_ffn_fn));
        let f_sc = const_buf(c, bytemuck::cast_slice(s.hc_ffn_scale));
        let f_bs = const_buf(c, bytemuck::cast_slice(s.hc_ffn_base));
        let f_nw = const_buf(c, bytemuck::cast_slice(s.ffn_norm));
        let qnw = const_buf(c, bytemuck::cast_slice(s.q_norm));
        let knw = const_buf(c, bytemuck::cast_slice(s.kv_norm));
        let sink = const_buf(c, bytemuck::cast_slice(s.sink));
        let router = const_buf(c, bytemuck::cast_slice(s.router));
        let mask_b = const_buf(c, bytemuck::cast_slice(s.mask_u32));
        let map_b = const_buf(c, bytemuck::cast_slice(s.map_u32));
        let n_res = if global_bufs.is_some() {
            s.mask_u32.iter().filter(|&&v| v != 0).count()
        } else {
            s.experts.len() - 1
        };

        let mixfold = |pass: &mut wgpu::ComputePass<'_>,
                       t1: u8,
                       t2: u8,
                       st: &wgpu::Buffer,
                       fnw: &wgpu::Buffer,
                       sc: &wgpu::Buffer,
                       bs: &wgpu::Buffer,
                       nw: &wgpu::Buffer| {
            let bind = cached_bind(c, bk(t1), || {
                let p = uniform_u32x4(c, [(hc * dim) as u32, mix_hc as u32, 0, 0]);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_f32_matvec_x.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, fnw),
                        bind_buf(1, st),
                        bind_buf(2, &mixes_bt),
                        bind_buf(3, &p),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_f32_matvec_x);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(mix_hc as u32, block as u32, 1);
            let bind = cached_bind(c, bk(t2), || {
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_hc_pre_fold.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, st),
                        bind_buf(1, &mixes_bt),
                        bind_buf(2, sc),
                        bind_buf(3, bs),
                        bind_buf(4, &fold_bt),
                        bind_buf(5, &hpost_bt),
                        bind_buf(6, &hcomb_bt),
                        bind_buf(7, &hcp_n),
                        bind_buf(8, nw),
                        bind_buf(9, &x2_bt),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_hc_pre_fold);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(block as u32, 1, 1);
        };
        let expand = |pass: &mut wgpu::ComputePass<'_>,
                      tag: u8,
                      x: &wgpu::Buffer,
                      res: &wgpu::Buffer,
                      out: &wgpu::Buffer| {
            let bind = cached_bind(c, bk(tag), || {
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_hc_post_expand.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, x),
                        bind_buf(1, res),
                        bind_buf(2, &hpost_bt),
                        bind_buf(3, &hcomb_bt),
                        bind_buf(4, out),
                        bind_buf(5, &hcp),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_hc_post_expand);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(((hc * dim) as u32).div_ceil(256), block as u32, 1);
        };
        let rms_b = |pass: &mut wgpu::ComputePass<'_>,
                     tag: u8,
                     x: &wgpu::Buffer,
                     w: &wgpu::Buffer,
                     o: &wgpu::Buffer,
                     n: usize| {
            let bind = cached_bind(c, bk(tag), || {
                let p = uniform_u32x4(c, [n as u32, 0, g.eps.to_bits(), 0]);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.rmsnorm_b.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, x),
                        bind_buf(1, w),
                        bind_buf(2, o),
                        bind_buf(3, &p),
                    ],
                })
            });
            pass.set_pipeline(&c.rmsnorm_b);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(block as u32, 1, 1);
        };
        let rope_bt = |pass: &mut wgpu::ComputePass<'_>,
                       tag: u8,
                       x: &wgpu::Buffer,
                       nh: usize,
                       hd: usize,
                       flags: u32| {
            let bind = cached_bind(c, bk(tag), || {
                let p = uniform_u32x4(c, [nh as u32, hd as u32, g.rd as u32, flags]);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_rope_heads.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, x),
                        bind_buf(1, &freq),
                        bind_buf(2, &p),
                        bind_buf(3, &rope_meta),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_rope_heads);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(nh as u32, block as u32, 1);
        };

        // attention half's fold
        {
            let mut pass = begin_pass(&mut enc);
            mixfold(&mut pass, 236, 237, &state_bt, &a_fn, &a_sc, &a_bs, &a_nw);
        }
        // the block's own keys into rows [window, window+block)
        let kvfull = frame_buf_t(c, BT_DSPARK_KV, 8, block * kvw * 4, false);
        let kvnorm = frame_buf_t(c, BT_DSPARK_KV, 16, block * kvw * 4, false);
        if !encode_q4tp_mv4_b(c, &mut enc, &wkv_w, &x2_bt, &kvfull, kvw, dim, block) {
            return false;
        }
        {
            let mut pass = begin_pass(&mut enc);
            rms_b(&mut pass, 238, &kvfull, &knw, &kvnorm, kvw);
            rope_bt(&mut pass, 239, &kvnorm, 1, kvw, 0);
            for i in 0..block {
                encode_blit_p(
                    &mut pass,
                    c,
                    &kvnorm,
                    &cache,
                    g.hd,
                    i * kvw + kvw - g.hd,
                    (g.window + i) * g.hd,
                    None,
                );
            }
        }
        // q
        // upload=true: the trunk's batch shares this (tag, len) entry and
        // seeds it with write_buffer — first creation decides the usage.
        let qn_bt = fb(BT_QN, g.q_lora, true);
        if !encode_q4tp_mv4_b(c, &mut enc, &wq_a, &x2_bt, &qr_bt, g.q_lora, dim, block) {
            return false;
        }
        {
            let mut pass = begin_pass(&mut enc);
            rms_b(&mut pass, 240, &qr_bt, &qnw, &qn_bt, g.q_lora);
        }
        if !encode_q4tp_mv4_b(
            c,
            &mut enc,
            &wq_b,
            &qn_bt,
            &q_bt,
            g.nh * g.hd,
            g.q_lora,
            block,
        ) {
            return false;
        }
        {
            let mut pass = begin_pass(&mut enc);
            rope_bt(&mut pass, 241, &q_bt, g.nh, g.hd, 1);
            // attend: one shared list (stride 0), every position sees the
            // ring and the whole block.
            let bind = cached_bind(c, bk(242), || {
                let p = uniform_mixed(c, [g.nh as u32, g.hd as u32, 0], g.scale);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_sparse_attend.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &q_bt),
                        bind_buf(1, &cache),
                        bind_buf(2, &idx_b),
                        bind_buf(3, &sink),
                        bind_buf(4, &attn_bt),
                        bind_buf(5, &p),
                        bind_buf(6, &sa_m),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_sparse_attend);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(g.nh as u32, block as u32, 1);
            rope_bt(&mut pass, 243, &attn_bt, g.nh, g.hd, 2);
            // o_project
            let o_rows = g.o_groups * g.o_lora;
            let o_cols = g.nh * g.hd / g.o_groups;
            let p_ol = if g.o_lora % 4 == 0 && ol_a4() {
                &c.bt_o_lora_a4
            } else {
                &c.bt_o_lora_a
            };
            let a4 = std::ptr::eq(
                p_ol as *const wgpu::ComputePipeline,
                &c.bt_o_lora_a4 as *const _,
            );
            let bind = cached_bind(c, bk(244), || {
                let p = uniform_u32x4(c, [(o_cols / 32) as u32, o_rows as u32, g.o_lora as u32, 0]);
                let mut entries = vec![bind_buf(0, &wo_a)];
                if !a4 {
                    entries.push(bind_buf(1, &attn_bt));
                }
                entries.push(bind_buf(2, &mid_bt));
                entries.push(bind_buf(3, &p));
                entries.push(bind_buf(4, &wo_a));
                if a4 {
                    entries.push(bind_buf(5, &attn_bt));
                }
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &p_ol.get_bind_group_layout(0),
                    entries: &entries,
                })
            });
            pass.set_pipeline(p_ol);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups((o_rows as u32).div_ceil(4), block as u32, 1);
        }
        if !encode_q4tp_mv4_b(
            c,
            &mut enc,
            &wo_b,
            &mid_bt,
            &ao_bt,
            dim,
            g.o_groups * g.o_lora,
            block,
        ) {
            return false;
        }
        // glue + MoE
        {
            let mut pass = begin_pass(&mut enc);
            expand(&mut pass, 245, &ao_bt, &state_bt, &state2_bt);
            mixfold(&mut pass, 246, 247, &state2_bt, &f_fn, &f_sc, &f_bs, &f_nw);
            let pipe = if g.n_experts < 64 {
                &c.bt_f32_matvec_x
            } else {
                &c.bt_f32_matvec_w
            };
            let bind = cached_bind(c, bk(248), || {
                let p = uniform_u32x4(c, [dim as u32, g.n_experts as u32, 0, 0]);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &pipe.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &router),
                        bind_buf(1, &x2_bt),
                        bind_buf(2, &logit_bt),
                        bind_buf(3, &p),
                    ],
                })
            });
            pass.set_pipeline(pipe);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(g.n_experts as u32, block as u32, 1);
            let bias_b = match s.bias {
                Some(b) => const_buf(c, bytemuck::cast_slice(b)),
                None => logit_bt.clone(),
            };
            let rflags = (s.bias.is_some() as u32)
                | 2  // mask: only resident experts are selectable
                | 8
                | 16 // subset: winners arrive as pack slots via the map
                | (s.global.map_or(n_res as u32, |gl| gl.shared_slot) << 8);
            let forced_dummy = frame_buf_t(c, BT_FORCED, dk(si), block * g.top_k * 4, true);
            let bind = cached_bind(c, bk(249), || {
                let rp = uniform_mixed(
                    c,
                    [g.n_experts as u32, g.top_k as u32, rflags],
                    g.route_scale,
                );
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_moe_route.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &logit_bt),
                        bind_buf(1, &bias_b),
                        bind_buf(2, &mask_b),
                        bind_buf(3, &forced_dummy),
                        bind_buf(4, &msel_bt),
                        bind_buf(5, &mwt_bt),
                        bind_buf(6, &mcnt_bt),
                        bind_buf(7, &rp),
                        bind_buf(8, &map_b),
                        bind_buf(9, &cold_bt),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_moe_route);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(block as u32, 1, 1);
            let stride16 = |rows: usize, cols: usize, q2: bool| -> u32 {
                let dt = if q2 {
                    cortiq_core::TensorDtype::Q2TiledP
                } else {
                    cortiq_core::TensorDtype::Q4TiledP
                };
                (cortiq_core::quant::expected_nbytes(dt, &[rows, cols]).unwrap_or(0) / 2) as u32
            };
            let gu_u = uniform_u32x8(
                c,
                [
                    (dim / 32) as u32,
                    g.inter as u32,
                    slots as u32,
                    stride16(g.inter, dim, g.gu_q2),
                    g.swiglu_limit.to_bits(),
                    s.global.map_or(0, |gl| gl.segment_slots),
                    0,
                    0,
                ],
            );
            let dn_u = uniform_u32x4(
                c,
                [
                    (g.inter / 32) as u32,
                    dim as u32,
                    slots as u32,
                    stride16(dim, g.inter, g.dn_q2),
                ],
            );
            if let Some(gb) = global_bufs.as_ref() {
                let Some((p_gu, p_dn)) = dsv4_global_moe_pipelines(c, g.gu_q2, gb.segments) else {
                    return false;
                };
                let dn_g = uniform_u32x8(
                    c,
                    [
                        (g.inter / 32) as u32,
                        dim as u32,
                        slots as u32,
                        stride16(dim, g.inter, false),
                        s.global.unwrap().segment_slots,
                        0,
                        0,
                        0,
                    ],
                );
                let gate_bindings: Vec<_> = gb
                    .gate
                    .iter()
                    .map(wgpu::Buffer::as_entire_buffer_binding)
                    .collect();
                let up_bindings: Vec<_> = gb
                    .up
                    .iter()
                    .map(wgpu::Buffer::as_entire_buffer_binding)
                    .collect();
                let down_bindings: Vec<_> = gb
                    .down
                    .iter()
                    .map(wgpu::Buffer::as_entire_buffer_binding)
                    .collect();
                let bg_gu = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("dspark-global-gu"),
                    layout: &p_gu.get_bind_group_layout(0),
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: wgpu::BindingResource::BufferArray(&gate_bindings),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: wgpu::BindingResource::BufferArray(&up_bindings),
                        },
                        bind_buf(2, &x2_bt),
                        bind_buf(3, &msel_bt),
                        bind_buf(4, &mact_bt),
                        bind_buf(5, &mwt_bt),
                    ],
                });
                let bg_gu_p = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("dspark-global-gu-p"),
                    layout: &p_gu.get_bind_group_layout(1),
                    entries: &[bind_buf(0, &gu_u)],
                });
                pass.set_pipeline(p_gu);
                pass.set_bind_group(0, &bg_gu, &[]);
                pass.set_bind_group(1, &bg_gu_p, &[]);
                pass.dispatch_workgroups(g.inter as u32, slots as u32, block as u32);

                let bg_dn = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("dspark-global-dn"),
                    layout: &p_dn.get_bind_group_layout(0),
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: wgpu::BindingResource::BufferArray(&down_bindings),
                        },
                        bind_buf(1, &mact_bt),
                        bind_buf(2, &msel_bt),
                        bind_buf(3, &mwt_bt),
                        bind_buf(4, &mo_bt),
                    ],
                });
                let bg_dn_p = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("dspark-global-dn-p"),
                    layout: &p_dn.get_bind_group_layout(1),
                    entries: &[bind_buf(0, &dn_g)],
                });
                pass.set_pipeline(p_dn);
                pass.set_bind_group(0, &bg_dn, &[]);
                pass.set_bind_group(1, &bg_dn_p, &[]);
                pass.dispatch_workgroups(dim as u32, block as u32, 1);
            } else {
                let (gate_all, up_all, down_all) = local_bufs.as_ref().unwrap();
                let gu_r4 = g.inter % 4 == 0 && bt_gu_r4_on();
                let p_gu = if g.gu_q2 {
                    if gu_r4 {
                        &c.bt_moe_gate_up_q2tp_r4
                    } else {
                        &c.bt_moe_gate_up_q2tp
                    }
                } else if gu_r4 {
                    &c.moe_gate_up_q4tp_b_r4
                } else {
                    &c.moe_gate_up_q4tp_b
                };
                let bind = cached_bind(c, bk(250), || {
                    c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &p_gu.get_bind_group_layout(0),
                        entries: &[
                            bind_buf(0, gate_all),
                            bind_buf(1, up_all),
                            bind_buf(2, &x2_bt),
                            bind_buf(3, &msel_bt),
                            bind_buf(4, &mact_bt),
                            bind_buf(5, &gu_u),
                        ],
                    })
                });
                pass.set_pipeline(p_gu);
                pass.set_bind_group(0, &bind, &[]);
                let gx = if gu_r4 {
                    (g.inter as u32).div_ceil(4)
                } else {
                    g.inter as u32
                };
                pass.dispatch_workgroups(gx, slots as u32, block as u32);
                let p_dn = if g.dn_q2 {
                    &c.moe_down_q2tp_b
                } else {
                    &c.moe_down_q4tp_b
                };
                let bind = cached_bind(c, bk(251), || {
                    // The 2-bit kernel reads x through the vec4 view and never
                    // touches the scalar activation binding; the layouts differ.
                    let mut entries = vec![
                        bind_buf(0, down_all),
                        bind_buf(2, &msel_bt),
                        bind_buf(3, &mwt_bt),
                        bind_buf(4, &mo_bt),
                        bind_buf(5, &dn_u),
                    ];
                    if g.dn_q2 {
                        entries.push(bind_buf(7, &mact_bt));
                    } else {
                        entries.insert(1, bind_buf(1, &mact_bt));
                    }
                    c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &p_dn.get_bind_group_layout(0),
                        entries: &entries,
                    })
                });
                pass.set_pipeline(p_dn);
                pass.set_bind_group(0, &bind, &[]);
                pass.dispatch_workgroups(dim as u32, block as u32, 1);
            }
            expand(&mut pass, 252, &mo_bt, &state2_bt, &state_bt);
        }
    }
    // ── home: the final states, one fence ──
    let bytes_out = (block * hc * dim * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        bytes_out,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dspark-stage",
    );
    let ok = readback(
        c,
        enc,
        &state_bt,
        &stage,
        bytes_out,
        &mut states_out[..block * hc * dim],
    );
    drop(sc);
    ok
}

thread_local! {
    /// Armed by the speculative verify: how many layers' per-token hidden
    /// inputs the batch should retain on the device (0 = off), and which
    /// layers' post-layer hyper-connection states to photograph for the
    /// draft's capture. The retained hiddens feed the state replay after a
    /// partial acceptance; the captures feed the next draft.
    static SPEC_RETAIN: std::cell::RefCell<(usize, Vec<usize>)> =
        const { std::cell::RefCell::new((0, Vec::new())) };
}

/// Read a pooled frame buffer back, for parity debugging only: the tag
/// registry names what lives where.
/// Debug window into a layer's INDEX cache: `n` floats from float offset
/// `off`. The fold-in-verify investigation reads the entry a fused fold
/// just landed and compares it against the per-token path's.
pub fn dsv4_dbg_read_ix(kv_id: u64, li: usize, off: usize, n: usize) -> Option<Vec<f32>> {
    let c = ctx()?;
    let b = {
        let m = c.dsv4_ixkv.lock().unwrap();
        m.get(&(kv_id, li)).map(|(b, _)| b.clone())?
    };
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
    let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("dbg-ix-stage"),
        size: (n * 4) as u64,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&b, (off * 4) as u64, &stage, 0, (n * 4) as u64);
    submit(c, finish_enc(enc));
    let slice = stage.slice(..);
    slice.map_async(wgpu::MapMode::Read, |_| {});
    c.device.poll(wgpu::PollType::wait_indefinitely()).ok()?;
    let data = slice.get_mapped_range().ok()?;
    let v: Vec<f32> = bytemuck::cast_slice(&data).to_vec();
    drop(data);
    Some(v)
}

pub fn dsv4_dbg_read_tag(tag: u8, tok: usize, n: usize) -> Option<Vec<f32>> {
    let c = ctx()?;
    let b = {
        let m = c.dsv4_scratch.lock().unwrap();
        // The pool keys on (tag, tok, len); a debug reader doesn't know the
        // len, so take the first entry matching (tag, tok).
        m.iter()
            .find(|((t, tk, len), _)| *t == tag && *tk == tok && *len == n * 4)
            .or_else(|| m.iter().find(|((t, tk, _), _)| *t == tag && *tk == tok))
            .map(|(_, b)| b.clone())?
    };
    let bytes = ((n * 4) as u64).min(b.size());
    let mut out = vec![0.0f32; (bytes / 4) as usize];
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dbg-stage",
    );
    let enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("dbg") });
    let ok = readback(c, enc, &b, &stage, bytes, &mut out);
    drop(sc);
    ok.then_some(out)
}

/// Arm (n_layers > 0) or disarm (0) hidden-state retention for the next
/// batched chain call on this thread. `caps` lists the layer ordinals whose
/// post-layer state the draft's capture needs.
pub fn dsv4_spec_retain_arm(n_layers: usize, caps: &[usize]) {
    SPEC_RETAIN.with(|c| *c.borrow_mut() = (n_layers, caps.to_vec()));
}

const BT_CAP: u8 = 206;

/// Read the WHOLE capture photograph back in one fence:
/// `[n_caps, batch, hc*dim]`, host-sliced by the caller. The per-slot
/// variant cost a fence per (slot, token) — five of them a pass.
pub fn dsv4_spec_cap_read_all(batch: usize, n_caps: usize, hc_dim: usize, out: &mut [f32]) -> bool {
    if n_caps == 0 {
        return true;
    }
    let Some(c) = ctx() else { return false };
    let cap = frame_buf_t(c, BT_CAP, 0, n_caps * batch * hc_dim * 4, true);
    let bytes = (n_caps * batch * hc_dim * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dsv4-cap-stage",
    );
    let enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-cap-read-all"),
        });
    let ok = readback(
        c,
        enc,
        &cap,
        &stage,
        bytes,
        &mut out[..n_caps * batch * hc_dim],
    );
    drop(sc);
    ok
}

/// Read one photographed capture back: `slot` indexes the armed `caps`
/// list, `tok` the batch position. `out` takes hc*dim floats.
pub fn dsv4_spec_cap_read(
    slot: usize,
    tok: usize,
    batch: usize,
    n_caps: usize,
    hc_dim: usize,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    let cap = frame_buf_t(c, BT_CAP, 0, n_caps * batch * hc_dim * 4, true);
    let bytes = (hc_dim * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dsv4-cap-stage",
    );
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-cap-read"),
        });
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(
        &cap,
        ((slot * batch + tok) * hc_dim * 4) as u64,
        &stage,
        0,
        bytes,
    );
    let ok = {
        submit(c, finish_enc(enc));
        let slice = stage.slice(..bytes);
        let done = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
        let d2 = done.clone();
        slice.map_async(wgpu::MapMode::Read, move |_| {
            d2.store(true, std::sync::atomic::Ordering::Release);
        });
        if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
            false
        } else if let Ok(data) = slice.get_mapped_range() {
            out[..hc_dim].copy_from_slice(bytemuck::cast_slice(&data[..hc_dim * 4]));
            drop(data);
            stage.unmap();
            true
        } else {
            false
        }
    };
    drop(sc);
    ok
}

/// The device state a speculative pass damages and how to put it back.
///
/// Append-only regions (the compressed tail inside the KV cache, the
/// indexer's cache) roll back by COUNT and are not copied. What is copied:
/// the window rows the B appends will shift out (at most B per layer), and
/// the compressor streams (pending/previous, both kinds), which mutate on
/// every token.
pub struct Dsv4SpecShadow {
    kv_id: u64,
    batch: usize,
    layers: Vec<SpecLayerShadow>,
}
struct SpecLayerShadow {
    li: usize,
    hd: usize,
    window: usize,
    filled: usize,
    /// The first `dk` window rows, photographed before the pass.
    dk: usize,
    head: Option<wgpu::Buffer>,
    comp: Vec<(u8, [wgpu::Buffer; 4])>,
}

/// Photograph what the B-token pass will destroy. One submission.
pub fn dsv4_spec_shadow(
    kv_id: u64,
    metas: &[(usize, usize, usize, usize)], // (li, hd, window, filled)
    batch: usize,
) -> Option<Dsv4SpecShadow> {
    let c = ctx()?;
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-spec-shadow"),
        });
    let mut layers = Vec::with_capacity(metas.len());
    for &(li, hd, window, filled) in metas {
        // Staged batches freeze the window until commit, so there is
        // nothing to photograph there any more; only the compressor
        // streams still mutate mid-pass.
        let dk = 0usize;
        let head = if dk > 0 {
            let cache = {
                let map = c.dsv4_kv.lock().unwrap();
                map.get(&(kv_id, li)).map(|(b, _)| b.clone())?
            };
            let b = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("dsv4-spec-head"),
                size: (dk * hd * 4) as u64,
                usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            });
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(&cache, 0, &b, 0, (dk * hd * 4) as u64);
            Some(b)
        } else {
            None
        };
        let mut comp = Vec::new();
        for kind in [0u8, 1u8] {
            let live = {
                let map = c.dsv4_comp.lock().unwrap();
                map.get(&(kind, kv_id, li)).cloned()
            };
            if let Some(bufs) = live {
                let clones: [wgpu::Buffer; 4] = std::array::from_fn(|i| {
                    let b = c.device.create_buffer(&wgpu::BufferDescriptor {
                        label: Some("dsv4-spec-comp"),
                        size: bufs[i].size(),
                        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC,
                        mapped_at_creation: false,
                    });
                    flush_pass(&enc);
                    enc.copy_buffer_to_buffer(&bufs[i], 0, &b, 0, bufs[i].size());
                    b
                });
                comp.push((kind, clones));
            }
        }
        layers.push(SpecLayerShadow {
            li,
            hd,
            window,
            filled,
            dk,
            head,
            comp,
        });
    }
    submit(c, finish_enc(enc));
    Some(Dsv4SpecShadow {
        kv_id,
        batch,
        layers,
    })
}

/// Put the device back to the snapshot (as if NO speculative token ran):
/// drop the B appended window rows, un-shift, restore the shadowed head
/// rows, copy the compressor streams back. Counts are the host's business.
pub fn dsv4_spec_restore(sh: &Dsv4SpecShadow) -> bool {
    let Some(c) = ctx() else { return false };
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-spec-restore"),
        });
    for l in &sh.layers {
        if l.dk > 0 {
            let cache = {
                let map = c.dsv4_kv.lock().unwrap();
                match map.get(&(sh.kv_id, l.li)) {
                    Some((b, _)) => b.clone(),
                    None => return false,
                }
            };
            // Surviving original rows sit at [0, keep); slide them right by
            // dk and put the photographed head back in front.
            let keep = l.filled - l.dk;
            if keep > 0 {
                let scratch = frame_buf(c, BT_SPEC_SCRATCH, l.window * l.hd * 4, true);
                flush_pass(&enc);
                enc.copy_buffer_to_buffer(&cache, 0, &scratch, 0, (keep * l.hd * 4) as u64);
                flush_pass(&enc);
                enc.copy_buffer_to_buffer(
                    &scratch,
                    0,
                    &cache,
                    (l.dk * l.hd * 4) as u64,
                    (keep * l.hd * 4) as u64,
                );
            }
            if let Some(head) = &l.head {
                flush_pass(&enc);
                enc.copy_buffer_to_buffer(head, 0, &cache, 0, (l.dk * l.hd * 4) as u64);
            }
        }
        for (kind, clones) in &l.comp {
            let live = {
                let map = c.dsv4_comp.lock().unwrap();
                map.get(&(*kind, sh.kv_id, l.li)).cloned()
            };
            let Some(live) = live else { return false };
            for i in 0..4 {
                let n = live[i].size().min(clones[i].size());
                flush_pass(&enc);
                enc.copy_buffer_to_buffer(&clones[i], 0, &live[i], 0, n);
            }
        }
    }
    submit(c, finish_enc(enc));
    true
}

/// Land the accepted prefix's staged rows in every device layer's window —
/// the staged batch's deferred slide. Safe for any k in 0..=batch.
pub fn dsv4_spec_commit_windows(
    kv_id: u64,
    metas: &[(usize, usize, usize, usize)], // (li, filled, window, hd)
    batch: usize,
    k: usize,
) -> bool {
    let Some(c) = ctx() else { return false };
    if k == 0 {
        return true;
    }
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-spec-commit"),
        });
    for &(li, filled, window, hd) in metas {
        let (cache, cap) = {
            let map = c.dsv4_kv.lock().unwrap();
            match map.get(&(kv_id, li)) {
                Some((b, cap)) => (b.clone(), *cap),
                None => return false,
            }
        };
        let srow0 = (cap / hd).saturating_sub(batch + 1);
        encode_staged_commit(c, &mut enc, &cache, filled, window, hd, srow0, k);
    }
    submit(c, finish_enc(enc));
    true
}

/// Re-append the state of the ACCEPTED tokens after a restore: window rows,
/// compressor streams and folds, the indexer's compressor — from the hidden
/// inputs the batch retained per (layer, token). No attention, no scoring:
/// state only, exactly what a sequential walk of those k tokens would have
/// written.
#[allow(clippy::too_many_arguments)]
pub fn dsv4_spec_replay(
    model: &Arc<CmfModel>,
    layers: &[(usize, Dsv4Prep<'_>)], // (li, prep with counts AS OF the snapshot)
    kv_id: u64,
    pos0: usize,
    batch: usize,
    k: usize,
    inv_freq: &[&[f32]],
    hd: usize,
    dim: usize,
    rope_dim: usize,
    eps: f32,
    // The staged batch commits window rows by blit; the replay then owes
    // only the compressor streams and folds.
    skip_window: bool,
) -> bool {
    if k == 0 {
        return true;
    }
    macro_rules! rfail {
        ($($t:tt)*) => {{
            if std::env::var("CMF_DSV4_SPEC_DEBUG").is_ok() {
                eprintln!("spec_replay: {}", format_args!($($t)*));
            }
            return false;
        }};
    }
    let Some(c) = ctx() else {
        rfail!("нет контекста")
    };
    let retain = frame_buf_t(c, BT_RETAIN, 0, layers.len() * batch * dim * 4, false);
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-spec-replay"),
        });
    for (i, (li, p)) in layers.iter().enumerate() {
        for t in 0..k {
            let _salt = Dsv4FrameSalt::enter(t + 1);
            let x2_t = frame_buf_t(c, 45, t, dim * 4, true);
            {
                let mut pass = begin_pass(&mut enc);
                encode_blit_p(
                    &mut pass,
                    c,
                    &retain,
                    &x2_t,
                    dim,
                    (i * batch + t) * dim,
                    0,
                    None,
                );
            }
            let advanced = |ratio: usize| -> usize {
                if ratio == 0 {
                    return 0;
                }
                (0..t).filter(|j| (pos0 + j + 1) % ratio == 0).count()
            };
            let cache = {
                let map = c.dsv4_kv.lock().unwrap();
                match map.get(&(kv_id, *li)) {
                    Some((b, _)) => b.clone(),
                    None => rfail!("нет кеша слоя {li}"),
                }
            };
            if let Some((cw, cg)) = &p.comp {
                let ew = if cg.overlap { cg.width / 2 } else { cg.width };
                let n_comp = p.n_comp + advanced(cg.ratio);
                let off = p.comp_dst_off + (n_comp - p.n_comp) * ew;
                if dsv4_compressor_frame(
                    model,
                    cw,
                    *cg,
                    0,
                    kv_id,
                    *li,
                    &x2_t,
                    pos0 + t,
                    inv_freq[i],
                    &cache,
                    off,
                    &mut enc,
                )
                .is_none()
                {
                    // None means "no fold this token" — the pending append
                    // is encoded either way; the prep path reads it the
                    // same. A real refusal (missing weight) also lands here
                    // and surfaces as a wrong answer downstream, which the
                    // sequential-parity test is what catches.
                }
            }
            if let Some((iw, ig, _, ixg)) = &p.ix {
                let ew = if ig.overlap { ig.width / 2 } else { ig.width };
                let n_ix = p.n_ix + advanced(ig.ratio);
                let off = p.ix_dst_off + (n_ix - p.n_ix) * ew;
                let Some(ixkv) = dsv4_index_cache(kv_id, *li, ixg.idim * (n_ix + 2)) else {
                    rfail!("ix-кеш слоя {li}");
                };
                if dsv4_compressor_frame(
                    model,
                    iw,
                    *ig,
                    1,
                    kv_id,
                    *li,
                    &x2_t,
                    pos0 + t,
                    inv_freq[i],
                    &ixkv,
                    off,
                    &mut enc,
                )
                .is_none()
                {
                    // Same contract as above: None is "no fold".
                }
            }
            let filled = (p.filled + t).min(p.window);
            if !skip_window
                && dsv4_window_append(
                    model,
                    p.wkv,
                    p.kv_norm,
                    &x2_t,
                    &cache,
                    hd,
                    p.window,
                    filled,
                    dim,
                    rope_dim,
                    eps,
                    pos0 + t,
                    inv_freq[i],
                    kv_id,
                    *li,
                    &mut enc,
                )
                .is_none()
            {
                rfail!("окно слоя {li} токена {t}");
            }
        }
    }
    submit(c, finish_enc(enc));
    true
}

/// Slide the window by what a run of `k` appends would have slid it, and
/// land the first `k` staged rows as its newest — through scratch, because
/// neither a copy nor a blit may alias its own buffer. Three copies when it
/// slides, two when it does not.
#[allow(clippy::too_many_arguments)]
fn encode_staged_commit(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    cache: &wgpu::Buffer,
    filled: usize,
    window: usize,
    hd: usize,
    srow0: usize,
    k: usize,
) {
    if k == 0 {
        return;
    }
    let shift = (filled + k).saturating_sub(window);
    let keep = filled - shift.min(filled);
    let scratch = frame_buf(c, BT_SPEC_SCRATCH, window * hd * 4, true);
    if shift > 0 {
        if keep > 0 {
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(
                cache,
                (shift * hd * 4) as u64,
                &scratch,
                0,
                (keep * hd * 4) as u64,
            );
        }
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(
            cache,
            (srow0 * hd * 4) as u64,
            &scratch,
            (keep * hd * 4) as u64,
            (k * hd * 4) as u64,
        );
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&scratch, 0, cache, 0, ((keep + k) * hd * 4) as u64);
    } else {
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(
            cache,
            (srow0 * hd * 4) as u64,
            &scratch,
            0,
            (k * hd * 4) as u64,
        );
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(
            &scratch,
            0,
            cache,
            (filled * hd * 4) as u64,
            (k * hd * 4) as u64,
        );
    }
}

/// One layer of a whole BATCH, with the token on a grid axis: every
/// dispatch covers all B tokens, so the layer costs the dispatch count of a
/// single token. The arithmetic per token is the single frame's, term for
/// term; only the prep (window append, compressors, indexer) still encodes
/// per token — its writes are ordered by the pass, which is what lets token
/// t attend to the entries token t-1 just appended.
#[allow(clippy::too_many_arguments)]
fn dsv4_layer_frame_bt_enc(
    model: &Arc<CmfModel>,
    w: &Dsv4LayerW,
    g: Dsv4LayerGeom,
    kv_id: u64,
    li: usize,
    batch: usize,
    preps: &[Dsv4Prep],
    forced_rows: Option<&[Option<Vec<usize>>]>,
    inv_freq: &[f32],
    pos0: usize,
    enc: &mut wgpu::CommandEncoder,
) -> Option<(wgpu::Buffer, wgpu::Buffer)> {
    macro_rules! no {
        ($($t:tt)*) => {{
            if std::env::var("CMF_DSV4_FRAME_DEBUG").is_ok() {
                eprintln!("пакетный кадр слоя отклонён: {}", format_args!($($t)*));
            }
            return None;
        }};
    }
    let Some(c) = ctx() else {
        no!("нет контекста wgpu")
    };
    let a = g.attn;
    let m = g.moe;
    let (hc, dim) = (g.hc, a.dim);
    let mix_hc = (2 + hc) * hc;
    if preps.len() != batch {
        no!("подготовок {} на пакет {batch}", preps.len());
    }
    // The bind-group key: the layer, salted by the batch width so a 3-wide
    // chunk never reuses a 5-wide chunk's groups (their buffers differ by
    // length and therefore identity).
    let bk = |tag: u8| (tag, kv_id, li * FRAME_TOK_STRIDE + batch);

    // ── weights (same set and order as the single frame) ──
    let bytes = model.primary_bytes();
    let mut wb = Vec::with_capacity(5);
    for &idx in &[
        w.attn.wq_a,
        w.attn.wq_b,
        w.attn.wo_a,
        w.attn.wo_b,
        w.next_wq_a.unwrap_or(w.attn.wq_a),
    ] {
        let Some(e) = model.tensors.get(idx) else {
            no!("тензора {idx} нет");
        };
        if e.dtype != cortiq_core::TensorDtype::Q4TiledP {
            no!("{} не q4tp", e.name);
        }
        let (Some(abs), plen) = (model.entry_abs_offset(e), e.nbytes as usize) else {
            no!("{} без смещения", e.name);
        };
        let Some(b) = weight_buffer_l(
            c,
            (model.uid() as usize, idx),
            &bytes[abs..abs + plen],
            layer_of_name(&model.tensors[idx].name),
        ) else {
            no!("{} не влез в VRAM", e.name);
        };
        wb.push(b);
    }
    let Some((gate_all, up_all, down_all)) = moe_expert_bufs(
        c,
        model,
        w.moe.experts,
        m.inter,
        m.hidden,
        true,
        m.gu_q2,
        false,
    ) else {
        no!("эксперты не влезли в VRAM");
    };
    let cache = {
        let map = c.dsv4_kv.lock().unwrap();
        match map.get(&(kv_id, li)) {
            Some((b, _)) => b.clone(),
            None => no!("кеш ({kv_id}, {li}) не заведён"),
        }
    };

    // ── constants ──
    let sink = const_buf(c, bytemuck::cast_slice(&w.attn.sink[..a.nh]));
    let freq = const_buf(c, bytemuck::cast_slice(&inv_freq[..a.rd / 2]));
    let ffn_fn = const_buf(c, bytemuck::cast_slice(w.hc_ffn_fn));
    let ffn_sc = const_buf(c, bytemuck::cast_slice(w.hc_ffn_scale));
    let ffn_bs = const_buf(c, bytemuck::cast_slice(&w.hc_ffn_base[..mix_hc]));
    let ffn_nw = const_buf(c, bytemuck::cast_slice(&w.ffn_norm[..dim]));
    let n_pack = w.moe.experts.len().saturating_sub(1);
    let n_route = w.moe.remap.map_or(n_pack, |r| r.len().max(n_pack));
    if w.router.len() < m.hidden * n_route {
        no!("роутер короче {} × {}", n_route, m.hidden);
    }
    let router = const_buf(c, bytemuck::cast_slice(&w.router[..m.hidden * n_route]));
    let next_nw = const_buf(c, bytemuck::cast_slice(&w.next_norm[..dim]));
    let next_qn = const_buf(c, bytemuck::cast_slice(&w.next_q_norm[..a.q_lora]));

    // ── strided working buffers ──
    let slots = m.top_k + 1;
    let fb = |tag: u8, len: usize, upload: bool| frame_buf_t(c, tag, 0, batch * len * 4, upload);
    let state_bt = fb(BT_STATE, hc * dim, true);
    let x2_bt = fb(BT_X2, dim, true);
    let fold_bt = fb(BT_FOLD, dim, false);
    let qn_bt = fb(BT_QN, a.q_lora, true);
    let q_bt = fb(BT_Q, a.nh * a.hd, false);
    let attn_bt = fb(BT_ATTN, a.nh * a.hd, false);
    let mid_bt = fb(BT_MID, a.o_groups * a.o_lora, false);
    let ao_bt = fb(BT_AO, dim, false);
    let mixes_bt = fb(BT_MIXES, mix_hc, false);
    let hpost_bt = fb(BT_HPOST, hc, true);
    let hcomb_bt = fb(BT_HCOMB, hc * hc, true);
    let state2_bt = fb(BT_STATE2, hc * dim, false);
    let logit_bt = fb(BT_LOGIT, n_route, false);
    let msel_bt = fb(BT_MSEL, slots, false);
    let mwt_bt = fb(BT_MWT, slots, false);
    let mcnt_bt = fb(BT_MCNT, 1, false);
    let mact_bt = fb(BT_MACT, slots * m.inter, false);
    let mo_bt = fb(BT_MO, dim, false);
    let qr_bt = fb(BT_QR, a.q_lora, false);
    let rope_meta = fb(BT_ROPE_META, 2, true);
    // PER LAYER, not pooled: these are filled with `queue.write_buffer`,
    // and every such write lands before the run's single submit — a shared
    // buffer would hand every layer the LAST layer's contents. The same
    // trap once made two hash layers route with one list.
    let sa_m = frame_buf_t(c, BT_SA_M, li, batch * 4, true);
    let idx_bt = fb(BT_IDX, 1024, true);
    let forced_bt = frame_buf_t(c, BT_FORCED, li, batch * m.top_k * 4, true);

    let _ = (&sa_m, &idx_bt);

    // ── STAGED batch: the window never slides mid-pass. Every token's new
    //    key row lands in staging at the cache's tail; the attends read the
    //    frozen window plus each token's staged prefix through per-token
    //    index lists; the slide happens ONCE at commit. This is what lets
    //    the attends (and the indexer) run batched instead of interleaved —
    //    the interleave was the dispatch count, and the dispatch count was
    //    the pass. ──
    let metas: Vec<f32> = (0..batch)
        .flat_map(|t| [(pos0 + t) as f32, a.eps])
        .collect();
    c.queue
        .write_buffer(&rope_meta, 0, bytemuck::cast_slice(&metas));
    let p0 = &preps[0];
    let (cache_cap, kvw) = {
        let map = c.dsv4_kv.lock().unwrap();
        let cap = match map.get(&(kv_id, li)) {
            Some((_, cap)) => *cap,
            None => no!("кеш ({kv_id}, {li}) не заведён"),
        };
        let Some(e) = model.tensors.get(p0.wkv) else {
            no!("wkv без записи");
        };
        (cap, e.shape[0])
    };
    let srow0 = (cache_cap / a.hd).saturating_sub(batch + 1);
    {
        // Staging must sit past everything the pass appends.
        let ew_c = p0.comp.as_ref().map_or(
            0,
            |(_, cg)| {
                if cg.overlap {
                    cg.width / 2
                } else {
                    cg.width
                }
            },
        );
        let comp_top = p0.window * a.hd + (p0.n_comp + batch.div_ceil(4).max(1) + 2) * ew_c.max(1);
        if srow0 * a.hd < comp_top {
            no!("нет места под staging: srow0 {srow0}, comp_top {comp_top}");
        }
    }

    // Sampled-layer GPU profile: every pass opened between two marks lands
    // under the earlier mark's stage.
    let ts_armed = c.ts_query.is_some() && bt_ts_lis().contains(&li);
    let mark = |s: usize| {
        if ts_armed {
            bt_ts(s);
        }
    };
    mark(1);
    // q for the whole batch — touches no cache.
    if !encode_q4tp_mv4_b(c, enc, &wb[1], &qn_bt, &q_bt, a.nh * a.hd, a.q_lora, batch) {
        no!("q-проекция пакетом не закодировалась");
    }

    // Compressor streams: the two projections batch over the tokens (they
    // read only each token's hidden), the state step stays per token in
    // causal order — its pending/prev shuffle is order-dependent. What it
    // appends goes PAST the logical extents, so the batched attends below
    // stay safe.
    let mut comp_jobs: Vec<(u8, Dsv4CompW, Dsv4CompGeom)> = Vec::new();
    if let Some((cw, cg)) = &p0.comp {
        comp_jobs.push((0, cw.clone(), *cg));
    }
    if let Some((iw, ig, _, ixg0)) = &p0.ix {
        if dsv4_index_cache(kv_id, li, ixg0.idim * (p0.n_ix + batch + 2)).is_none() {
            no!("ix-кеш слоя {li}");
        }
        comp_jobs.push((1, iw.clone(), *ig));
    }
    mark(2);
    if !dsv4_skip("comp") {
        for (kind, cw, cg) in &comp_jobs {
            let ckv_bt = frame_buf_t(
                c,
                BT_DSPARK_KV,
                26 + *kind as usize,
                batch * cg.width * 4,
                false,
            );
            let csc_bt = frame_buf_t(
                c,
                BT_DSPARK_KV,
                28 + *kind as usize,
                batch * cg.width * 4,
                false,
            );
            for (wi, out) in [(cw.wkv, &ckv_bt), (cw.wgate, &csc_bt)] {
                let Some(e) = model.tensors.get(wi) else {
                    no!("компрессор без тензора")
                };
                let (Some(abs), plen) = (model.entry_abs_offset(e), e.nbytes as usize) else {
                    no!("компрессор без смещения");
                };
                let Some(wbuf) =
                    weight_buffer(c, (model.uid() as usize, wi), &bytes[abs..abs + plen])
                else {
                    no!("компрессор не влез");
                };
                if !encode_q4tp_mv4_b(c, enc, &wbuf, &x2_bt, out, cg.width, cg.hidden, batch) {
                    no!("проекция компрессора пакетом");
                }
            }
            let dst = if *kind == 0 {
                cache.clone()
            } else {
                let m2 = c.dsv4_ixkv.lock().unwrap();
                match m2.get(&(kv_id, li)) {
                    Some((b, _)) => b.clone(),
                    None => no!("ix-кеш ({kv_id},{li}) не заведён"),
                }
            };
            // Fold-free tokens append in ONE dispatch per segment (their
            // slots are all distinct between folds, so the writes commute);
            // only the token that CLOSES a window walks the full per-token
            // state step, folds included.
            let streams = {
                let mut m = c.dsv4_comp.lock().unwrap();
                m.entry((*kind, kv_id, li))
                    .or_insert_with(|| {
                        let span = cg.ratio * cg.width;
                        let mk = || {
                            c.device.create_buffer(&wgpu::BufferDescriptor {
                                label: Some("dsv4-comp-stream"),
                                size: (span * 4).max(4) as u64,
                                usage: wgpu::BufferUsages::STORAGE
                                    | wgpu::BufferUsages::COPY_SRC
                                    | wgpu::BufferUsages::COPY_DST,
                                mapped_at_creation: false,
                            })
                        };
                        [mk(), mk(), mk(), mk()]
                    })
                    .clone()
            };
            let ape_all = const_buf(c, bytemuck::cast_slice(cw.ape));
            let slots: Vec<u32> = (0..batch).map(|t| ((pos0 + t) % cg.ratio) as u32).collect();
            let slot_b = frame_buf_t(
                c,
                BT_DSPARK_META,
                2 * FRAME_TOK_STRIDE + li * 2 + *kind as usize,
                batch * 4,
                true,
            );
            c.queue
                .write_buffer(&slot_b, 0, bytemuck::cast_slice(&slots));
            let append_seg = |enc: &mut wgpu::CommandEncoder, t0: usize, n: usize| {
                if n == 0 {
                    return;
                }
                let tag = 231 + kind;
                // The uniform carries (t0, n); the fold position slides with
                // pos0, so the SAME t0 recurs with a DIFFERENT n — n must be
                // in the key or a stale shorter uniform silently drops the
                // tail tokens' appends.
                let bkey = ((li * FRAME_TOK_STRIDE + batch) * FRAME_TOK_STRIDE + t0)
                    * FRAME_TOK_STRIDE
                    + n;
                let bind = cached_bind(c, (tag, kv_id, bkey), || {
                    let flags = cg.overlap as u32;
                    let pu = uniform_u32x4(c, [cg.width as u32, t0 as u32, n as u32, flags]);
                    c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                        label: None,
                        layout: &c.bt_comp_append.get_bind_group_layout(0),
                        entries: &[
                            bind_buf(0, &ckv_bt),
                            bind_buf(1, &csc_bt),
                            bind_buf(2, &ape_all),
                            bind_buf(3, &streams[0]),
                            bind_buf(4, &streams[1]),
                            bind_buf(5, &pu),
                            bind_buf(6, &slot_b),
                        ],
                    })
                });
                let mut pass = begin_pass(enc);
                pass.set_pipeline(&c.bt_comp_append);
                pass.set_bind_group(0, &bind, &[]);
                pass.dispatch_workgroups((cg.width as u32).div_ceil(256), n as u32, 1);
            };
            // Which kinds batch their appends. Both are walk-exact now that
            // the uniform's (t0, n) pair is part of the bind key — the old
            // kind-0 "divergence" was a stale shorter n from a previous
            // pass's cached uniform dropping tail appends. CMF_DSV4_SEGAPP
            // overrides for a bisect.
            let seg_on = std::env::var("CMF_DSV4_SEGAPP").unwrap_or_else(|_| "01".into());
            let seg_this = seg_on.contains(char::from(b'0' + *kind));
            let mut seg0 = 0usize;
            for (t, p) in preps.iter().enumerate() {
                let folds = (pos0 + t + 1) % cg.ratio == 0;
                if !folds && seg_this {
                    continue;
                }
                if seg_this {
                    append_seg(enc, seg0, t - seg0);
                }
                seg0 = t + 1;
                let ew_f = if cg.overlap { cg.width / 2 } else { cg.width };
                if bt_comp_fold_on() && ew_f <= 512 {
                    // Append + pool + norm + rope + land + shift, one link.
                    let off = if *kind == 0 {
                        p.comp_dst_off
                    } else {
                        p.ix_dst_off
                    };
                    let pos = pos0 + t;
                    let have_prev = pos + 1 > cg.ratio;
                    let flags = (cg.overlap as u32) | ((have_prev as u32) << 1);
                    let posb = (pos + 1 - cg.ratio) as f32;
                    // PER (layer, kind, token): two folds can share one
                    // pass at B=5, and every queue.write_buffer lands
                    // before the submit — a shared buffer would hand the
                    // first fold the second fold's uniform.
                    let pu = frame_buf_t(
                        c,
                        BT_DSPARK_META,
                        8 * FRAME_TOK_STRIDE + (li * FRAME_TOK_STRIDE + t) * 2 + *kind as usize,
                        32,
                        true,
                    );
                    c.queue.write_buffer(
                        &pu,
                        0,
                        bytemuck::cast_slice(&[
                            cg.width as u32,
                            cg.ratio as u32,
                            cg.rope_dim as u32,
                            flags,
                            off as u32,
                            posb.to_bits(),
                            cg.eps.to_bits(),
                            t as u32,
                        ]),
                    );
                    let nw_f = const_buf(c, bytemuck::cast_slice(&cw.norm[..ew_f]));
                    let fr_f = const_buf(c, bytemuck::cast_slice(&inv_freq[..cg.rope_dim / 2]));
                    let bind =
                        cached_bind(c, ((253 + *kind), kv_id, li * FRAME_TOK_STRIDE + t), || {
                            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                                label: None,
                                layout: &c.bt_comp_fold.get_bind_group_layout(0),
                                entries: &[
                                    bind_buf(0, &ckv_bt),
                                    bind_buf(1, &csc_bt),
                                    bind_buf(2, &ape_all),
                                    bind_buf(3, &streams[0]),
                                    bind_buf(4, &streams[1]),
                                    bind_buf(5, &streams[2]),
                                    bind_buf(6, &streams[3]),
                                    bind_buf(7, &nw_f),
                                    bind_buf(8, &fr_f),
                                    bind_buf(9, &dst),
                                    bind_buf(10, &pu),
                                ],
                            })
                        });
                    let mut pass = begin_pass(enc);
                    pass.set_pipeline(&c.bt_comp_fold);
                    pass.set_bind_group(0, &bind, &[]);
                    pass.dispatch_workgroups(1, 1, 1);
                } else {
                    let _salt = Dsv4FrameSalt::enter(t + 1);
                    let ckv_t = frame_buf(c, 70 + kind, cg.width * 4, false);
                    let csc_t = frame_buf(c, 72 + kind, cg.width * 4, false);
                    {
                        let mut pass = begin_pass(enc);
                        encode_blit_p(
                            &mut pass,
                            c,
                            &ckv_bt,
                            &ckv_t,
                            cg.width,
                            t * cg.width,
                            0,
                            None,
                        );
                        encode_blit_p(
                            &mut pass,
                            c,
                            &csc_bt,
                            &csc_t,
                            cg.width,
                            t * cg.width,
                            0,
                            None,
                        );
                    }
                    let off = if *kind == 0 {
                        p.comp_dst_off
                    } else {
                        p.ix_dst_off
                    };
                    let _ = comp_state_step(
                        c,
                        cw,
                        *cg,
                        *kind,
                        kv_id,
                        li,
                        &ckv_t,
                        &csc_t,
                        pos0 + t,
                        inv_freq,
                        &dst,
                        off,
                        enc,
                        None,
                    );
                }
            }
            if seg_this {
                append_seg(enc, seg0, batch - seg0);
            }
        }
    }

    mark(3);
    // The staged key rows, batched: project, norm, rotate, park at the tail.
    let kvfull = frame_buf_t(c, BT_DSPARK_KV, 20, batch * kvw * 4, false);
    let kvnorm = frame_buf_t(c, BT_DSPARK_KV, 21, batch * kvw * 4, false);
    if !dsv4_skip("win") {
        let Some(e) = model.tensors.get(p0.wkv) else {
            no!("wkv без записи");
        };
        let (Some(abs), plen) = (model.entry_abs_offset(e), e.nbytes as usize) else {
            no!("wkv без смещения");
        };
        let Some(wkv_w) = weight_buffer(c, (model.uid() as usize, p0.wkv), &bytes[abs..abs + plen])
        else {
            no!("wkv не влез");
        };
        if !encode_q4tp_mv4_b(c, enc, &wkv_w, &x2_bt, &kvfull, kvw, dim, batch) {
            no!("wkv пакетом не закодировался");
        }
        let knw = const_buf(c, bytemuck::cast_slice(&p0.kv_norm[..kvw]));
        let mut pass = begin_pass(enc);
        {
            let bind = cached_bind(c, bk(231), || {
                let pu = uniform_u32x4(c, [kvw as u32, 0, a.eps.to_bits(), 0]);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.rmsnorm_b.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &kvfull),
                        bind_buf(1, &knw),
                        bind_buf(2, &kvnorm),
                        bind_buf(3, &pu),
                    ],
                })
            });
            pass.set_pipeline(&c.rmsnorm_b);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(batch as u32, 1, 1);
        }
        {
            let bind = cached_bind(c, bk(232), || {
                let pu = uniform_u32x4(c, [1, kvw as u32, a.rd as u32, 0]);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_rope_heads.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &kvnorm),
                        bind_buf(1, &freq),
                        bind_buf(2, &pu),
                        bind_buf(3, &rope_meta),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_rope_heads);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(1, batch as u32, 1);
        }
        for t in 0..batch {
            encode_blit_p(
                &mut pass,
                c,
                &kvnorm,
                &cache,
                a.hd,
                t * kvw + kvw - a.hd,
                (srow0 + t) * a.hd,
                None,
            );
        }
    }

    // The indexer, batched: queries and head weights for every token in two
    // B-axis projections, then scores, top-k and the staged lists.
    let kmax = p0.idx_cap.saturating_sub(p0.window).max(1);
    let pick_bt = fb(203, kmax, false);
    let lim_bt = frame_buf_t(c, BT_SA_M, FRAME_TOK_STRIDE + li, batch * 4, true);
    let meta_bt = frame_buf_t(c, BT_FORCED, FRAME_TOK_STRIDE + li, batch * 4 * 4, true);
    let is_ix = p0.ix.is_some();
    let mut ms = vec![0u32; batch];
    let mut metas_ix = vec![0u32; batch * 4];
    let mut lims = vec![0u32; batch];
    for (t, p) in preps.iter().enumerate() {
        let adv_c = p.comp.as_ref().map_or(0, |(_, cg)| {
            usize::from(cg.ratio > 0 && (pos0 + t + 1) % cg.ratio.max(1) == 0)
        });
        let n_comp_t = p.n_comp + adv_c;
        let adv_i = p.ix.as_ref().map_or(0, |(_, cg, _, _)| {
            usize::from(cg.ratio > 0 && (pos0 + t + 1) % cg.ratio.max(1) == 0)
        });
        let n_ix_t = p.n_ix + adv_i;
        let old_vis = p0.filled.min(p0.window.saturating_sub(t + 1));
        let win_start = p0.filled - old_vis;
        let staged_n = t + 1;
        let k_t = if is_ix {
            let limit = n_ix_t.min(n_comp_t);
            lims[t] = limit as u32;
            if limit == 0 {
                0
            } else {
                p0.idx_cap
                    .saturating_sub(p0.window)
                    .min(4096)
                    .min({
                        let Some((_, _, _, ixg)) = p.ix.as_ref() else {
                            unreachable!()
                        };
                        ixg.top_k
                    })
                    .min(limit)
            }
        } else {
            n_comp_t
        };
        metas_ix[t * 4] = win_start as u32;
        metas_ix[t * 4 + 1] = old_vis as u32;
        metas_ix[t * 4 + 2] = staged_n as u32;
        metas_ix[t * 4 + 3] = k_t as u32;
        ms[t] = (old_vis + staged_n + k_t) as u32;
        if ms[t] > 1024 {
            no!("список токена {t} длиной {}", ms[t]);
        }
    }
    c.queue
        .write_buffer(&meta_bt, 0, bytemuck::cast_slice(&metas_ix));
    c.queue
        .write_buffer(&lim_bt, 0, bytemuck::cast_slice(&lims));
    let sa_m_li = frame_buf_t(c, BT_SA_M, li, batch * 4, true);
    c.queue.write_buffer(&sa_m_li, 0, bytemuck::cast_slice(&ms));
    if is_ix && !dsv4_skip("ix") {
        let Some((_, _, ixw, ixg)) = p0.ix.as_ref() else {
            unreachable!()
        };
        let mut ixb = Vec::with_capacity(2);
        for &idx in &[ixw.wq_b, ixw.weights_proj] {
            let Some(e) = model.tensors.get(idx) else {
                no!("индексер без тензора")
            };
            let (Some(abs), plen) = (model.entry_abs_offset(e), e.nbytes as usize) else {
                no!("индексер без смещения");
            };
            let Some(b) = weight_buffer_l(
                c,
                (model.uid() as usize, idx),
                &bytes[abs..abs + plen],
                layer_of_name(&model.tensors[idx].name),
            ) else {
                no!("индексер не влез");
            };
            ixb.push(b);
        }
        let ixkv = {
            let m2 = c.dsv4_ixkv.lock().unwrap();
            match m2.get(&(kv_id, li)) {
                Some((b, _)) => b.clone(),
                None => no!("ix-кеш ({kv_id},{li}) не заведён"),
            }
        };
        let qi_bt = frame_buf_t(c, BT_DSPARK_KV, 22, batch * ixg.ih * ixg.idim * 4, false);
        let hw_bt = frame_buf_t(c, BT_DSPARK_KV, 23, batch * ixg.ih * 4, false);
        let sc_bt = frame_buf_t(c, BT_DSPARK_KV, 24, batch * 4096 * 4, false);
        let cnt_bt = frame_buf_t(c, BT_DSPARK_KV, 25, batch * 4, false);
        mark(4);
        if !encode_q4tp_mv4_b(
            c,
            enc,
            &ixb[0],
            &qn_bt,
            &qi_bt,
            ixg.ih * ixg.idim,
            ixg.q_lora,
            batch,
        ) {
            no!("индексер q пакетом");
        }
        if !encode_q4tp_mv4_b(c, enc, &ixb[1], &x2_bt, &hw_bt, ixg.ih, ixg.hidden, batch) {
            no!("индексер веса пакетом");
        }
        mark(5);
        let sc_factor = (ixg.idim as f32).powf(-0.5) * (ixg.ih as f32).powf(-0.5);
        let n_pos = lims.iter().copied().max().unwrap_or(0) as usize;
        let mut pass = begin_pass(enc);
        {
            let bind = cached_bind(c, bk(233), || {
                let pu = uniform_u32x4(c, [ixg.ih as u32, ixg.idim as u32, ixg.rope_dim as u32, 0]);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_rope_heads.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &qi_bt),
                        bind_buf(1, &freq),
                        bind_buf(2, &pu),
                        bind_buf(3, &rope_meta),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_rope_heads);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(ixg.ih as u32, batch as u32, 1);
        }
        if n_pos > 0 {
            let bind = cached_bind(c, bk(234), || {
                let pu = uniform_mixed(c, [ixg.ih as u32, ixg.idim as u32, 4096], sc_factor);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_index_scores.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &qi_bt),
                        bind_buf(1, &ixkv),
                        bind_buf(2, &hw_bt),
                        bind_buf(3, &sc_bt),
                        bind_buf(4, &pu),
                        bind_buf(5, &lim_bt),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_index_scores);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(n_pos as u32, batch as u32, 1);
            let bind = cached_bind(c, bk(235), || {
                let pu = uniform_u32x4(c, [kmax as u32, 0, 0, 0]);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_top_k.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &sc_bt),
                        bind_buf(1, &pick_bt),
                        bind_buf(2, &cnt_bt),
                        bind_buf(3, &pu),
                        bind_buf(4, &lim_bt),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_top_k);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(batch as u32, 1, 1);
        }
    }
    mark(6);
    {
        let mut pass = begin_pass(enc);
        let bind = cached_bind(c, bk(236), || {
            let pu = uniform_u32x4(
                c,
                [p0.window as u32, kmax as u32, (!is_ix) as u32, srow0 as u32],
            );
            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &c.bt_idx_build_staged.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, &pick_bt),
                    bind_buf(1, &idx_bt),
                    bind_buf(2, &pu),
                    bind_buf(3, &meta_bt),
                ],
            })
        });
        pass.set_pipeline(&c.bt_idx_build_staged);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(4, batch as u32, 1);
        // The attends, all tokens at once, over the frozen cache.
        if !dsv4_skip("attn") && !dsv4_skip("sa") {
            let bind = cached_bind(c, bk(228), || {
                let pu = uniform_u32x4(c, [a.nh as u32, a.hd as u32, a.rd as u32, 1]);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_rope_heads.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &q_bt),
                        bind_buf(1, &freq),
                        bind_buf(2, &pu),
                        bind_buf(3, &rope_meta),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_rope_heads);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(a.nh as u32, batch as u32, 1);
            let bind = cached_bind(c, bk(229), || {
                let pu = uniform_mixed(c, [a.nh as u32, a.hd as u32, 1024], a.scale);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_sparse_attend.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &q_bt),
                        bind_buf(1, &cache),
                        bind_buf(2, &idx_bt),
                        bind_buf(3, &sink),
                        bind_buf(4, &attn_bt),
                        bind_buf(5, &pu),
                        bind_buf(6, &sa_m_li),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_sparse_attend);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(a.nh as u32, batch as u32, 1);
            let bind = cached_bind(c, bk(230), || {
                let pu = uniform_u32x4(c, [a.nh as u32, a.hd as u32, a.rd as u32, 2]);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_rope_heads.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &attn_bt),
                        bind_buf(1, &freq),
                        bind_buf(2, &pu),
                        bind_buf(3, &rope_meta),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_rope_heads);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(a.nh as u32, batch as u32, 1);
        }
    }
    {
        let _ = ();
        // Grouped output projection, one dispatch for the whole batch. The
        // per-token walk-exact loop cost five dependent links of the pass's
        // critical path per layer; the twin reduces in a different tree,
        // which the speculative mode's round-off contract already covers.
        if !dsv4_skip("olora") && !dsv4_skip("oproj") {
            let o_rows = a.o_groups * a.o_lora;
            let o_cols = a.nh * a.hd / a.o_groups;
            if ts_armed {
                // Attribution probe: an empty pass between the attends and
                // this dispatch absorbs the prior pass's drain into its own
                // window, so `olora` reads the dispatch alone.
                bt_ts(10);
                let _p = begin_pass(enc);
            }
            mark(7);
            // The staged twin shares one o-group's x span across its four
            // sub-rows; shapes that straddle a group keep the direct kernel.
            // Staging the span bought nothing on the release (0.41→0.45:
            // the weights, not x, are what this dispatch waits on) — the
            // twin stays for the next investigation, off by default.
            let ol_stage = {
                static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
                *ON.get_or_init(|| std::env::var("CMF_DSV4_OLORA_STAGE").is_ok_and(|v| v != "0"))
            };
            let p_ol = match c.bt_o_lora_a2.as_ref() {
                Some(p2) if ol_stage && o_cols <= 4608 && a.o_lora % 4 == 0 => p2,
                _ if a.o_lora % 4 == 0 && ol_a4() => &c.bt_o_lora_a4,
                _ => &c.bt_o_lora_a,
            };
            let bind = cached_bind(c, bk(213), || {
                let p = uniform_u32x4(c, [(o_cols / 32) as u32, o_rows as u32, a.o_lora as u32, 0]);
                // Entries follow the chosen pipeline's auto layout: the a4
                // twin reads x only through the vec4 view and drops the
                // scalar binding; the staged twin has neither vec4 view.
                let a4 = std::ptr::eq(
                    p_ol as *const wgpu::ComputePipeline,
                    &c.bt_o_lora_a4 as *const _,
                );
                let a2 = std::ptr::eq(
                    p_ol as *const wgpu::ComputePipeline,
                    c.bt_o_lora_a2
                        .as_ref()
                        .map_or(std::ptr::null(), |x| x as *const _),
                );
                let mut entries = vec![bind_buf(0, &wb[2])];
                if !a4 {
                    entries.push(bind_buf(1, &attn_bt));
                }
                entries.push(bind_buf(2, &mid_bt));
                entries.push(bind_buf(3, &p));
                if !a2 {
                    entries.push(bind_buf(4, &wb[2]));
                }
                if a4 {
                    entries.push(bind_buf(5, &attn_bt));
                }
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &p_ol.get_bind_group_layout(0),
                    entries: &entries,
                })
            });
            let mut pass = begin_pass(enc);
            pass.set_pipeline(p_ol);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups((o_rows as u32).div_ceil(4), batch as u32, 1);
        }
    }
    mark(15);
    if !dsv4_skip("wob")
        && !dsv4_skip("oproj")
        && !encode_q4tp_mv4_b(
            c,
            enc,
            &wb[3],
            &mid_bt,
            &ao_bt,
            dim,
            a.o_groups * a.o_lora,
            batch,
        )
    {
        no!("wo_b пакетом не закодировался");
    }

    mark(8);
    // ── glue and MoE, one pass ──
    let hcp = uniform_u32x8(
        c,
        [
            hc as u32,
            dim as u32,
            g.sinkhorn_iters as u32,
            g.hc_eps.to_bits(),
            0,
            0,
            0,
            0,
        ],
    );
    let hcp_n = uniform_u32x8(
        c,
        [
            hc as u32,
            dim as u32,
            g.sinkhorn_iters as u32,
            g.hc_eps.to_bits(),
            1,
            mix_hc as u32,
            0,
            0,
        ],
    );
    // The routing description: per-token forced rows when the layer hashes.
    let has_forced = forced_rows.is_some_and(|f| f.iter().any(|r| r.is_some()));
    if has_forced {
        let mut rows = vec![0u32; batch * m.top_k];
        for (t, r) in forced_rows.unwrap().iter().enumerate() {
            let Some(r) = r else {
                no!("хэш-слой без строки токена {t}");
            };
            if r.len() < m.top_k {
                no!("хэш-строка токена {t} короче top_k");
            }
            for (i, &e) in r[..m.top_k].iter().enumerate() {
                rows[t * m.top_k + i] = e as u32;
            }
        }
        c.queue
            .write_buffer(&forced_bt, 0, bytemuck::cast_slice(&rows));
    }
    {
        let mut pass = begin_pass(enc);
        let expand = |pass: &mut wgpu::ComputePass<'_>,
                      tag: u8,
                      x: &wgpu::Buffer,
                      res: &wgpu::Buffer,
                      out: &wgpu::Buffer| {
            let bind = cached_bind(c, bk(tag), || {
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_hc_post_expand.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, x),
                        bind_buf(1, res),
                        bind_buf(2, &hpost_bt),
                        bind_buf(3, &hcomb_bt),
                        bind_buf(4, out),
                        bind_buf(5, &hcp),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_hc_post_expand);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(((hc * dim) as u32).div_ceil(256), batch as u32, 1);
        };
        let mix = |pass: &mut wgpu::ComputePass<'_>,
                   tag: u8,
                   fnw: &wgpu::Buffer,
                   state: &wgpu::Buffer| {
            // The walk picks the 1024-thread kernel below 64 rows; the twin
            // must reduce in the same tree.
            let pipe = if mix_hc < 64 {
                &c.bt_f32_matvec_x
            } else {
                &c.bt_f32_matvec_w
            };
            let bind = cached_bind(c, bk(tag), || {
                let p = uniform_u32x4(c, [(hc * dim) as u32, mix_hc as u32, 0, 0]);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &pipe.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, fnw),
                        bind_buf(1, state),
                        bind_buf(2, &mixes_bt),
                        bind_buf(3, &p),
                    ],
                })
            });
            pass.set_pipeline(pipe);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(mix_hc as u32, batch as u32, 1);
        };
        let fold = |pass: &mut wgpu::ComputePass<'_>,
                    tag: u8,
                    state: &wgpu::Buffer,
                    sc: &wgpu::Buffer,
                    bs: &wgpu::Buffer,
                    nw: &wgpu::Buffer| {
            let bind = cached_bind(c, bk(tag), || {
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_hc_pre_fold.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, state),
                        bind_buf(1, &mixes_bt),
                        bind_buf(2, sc),
                        bind_buf(3, bs),
                        bind_buf(4, &fold_bt),
                        bind_buf(5, &hpost_bt),
                        bind_buf(6, &hcomb_bt),
                        bind_buf(7, &hcp_n),
                        bind_buf(8, nw),
                        bind_buf(9, &x2_bt),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_hc_pre_fold);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(batch as u32, 1, 1);
        };
        // The whole FFN-half join — expand, mix, Sinkhorn fold, norm — as
        // ONE link of the critical path.
        let fused = |pass: &mut wgpu::ComputePass<'_>,
                     tag: u8,
                     x: &wgpu::Buffer,
                     res: &wgpu::Buffer,
                     state_out: &wgpu::Buffer,
                     mixw: &wgpu::Buffer,
                     sc: &wgpu::Buffer,
                     bs: &wgpu::Buffer,
                     nw: &wgpu::Buffer| {
            let bind = cached_bind(c, bk(tag), || {
                let p = uniform_u32x8(
                    c,
                    [
                        hc as u32,
                        dim as u32,
                        g.sinkhorn_iters as u32,
                        g.hc_eps.to_bits(),
                        mix_hc as u32,
                        0,
                        0,
                        0,
                    ],
                );
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_hc_block.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, x),
                        bind_buf(1, res),
                        bind_buf(2, &hpost_bt),
                        bind_buf(3, &hcomb_bt),
                        bind_buf(4, mixw),
                        bind_buf(5, sc),
                        bind_buf(6, bs),
                        bind_buf(7, nw),
                        bind_buf(8, state_out),
                        bind_buf(9, &fold_bt),
                        bind_buf(10, &x2_bt),
                        bind_buf(11, &p),
                    ],
                })
            });
            pass.set_pipeline(&c.bt_hc_block);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(batch as u32, 1, 1);
        };
        if bt_hc_split() {
            expand(&mut pass, 240, &ao_bt, &state_bt, &state2_bt);
            mix(&mut pass, 241, &ffn_fn, &state2_bt);
            fold(&mut pass, 242, &state2_bt, &ffn_sc, &ffn_bs, &ffn_nw);
        } else {
            fused(
                &mut pass, 214, &ao_bt, &state_bt, &state2_bt, &ffn_fn, &ffn_sc, &ffn_bs, &ffn_nw,
            );
        }
        // ...router logits over the normed fold...
        if ts_armed {
            drop(pass);
            bt_ts(11);
            pass = begin_pass(enc);
        }
        {
            let pipe = if n_route < 64 {
                &c.bt_f32_matvec_x
            } else {
                &c.bt_f32_matvec_w
            };
            let bind = cached_bind(c, bk(217), || {
                let p = uniform_u32x4(c, [m.hidden as u32, n_route as u32, 0, 0]);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &pipe.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &router),
                        bind_buf(1, &x2_bt),
                        bind_buf(2, &logit_bt),
                        bind_buf(3, &p),
                    ],
                })
            });
            pass.set_pipeline(pipe);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(n_route as u32, batch as u32, 1);
        }
        // ...route, gate/up, down...
        let bs = match w.moe.bias {
            Some(b) if b.len() >= n_route => const_buf(c, bytemuck::cast_slice(&b[..n_route])),
            _ => logit_bt.clone(),
        };
        let subset = w.moe.remap.is_some();
        let rflags = (w.moe.bias.is_some_and(|b| b.len() >= n_route) as u32)
            | ((w.moe.mask.is_some_and(|m| m.len() >= n_route) as u32) << 1)
            | ((has_forced as u32) << 2)
            | 8
            | ((subset as u32) << 4)
            | ((n_pack as u32) << 8);
        let mk = match w.moe.mask {
            Some(mask) if mask.len() >= n_route => {
                const_buf(c, bytemuck::cast_slice(&mask[..n_route]))
            }
            _ => frame_buf(c, 17, n_route.max(1) * 4, true),
        };
        let rmap = match w.moe.remap {
            Some(remap) if remap.len() >= n_route => {
                const_buf(c, bytemuck::cast_slice(&remap[..n_route]))
            }
            _ => frame_buf(c, 26, n_route.max(1) * 4, true),
        };
        let cold = fb(203, 4 * m.top_k, false);
        {
            let bind = cached_bind(c, bk(218), || {
                let rp = uniform_mixed(c, [n_route as u32, m.top_k as u32, rflags], m.route_scale);
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.bt_moe_route.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &logit_bt),
                        bind_buf(1, &bs),
                        bind_buf(2, &mk),
                        bind_buf(3, &forced_bt),
                        bind_buf(4, &msel_bt),
                        bind_buf(5, &mwt_bt),
                        bind_buf(6, &mcnt_bt),
                        bind_buf(7, &rp),
                        bind_buf(8, &rmap),
                        bind_buf(9, &cold),
                    ],
                })
            });
            if !dsv4_skip("route") && !dsv4_skip("moe") {
                pass.set_pipeline(&c.bt_moe_route);
                pass.set_bind_group(0, &bind, &[]);
                pass.dispatch_workgroups(batch as u32, 1, 1);
            }
        }
        let stride16 = |rows: usize, cols: usize, q2: bool| -> u32 {
            let dt = if q2 {
                cortiq_core::TensorDtype::Q2TiledP
            } else {
                cortiq_core::TensorDtype::Q4TiledP
            };
            (cortiq_core::quant::expected_nbytes(dt, &[rows, cols]).unwrap_or(0) / 2) as u32
        };
        let gu_u = uniform_u32x8(
            c,
            [
                (m.hidden / 32) as u32,
                m.inter as u32,
                slots as u32,
                stride16(m.inter, m.hidden, m.gu_q2),
                m.swiglu_limit.to_bits(),
                0,
                0,
                0,
            ],
        );
        let dn_u = uniform_u32x4(
            c,
            [
                (m.inter / 32) as u32,
                m.hidden as u32,
                slots as u32,
                stride16(m.hidden, m.inter, false),
            ],
        );
        let gu_r4 = m.inter % 4 == 0 && bt_gu_r4_on();
        let p_gu = if m.gu_q2 {
            if gu_r4 {
                &c.bt_moe_gate_up_q2tp_r4
            } else {
                &c.bt_moe_gate_up_q2tp
            }
        } else if gu_r4 {
            &c.moe_gate_up_q4tp_b_r4
        } else {
            &c.moe_gate_up_q4tp_b
        };
        if ts_armed {
            drop(pass);
            bt_ts(12);
            pass = begin_pass(enc);
        }
        let bind = cached_bind(c, bk(219), || {
            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &p_gu.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, &gate_all),
                    bind_buf(1, &up_all),
                    bind_buf(2, &x2_bt),
                    bind_buf(3, &msel_bt),
                    bind_buf(4, &mact_bt),
                    bind_buf(5, &gu_u),
                ],
            })
        });
        if !dsv4_skip("gu") && !dsv4_skip("moe") {
            pass.set_pipeline(p_gu);
            pass.set_bind_group(0, &bind, &[]);
            let gx = if gu_r4 {
                (m.inter as u32).div_ceil(4)
            } else {
                m.inter as u32
            };
            pass.dispatch_workgroups(gx, slots as u32, batch as u32);
        }
        if ts_armed {
            drop(pass);
            bt_ts(13);
            pass = begin_pass(enc);
        }
        if bt_dn_mode() == 3 && stride16(m.hidden, m.inter, false) % 8 == 0 {
            let bind = cached_bind(c, bk(248), || {
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.moe_down_q4tp_b4.get_bind_group_layout(0),
                    // No entry for binding 1: the register-vec4 kernel
                    // reads x only through the vec4 view, and the auto
                    // layout drops what the entry point never touches.
                    entries: &[
                        bind_buf(0, &down_all),
                        bind_buf(2, &msel_bt),
                        bind_buf(3, &mwt_bt),
                        bind_buf(4, &mo_bt),
                        bind_buf(5, &dn_u),
                        bind_buf(6, &down_all),
                        bind_buf(7, &mact_bt),
                    ],
                })
            });
            if !dsv4_skip("dn") && !dsv4_skip("moe") {
                pass.set_pipeline(&c.moe_down_q4tp_b4);
                pass.set_bind_group(0, &bind, &[]);
                pass.dispatch_workgroups((m.hidden as u32).div_ceil(4), batch as u32, 1);
            }
        } else if bt_dn_mode() == 2 {
            // Per-slot partials, then the ascending-slot sum.
            let dpart = frame_buf_t(c, BT_DSPARK_KV, 30, batch * slots * m.hidden * 4, false);
            let bindp = cached_bind(c, bk(246), || {
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.moe_down_q4tp_part.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &down_all),
                        bind_buf(1, &mact_bt),
                        bind_buf(2, &msel_bt),
                        bind_buf(3, &mwt_bt),
                        bind_buf(4, &dpart),
                        bind_buf(5, &dn_u),
                    ],
                })
            });
            let bindr = cached_bind(c, bk(247), || {
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.moe_down_q4tp_red.get_bind_group_layout(0),
                    entries: &[bind_buf(0, &dpart), bind_buf(1, &mo_bt), bind_buf(2, &dn_u)],
                })
            });
            if !dsv4_skip("dn") && !dsv4_skip("moe") {
                pass.set_pipeline(&c.moe_down_q4tp_part);
                pass.set_bind_group(0, &bindp, &[]);
                pass.dispatch_workgroups(m.hidden as u32, slots as u32, batch as u32);
                pass.set_pipeline(&c.moe_down_q4tp_red);
                pass.set_bind_group(0, &bindr, &[]);
                pass.dispatch_workgroups((m.hidden as u32).div_ceil(256), batch as u32, 1);
            }
        } else {
            // The direct-load twin needs an even u16 stride so a tile's four
            // words ARE four words; every published layout satisfies it, the
            // check keeps an exotic one honest.
            let p_dn = if bt_dn_mode() == 1 && stride16(m.hidden, m.inter, false) % 2 == 0 {
                &c.moe_down_q4tp_b2
            } else {
                &c.moe_down_q4tp_b
            };
            let bind = cached_bind(c, bk(220), || {
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &p_dn.get_bind_group_layout(0),
                    entries: &[
                        bind_buf(0, &down_all),
                        bind_buf(1, &mact_bt),
                        bind_buf(2, &msel_bt),
                        bind_buf(3, &mwt_bt),
                        bind_buf(4, &mo_bt),
                        bind_buf(5, &dn_u),
                    ],
                })
            });
            if !dsv4_skip("dn") && !dsv4_skip("moe") {
                pass.set_pipeline(p_dn);
                pass.set_bind_group(0, &bind, &[]);
                pass.dispatch_workgroups(m.hidden as u32, batch as u32, 1);
            }
        }
        // ...the MoE half's join and the NEXT layer's opening, fused too.
        if ts_armed {
            drop(pass);
            bt_ts(14);
            pass = begin_pass(enc);
        }
        if let Some(nf) = w.hc_next_fn {
            let nfn = const_buf(c, bytemuck::cast_slice(nf));
            let nsc = const_buf(c, bytemuck::cast_slice(w.hc_next_scale));
            let nbs = const_buf(c, bytemuck::cast_slice(&w.hc_next_base[..mix_hc]));
            if bt_hc_split() {
                expand(&mut pass, 243, &mo_bt, &state2_bt, &state_bt);
                mix(&mut pass, 244, &nfn, &state_bt);
                fold(&mut pass, 245, &state_bt, &nsc, &nbs, &next_nw);
            } else {
                fused(
                    &mut pass, 221, &mo_bt, &state2_bt, &state_bt, &nfn, &nsc, &nbs, &next_nw,
                );
            }
        } else {
            expand(&mut pass, 221, &mo_bt, &state2_bt, &state_bt);
        }
    }
    mark(9);
    if w.hc_next_fn.is_some() && !dsv4_skip("nextq") {
        if !encode_q4tp_mv4_b(c, enc, &wb[4], &x2_bt, &qr_bt, a.q_lora, dim, batch) {
            no!("next-q пакетом не закодировался");
        }
        // One row-per-workgroup dispatch for every token's LoRA norm: the
        // per-token walk tree cost five dependent links; round-off class.
        let mut pass = begin_pass(enc);
        let bind = cached_bind(c, bk(227), || {
            let p = uniform_u32x4(c, [a.q_lora as u32, 0, a.eps.to_bits(), 0]);
            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &c.rmsnorm_b.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, &qr_bt),
                    bind_buf(1, &next_qn),
                    bind_buf(2, &qn_bt),
                    bind_buf(3, &p),
                ],
            })
        });
        pass.set_pipeline(&c.rmsnorm_b);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(batch as u32, 1, 1);
    }
    // ── the slide, once: outside a speculative verify the batch commits
    //    wholesale — the staged rows become the window's newest, the oldest
    //    leave, exactly what B walked appends would have left. A verify
    //    (retention armed) defers this to dsv4_spec_finish, which commits
    //    only the accepted prefix. ──
    mark(10);
    if SPEC_RETAIN.with(|v| v.borrow().0) == 0 && !dsv4_skip("win") {
        encode_staged_commit(c, enc, &cache, p0.filled, p0.window, a.hd, srow0, batch);
    }
    let src = if w.hc_next_fn.is_some() {
        x2_bt.clone()
    } else {
        ao_bt.clone()
    };
    if ts_armed {
        bt_ts(0);
    }
    Some((src, state_bt))
}

pub fn dsv4_layer_chain(
    model: &Arc<CmfModel>,
    layers: &[(Dsv4LayerW<'_>, Dsv4LayerGeom, Dsv4Prep<'_>)],
    kv_id: u64,
    first_li: usize,
    inv_freq: &[&[f32]],
    pos: usize,
    folded_out: &mut [f32],
    // When present, the hyper-connection state comes back in the SAME
    // submission — the caller then skips `dsv4_state_read` entirely.
    state_out: Option<&mut [f32]>,
) -> bool {
    let Some(c) = ctx() else { return false };
    let Some((_, g0, _)) = layers.first() else {
        return false;
    };
    let dim = g0.attn.dim;
    if folded_out.len() < dim || inv_freq.len() != layers.len() {
        return false;
    }
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-chain"),
        });
    let mut last = None;
    let t_enc = std::time::Instant::now();
    for (i, (w, g, p)) in layers.iter().enumerate() {
        // `qn = None` throughout: the first layer's LoRA vector was left on
        // the card by the caller's seed, every later one by the frame before.
        let Some(b) = dsv4_layer_frame_enc(
            model,
            w,
            *g,
            kv_id,
            first_li + i,
            0,
            false,
            false,
            false,
            None,
            &[],
            Some(p),
            inv_freq[i],
            pos,
            &mut enc,
        ) else {
            // Nothing has been submitted, so the token can still be run the
            // old way — but the caches this chain advanced have NOT been
            // touched either, because every write went into this encoder.
            return false;
        };
        last = Some(b);
    }
    let Some(src) = last else { return false };
    // The whole run's HOST encode time, before the fence: the number that
    // says whether the chain is submit-bound or encode-bound.
    CHAIN_ENC_NS.fetch_add(
        t_enc.elapsed().as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    CHAIN_LAYERS.fetch_add(layers.len() as u64, std::sync::atomic::Ordering::Relaxed);
    CHAIN_RUNS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    // From here to the fence is the wait the counter below reports.
    let t_wait = std::time::Instant::now();

    let ok = match state_out {
        // The state comes back in the SAME submission as the folded vector.
        // Asking for it afterwards cost a second fence on a token that is
        // otherwise one submission — a third of the round trips, for a copy
        // of a few kilobytes.
        Some(st) => {
            let sb = frame_buf(c, 40, st.len() * 4, true);
            readback2(c, enc, (&src, &mut folded_out[..dim]), (&sb, st))
        }
        None => {
            let mut sc = c.scratch.lock().unwrap();
            let stage = Scratch::ensure(
                &c.device,
                &mut sc.stage,
                (dim * 4) as u64,
                wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                "dsv4-chain-stage",
            );
            let ok = readback(
                c,
                enc,
                &src,
                &stage,
                (dim * 4) as u64,
                &mut folded_out[..dim],
            );
            drop(sc);
            ok
        }
    };
    CHAIN_WAIT_NS.fetch_add(
        t_wait.elapsed().as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    ok
}

/// Host encode and fence-wait of the chain, per run, with the layer count —
/// the split that decides where its 12.5 tok/s goes.
pub static CHAIN_ENC_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static CHAIN_WAIT_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static CHAIN_LAYERS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
/// Submissions per token — one fence each. The number that says whether a
/// layer kind is still breaking the chain into pieces.
pub static CHAIN_RUNS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// Exact hand-off from one token-axis partial layer.  Resident experts have
/// already been expanded into `states`; `cold` names only the routed winners
/// absent from the pack, and `cold_x` is the corresponding normalized FFN
/// input.  The host computes those few experts, adds `posts[j] * cold`, then
/// seeds the next layer.  No routing restriction is involved.
pub struct Dsv4PartialBatchOut {
    pub states: Vec<f32>,
    pub cold_x: Vec<f32>,
    pub posts: Vec<f32>,
    pub cold: Vec<Vec<(usize, f32)>>,
    /// Exact global top-k selected by the device, per token.  The caller uses
    /// the guaranteed-accepted first row to keep the live LRU pack in step
    /// with ordinary sequential decode.
    pub routed: Vec<Vec<usize>>,
}

/// Run one partial layer over the token axis and return everything needed to
/// make its cold correction before the next dependent layer.  This is the
/// missing middle ground between a complete fused chain and the old all-host
/// tail: attention and resident experts stay batched on the device, while
/// exact cold experts are completed from the checkpoint on the host.
#[allow(clippy::too_many_arguments)]
pub fn dsv4_layer_batch_partial(
    model: &Arc<CmfModel>,
    w: &Dsv4LayerW,
    g: Dsv4LayerGeom,
    kv_id: u64,
    li: usize,
    batch: usize,
    preps: &[Dsv4Prep],
    forced_rows: Option<&[Option<Vec<usize>>]>,
    inv_freq: &[f32],
    pos0: usize,
) -> Option<Dsv4PartialBatchOut> {
    let c = ctx()?;
    if batch == 0 || batch > FRAME_TOK_STRIDE || w.moe.remap.is_none() || w.hc_next_fn.is_some() {
        return None;
    }
    let (hc, dim, top_k) = (g.hc, g.attn.dim, g.moe.top_k);
    if let Some((_, _, _, ixg)) = preps.first()?.ix.as_ref() {
        let need = preps
            .iter()
            .map(|p| p.n_ix)
            .max()
            .unwrap_or(0)
            .saturating_add(2);
        dsv4_index_cache(kv_id, li, ixg.idim * need)?;
    }
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-partial-batch"),
        });

    // Replay after a rejected speculative suffix needs the exact input of
    // every layer.  It is live in BT_X2 before this frame overwrites that
    // buffer with the FFN-normalized input used by the cold experts.
    let (retain_n, _) = SPEC_RETAIN.with(|v| v.borrow().clone());
    if retain_n > li {
        let x2_bt = frame_buf_t(c, BT_X2, 0, batch * dim * 4, true);
        let retain = frame_buf_t(c, BT_RETAIN, 0, retain_n * batch * dim * 4, false);
        let mut pass = begin_pass(&mut enc);
        encode_blit_p(
            &mut pass,
            c,
            &x2_bt,
            &retain,
            batch * dim,
            0,
            li * batch * dim,
            None,
        );
    }

    let (_, state_bt) = dsv4_layer_frame_bt_enc(
        model,
        w,
        g,
        kv_id,
        li,
        batch,
        preps,
        forced_rows,
        inv_freq,
        pos0,
        &mut enc,
    )?;
    let x2_bt = frame_buf_t(c, BT_X2, 0, batch * dim * 4, true);
    let post_bt = frame_buf_t(c, BT_HPOST, 0, batch * hc * 4, true);
    let cold_bt = frame_buf_t(c, 203, 0, batch * 4 * top_k * 4, false);

    let state_bytes = batch * hc * dim * 4;
    let x_bytes = batch * dim * 4;
    let post_bytes = batch * hc * 4;
    let cold_bytes = batch * 4 * top_k * 4;
    let total = state_bytes + x_bytes + post_bytes + cold_bytes;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        total as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dsv4-partial-batch-stage",
    );
    flush_pass(&enc);
    let mut off = 0u64;
    for (src, n) in [
        (&state_bt, state_bytes),
        (&x2_bt, x_bytes),
        (&post_bt, post_bytes),
        (&cold_bt, cold_bytes),
    ] {
        enc.copy_buffer_to_buffer(src, 0, &stage, off, n as u64);
        off += n as u64;
    }
    submit(c, finish_enc(enc));
    let slice = stage.slice(..total as u64);
    slice.map_async(wgpu::MapMode::Read, |_| {});
    c.device.poll(wgpu::PollType::wait_indefinitely()).ok()?;
    let data = slice.get_mapped_range().ok()?;
    let mut at = 0usize;
    let mut take_f32 = |n: usize| {
        let end = at + n * 4;
        let out: Vec<f32> = bytemuck::cast_slice(&data[at..end]).to_vec();
        at = end;
        out
    };
    let states = take_f32(batch * hc * dim);
    let cold_x = take_f32(batch * dim);
    let posts = take_f32(batch * hc);
    let cold_words: &[u32] = bytemuck::cast_slice(&data[at..at + cold_bytes]);
    let mut cold = vec![Vec::new(); batch];
    let mut routed = vec![Vec::new(); batch];
    for t in 0..batch {
        let row = &cold_words[t * 4 * top_k..(t + 1) * 4 * top_k];
        for slot in 0..top_k {
            let ei = row[2 * slot];
            if ei != u32::MAX {
                cold[t].push((ei as usize, f32::from_bits(row[2 * slot + 1])));
            }
            let global = row[2 * top_k + 2 * slot];
            if global != u32::MAX {
                routed[t].push(global as usize);
            }
        }
    }
    drop(data);
    stage.unmap();
    drop(sc);
    Some(Dsv4PartialBatchOut {
        states,
        cold_x,
        posts,
        cold,
        routed,
    })
}

/// Replace one armed capture photograph with an exact host-corrected batch.
/// Queue writes are ordered after the partial layer's readback fence and
/// before the eventual capture read, so no extra submission is needed.
pub fn dsv4_spec_cap_write_host(li: usize, batch: usize, hc_dim: usize, states: &[f32]) -> bool {
    let Some(c) = ctx() else { return false };
    let (_, caps) = SPEC_RETAIN.with(|v| v.borrow().clone());
    let Some(slot) = caps.iter().position(|&x| x == li) else {
        return true;
    };
    if states.len() < batch * hc_dim {
        return false;
    }
    let cap = frame_buf_t(c, BT_CAP, 0, caps.len() * batch * hc_dim * 4, true);
    c.queue.write_buffer(
        &cap,
        ((slot * batch * hc_dim) * 4) as u64,
        bytemuck::cast_slice(&states[..batch * hc_dim]),
    );
    true
}

/// One layer, submitted on its own and read back — the shape the two-frame
/// path and the current layer loop use.
#[allow(clippy::too_many_arguments)]
pub fn dsv4_layer_frame(
    model: &Arc<CmfModel>,
    w: &Dsv4LayerW,
    g: Dsv4LayerGeom,
    kv_id: u64,
    li: usize,
    qn: Option<&[f32]>,
    idxs: &[u32],
    inv_freq: &[f32],
    pos: usize,
    folded_next: &mut [f32],
    // Subset packs (live remap in `w.moe.remap`) return the winners the
    // slots do not hold; the caller completes them on the host and owes
    // the state the correction (`dsv4_state_add_cold`). None keeps the
    // full-pack contract: every winner resident, nothing to return.
    mut cold_out: Option<&mut Vec<(usize, f32)>>,
    // The route shader already mirrors every GLOBAL winner and its normalized
    // weight into the second half of its cold buffer. Stats consume that
    // mirror from the same readback/fence; there is no diagnostic submission
    // on the hot path. Keeping the weight matters: mask `cover` is output
    // mass, not a vote count where the weakest and strongest top-k routes are
    // treated as equal.
    mut route_out: Option<&mut Vec<(usize, f32)>>,
    cold_x_out: &mut Vec<f32>,
) -> bool {
    let Some(c) = ctx() else { return false };
    let dim = g.attn.dim;
    if folded_next.len() < dim {
        return false;
    }
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-layer"),
        });
    let Some(folded) = dsv4_layer_frame_enc(
        model, w, g, kv_id, li, 0, false, false, false, qn, idxs, None, inv_freq, pos, &mut enc,
    ) else {
        return false;
    };
    let cold_bytes = (4 * g.moe.top_k * 4) as u64;
    let read_route = cold_out.is_some() || route_out.is_some();
    let x_bytes = if cold_out.is_some() {
        (dim * 4) as u64
    } else {
        0
    };
    let total = (dim * 4) as u64 + if read_route { cold_bytes + x_bytes } else { 0 };
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        total,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dsv4-layer-stage",
    );
    if read_route {
        // The cold list rides the same staging and the same fence as the
        // fold — one barrier pays for both, exactly like the two-frame path.
        enc.copy_buffer_to_buffer(&folded, 0, &stage, 0, (dim * 4) as u64);
        let cold_src = if w.moe.remap.is_some() {
            store_slot(c, 27, kv_id, li, &[])
        } else {
            frame_buf(c, 27, 4 * g.moe.top_k * 4, false)
        };
        enc.copy_buffer_to_buffer(&cold_src, 0, &stage, (dim * 4) as u64, cold_bytes);
        if x_bytes > 0 {
            enc.copy_buffer_to_buffer(
                &frame_buf_t(c, 45, 0, dim * 4, true),
                0,
                &stage,
                (dim * 4) as u64 + cold_bytes,
                x_bytes,
            );
        }
        submit(c, finish_enc(enc));
        let slice = stage.slice(..total);
        slice.map_async(wgpu::MapMode::Read, |_| {});
        if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
            return false;
        }
        let mut ok = false;
        if let Ok(data) = slice.get_mapped_range() {
            folded_next[..dim].copy_from_slice(bytemuck::cast_slice(&data[..dim * 4]));
            let ct = dim * 4 + cold_bytes as usize;
            let tail: &[u32] = bytemuck::cast_slice(&data[dim * 4..ct]);
            if let Some(cold_out) = cold_out.as_deref_mut() {
                cold_out.clear();
                for t in 0..g.moe.top_k {
                    if tail[2 * t] != u32::MAX {
                        cold_out.push((tail[2 * t] as usize, f32::from_bits(tail[2 * t + 1])));
                    }
                }
                cold_x_out.clear();
                cold_x_out.extend_from_slice(bytemuck::cast_slice(&data[ct..total as usize]));
            }
            if let Some(route_out) = route_out.as_deref_mut() {
                route_out.clear();
                for t in 0..g.moe.top_k {
                    let e = tail[2 * g.moe.top_k + 2 * t];
                    if e != u32::MAX {
                        route_out.push((
                            e as usize,
                            f32::from_bits(tail[2 * g.moe.top_k + 2 * t + 1]),
                        ));
                    }
                }
            }
            ok = true;
        }
        stage.unmap();
        drop(sc);
        return ok;
    }
    let ok = readback(
        c,
        enc,
        &folded,
        &stage,
        (dim * 4) as u64,
        &mut folded_next[..dim],
    );
    drop(sc);
    ok
}

/// Is this tensor resident, or can it be made so? Uploads it if it can.
pub fn dsv4_weight_ready(model: &Arc<CmfModel>, idx: usize) -> bool {
    let Some(c) = ctx() else { return false };
    let Some(e) = model.tensors.get(idx) else {
        return false;
    };
    let Some(abs) = model.entry_abs_offset(e) else {
        return false;
    };
    let bytes = model.primary_bytes();
    let plen = e.nbytes as usize;
    if abs + plen > bytes.len() {
        return false;
    }
    weight_buffer_l(
        c,
        (model.uid() as usize, idx),
        &bytes[abs..abs + plen],
        layer_of_name(&model.tensors[idx].name),
    )
    .is_some()
}

/// Whether this adapter can execute the segmented global mixed Q2TP/Q4TP MoE cache.
/// Descriptor-indexing is optional in wgpu, so unsupported platforms keep
/// using the parity-proven per-layer slot banks.
pub fn dsv4_global_moe_supported() -> bool {
    ctx().is_some_and(|c| {
        c.dsv4_global_gu.is_some() && c.dsv4_global_gu_q2.is_some() && c.dsv4_global_dn.is_some()
    }) && std::env::var("CMF_DSV4_GLOBAL_POOL").as_deref() != Ok("0")
}

/// Select the pipeline family matching one cache's descriptor-array width.
/// Keeping this selection beside the allocator prevents an S16 buffer bank
/// from ever being submitted with the S8 bind-group layout.
fn dsv4_global_moe_pipelines(
    c: &Ctx,
    gu_q2: bool,
    segments: usize,
) -> Option<(&wgpu::ComputePipeline, &wgpu::ComputePipeline)> {
    match segments {
        DSV4_GLOBAL_MOE_SEGMENTS => Some((
            if gu_q2 {
                c.dsv4_global_gu_q2.as_ref()?
            } else {
                c.dsv4_global_gu.as_ref()?
            },
            c.dsv4_global_dn.as_ref()?,
        )),
        DSV4_GLOBAL_MOE_SEGMENTS_S16 => Some((
            if gu_q2 {
                c.dsv4_global_gu_q2_s16.as_ref()?
            } else {
                c.dsv4_global_gu_s16.as_ref()?
            },
            c.dsv4_global_dn_s16.as_ref()?,
        )),
        _ => None,
    }
}

fn dsv4_global_moe_segments_for_dsv41(c: &Ctx) -> usize {
    if std::env::var("CMF_DSV41_GLOBAL_SEGMENTS").as_deref() == Ok("16")
        && c.dsv4_global_gu_s16.is_some()
        && c.dsv4_global_gu_q2_s16.is_some()
        && c.dsv4_global_dn_s16.is_some()
    {
        DSV4_GLOBAL_MOE_SEGMENTS_S16
    } else {
        DSV4_GLOBAL_MOE_SEGMENTS
    }
}

/// Pure capacity calculation shared by the allocator and its bounded geometry
/// tests. `workspace` is reserved in units of bytes before segment rounding.
fn dsv4_global_moe_capacity(
    requested: usize,
    per: usize,
    max_len: usize,
    range: u64,
    workspace: u64,
    segments: usize,
) -> Option<(usize, usize)> {
    if per == 0 || max_len == 0 || segments == 0 {
        return None;
    }
    let max_seg = (range / max_len as u64) as usize;
    let requested = requested.saturating_sub((workspace / per as u64) as usize);
    let capacity = requested.min(max_seg.saturating_mul(segments)) / segments * segments;
    (capacity >= segments).then_some((capacity, capacity / segments))
}

/// Allocate the single model-wide mixed-Q2TP/Q4TP bank cache using the
/// parity-proven S8 geometry. Generic models retain this entry point.
pub fn dsv4_global_moe_create(
    model: &Arc<CmfModel>,
    requested: usize,
    inter: usize,
    hidden: usize,
    gu_q2: bool,
) -> Option<(usize, usize)> {
    dsv4_global_moe_create_with_segments(
        model,
        requested,
        inter,
        hidden,
        gu_q2,
        DSV4_GLOBAL_MOE_SEGMENTS,
    )
}

/// Allocate the V4.1 global bank, selecting S16 only when the explicit
/// `CMF_DSV41_GLOBAL_SEGMENTS=16` request was accepted during device init.
/// If the adapter or shader family cannot provide S16, this safely falls back
/// to the existing S8 pool.
pub fn dsv4_global_moe_create_for_dsv41(
    model: &Arc<CmfModel>,
    requested: usize,
    inter: usize,
    hidden: usize,
    gu_q2: bool,
) -> Option<(usize, usize)> {
    let c = ctx()?;
    let segments = dsv4_global_moe_segments_for_dsv41(&c);
    dsv4_global_moe_create_with_segments(model, requested, inter, hidden, gu_q2, segments)
}

fn dsv4_global_moe_create_with_segments(
    model: &Arc<CmfModel>,
    requested: usize,
    inter: usize,
    hidden: usize,
    gu_q2: bool,
    segments: usize,
) -> Option<(usize, usize)> {
    use std::sync::atomic::Ordering;
    let c = ctx()?;
    if dsv4_global_moe_pipelines(&c, gu_q2, segments).is_none() {
        return None;
    }
    if let Some(b) = c.dsv4_global_moe.lock().unwrap().get(&model.uid()).cloned() {
        return (b.gu_q2 == gu_q2 && b.segments == segments)
            .then_some((b.capacity, b.segment_slots));
    }
    let gu_len = cortiq_core::quant::expected_nbytes(
        if gu_q2 {
            cortiq_core::TensorDtype::Q2TiledP
        } else {
            cortiq_core::TensorDtype::Q4TiledP
        },
        &[inter, hidden],
    )?;
    let d_len =
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[hidden, inter])?;
    let per = 2usize.checked_mul(gu_len)?.checked_add(d_len)?;
    let range = c
        .device
        .limits()
        .max_storage_buffer_binding_size
        .min(c.device.limits().max_buffer_size);
    // Binding-array descriptors, per-layer activation/cold-readback buffers,
    // KV growth and queue staging are physical VRAM too but are not counted
    // as resident weights. Reserve 2-4 GiB before rounding the logical bank.
    let gib = 1024 * 1024 * 1024u64;
    let workspace = (c.vram_budget / 10).clamp(2 * gib, 4 * gib);
    let (capacity, segment_slots) = dsv4_global_moe_capacity(
        requested,
        per,
        gu_len.max(d_len),
        range,
        workspace,
        segments,
    )?;
    let total = (capacity as u64).checked_mul(per as u64)?;
    if c.resident.load(Ordering::Relaxed).saturating_add(total) > c.vram_budget {
        return None;
    }
    let mk = |label: &'static str, plen: usize| -> Vec<wgpu::Buffer> {
        (0..segments)
            .map(|_| {
                c.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some(label),
                    size: (segment_slots * plen) as u64,
                    usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                })
            })
            .collect()
    };
    let bufs = Arc::new(Dsv4GlobalMoeBufs {
        gate: mk("dsv4-global-gate", gu_len),
        up: mk("dsv4-global-up", gu_len),
        down: mk("dsv4-global-down", d_len),
        capacity,
        segment_slots,
        gu_len,
        d_len,
        gu_q2,
        segments,
    });
    c.resident.fetch_add(total, Ordering::Relaxed);
    note_resident_peak(c);
    c.dsv4_global_moe.lock().unwrap().insert(model.uid(), bufs);
    tracing::info!(
        "DSV4 unified global pool: {} slots, {} segments × {}, {} MB, gate/up {}, requested {}",
        capacity,
        segments,
        segment_slots,
        total / 1024 / 1024,
        if gu_q2 { "q2tp" } else { "q4tp" },
        requested,
    );
    Some((capacity, segment_slots))
}

#[cfg(test)]
mod global_moe_capacity_tests {
    use super::{dsv4_global_moe_capacity, host_tier_cache_admissible};

    const RANGE: u64 = (1u64 << 31) - 1;
    const PER: usize = 18_424_832;
    const MAX_LEN: usize = 6_149_120;
    const GIB: u64 = 1024 * 1024 * 1024;

    #[test]
    fn disabled_host_tier_rejects_cache_admission() {
        // The test command runs with CMF_RAM_TIER_MB unset. This is the
        // production default that previously paid for a discarded clone.
        assert!(!host_tier_cache_admissible(6_149_120));
    }

    #[test]
    fn s8_capacity_matches_existing_geometry() {
        assert_eq!(
            dsv4_global_moe_capacity(3146, PER, MAX_LEN, RANGE, 4 * GIB, 8),
            Some((2792, 349))
        );
    }

    #[test]
    fn s16_capacity_rounds_to_whole_segments() {
        assert_eq!(
            dsv4_global_moe_capacity(3146, PER, MAX_LEN, RANGE, 4 * GIB, 16),
            Some((2912, 182))
        );
    }

    #[test]
    fn s16_small_budget_stays_inside_allocator_envelope() {
        let budget = 16 * GIB;
        let requested = (budget as usize * 75 / 100) / PER;
        let workspace = (budget / 10).clamp(2 * GIB, 4 * GIB);
        let (capacity, slots) =
            dsv4_global_moe_capacity(requested, PER, MAX_LEN, RANGE, workspace, 16)
                .expect("16 GiB profile should retain a usable S16 bank");
        assert_eq!(capacity, slots * 16);
        assert!(capacity as u64 * PER as u64 + workspace <= budget);
    }
}

pub fn dsv4_global_moe_ready(model: &Arc<CmfModel>) -> bool {
    ctx().is_some_and(|c| c.dsv4_global_moe.lock().unwrap().contains_key(&model.uid()))
}

/// Install one `(layer, expert)` triple into a flat global slot. Slot metadata
/// is committed by the caller only after all three ordered queue writes
/// succeed, so a remap never exposes partially replaced weights.
pub fn dsv4_global_slot_fill(model: &Arc<CmfModel>, slot: usize, t: (usize, usize, usize)) -> bool {
    use std::sync::atomic::Ordering;
    let t_fill = dsv4_fill_profile_on().then(std::time::Instant::now);
    let finish = |ok: bool| {
        if let Some(started) = t_fill.as_ref() {
            DSV4_FILL_NS.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed);
        }
        ok
    };
    let Some(c) = ctx() else { return finish(false) };
    let Some(b) = c.dsv4_global_moe.lock().unwrap().get(&model.uid()).cloned() else {
        return finish(false);
    };
    if slot >= b.capacity {
        return finish(false);
    }
    let seg = slot / b.segment_slots;
    let local = slot % b.segment_slots;
    let bytes = model.primary_bytes();
    let put = |buf: &wgpu::Buffer, idx: usize, plen: usize| -> bool {
        let Some(e) = model.tensors.get(idx) else {
            return false;
        };
        if e.nbytes as usize != plen {
            return false;
        }
        let Some(abs) = model.entry_abs_offset(e) else {
            return false;
        };
        let off = local.saturating_mul(plen);
        if off.saturating_add(plen) as u64 > buf.size() {
            return false;
        }
        let key = (model.uid() as usize, idx);
        let tier = host_tier_get(key);
        let src: &[u8] = if let Some(v) = tier.as_deref() {
            if v.len() != plen {
                return false;
            }
            v
        } else if let Some(v) = pread_range(model, abs, plen) {
            c.queue.write_buffer(buf, off as u64, &v);
            host_tier_put(key, Arc::new(v));
            return true;
        } else {
            let Some(src) = bytes.get(abs..abs + plen) else {
                return false;
            };
            src
        };
        c.queue.write_buffer(buf, off as u64, src);
        // `host_tier_put` is a no-op when CMF_RAM_TIER_MB is unset/zero.
        // Check that admission is possible before cloning the full mapped
        // tensor; this path runs for every global-bank miss and otherwise
        // discarded one allocation per projection. A full tier may evict
        // sampled entries in `host_tier_put`, so only reject the intrinsic
        // no-tier and over-budget cases here.
        if tier.is_none() && host_tier_cache_admissible(plen) {
            host_tier_put(key, Arc::new(src.to_vec()));
        }
        true
    };
    let ok = put(&b.gate[seg], t.0, b.gu_len)
        && put(&b.up[seg], t.1, b.gu_len)
        && put(&b.down[seg], t.2, b.d_len);
    if ok {
        DSV4_FILLS.fetch_add(1, Ordering::Relaxed);
        DSV4_FILL_BYTES.fetch_add((2 * b.gu_len + b.d_len) as u64, Ordering::Relaxed);
    }
    finish(ok)
}

/// How many experts of this shape still fit on the card. The caller packs
/// that many and leaves the rest to the host — per EXPERT, so no layer ever
/// has to leave the device wholesale.
pub fn dsv4_experts_fit(inter: usize, hidden: usize, gu_q2: bool, dn_q2: bool) -> usize {
    use std::sync::atomic::Ordering;
    let Some(c) = ctx() else { return 0 };
    let gu = if gu_q2 {
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q2TiledP, &[inter, hidden])
    } else {
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[inter, hidden])
    }
    .unwrap_or(0);
    let dn = cortiq_core::quant::expected_nbytes(
        if dn_q2 {
            cortiq_core::TensorDtype::Q2TiledP
        } else {
            cortiq_core::TensorDtype::Q4TiledP
        },
        &[hidden, inter],
    )
    .unwrap_or(0);
    let per = (2 * gu + dn) as u64;
    if per == 0 {
        return 0;
    }
    let used = c.resident.load(Ordering::Relaxed);
    // A weight budget is not the physical allocation ceiling: command
    // staging, KV/state buffers and the driver's own bookkeeping live beside
    // it. Keep a small geometry-independent reserve that scales down on
    // small cards and caps out on large ones. The value is intentionally a
    // function of the configured budget, never a checkpoint/layer cutoff.
    let mib = 1024 * 1024u64;
    let reserve = std::env::var("CMF_GPU_WORKSPACE_MB")
        .ok()
        .and_then(|v| v.parse::<u64>().ok())
        .map(|v| v.saturating_mul(mib))
        .unwrap_or_else(|| {
            // The base covers command staging, KV growth and driver
            // bookkeeping; the draft term carves out what the speculative
            // draft's own resident pack will take, so the trunk's greedy
            // packing stops before the draft has nowhere to live. Packing
            // against the physical ceiling fails NONDETERMINISTICALLY as
            // the KV grows mid-run, so the base rides twice when a draft
            // is coming.
            let base = (c.vram_budget / 384).clamp(256 * mib, 512 * mib);
            let draft = DRAFT_RESERVE.load(Ordering::Relaxed);
            if draft > 0 {
                2 * base + draft
            } else {
                base
            }
        });
    let usable = c.vram_budget.saturating_sub(reserve);
    ((usable.saturating_sub(used)) / per) as usize
}

/// What the speculative draft's device pack will need, set at load when the
/// file carries an MTP stack and speculation is not disabled. Zero means no
/// draft is coming and the trunk may pack into the whole budget.
pub static DRAFT_RESERVE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
/// Expert-bank part of `DRAFT_RESERVE`.  The draft builder may reclaim this
/// part when asking how many experts fit; graph workspace must stay reserved.
pub static DRAFT_PACK_RESERVE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// f32 GEMM for the skill-bake trainer: y[n,m] = x[n,k] · w[m,k]ᵀ on the
/// card, riding the existing batched f32 matvec kernel. The weight side
/// is the bake's long-lived f32 replica, cached device-side by address —
/// uploaded once, reused for every one of the run's hundreds of steps.
/// False = caller keeps its CPU path (no device, tiny job, or oversize
/// dispatch). Bit-parity is NOT claimed (GPU sum order differs); the
/// trainer's loss landscape does not care, and the quality gate at the
/// end is the arbiter.
use crate::gpu::fp_bytes;

/// One bake-GPU op at a time: gemm_nt / gemm_dx / the FFN chain share the
/// pooled slots (bx/by/bst/bb/ba), and two callers interleaving dispatch
/// and readback on the same buffers return each other's numbers. The bake
/// is sequential in practice; this makes the API safe rather than lucky.
static BAKE_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());

/// Strict-f32 mode for the bake GEMMs: the coop arms and device chains
/// stand down and everything runs the scalar f32 kernels. Phase A turns
/// this ON for its training steps: the mask SELECTS neurons by a
/// gradient signal, and f16 operand rounding (~1e-3) on that signal
/// compounds over ~90 steps into closing the wrong ones — measured as
/// hard-PPL 5.207 vs 4.293 at the same 2.56% sparsity, twice, while an
/// f32 run reproduces the reference trajectory to the third decimal.
/// Forward/eval sweeps stay on tensor cores everywhere: their PPL
/// matches the f32 path exactly at print precision.
static BAKE_F32_STRICT: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);

pub fn bake_precision_strict(on: bool) {
    BAKE_F32_STRICT.store(on, std::sync::atomic::Ordering::Relaxed);
}

#[inline]
fn f32_strict() -> bool {
    BAKE_F32_STRICT.load(std::sync::atomic::Ordering::Relaxed)
}

/// A bake weight on the card — but RESIDENT only once it has proven it
/// recurs. Phase B hands this function a fresh activation matrix per dw
/// step; a cache that adopts everything it sees grows by 250 MB a step
/// and killed a 90-step phase B at step ~40 with a silent device OOM.
/// The ledger: first sighting of (ptr, fp) uploads a TRANSIENT buffer
/// (freed with the submission) and only records the fingerprint; the
/// second sighting with the same fp promotes to resident. An in-place
/// update (Adam masters: same ptr, new fp on a resident entry) refreshes
/// the resident buffer rather than demoting it. Transients never repeat
/// a fingerprint, so they never occupy a byte past their own call.
/// Do all of these fit as single storage buffers?
///
/// wgpu caps one buffer — 4 292 870 144 bytes on an A100 — and
/// `create_buffer` treats a request over it as a FATAL validation error,
/// not a recoverable one. A 248 320-vocab head in f32 is 5 085 593 600
/// bytes, so baking a 27B panicked in the middle of phase A instead of
/// declining to the host arm that was sitting right there.
pub(crate) fn buffers_fit(max_buffer_size: u64, sizes: &[u64]) -> bool {
    sizes.iter().all(|&s| s <= max_buffer_size)
}

fn bake_weight(c: &Ctx, w: &[f32], label: &'static str) -> wgpu::Buffer {
    let key = (w.as_ptr() as usize, w.len());
    let fp = fp_bytes(bytemuck::cast_slice(w));
    let fresh = |note: bool| {
        if note {
            crate::gpu::probe_note_cold();
        }
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size: (w.len() * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        c.queue.write_buffer(&b, 0, bytemuck::cast_slice(w));
        b
    };
    let mut cb = c.gemm_w_bufs.lock().unwrap();
    match cb.get_mut(&key) {
        Some((Some(b), f)) if *f == fp => b.clone(),
        Some((slot @ Some(_), f)) => {
            // Resident, contents moved on: an in-place weight update.
            crate::gpu::probe_note_cold();
            *f = fp;
            let b = slot.as_ref().unwrap().clone();
            c.queue.write_buffer(&b, 0, bytemuck::cast_slice(w));
            b
        }
        Some((slot @ None, f)) if *f == fp => {
            // Second sighting of the same bytes — it recurs; promote.
            let b = fresh(true);
            *slot = Some(b.clone());
            b
        }
        Some((None, f)) => {
            // Same address, different transient — stay transient.
            *f = fp;
            fresh(false)
        }
        None => {
            cb.insert(key, (None, fp));
            fresh(true)
        }
    }
}

/// A 2-D convolution as a GEMM on the matrix units: im2col ON THE CARD
/// into a pixel tile, then `gemm_nt_coop` against the weight, which is
/// already `[oc, ic·k·k]` in the order the patch column is built — so
/// nothing is repacked. The scalar `vae_conv` kernel this replaces runs
/// at about 220 GFLOP/s; a 512×512×128×128 3×3 conv is 77 GFLOP of it.
///
/// Tiled over pixels because the column matrix does not fit: at that
/// size it would be 1.2 GB whole. One tile of 16k pixels is ~75 MB.
#[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 {
    let Some(c) = ctx() else { return false };
    let (Some(im2col), Some(gemm)) = (c.vae_im2col.as_ref(), c.gemm_nt_coop.as_ref()) else {
        return false;
    };
    let kdim = ic * k * k;
    let hw = h * wi;
    // The coop kernel reads the reduction four wide, and its grid is
    // one workgroup per 64×64 output tile.
    if kdim % 4 != 0
        || k % 2 == 0
        || w.len() < oc * kdim
        || x.len() < ic * hw
        || out.len() < oc * hw
    {
        return false;
    }
    let pad = k / 2;
    let _gate = c.mm_gate.lock().unwrap();
    let tile = (1usize << 24) / kdim.max(1); // ~64 MB of columns
    let tile = tile.clamp(64, hw).min(hw);
    let wbuf = bake_weight(c, &w[..oc * kdim], "vae-conv-w");
    let st = wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST;
    let (xb, colb, yb, stage) = {
        let mut sc = c.scratch.lock().unwrap();
        (
            Scratch::ensure(
                &c.device,
                &mut sc.vcx,
                (ic * hw * 4) as u64,
                st,
                "vae-conv-x",
            ),
            Scratch::ensure(
                &c.device,
                &mut sc.vcc,
                (tile * kdim * 4) as u64,
                st,
                "vae-conv-col",
            ),
            Scratch::ensure(
                &c.device,
                &mut sc.vcy,
                (tile * oc * 4) as u64,
                st | wgpu::BufferUsages::COPY_SRC,
                "vae-conv-y",
            ),
            Scratch::ensure(
                &c.device,
                &mut sc.vcs,
                (hw * oc * 4) as u64,
                wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
                "vae-conv-stage",
            ),
        )
    };
    c.queue
        .write_buffer(&xb, 0, bytemuck::cast_slice(&x[..ic * hw]));
    // Every tile lands in ONE device buffer and comes home once. Reading
    // each tile back instead drains the queue per tile — eighteen stalls
    // in a 512×512 conv, which is why the first version of this was no
    // faster than the scalar kernel it replaced.
    let ybig = {
        let mut sc = c.scratch.lock().unwrap();
        Scratch::ensure(
            &c.device,
            &mut sc.vcb,
            (hw * oc * 4) as u64,
            wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_DST
                | wgpu::BufferUsages::COPY_SRC,
            "vae-conv-all",
        )
    };
    let mut yt = vec![0f32; hw * oc];
    let mut p0 = 0usize;
    while p0 < hw {
        let t = tile.min(hw - p0);
        let up = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&[
                    ic as u32, h as u32, wi as u32, k as u32, pad as u32, p0 as u32, t as u32, 0u32,
                ]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let bg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &im2col.get_bind_group_layout(0),
            entries: &[bind_buf(0, &xb), bind_buf(1, &colb), bind_buf(2, &up)],
        });
        let ug = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&[(kdim / 4) as u32, oc as u32, t as u32, 0u32]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let bgg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &gemm.get_bind_group_layout(0),
            entries: [
                bind_buf(0, &wbuf),
                bind_buf(1, &colb),
                bind_buf(2, &yb),
                bind_buf(3, &ug),
            ]
            .as_slice(),
        });
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("vae-conv"),
            });
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(im2col);
            pass.set_bind_group(0, &bg, &[]);
            let wgs = ((t * kdim) as u32).div_ceil(256);
            pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
        }
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(gemm);
            pass.set_bind_group(0, &bgg, &[]);
            pass.dispatch_workgroups((oc as u32).div_ceil(64), (t as u32).div_ceil(64), 1);
        }
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&yb, 0, &ybig, (p0 * oc * 4) as u64, (t * oc * 4) as u64);
        submit(c, finish_enc(enc));
        p0 += t;
    }
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("vae-conv-out"),
        });
    if !readback(
        c,
        enc_take(&mut enc),
        &ybig,
        &stage,
        (hw * oc * 4) as u64,
        &mut yt,
    ) {
        return false;
    }
    // [hw, oc] → NCHW, bias on the way, across the pool.
    let po = SendPtrC(out.as_mut_ptr());
    let work = |lo: usize, hi: usize| {
        for o in lo..hi {
            let b = bias.map_or(0.0, |bb| bb[o]);
            // SAFETY: one output channel per worker, disjoint.
            let dst = unsafe { std::slice::from_raw_parts_mut(po.0.add(o * hw), hw) };
            for (p, d) in dst.iter_mut().enumerate() {
                *d = yt[p * oc + o] + b;
            }
        }
    };
    // No pool handle down here, and the scatter is memory-bound anyway.
    work(0, oc);
    true
}

struct SendPtrC(*mut f32);
unsafe impl Send for SendPtrC {}
unsafe impl Sync for SendPtrC {}

fn enc_take(e: &mut wgpu::CommandEncoder) -> wgpu::CommandEncoder {
    flush_pass(&e);
    std::mem::replace(
        e,
        ctx()
            .unwrap()
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }),
    )
}

/// A 1D convolution as an NT GEMM whose A matrix never exists on the
/// host. `yt` comes back as `[out_n x oc]` — the same layout
/// `gemm_nt_f32` produces, so the caller's epilogue is unchanged.
///
/// Measured on a 3090: the host arm of this convolution spent its time
/// building a `kk x out_n` column buffer, transposing it into a second
/// buffer of equal size, and uploading that — up to 2.37 GB of traffic
/// for a 20-second song. Only `x` (k times smaller) crosses the bus now.
#[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 {
    if std::env::var("CMF_BAKE_GPU").as_deref() == Ok("0") || f32_strict() {
        return false;
    }
    let Some(c) = ctx() else { return false };
    let (Some(ic_pipe), Some(gemm)) = (c.conv1d_im2col.as_ref(), c.gemm_nt_coop.as_ref()) else {
        return false;
    };
    let kk = ic * k;
    // The coop kernel reads its reduction four wide, and a job this small
    // loses to the round trip either way.
    if kk % 4 != 0 || out_n == 0 || out_n * kk * oc < (1 << 22) {
        return false;
    }
    if x.len() < ic * n || w.len() < oc * kk || yt.len() < out_n * oc {
        return false;
    }
    if oc.div_ceil(64) > 65_535 {
        return false;
    }
    // One column tile stays well under Vulkan's 2 GiB per-binding limit.
    const MAX_COL_FLOATS: usize = 128 << 20; // 512 MB
    let span = (MAX_COL_FLOATS / kk).max(1).min(out_n);
    if span.div_ceil(64) > 65_535 {
        return false;
    }
    let _bake = BAKE_LOCK.lock().unwrap();
    let wbuf = bake_weight(c, &w[..oc * kk], "conv1d-w");
    let xbuf = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("conv1d-x"),
            contents: bytemuck::cast_slice(&x[..ic * n]),
            usage: wgpu::BufferUsages::STORAGE,
        });
    let mut sc = c.scratch.lock().unwrap();
    let colbuf = Scratch::ensure(
        &c.device,
        &mut sc.bx,
        (span * kk * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "conv1d-col",
    );
    let ybuf = Scratch::ensure(
        &c.device,
        &mut sc.by,
        (span * oc * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "conv1d-y",
    );
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.bst,
        (span * oc * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "conv1d-stage",
    );
    drop(sc);
    let mut p0 = 0usize;
    while p0 < out_n {
        let tile = span.min(out_n - p0);
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("conv1d"),
            });
        let cu = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&[
                    ic as u32,
                    n as u32,
                    k as u32,
                    dil as u32,
                    pad as u32,
                    p0 as u32,
                    tile as u32,
                    0u32,
                ]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let cb = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &ic_pipe.get_bind_group_layout(0),
            entries: &[bind_buf(0, &xbuf), bind_buf(1, &colbuf), bind_buf(2, &cu)],
        });
        // The grid folds into y past 65535 workgroups, matching the
        // shader's `gid.y * (65535 * 256)` unpack.
        let groups = (tile * kk).div_ceil(256) as u32;
        let gx = groups.min(65_535);
        let gy = groups.div_ceil(65_535);
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(ic_pipe);
            pass.set_bind_group(0, &cb, &[]);
            pass.dispatch_workgroups(gx, gy, 1);
        }
        let gu = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&[(kk / 4) as u32, oc as u32, tile as u32, 0u32]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let gb = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &gemm.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wbuf),
                bind_buf(1, &colbuf),
                bind_buf(2, &ybuf),
                bind_buf(3, &gu),
            ],
        });
        {
            let mut pass = begin_pass(&mut enc);
            pass.set_pipeline(gemm);
            pass.set_bind_group(0, &gb, &[]);
            pass.dispatch_workgroups((oc as u32).div_ceil(64), (tile as u32).div_ceil(64), 1);
        }
        let bytes = (tile * oc * 4) as u64;
        if !readback(
            c,
            enc,
            &ybuf,
            &stage,
            bytes,
            &mut yt[p0 * oc..(p0 + tile) * oc],
        ) {
            return false;
        }
        p0 += tile;
    }
    true
}

/// `gemm_nt_f32` for operands that change every call: `w` goes to a
/// reused scratch buffer instead of the resident ledger, and its bytes
/// are never fingerprinted. This is the shape the refit accumulation
/// needs — `C[n, m] = X[n, k]·Wᵀ[k, m]` where both operands are fresh
/// activations, not weights.
pub fn gemm_nt_f32_transient(
    x: &[f32],
    w: &[f32],
    y: &mut [f32],
    n: usize,
    k: usize,
    m: usize,
) -> bool {
    TRANSIENT_W.with(|t| t.set(true));
    let r = gemm_nt_f32(x, w, y, n, k, m);
    TRANSIENT_W.with(|t| t.set(false));
    r
}

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

pub fn gemm_nt_f32(x: &[f32], w: &[f32], y: &mut [f32], n: usize, k: usize, m: usize) -> bool {
    if std::env::var("CMF_BAKE_GPU").as_deref() == Ok("0") || f32_strict() {
        return false;
    }
    let Some(c) = ctx() else { return false };
    // Small jobs lose to the round trip.
    if n * k * m < (1 << 22) {
        return false;
    }
    // Tensor cores when the device brought them up: 64x64 tiles, so the
    // dispatch grid is m/64 x n/64 and the vocabulary head (m = 166k, which
    // the scalar arm's one-workgroup-per-element grid could never dispatch)
    // becomes an ordinary call. Held to discrete cards for the big-m case:
    // on unified memory the head weight's device copy would be 2 GB of the
    // laptop's RAM that the CPU path does not spend.
    let coop = c
        .gemm_nt_coop
        .as_ref()
        .filter(|_| k % 4 == 0 && m.div_ceil(64) <= 65_535 && n.div_ceil(64) <= 65_535)
        .filter(|_| c.discrete || (m <= 65_000 && n <= 65_000));
    // The scalar arm's dispatch is (m, n) workgroups — huge dims are
    // illegal, and a huge PRODUCT is a hang: a phase-B dw shape
    // (3072×21504 = 66M workgroups, four billion threads) ran minutes
    // per call on an M4 and looked like a stuck bake. Matvec-style
    // kernels are for matvec-style grids; anything bigger is the CPU's.
    if coop.is_none() && (m > 65_000 || n > 65_000 || m.saturating_mul(n) > (1 << 22)) {
        return false;
    }
    if x.len() < n * k || w.len() < m * k || y.len() < n * m {
        return false;
    }
    // Decline rather than die: see `buffers_fit`.
    if !buffers_fit(
        c.device.limits().max_buffer_size,
        &[(w.len() * 4) as u64, (n * k * 4) as u64, (n * m * 4) as u64],
    ) {
        return false;
    }
    let _bake = BAKE_LOCK.lock().unwrap();
    let transient = TRANSIENT_W.with(|t| t.get());
    let wbuf = if transient {
        let mut sc = c.scratch.lock().unwrap();
        let b = Scratch::ensure(
            &c.device,
            &mut sc.bwt,
            (w.len() * 4) as u64,
            wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            "gemm-wt",
        );
        drop(sc);
        c.queue.write_buffer(&b, 0, bytemuck::cast_slice(w));
        b
    } else {
        bake_weight(c, w, "bake-w")
    };
    let mut sc = c.scratch.lock().unwrap();
    let xbuf = Scratch::ensure(
        &c.device,
        &mut sc.bx,
        (n * k * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "bake-x",
    );
    let ybuf = Scratch::ensure(
        &c.device,
        &mut sc.by,
        (n * m * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "bake-y",
    );
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.bst,
        (n * m * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "bake-stage",
    );
    drop(sc);
    c.queue
        .write_buffer(&xbuf, 0, bytemuck::cast_slice(&x[..n * k]));
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("bake-gemm"),
        });
    if let Some(pipe) = coop {
        // MmP { cols4, rows, nb } in the coop kernel's terms: rows = weight
        // rows (m), nb = tokens (n).
        let u = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&[(k / 4) as u32, m as u32, n as u32, 0u32]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &pipe.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wbuf),
                bind_buf(1, &xbuf),
                bind_buf(2, &ybuf),
                bind_buf(3, &u),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(pipe);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((m as u32).div_ceil(64), (n as u32).div_ceil(64), 1);
    } else {
        let u = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&[k as u32, m as u32, 0u32, 0u32]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.f32_matvec_b.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wbuf),
                bind_buf(1, &xbuf),
                bind_buf(2, &ybuf),
                bind_buf(3, &u),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.f32_matvec_b);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(m as u32, n as u32, 1);
    }
    let bytes = (n * m * 4) as u64;
    readback(c, enc, &ybuf, &stage, bytes, &mut y[..n * m])
}

/// Backward twin of `gemm_nt_f32`: dx[n,k] = dy[n,m] · w[m,k], the
/// gradient the bake pushes back through a frozen projection. Phase A
/// calls it three times a layer and nothing else touches the card
/// there, so this is what decides whether a bake is GPU work or not.
pub fn gemm_dx_f32(dy: &[f32], w: &[f32], dx: &mut [f32], n: usize, k: usize, m: usize) -> bool {
    if std::env::var("CMF_BAKE_GPU").as_deref() == Ok("0") || f32_strict() {
        return false;
    }
    let Some(c) = ctx() else { return false };
    if n * k * m < (1 << 22) || k > 65_000 || n > 65_000 {
        return false;
    }
    if dy.len() < n * m || w.len() < m * k || dx.len() < n * k {
        return false;
    }
    let wbuf = {
        // The streaming replica re-dequantizes evicted layers into recycled
        // Vecs — same address, different matrix — so the hit must prove the
        // contents too, or a narrow residency window bakes with the wrong
        // frozen weights.
        let key = (w.as_ptr() as usize, w.len());
        let fp = fp_bytes(bytemuck::cast_slice(w));
        let mut cb = c.const_bufs.lock().unwrap();
        if let Some((b, f)) = cb.get_mut(&key) {
            if *f != fp {
                crate::gpu::probe_note_cold();
                c.queue.write_buffer(b, 0, bytemuck::cast_slice(w));
                *f = fp;
            }
            b.clone()
        } else {
            crate::gpu::probe_note_cold(); // first touch = upload, not a steady sample
            let b = c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("bake-dxw"),
                size: (w.len() * 4) as u64,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            c.queue.write_buffer(&b, 0, bytemuck::cast_slice(w));
            cb.insert(key, (b.clone(), fp));
            b
        }
    };
    let _bake = BAKE_LOCK.lock().unwrap();
    // Shares the forward GEMM's pooled slots under the same bake lock.
    let mut sc = c.scratch.lock().unwrap();
    let dyb = Scratch::ensure(
        &c.device,
        &mut sc.bx,
        (n * m * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "bake-x",
    );
    let dxb = Scratch::ensure(
        &c.device,
        &mut sc.by,
        (n * k * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "bake-y",
    );
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.bst,
        (n * k * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "bake-stage",
    );
    drop(sc);
    c.queue
        .write_buffer(&dyb, 0, bytemuck::cast_slice(&dy[..n * m]));
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("bake-dx"),
        });
    let coop = c
        .gemm_nn_coop
        .as_ref()
        .filter(|_| m % 4 == 0 && k.div_ceil(64) <= 65_535 && n.div_ceil(64) <= 65_535);
    // The scalar dx arm has NO probe arbitrating it — the call site takes
    // whatever this function accepts, shadowing the Accelerate arm on
    // macOS. So without cooperative matrices it accepts only genuinely
    // matvec-sized grids; a phase-B dx (3072×256 workgroups) ground a
    // Mac bake to 12% CPU while the BLAS arm sat right below it.
    if coop.is_none() && k.saturating_mul(n) > (1 << 16) {
        return false;
    }
    if let Some(pipe) = coop {
        // MmP for the NN twin: cols4 = reduction (m) / 4, rows = dx's
        // width (k), nb = tokens (n).
        let u = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&[(m / 4) as u32, k as u32, n as u32, 0u32]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &pipe.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wbuf),
                bind_buf(1, &dyb),
                bind_buf(2, &dxb),
                bind_buf(3, &u),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(pipe);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((k as u32).div_ceil(64), (n as u32).div_ceil(64), 1);
    } else {
        let u = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&[m as u32, k as u32, 0u32, 0u32]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.f32_gemm_dx.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wbuf),
                bind_buf(1, &dyb),
                bind_buf(2, &dxb),
                bind_buf(3, &u),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.f32_gemm_dx);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(k as u32, n as u32, 1);
    }
    let bytes = (n * k * 4) as u64;
    readback(c, enc, &dxb, &stage, bytes, &mut dx[..n * k])
}

/// The bake's frozen FFN as ONE device chain: both = n2·guᵀ, act =
/// silu(g)·u·(scale), ffn = act·downᵀ — one submit, one map. The gate+up
/// plane (n×2·inter, tens of MB) stops crossing PCIe at all on eval calls
/// — most of a bake — and a training call reads it back beside ffn in the
/// same map for the backward pass. Numerics are the tensor-core GEMM's
/// (f16 operands, f32 accumulator) with silu in plain f32 on device.
/// Returns false untouched wherever the cooperative pipeline is absent or
/// a shape steps outside the tiles — the caller keeps its scalar path.
#[allow(clippy::too_many_arguments)]
pub fn ffn_chain_f32(
    n2: &[f32],
    gu: &[f32],
    down: &[f32],
    scale: Option<&[f32]>,
    both_out: Option<&mut [f32]>,
    plane: Option<usize>,
    ffn: &mut [f32],
    n: usize,
    hsz: usize,
    inter: usize,
) -> bool {
    if std::env::var("CMF_BAKE_GPU").as_deref() == Ok("0") {
        return false;
    }
    let Some(c) = ctx() else { return false };
    let Some(nt) = c.gemm_nt_coop.as_ref() else {
        return false;
    };
    if f32_strict() || hsz % 4 != 0 || inter % 4 != 0 || n * hsz * 2 * inter < (1 << 22) {
        return false;
    }
    if (2 * inter).div_ceil(64) > 65_535 || n.div_ceil(64) > 65_535 || hsz.div_ceil(64) > 65_535 {
        return false;
    }
    if n2.len() < n * hsz
        || gu.len() < 2 * inter * hsz
        || down.len() < hsz * inter
        || ffn.len() < n * hsz
        || scale.is_some_and(|s| s.len() < inter)
        || both_out.as_ref().is_some_and(|b| b.len() < n * 2 * inter)
    {
        return false;
    }
    let _bake = BAKE_LOCK.lock().unwrap();
    let wgu = bake_weight(c, gu, "bake-gu");
    let wdn = bake_weight(c, down, "bake-dn");
    let ffn_bytes = (n * hsz * 4) as u64;
    let both_bytes = (n * 2 * inter * 4) as u64;
    // ffn first, both after — the second slice starts 16-aligned.
    let off = ffn_bytes.div_ceil(16) * 16;
    let stage_need = if both_out.is_some() {
        off + both_bytes
    } else {
        ffn_bytes
    };
    let mut sc = c.scratch.lock().unwrap();
    let xbuf = Scratch::ensure(
        &c.device,
        &mut sc.bx,
        (n * hsz * 4) as u64,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "bake-x",
    );
    // A training call parks the plane per LAYER so the backward chain can
    // consume it without the round trip; anything else shares one slot.
    // Discrete cards only — on unified memory those 22 planes are RAM.
    let both = match plane.filter(|_| c.discrete) {
        Some(li) => {
            let mut planes = c.bake_planes.lock().unwrap();
            let slot = planes.entry(li).or_insert((
                c.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("bake-plane"),
                    size: both_bytes.next_power_of_two(),
                    usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                    mapped_at_creation: false,
                }),
                both_bytes.next_power_of_two(),
            ));
            if slot.1 < both_bytes {
                crate::gpu::probe_note_cold();
                slot.0 = c.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("bake-plane"),
                    size: both_bytes.next_power_of_two(),
                    usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
                    mapped_at_creation: false,
                });
                slot.1 = both_bytes.next_power_of_two();
            }
            slot.0.clone()
        }
        None => Scratch::ensure(
            &c.device,
            &mut sc.bb,
            both_bytes,
            wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            "bake-both",
        ),
    };
    let act = Scratch::ensure(
        &c.device,
        &mut sc.ba,
        (n * inter * 4) as u64,
        wgpu::BufferUsages::STORAGE,
        "bake-act",
    );
    let ybuf = Scratch::ensure(
        &c.device,
        &mut sc.by,
        ffn_bytes,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "bake-y",
    );
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.bst,
        stage_need,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "bake-stage",
    );
    drop(sc);
    c.queue
        .write_buffer(&xbuf, 0, bytemuck::cast_slice(&n2[..n * hsz]));
    // The mask gate rides a tiny upload; a static zero word keeps the
    // binding well-formed when no scale is active.
    static NO_SCALE: [u8; 4] = [0; 4];
    let sbuf = match scale {
        Some(s) => const_buf(c, bytemuck::cast_slice(&s[..inter])),
        None => const_buf(c, &NO_SCALE),
    };
    let uni = |v: [u32; 4]| {
        c.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&v),
                usage: wgpu::BufferUsages::UNIFORM,
            })
    };
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("bake-ffn"),
        });
    {
        let u1 = uni([(hsz / 4) as u32, (2 * inter) as u32, n as u32, 0]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &nt.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wgu),
                bind_buf(1, &xbuf),
                bind_buf(2, &both),
                bind_buf(3, &u1),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(nt);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((2 * inter as u32).div_ceil(64), (n as u32).div_ceil(64), 1);
    }
    {
        let total = (n * inter) as u32;
        let wgs = total.div_ceil(256);
        let gx = wgs.min(32_768);
        let gy = wgs.div_ceil(32_768);
        let u2 = uni([inter as u32, total, scale.is_some() as u32, gx * 256]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.bake_silu.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &both),
                bind_buf(1, &sbuf),
                bind_buf(2, &act),
                bind_buf(3, &u2),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.bake_silu);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(gx, gy, 1);
    }
    {
        let u3 = uni([(inter / 4) as u32, hsz as u32, n as u32, 0]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &nt.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wdn),
                bind_buf(1, &act),
                bind_buf(2, &ybuf),
                bind_buf(3, &u3),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(nt);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((hsz as u32).div_ceil(64), (n as u32).div_ceil(64), 1);
    }
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&ybuf, 0, &stage, 0, ffn_bytes);
    if both_out.is_some() {
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(&both, 0, &stage, off, both_bytes);
    }
    submit(c, finish_enc(enc));
    let slice = stage.slice(..stage_need);
    let done = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
    let d2 = done.clone();
    slice.map_async(wgpu::MapMode::Read, move |_| {
        d2.store(true, std::sync::atomic::Ordering::Release);
    });
    if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        return false;
    }
    if !done.load(std::sync::atomic::Ordering::Acquire) {
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
    }
    {
        let Ok(data) = slice.get_mapped_range() else {
            return false;
        };
        par_copy(
            &mut ffn[..n * hsz],
            bytemuck::cast_slice(&data[..ffn_bytes as usize]),
        );
        if let Some(bo) = both_out {
            par_copy(
                &mut bo[..n * 2 * inter],
                bytemuck::cast_slice(&data[off as usize..(off + both_bytes) as usize]),
            );
        }
    }
    stage.unmap();
    true
}

/// Release everything the bake parked on the card: resident f32 weights
/// (~14 GB on a 4 B model), per-layer planes, pooled activations and
/// staging. Called when a bake hands the process to something else — the
/// runtime gate opens the model through the ordinary engine, and on a
/// 32 GB card the two residencies do not fit together (measured: the
/// gate OOM-panicked the moment the planes joined the weights).
pub fn bake_release() {
    let Some(c) = ctx() else { return };
    let _bake = BAKE_LOCK.lock().unwrap();
    c.bake_planes.lock().unwrap().clear();
    c.gemm_w_bufs.lock().unwrap().clear();
    let mut sc = c.scratch.lock().unwrap();
    sc.bx = None;
    sc.by = None;
    sc.bst = None;
    sc.bb = None;
    sc.ba = None;
}

/// Everything the trainer's backward needs from a device attention
/// forward, read back in one map. `qkv_plane` is the RAW fused
/// projection (bias not yet applied) — the host derives qpre/kpre/
/// vproj/gate_pre from it exactly as the host path's split does.
pub struct AttnChainActs {
    pub qkv_plane: Vec<f32>,
    pub qrot: Vec<f32>,
    pub krot: Vec<f32>,
    pub qinv: Vec<f32>,
    pub kinv: Vec<f32>,
    pub ao: Vec<f32>,
}

/// Configuration mirror of the replica's `full_attn_fwd` — every field
/// the host math reads, so the chain can reproduce it bit-honestly.
pub struct AttnChainCfg<'a> {
    pub wqkv: &'a [f32],
    pub wo: &'a [f32],
    pub q_norm: Option<&'a [f32]>,
    pub k_norm: Option<&'a [f32]>,
    /// Concatenated (bq | bk | bv) or None.
    pub bias: Option<&'a [f32]>,
    pub output_gate: bool,
    pub gemma: bool,
    pub eps: f32,
    pub rotary_half: usize,
    /// cos/sin table [t × rotary_half × 2], f32 from HOST f64 trig.
    pub rope: &'a [f32],
    pub b: usize,
    pub t: usize,
    pub nh: usize,
    pub nkv: usize,
    pub hd: usize,
    pub hsz: usize,
}

/// The bake's attention forward as ONE device chain: qkv GEMM (tensor
/// cores) → bias+split+qk-norm+RoPE → causal softmax·V per head →
/// output gate → wo GEMM (tensor cores) → attn_out. An eval call reads
/// back 3 MB; a training call adds the activation planes in the same
/// map. Declines (false) in strict-f32 mode and wherever a shape steps
/// outside what the kernels take — the caller keeps the host path.
pub fn attn_chain_f32(
    n1: &[f32],
    cfg: &AttnChainCfg,
    attn_out: &mut [f32],
    want_acts: bool,
) -> Option<AttnChainActs> {
    if std::env::var("CMF_BAKE_GPU").as_deref() == Ok("0") || f32_strict() {
        return None;
    }
    let c = ctx()?;
    let nt = c.gemm_nt_coop.as_ref()?;
    let (b, t, nh, nkv, hd, hsz) = (cfg.b, cfg.t, cfg.nh, cfg.nkv, cfg.hd, cfg.hsz);
    let n = b * t;
    let qdim = nh * hd;
    let kvdim = nkv * hd;
    let qrows = if cfg.output_gate { 2 * qdim } else { qdim };
    let fused = qrows + 2 * kvdim;
    if hd > 128
        || t > 1024
        || hsz % 4 != 0
        || nh % nkv != 0
        || n * hsz * fused < (1 << 22)
        || cfg.rotary_half * 2 > hd
        || cfg.rope.len() < t * cfg.rotary_half * 2
        || n1.len() < n * hsz
        || cfg.wqkv.len() < fused * hsz
        || cfg.wo.len() < hsz * qdim
        || attn_out.len() < n * hsz
    {
        return None;
    }
    // Same ceiling as the GEMM path — decline, do not die.
    if !buffers_fit(
        c.device.limits().max_buffer_size,
        &[(cfg.wqkv.len() * 4) as u64, (cfg.wo.len() * 4) as u64],
    ) {
        return None;
    }
    let _bake = BAKE_LOCK.lock().unwrap();
    let wqkv = bake_weight(c, cfg.wqkv, "bake-wqkv");
    let wo = bake_weight(c, cfg.wo, "bake-wo");
    let f4 = |x: usize| (x * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let xbuf = Scratch::ensure(
        &c.device,
        &mut sc.bx,
        f4(n * hsz),
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "bake-x",
    );
    let plane = Scratch::ensure(
        &c.device,
        &mut sc.bb,
        f4(n * fused),
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "bake-both",
    );
    let qrot = Scratch::ensure(
        &c.device,
        &mut sc.bqr,
        f4(n * qdim),
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "bake-qrot",
    );
    let krot = Scratch::ensure(
        &c.device,
        &mut sc.bkr,
        f4(n * kvdim),
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "bake-krot",
    );
    let vproj = Scratch::ensure(
        &c.device,
        &mut sc.bvp,
        f4(n * kvdim),
        wgpu::BufferUsages::STORAGE,
        "bake-vproj",
    );
    let gate = Scratch::ensure(
        &c.device,
        &mut sc.bgp,
        f4((n * qdim).max(1)),
        wgpu::BufferUsages::STORAGE,
        "bake-gate",
    );
    let ao = Scratch::ensure(
        &c.device,
        &mut sc.bao,
        f4(n * qdim),
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "bake-ao",
    );
    let ao_eff = Scratch::ensure(
        &c.device,
        &mut sc.bae,
        f4(n * qdim),
        wgpu::BufferUsages::STORAGE,
        "bake-aoeff",
    );
    let qinv = Scratch::ensure(
        &c.device,
        &mut sc.biq,
        f4(n * nh),
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "bake-qinv",
    );
    let kinv = Scratch::ensure(
        &c.device,
        &mut sc.bik,
        f4(n * nkv),
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "bake-kinv",
    );
    let ybuf = Scratch::ensure(
        &c.device,
        &mut sc.by,
        f4(n * hsz),
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "bake-y",
    );
    // Readback plan: attn_out first, then (training) the act planes,
    // every slice 16-aligned into one staging buffer.
    let mut offs = vec![(0u64, f4(n * hsz))];
    if want_acts {
        for sz in [n * fused, n * qdim, n * kvdim, n * nh, n * nkv, n * qdim] {
            let start = offs.last().map(|(o, l)| (o + l).div_ceil(16) * 16).unwrap();
            offs.push((start, f4(sz)));
        }
    }
    let need = offs.last().map(|(o, l)| o + l).unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.bst,
        need,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "bake-stage",
    );
    drop(sc);
    c.queue
        .write_buffer(&xbuf, 0, bytemuck::cast_slice(&n1[..n * hsz]));
    // Small constants ride the fingerprinted cache; absent ones bind a
    // zero plane so the kernel adds nothing.
    static ONE_F32: [u8; 4] = [0, 0, 0, 0];
    // The zero-bias plane lives under the graph's (0, len) sentinel key —
    // a fresh Vec per call would mint a new cache entry per call.
    let bias_buf = match cfg.bias {
        Some(bs) => const_buf(c, bytemuck::cast_slice(&bs[..fused])),
        None => {
            let key = (0usize, fused * 4);
            let mut cb = c.const_bufs.lock().unwrap();
            match cb.get(&key) {
                Some((b, _)) => b.clone(),
                None => {
                    let b = c.device.create_buffer(&wgpu::BufferDescriptor {
                        label: Some("g-zero"),
                        size: (fused * 4) as u64,
                        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                        mapped_at_creation: false,
                    });
                    c.queue.write_buffer(&b, 0, &vec![0u8; fused * 4]);
                    cb.insert(key, (b.clone(), 0));
                    b
                }
            }
        }
    };
    let qnorm_buf = match cfg.q_norm {
        Some(w) => const_buf(c, bytemuck::cast_slice(&w[..hd])),
        None => const_buf(c, &ONE_F32),
    };
    let knorm_buf = match cfg.k_norm {
        Some(w) => const_buf(c, bytemuck::cast_slice(&w[..hd])),
        None => const_buf(c, &ONE_F32),
    };
    let rope_buf = const_buf(
        c,
        bytemuck::cast_slice(&cfg.rope[..t * cfg.rotary_half * 2]),
    );
    let uni = |v: [u32; 4]| {
        c.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&v),
                usage: wgpu::BufferUsages::UNIFORM,
            })
    };
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("bake-attn"),
        });
    {
        // qkv = n1 · wqkvᵀ on the matrix units.
        let u = uni([(hsz / 4) as u32, fused as u32, n as u32, 0]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &nt.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wqkv),
                bind_buf(1, &xbuf),
                bind_buf(2, &plane),
                bind_buf(3, &u),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(nt);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((fused as u32).div_ceil(64), (n as u32).div_ceil(64), 1);
    }
    {
        let flags = (cfg.output_gate as u32)
            | ((cfg.q_norm.is_some() as u32) << 1)
            | ((cfg.k_norm.is_some() as u32) << 2)
            | ((cfg.gemma as u32) << 3);
        let params: [u32; 12] = [
            n as u32,
            t as u32,
            nh as u32,
            nkv as u32,
            hd as u32,
            qrows as u32,
            cfg.rotary_half as u32,
            flags,
            cfg.eps.to_bits(),
            0,
            0,
            0,
        ];
        let u = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&params),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.bake_qkr.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &plane),
                bind_buf(1, &qnorm_buf),
                bind_buf(2, &knorm_buf),
                bind_buf(3, &bias_buf),
                bind_buf(4, &rope_buf),
                bind_buf(5, &qrot),
                bind_buf(6, &krot),
                bind_buf(7, &vproj),
                bind_buf(8, &gate),
                bind_buf(9, &qinv),
                bind_buf(10, &kinv),
                bind_buf(11, &u),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.bake_qkr);
        pass.set_bind_group(0, &bind, &[]);
        let units = (n * (nh + 2 * nkv)) as u32;
        pass.dispatch_workgroups(units.div_ceil(64), 1, 1);
    }
    {
        let u2 = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&[
                    t as u32, nh as u32, nkv as u32, hd as u32, b as u32, 0, 0, 0,
                ]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.bake_attn_head.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &qrot),
                bind_buf(1, &krot),
                bind_buf(2, &vproj),
                bind_buf(3, &ao),
                bind_buf(4, &u2),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.bake_attn_head);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(t as u32, (b * nh) as u32, 1);
    }
    let wo_input = if cfg.output_gate {
        let total = (n * qdim) as u32;
        let wgs = total.div_ceil(256);
        let gx = wgs.min(32_768);
        let gy = wgs.div_ceil(32_768);
        let u = uni([total, gx * 256, 0, 0]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.bake_attn_gate.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &ao),
                bind_buf(1, &gate),
                bind_buf(2, &ao_eff),
                bind_buf(3, &u),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.bake_attn_gate);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(gx, gy, 1);
        &ao_eff
    } else {
        &ao
    };
    {
        // attn_out = ao_eff · woᵀ on the matrix units.
        let u = uni([(qdim / 4) as u32, hsz as u32, n as u32, 0]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &nt.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wo),
                bind_buf(1, wo_input),
                bind_buf(2, &ybuf),
                bind_buf(3, &u),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(nt);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((hsz as u32).div_ceil(64), (n as u32).div_ceil(64), 1);
    }
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&ybuf, 0, &stage, offs[0].0, offs[0].1);
    if want_acts {
        for (i, buf) in [&plane, &qrot, &krot, &qinv, &kinv, &ao].iter().enumerate() {
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(buf, 0, &stage, offs[i + 1].0, offs[i + 1].1);
        }
    }
    submit(c, finish_enc(enc));
    let slice = stage.slice(..need);
    let done = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
    let d2 = done.clone();
    slice.map_async(wgpu::MapMode::Read, move |_| {
        d2.store(true, std::sync::atomic::Ordering::Release);
    });
    if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        return None;
    }
    let mut acts = None;
    {
        let Ok(data) = slice.get_mapped_range() else {
            stage.unmap();
            return None;
        };
        let read = |o: u64, l: u64| -> Vec<f32> {
            bytemuck::cast_slice(&data[o as usize..(o + l) as usize]).to_vec()
        };
        par_copy(
            &mut attn_out[..n * hsz],
            bytemuck::cast_slice(&data[offs[0].0 as usize..(offs[0].0 + offs[0].1) as usize]),
        );
        if want_acts {
            acts = Some(AttnChainActs {
                qkv_plane: read(offs[1].0, offs[1].1),
                qrot: read(offs[2].0, offs[2].1),
                krot: read(offs[3].0, offs[3].1),
                qinv: read(offs[4].0, offs[4].1),
                kinv: read(offs[5].0, offs[5].1),
                ao: read(offs[6].0, offs[6].1),
            });
        }
    }
    stage.unmap();
    if !want_acts {
        return Some(AttnChainActs {
            qkv_plane: Vec::new(),
            qrot: Vec::new(),
            krot: Vec::new(),
            qinv: Vec::new(),
            kinv: Vec::new(),
            ao: Vec::new(),
        });
    }
    acts
}

/// The frozen FFN's backward as one device chain, fed by the plane the
/// forward parked: dact = dh2·down, dgu = silu·mul-backward(plane, dact),
/// dn2 = dgu·gu — one submit, and the only PCIe traffic is dh2 down
/// (3 MB) and dn2 back (3 MB) where the host path uploaded the 80 MB
/// concatenated dgu every layer. Declines (false) when the plane for
/// `li` is absent or undersized — the caller keeps its host path.
#[allow(clippy::too_many_arguments)]
pub fn ffn_bwd_chain_f32(
    dh2: &[f32],
    down: &[f32],
    gu: &[f32],
    li: usize,
    dn2: &mut [f32],
    n: usize,
    hsz: usize,
    inter: usize,
) -> bool {
    if std::env::var("CMF_BAKE_GPU").as_deref() == Ok("0") {
        return false;
    }
    let Some(c) = ctx() else { return false };
    let Some(nn) = c.gemm_nn_coop.as_ref() else {
        return false;
    };
    if f32_strict()
        || !c.discrete
        || hsz % 4 != 0
        || inter % 2 != 0
        || n * hsz * 2 * inter < (1 << 22)
    {
        return false;
    }
    if inter.div_ceil(64) > 65_535 || hsz.div_ceil(64) > 65_535 || n.div_ceil(64) > 65_535 {
        return false;
    }
    if dh2.len() < n * hsz
        || down.len() < hsz * inter
        || gu.len() < 2 * inter * hsz
        || dn2.len() < n * hsz
    {
        return false;
    }
    let both_bytes = (n * 2 * inter * 4) as u64;
    let _bake = BAKE_LOCK.lock().unwrap();
    let plane = {
        let planes = c.bake_planes.lock().unwrap();
        match planes.get(&li) {
            Some((b, cap)) if *cap >= both_bytes => b.clone(),
            _ => return false,
        }
    };
    let wdn = bake_weight(c, down, "bake-dn");
    let wgu = bake_weight(c, gu, "bake-gu");
    let dn2_bytes = (n * hsz * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let dhb = Scratch::ensure(
        &c.device,
        &mut sc.bx,
        dn2_bytes,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        "bake-x",
    );
    let dact = Scratch::ensure(
        &c.device,
        &mut sc.ba,
        (n * inter * 4) as u64,
        wgpu::BufferUsages::STORAGE,
        "bake-act",
    );
    let dgu = Scratch::ensure(
        &c.device,
        &mut sc.bb,
        both_bytes,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "bake-both",
    );
    let ybuf = Scratch::ensure(
        &c.device,
        &mut sc.by,
        dn2_bytes,
        wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        "bake-y",
    );
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.bst,
        dn2_bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "bake-stage",
    );
    drop(sc);
    c.queue
        .write_buffer(&dhb, 0, bytemuck::cast_slice(&dh2[..n * hsz]));
    let uni = |v: [u32; 4]| {
        c.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&v),
                usage: wgpu::BufferUsages::UNIFORM,
            })
    };
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("bake-ffn-bwd"),
        });
    {
        // dact[n, inter] = dh2[n, hsz] · down[hsz, inter]
        let u = uni([(hsz / 4) as u32, inter as u32, n as u32, 0]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &nn.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wdn),
                bind_buf(1, &dhb),
                bind_buf(2, &dact),
                bind_buf(3, &u),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(nn);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((inter as u32).div_ceil(64), (n as u32).div_ceil(64), 1);
    }
    {
        let total = (n * inter) as u32;
        let wgs = total.div_ceil(256);
        let gx = wgs.min(32_768);
        let gy = wgs.div_ceil(32_768);
        let u = uni([inter as u32, total, 0, gx * 256]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.bake_silu_bwd.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &plane),
                bind_buf(1, &dact),
                bind_buf(2, &dgu),
                bind_buf(3, &u),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.bake_silu_bwd);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(gx, gy, 1);
    }
    {
        // dn2[n, hsz] = dgu[n, 2·inter] · gu[2·inter, hsz]
        let u = uni([(2 * inter / 4) as u32, hsz as u32, n as u32, 0]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &nn.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wgu),
                bind_buf(1, &dgu),
                bind_buf(2, &ybuf),
                bind_buf(3, &u),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(nn);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((hsz as u32).div_ceil(64), (n as u32).div_ceil(64), 1);
    }
    readback(c, enc, &ybuf, &stage, dn2_bytes, &mut dn2[..n * hsz])
}

/// VAE conv2d on the card: same padding, stride 1, direct (no im2col).
/// The weights are cached by address — a decoder calls the same convs
/// once per image, and a server renders many.
#[allow(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 {
    vae_conv_impl(w, bias, x, ic, oc, h, w_img, k, false, out)
}

/// Nearest-2× upsample fused with the conv that follows it: the source
/// is the SMALL image (h/2 × w/2), so only it crosses the boundary.
#[allow(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 {
    vae_conv_impl(w, bias, x, ic, oc, h, w_img, k, true, out)
}

#[allow(clippy::too_many_arguments)]
fn vae_conv_impl(
    w: &[f32],
    bias: &[f32],
    x: &[f32],
    ic: usize,
    oc: usize,
    h: usize,
    w_img: usize,
    k: usize,
    up2: bool,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    // `h`/`w_img` are the SOURCE dims; an upsampling conv writes 2x.
    let (oh, ow) = if up2 { (h * 2, w_img * 2) } else { (h, w_img) };
    let ick2 = ic * k * k;
    if w.len() != oc * ick2 || x.len() != ic * h * w_img || out.len() != oc * oh * ow {
        return false;
    }
    if oc as u32 > 65_000 || (oh * ow) as u32 > 65_000 * 64 {
        return false;
    }
    let cache = |data: &[f32], label: &'static str| -> wgpu::Buffer {
        let key = (data.as_ptr() as usize, data.len());
        let fp = fp_bytes(bytemuck::cast_slice(data));
        let mut cb = c.const_bufs.lock().unwrap();
        if let Some((b, f)) = cb.get_mut(&key) {
            if *f != fp {
                c.queue.write_buffer(b, 0, bytemuck::cast_slice(data));
                *f = fp;
            }
            return b.clone();
        }
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size: (data.len() * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        c.queue.write_buffer(&b, 0, bytemuck::cast_slice(data));
        cb.insert(key, (b.clone(), fp));
        b
    };
    let wb = cache(w, "vae-w");
    let bb = cache(bias, "vae-b");
    let xb = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("vae-x"),
        size: (x.len() * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    c.queue.write_buffer(&xb, 0, bytemuck::cast_slice(x));
    let yb = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("vae-y"),
        size: (out.len() * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        mapped_at_creation: false,
    });
    let u = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[
                ic as u32,
                oc as u32,
                oh as u32,
                ow as u32,
                k as u32,
                up2 as u32,
                h as u32,
                w_img as u32,
            ]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("vae-conv"),
        });
    {
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.vae_conv.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wb),
                bind_buf(1, &bb),
                bind_buf(2, &xb),
                bind_buf(3, &yb),
                bind_buf(4, &u),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.vae_conv);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(((oh * ow) as u32).div_ceil(64), oc as u32, 1);
    }
    let bytes = (out.len() * 4) as u64;
    let stage = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("vae-stage"),
        size: bytes,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    readback(c, enc, &yb, &stage, bytes, out)
}

/// Three projections of ONE input in one submission: the DiT's q, k and
/// v. Each was its own round trip — upload x, compute, read back — and
/// x is identical for all three, so two uploads and two waits per block
/// were pure ceremony. Weights stay resident; only x in and the three
/// panels out cross the boundary.
#[allow(clippy::too_many_arguments)]
pub fn q4tp_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 {
    let Some(c) = ctx() else { return false };
    if hidden % 32 != 0 || b == 0 || xs.len() < b * hidden {
        return false;
    }
    if q_out.len() < b * qrows || k_out.len() < b * kvrows || v_out.len() < b * kvrows {
        return false;
    }
    let q4tp_weight = |idx: usize, rows: usize| -> Option<wgpu::Buffer> {
        let entry = model.tensors.get(idx)?;
        if entry.dtype != cortiq_core::TensorDtype::Q4TiledP
            || entry.shape.len() != 2
            || entry.shape[0] != rows
            || entry.shape[1] != hidden
        {
            return None;
        }
        let payload = cortiq_core::quant::expected_nbytes(
            cortiq_core::TensorDtype::Q4TiledP,
            &[rows, hidden],
        )?;
        tensor_weight_sized(c, model, idx, rows, payload)
    };
    let (Some(q1), Some(q2), Some(q3)) = (
        q4tp_weight(wq, qrows),
        q4tp_weight(wk, kvrows),
        q4tp_weight(wv, kvrows),
    ) else {
        return false;
    };
    let xb = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("qkv-x"),
        size: (b * hidden * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    c.queue
        .write_buffer(&xb, 0, bytemuck::cast_slice(&xs[..b * hidden]));
    let mk = |n: usize, label: &'static str| {
        c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size: (n * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        })
    };
    let qy = mk(b * qrows, "qkv-q");
    let ky = mk(b * kvrows, "qkv-k");
    let vy = mk(b * kvrows, "qkv-v");
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dit-qkv"),
        });
    encode_q4_tile_mm(c, &mut enc, &c.q4tp_mm, &q1, &xb, &qy, qrows, hidden, b);
    encode_q4_tile_mm(c, &mut enc, &c.q4tp_mm, &q2, &xb, &ky, kvrows, hidden, b);
    encode_q4_tile_mm(c, &mut enc, &c.q4tp_mm, &q3, &xb, &vy, kvrows, hidden, b);
    // One submission, one wait: the three readbacks share the fence.
    let qs = (b * qrows * 4) as u64;
    let ks = (b * kvrows * 4) as u64;
    let stage_q = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("qkv-sq"),
        size: qs,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    let stage_k = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("qkv-sk"),
        size: ks,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    let stage_v = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("qkv-sv"),
        size: ks,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&qy, 0, &stage_q, 0, qs);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&ky, 0, &stage_k, 0, ks);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&vy, 0, &stage_v, 0, ks);
    submit(c, finish_enc(enc));
    let read = |stage: &wgpu::Buffer, bytes: u64, dst: &mut [f32]| -> bool {
        let (tx, rx) = std::sync::mpsc::channel();
        stage.map_async(wgpu::MapMode::Read, ..bytes, move |r| {
            let _ = tx.send(r);
        });
        if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
            return false;
        }
        if !rx.recv().map(|r| r.is_ok()).unwrap_or(false) {
            return false;
        }
        let ok = match stage.get_mapped_range(..bytes) {
            Ok(raw) => {
                dst.copy_from_slice(bytemuck::cast_slice(&raw));
                drop(raw);
                true
            }
            Err(_) => false,
        };
        stage.unmap();
        ok
    };
    read(&stage_q, qs, &mut q_out[..b * qrows])
        && read(&stage_k, ks, &mut k_out[..b * kvrows])
        && read(&stage_v, ks, &mut v_out[..b * kvrows])
}

/// Return whether a scratch/storage binding remains valid after the
/// grow-only allocator rounds it to a power of two.  The real Qwen Image
/// prefill intermediate is roughly 240 MiB at 5120×12288, below the RTX
/// 3090's limit, but the refusal is required before `create_buffer` for
/// smaller/older adapters whose storage binding wall is lower.
fn qwen_storage_binding_fit(c: &Ctx, bytes: usize, label: &str) -> bool {
    let Some(need): Option<u64> = bytes.max(4096).try_into().ok() else {
        return false;
    };
    let rounded = need.checked_next_power_of_two().unwrap_or(u64::MAX);
    let limits = c.device.limits();
    let max = limits
        .max_storage_buffer_binding_size
        .min(limits.max_buffer_size);
    if rounded <= max {
        return true;
    }
    if std::env::var("CMF_GPU_DEBUG").is_ok() {
        eprintln!(
            "qwen q4tp GELU FFN refused: {label} binding {bytes} B rounds to {rounded} B, limit {max} B"
        );
    }
    false
}

/// Fetch one exact Q4TP projection without the generic dense-byte-size
/// assumption.  Q4TP's compressed payload is shorter than rows×cols; using
/// that raw element count here would bind following tensors as weight bytes.
fn qwen_q4tp_weight(
    c: &Ctx,
    model: &Arc<CmfModel>,
    idx: usize,
    rows: usize,
    cols: usize,
) -> Option<wgpu::Buffer> {
    let entry = model.tensors.get(idx)?;
    if entry.dtype != cortiq_core::TensorDtype::Q4TiledP
        || entry.shape.len() != 2
        || entry.shape[0] != rows
        || entry.shape[1] != cols
    {
        return None;
    }
    let payload =
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])?;
    if !qwen_storage_binding_fit(c, payload, "Q4TP weight") {
        return None;
    }
    tensor_weight_sized(c, model, idx, rows, payload)
}

/// Cache a small f32 vector in the existing constant arena.  Biases are
/// stable across diffusion steps, and using the same fingerprinted cache as
/// norms prevents one component's recycled mmap address from serving stale
/// values in a later component.
fn qwen_f32_const(c: &Ctx, data: &[f32], label: &'static str) -> wgpu::Buffer {
    let key = (data.as_ptr() as usize, data.len());
    let fp = fp_bytes(bytemuck::cast_slice(data));
    let mut cb = c.const_bufs.lock().unwrap();
    if let Some((b, f)) = cb.get(&key) {
        if *f == fp {
            return b.clone();
        }
        // A command encoder may still hold a bind group referring to the
        // old buffer. Updating that buffer here races commands recorded
        // before this call, so the unchanged-fingerprint fast path above is
        // the only case that reuses the cached handle. Fall through to
        // allocate a replacement and replace the cache entry.
    }
    let b = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some(label),
        size: (data.len().max(1) * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    c.queue.write_buffer(&b, 0, bytemuck::cast_slice(data));
    cb.insert(key, (b.clone(), fp));
    b
}

/// Encode the Qwen bias/GELU epilogue in the same command buffer as its
/// surrounding GEMMs.  `gelu=true` applies the input bias followed by the
/// exact tanh approximation; `false` only adds the output projection bias.
fn encode_qwen_gelu_bias(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    values: &wgpu::Buffer,
    bias: &wgpu::Buffer,
    n: usize,
    width: usize,
    gelu: bool,
) {
    let p = uniform_u32x4(c, [n as u32, width as u32, u32::from(gelu), 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("qwen-gelu-bias-bg"),
        layout: &c.qwen_gelu_bias.get_bind_group_layout(0),
        entries: &[bind_buf(0, values), bind_buf(1, bias), bind_buf(2, &p)],
    });
    let mut pass = begin_pass_with(enc, Some("qwen-gelu-bias"), None);
    pass.set_pipeline(&c.qwen_gelu_bias);
    pass.set_bind_group(0, &bind, &[]);
    let wgs = (n as u32).div_ceil(256);
    pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
}

/// Encode Qwen's affine-free LayerNorm and modulation into a caller-owned
/// device buffer.  The submission is ordered before the next queue submit,
/// so the following resident GEMM can consume the result without a fence or
/// host copy.
fn qwen_layernorm_mod_keep(
    c: &Ctx,
    src: &wgpu::Buffer,
    modulation: &[f32],
    dst: &wgpu::Buffer,
    batch: usize,
    hidden: usize,
    eps: f32,
) -> bool {
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("qwen-layernorm-mod"),
        });
    if !encode_qwen_layernorm_mod(c, &mut enc, src, modulation, dst, batch, hidden, eps) {
        return false;
    }
    submit(c, finish_enc(enc));
    true
}

/// Encode Qwen's affine-free LayerNorm and shift/scale modulation into a
/// caller-owned buffer without submitting.  The full Qwen block uses this
/// form so attention, residual, and MLP remain in one ordered encoder.
fn encode_qwen_layernorm_mod(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    src: &wgpu::Buffer,
    modulation: &[f32],
    dst: &wgpu::Buffer,
    batch: usize,
    hidden: usize,
    eps: f32,
) -> bool {
    if modulation.len() != hidden.saturating_mul(2) {
        return false;
    }
    let mod_buf = qwen_f32_const(c, modulation, "qwen-block-mod");
    let p = uniform_u32x4(c, [batch as u32, hidden as u32, eps.to_bits(), 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("qwen-layernorm-mod-bg"),
        layout: &c.qwen_layernorm_mod.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, src),
            bind_buf(1, &mod_buf),
            bind_buf(2, dst),
            bind_buf(3, &p),
        ],
    });
    {
        let mut pass = begin_pass_with(enc, Some("qwen-layernorm-mod"), None);
        pass.set_pipeline(&c.qwen_layernorm_mod);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(batch as u32, 1, 1);
    }
    true
}

/// Encode the Qwen residual `base + gate * delta` in place in `delta`.
fn encode_qwen_gated_residual(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    base: &wgpu::Buffer,
    delta: &wgpu::Buffer,
    gate: &wgpu::Buffer,
    batch: usize,
    hidden: usize,
) {
    let p = uniform_u32x4(c, [batch as u32, hidden as u32, 0, 0]);
    let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("qwen-gated-residual-bg"),
        layout: &c.qwen_gated_residual.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, base),
            bind_buf(1, delta),
            bind_buf(2, gate),
            bind_buf(3, &p),
        ],
    });
    let mut pass = begin_pass_with(enc, Some("qwen-gated-residual"), None);
    pass.set_pipeline(&c.qwen_gated_residual);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups_flat(((batch * hidden) as u32).div_ceil(256));
}

/// Qwen Image's second sub-block with all position-wise work resident on the
/// device: LayerNorm/modulation, both Q4TP projections, exact tanh-GELU,
/// output bias and gated residual.  The input and final state cross the host
/// boundary once each; all intermediate panels are pooled and bounded.
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 {
    let Some(c) = ctx() else { return false };
    if batch < 32
        || hidden == 0
        || inter == 0
        || hidden % 32 != 0
        || inter % 32 != 0
        || data.len() < batch.saturating_mul(hidden)
        || bias_in.len() != inter
        || bias_out.len() != hidden
        || modulation.len() != hidden.saturating_mul(2)
        || gate.len() != hidden
    {
        return false;
    }
    let Some(x_len) = batch.checked_mul(hidden) else {
        return false;
    };
    let Some(mid_len) = batch.checked_mul(inter) else {
        return false;
    };
    let Some(x_bytes) = x_len.checked_mul(4) else {
        return false;
    };
    let Some(mid_bytes) = mid_len.checked_mul(4) else {
        return false;
    };
    if !qwen_storage_binding_fit(c, x_bytes, "block input")
        || !qwen_storage_binding_fit(c, mid_bytes, "block intermediate")
        || !qwen_storage_binding_fit(c, x_bytes, "block output")
    {
        return false;
    }
    let Some(w_in_buf) = qwen_q4tp_weight(c, model, w_in, inter, hidden) else {
        return false;
    };
    let Some(w_out_buf) = qwen_q4tp_weight(c, model, w_out, hidden, inter) else {
        return false;
    };

    // Use the existing grow-only DiT pool so repeated blocks reuse their
    // allocations, while unique labels keep this chain independent from the
    // attention half's q/k/v scratch slots and later VAE stages.
    let pooled = |els: usize, usage: wgpu::BufferUsages, label: &'static str| {
        let want = (els.max(1) * 4) as u64;
        let mut pool = c.dit_pool.lock().unwrap();
        if let Some((buf, cap)) = pool.get(label) {
            if *cap >= want {
                return buf.clone();
            }
        }
        let size = want.next_power_of_two();
        let buf = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size,
            usage,
            mapped_at_creation: false,
        });
        pool.insert(label, (buf.clone(), size));
        buf
    };
    let st = wgpu::BufferUsages::STORAGE;
    let x_buf = pooled(
        x_len,
        st | wgpu::BufferUsages::COPY_DST,
        "qwen-block-x",
    );
    let norm_buf = pooled(mid_len.max(x_len), st, "qwen-block-norm");
    let mid_buf = pooled(mid_len, st | wgpu::BufferUsages::COPY_DST, "qwen-block-mid");
    let y_buf = pooled(
        x_len,
        st | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
        "qwen-block-y",
    );
    let stage = pooled(
        x_len,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "qwen-block-stage",
    );
    c.queue
        .write_buffer(&x_buf, 0, bytemuck::cast_slice(&data[..x_len]));

    if !qwen_layernorm_mod_keep(c, &x_buf, modulation, &norm_buf, batch, hidden, 1.0e-6) {
        return false;
    }
    if !qwen_q4tp_gemm_keep(model, w_in, &norm_buf, &mid_buf, batch, inter, hidden) {
        return false;
    }
    let bias_in_buf = qwen_f32_const(c, bias_in, "qwen-block-bias-in");
    // Keep the command construction explicit so the queue ordering is
    // visible and no temporary encoder is dropped before finish.
    let mut gelu_enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("qwen-block-gelu"),
        });
    encode_qwen_gelu_bias(c, &mut gelu_enc, &mid_buf, &bias_in_buf, mid_len, inter, true);
    submit(c, finish_enc(gelu_enc));
    if !qwen_q4tp_gemm_keep(model, w_out, &mid_buf, &y_buf, batch, hidden, inter) {
        return false;
    }
    let bias_out_buf = qwen_f32_const(c, bias_out, "qwen-block-bias-out");
    let mut tail_enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("qwen-block-tail"),
        });
    encode_qwen_gelu_bias(c, &mut tail_enc, &y_buf, &bias_out_buf, x_len, hidden, false);
    let gate_buf = qwen_f32_const(c, gate, "qwen-block-gate");
    encode_qwen_gated_residual(c, &mut tail_enc, &x_buf, &y_buf, &gate_buf, batch, hidden);
    let ok = readback(c, tail_enc, &y_buf, &stage, x_bytes as u64, &mut data[..x_len]);
    ok
}

/// Qwen Image's two-projection tanh-GELU FFN on WGPU.  The input projection,
/// exact bias+GELU, output projection and output bias share one command
/// buffer; only the source and final output cross the host/device boundary.
/// The cooperative f16 GEMM is reused when available, with a device max
/// reduction for the resident second operand.  The scalar Q4TP GEMM remains
/// the exact fallback when cooperative matrices, reductions, or binding
/// limits are unavailable.
#[allow(clippy::too_many_arguments)]
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 {
    let Some(c) = ctx() else { return false };
    let _gate = c.mm_gate.lock().unwrap();
    if b < 32
        || hidden == 0
        || inter == 0
        || hidden % 32 != 0
        || inter % 32 != 0
        || bias_in.len() != inter
        || bias_out.len() != hidden
    {
        return false;
    }
    let Some(x_len) = b.checked_mul(hidden) else {
        return false;
    };
    let Some(mid_len) = b.checked_mul(inter) else {
        return false;
    };
    if xs.len() < x_len
        || out.len() < x_len
        || x_len > u32::MAX as usize
        || mid_len > u32::MAX as usize
    {
        return false;
    }
    let Some(x_bytes) = x_len.checked_mul(4) else {
        return false;
    };
    let Some(mid_bytes) = mid_len.checked_mul(4) else {
        return false;
    };
    // Scratch::ensure rounds every slot to a power of two.  Check all three
    // live panels before touching the weight cache or creating any buffers.
    if !qwen_storage_binding_fit(c, x_bytes, "input")
        || !qwen_storage_binding_fit(c, mid_bytes, "intermediate")
        || !qwen_storage_binding_fit(c, x_bytes, "output")
    {
        return false;
    }
    let Some(w_in_buf) = qwen_q4tp_weight(c, model, w_in, inter, hidden) else {
        return false;
    };
    let Some(w_out_buf) = qwen_q4tp_weight(c, model, w_out, hidden, inter) else {
        return false;
    };
    let bias_in_buf = qwen_f32_const(c, bias_in, "qwen-mlp-bias-in");
    let bias_out_buf = qwen_f32_const(c, bias_out, "qwen-mlp-bias-out");

    let st = wgpu::BufferUsages::STORAGE;
    let (x_buf, mid_buf, y_buf, stage, planes, amax_parts) = {
        let mut sc = c.scratch.lock().unwrap();
        let x_buf = Scratch::ensure(
            &c.device,
            &mut sc.xs,
            x_bytes as u64,
            st | wgpu::BufferUsages::COPY_DST,
            "qwen-mlp-x",
        );
        let mid_buf = Scratch::ensure(&c.device, &mut sc.g, mid_bytes as u64, st, "qwen-mlp-mid");
        let y_buf = Scratch::ensure(
            &c.device,
            &mut sc.y,
            x_bytes as u64,
            st | wgpu::BufferUsages::COPY_SRC,
            "qwen-mlp-y",
        );
        let stage = Scratch::ensure(
            &c.device,
            &mut sc.stage,
            x_bytes as u64,
            wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            "qwen-mlp-stage",
        );
        let coop = c.q4tp_dq_f16.is_some()
            && c.q4tp_mm_coop_f16.is_some()
            && (c.act_absmax.is_some() || (c.act_amax_part.is_some() && c.act_amax_fold.is_some()))
            && std::env::var("CMF_QWEN_IMAGE_FUSED_MLP_COOP").as_deref() != Ok("0");
        let planes = if coop {
            let p1 = dq_f16_plane_slot(c, &mut sc, &w_in_buf, inter, hidden, false);
            let p2 = dq_f16_plane_slot(c, &mut sc, &w_out_buf, hidden, inter, true);
            match (p1, p2) {
                (Some(p1), Some(p2)) => Some((p1, p2)),
                _ => None,
            }
        } else {
            None
        };
        // `encode_act_absmax_with` can reuse these partials without trying to
        // lock `scratch` recursively while this scope owns it.
        let amax_parts =
            if planes.is_some() && c.act_amax_part.is_some() && c.act_amax_fold.is_some() {
                Some(Scratch::ensure(
                    &c.device,
                    &mut sc.amaxp,
                    (512 * 4) as u64,
                    st,
                    "qwen-mlp-amax-parts",
                ))
            } else {
                None
            };
        c.queue
            .write_buffer(&x_buf, 0, bytemuck::cast_slice(&xs[..x_len]));
        (x_buf, mid_buf, y_buf, stage, planes, amax_parts)
    };

    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("qwen-gelu-ffn"),
        });
    let use_coop = planes.is_some();
    if let Some(((plane1, bind_dq1), (plane2, bind_dq2))) = planes.as_ref() {
        let dq = c.q4tp_dq_f16.as_ref().unwrap();
        for (plane_bind, rows, cols) in [(bind_dq1, inter, hidden), (bind_dq2, hidden, inter)] {
            let mut pass = begin_pass_with(&mut enc, Some("qwen-q4tp-dequant"), None);
            pass.set_pipeline(dq);
            pass.set_bind_group(0, plane_bind, &[]);
            let Some(pairs) = rows.checked_mul(cols).and_then(|n| n.checked_div(2)) else {
                return false;
            };
            let Ok(pairs_u32) = u32::try_from(pairs) else {
                return false;
            };
            let wgs = pairs_u32.div_ceil(256);
            pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
        }
        let mx = xs[..x_len].iter().fold(
            0.0f32,
            |m, &v| if v.is_finite() { m.max(v.abs()) } else { m },
        );
        let ascale = if mx > 1000.0 { 1000.0 / mx } else { 1.0 };
        encode_q4_tile_mm_full(
            c,
            &mut enc,
            c.q4tp_mm_coop_f16.as_ref().unwrap(),
            plane1,
            &x_buf,
            &mid_buf,
            inter,
            hidden,
            b,
            ascale,
            None,
        );
        encode_qwen_gelu_bias(c, &mut enc, &mid_buf, &bias_in_buf, mid_len, inter, true);
        let asc = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("qwen-mlp-ascale"),
            size: 4,
            usage: st,
            mapped_at_creation: false,
        });
        if !encode_act_absmax_with(c, &mut enc, &mid_buf, mid_len, &asc, amax_parts.as_ref()) {
            return false;
        }
        encode_q4_tile_mm_full(
            c,
            &mut enc,
            c.q4tp_mm_coop_f16.as_ref().unwrap(),
            plane2,
            &mid_buf,
            &y_buf,
            hidden,
            inter,
            b,
            0.0,
            Some(&asc),
        );
    } else {
        encode_q4_tile_mm(
            c, &mut enc, &c.q4tp_mm, &w_in_buf, &x_buf, &mid_buf, inter, hidden, b,
        );
        encode_qwen_gelu_bias(c, &mut enc, &mid_buf, &bias_in_buf, mid_len, inter, true);
        encode_q4_tile_mm(
            c, &mut enc, &c.q4tp_mm, &w_out_buf, &mid_buf, &y_buf, hidden, inter, b,
        );
    }
    encode_qwen_gelu_bias(c, &mut enc, &y_buf, &bias_out_buf, x_len, hidden, false);
    let ok = readback(c, enc, &y_buf, &stage, x_bytes as u64, &mut out[..x_len]);
    if ok && std::env::var("CMF_GPU_DEBUG").is_ok() {
        static SEEN: std::sync::OnceLock<
            std::sync::Mutex<std::collections::HashSet<(usize, usize, bool)>>,
        > = std::sync::OnceLock::new();
        let seen = SEEN.get_or_init(|| std::sync::Mutex::new(std::collections::HashSet::new()));
        if seen.lock().unwrap().insert((hidden, inter, use_coop)) {
            let limits = c.device.limits();
            eprintln!(
                "qwen q4tp GELU FFN: fused {} b={b} hidden={hidden} inter={inter} mid_bytes={mid_bytes} storage_limit={} max_buffer={}",
                if use_coop { "coop_f16" } else { "scalar" },
                limits.max_storage_buffer_binding_size,
                limits.max_buffer_size,
            );
        }
    }
    ok
}

/// Encode Qwen's joint attention from already-packed head-major planes.
/// Unlike `dit_attention_inner`, this helper never submits or reads back:
/// the caller owns the encoder and can append both output projections and
/// the MLP before the one final fence. The layouts and kernels are the
/// established DiT attention path; only the ownership of the encoder and
/// the explicit head-major buffers differ.
#[allow(clippy::too_many_lines)]
fn qwen_attention_encode(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    q: &wgpu::Buffer,
    k: &wgpu::Buffer,
    v: &wgpu::Buffer,
    heads: usize,
    total: usize,
    head_dim: usize,
) -> Option<wgpu::Buffer> {
    if heads == 0 || total == 0 || head_dim == 0 || head_dim % 2 != 0 {
        return None;
    }
    let Some(hidden) = heads.checked_mul(head_dim) else {
        return None;
    };
    let Some(head_bytes) = total
        .checked_mul(head_dim)
        .and_then(|n| n.checked_mul(4))
    else {
        return None;
    };
    let Some(total_bytes) = total.checked_mul(hidden).and_then(|n| n.checked_mul(4)) else {
        return None;
    };
    let Some(score_bytes) = total.checked_mul(total).and_then(|n| n.checked_mul(4)) else {
        return None;
    };
    if !qwen_storage_binding_fit(c, head_bytes, "chain attention head")
        || !qwen_storage_binding_fit(c, total_bytes, "chain attention output")
        || !qwen_storage_binding_fit(c, score_bytes, "chain attention scores")
    {
        return None;
    }
    if (head_bytes as u64) > q.size()
        || (head_bytes as u64) > k.size()
        || (head_bytes as u64) > v.size()
    {
        return None;
    }

    let pooled = |bytes: usize, usage: wgpu::BufferUsages, label: &'static str| {
        let want = bytes.max(4) as u64;
        let mut pool = c.dit_pool.lock().unwrap();
        if let Some((buf, cap)) = pool.get(label) {
            if *cap >= want {
                return buf.clone();
            }
        }
        let size = want.next_power_of_two();
        let buf = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size,
            usage,
            mapped_at_creation: false,
        });
        pool.insert(label, (buf.clone(), size));
        buf
    };
    let st = wgpu::BufferUsages::STORAGE;
    let score = pooled(score_bytes, st, "qwen-chain-scores");
    let panel = pooled(total_bytes, st, "qwen-chain-panel");
    let output = pooled(
        total_bytes,
        st | wgpu::BufferUsages::COPY_SRC,
        "qwen-chain-attn",
    );
    let vt = if c.dit_v_transpose.is_some()
        && c.dit_gemm_coop.is_some()
        && std::env::var("CMF_DIT_ATTN_COOP").as_deref() != Ok("0")
        && std::env::var("CMF_DIT_PV_COOP").as_deref() != Ok("0")
        && head_dim % 4 == 0
    {
        Some(pooled(total_bytes, st, "qwen-chain-vt"))
    } else {
        None
    };

    let params = |m: u32, k: u32, n: u32, scale: f32| {
        uniform_u32x8(c, [m, k, n, scale.to_bits(), 0, 0, 0, 0])
    };
    let p_qk = params(total as u32, head_dim as u32, total as u32, 1.0 / (head_dim as f32).sqrt());
    let p_sm = params(total as u32, head_dim as u32, total as u32, 1.0);
    let p_pv = params(total as u32, total as u32, head_dim as u32, 1.0);
    let head = head_bytes as u64;
    let score_len = score_bytes as u64;

    // The transpose pipeline has a distinct auto layout and entry point.
    if let (Some(vt), Some(tp)) = (vt.as_ref(), c.dit_v_transpose.as_ref()) {
        let p = uniform_u32x4(c, [total as u32, heads as u32, head_dim as u32, 0]);
        let bg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("qwen-chain-vt-bg"),
            layout: &tp.get_bind_group_layout(0),
            entries: &[bind_buf(0, v), bind_buf(1, vt), bind_buf(2, &p)],
        });
        let mut pass = begin_pass_with(enc, Some("qwen-chain-vt"), None);
        pass.set_pipeline(tp);
        pass.set_bind_group(0, &bg, &[]);
        pass.dispatch_workgroups_flat(((heads * total * head_dim) as u32).div_ceil(256));
    }

    let bind = |pipe: &wgpu::ComputePipeline,
                a: &wgpu::Buffer,
                ao: u64,
                al: u64,
                b: &wgpu::Buffer,
                bo: u64,
                bl: u64,
                cc: &wgpu::Buffer,
                co: u64,
                cl: u64,
                pp: &wgpu::Buffer|
     -> wgpu::BindGroup {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("qwen-chain-attn-bg"),
            layout: &pipe.get_bind_group_layout(0),
            entries: &[
                bind_buf_off(0, a, ao, al),
                bind_buf_off(1, b, bo, bl),
                bind_buf_off(2, cc, co, cl),
                bind_buf(3, pp),
            ],
        })
    };

    for h in 0..heads {
        let hoff = h as u64 * head;
        let p_pv_coop = match (c.dit_gemm_coop.as_ref(), vt.as_ref()) {
            (Some(pipe), Some(vt)) => {
                let p = uniform_u32x8(
                    c,
                    [
                        (total / 4) as u32,
                        head_dim as u32,
                        total as u32,
                        1.0f32.to_bits(),
                        0,
                        0,
                        0,
                        total as u32,
                    ],
                );
                Some(c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("qwen-chain-pv-coop-bg"),
                    layout: &pipe.get_bind_group_layout(0),
                    entries: &[
                        bind_buf_off(0, vt, hoff, head),
                        bind_buf(1, &score),
                        bind_buf_off(2, &panel, hoff, head),
                        bind_buf(3, &p),
                    ],
                }))
            }
            _ => None,
        };
        let bg_qk = bind(
            &c.dit_qk,
            q,
            hoff,
            head,
            k,
            hoff,
            head,
            &score,
            0,
            score_len,
            &p_qk,
        );
        let bg_sm = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("qwen-chain-softmax-bg"),
            layout: &c.dit_softmax.get_bind_group_layout(0),
            entries: &[bind_buf(2, &score), bind_buf(3, &p_sm)],
        });
        let bg_pv = bind(
            &c.dit_pv,
            &score,
            0,
            score_len,
            v,
            hoff,
            head,
            &panel,
            hoff,
            head,
            &p_pv,
        );
        {
            // Cooperative QK converts activations to f16; keep this measured
            // accuracy-sensitive product on the scalar F32 path. PV can use
            // cooperative math independently after its odd-tail fix.
            let mut pass = begin_pass_with(enc, Some("qwen-chain-qk"), None);
            pass.set_pipeline(&c.dit_qk);
            pass.set_bind_group(0, &bg_qk, &[]);
            pass.dispatch_workgroups((total as u32).div_ceil(64), (total as u32).div_ceil(64), 1);
        }
        {
            let mut pass = begin_pass_with(enc, Some("qwen-chain-softmax"), None);
            pass.set_pipeline(&c.dit_softmax);
            pass.set_bind_group(0, &bg_sm, &[]);
            pass.dispatch_workgroups(total as u32, 1, 1);
        }
        {
            let mut pass = begin_pass_with(enc, Some("qwen-chain-pv"), None);
            if let (Some(pipe), Some(bg)) = (c.dit_gemm_coop.as_ref(), p_pv_coop.as_ref()) {
                pass.set_pipeline(pipe);
                pass.set_bind_group(0, bg, &[]);
            } else {
                pass.set_pipeline(&c.dit_pv);
                pass.set_bind_group(0, &bg_pv, &[]);
            }
            pass.dispatch_workgroups((head_dim as u32).div_ceil(64), (total as u32).div_ceil(64), 1);
        }
    }
    let p_un = uniform_u32x8(c, [heads as u32, total as u32, head_dim as u32, 1, 0, 0, 0, 0]);
    let bg_un = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("qwen-chain-unstack-bg"),
        layout: &c.dit_unstack.get_bind_group_layout(0),
        entries: &[bind_buf(0, &panel), bind_buf(2, &output), bind_buf(3, &p_un)],
    });
    let mut pass = begin_pass_with(enc, Some("qwen-chain-unstack"), None);
    pass.set_pipeline(&c.dit_unstack);
    pass.set_bind_group(0, &bg_un, &[]);
    pass.dispatch_workgroups_flat(((heads * total * head_dim) as u32).div_ceil(256));
    Some(output)
}

/// Keep one complete Qwen double-stream block in one command encoder.  The
/// input norm/mod panels are prepared by the native caller; from those
/// panels through QKV, RoPE/attention, output projections, both gated
/// residuals and both tanh-GELU MLPs, no activation is returned to the host.
/// A refusal happens before the caller mutates its state, preserving the
/// portable per-op path for unsupported devices/codecs.
#[allow(clippy::too_many_lines)]
pub fn qwen_image_block(
    model: &Arc<CmfModel>,
    a: &mut crate::gpu::QwenImageBlockArgs<'_>,
) -> bool {
    let Some(c) = ctx() else { return false };
    let hidden = a.heads.checked_mul(a.head_dim).unwrap_or(0);
    let total = a.image_tokens.checked_add(a.text_tokens).unwrap_or(0);
    let pairs = a.head_dim / 2;
    let image_len = a.image_tokens.checked_mul(hidden).unwrap_or(0);
    let text_len = a.text_tokens.checked_mul(hidden).unwrap_or(0);
    let image_bytes = image_len.checked_mul(4).unwrap_or(0);
    let text_bytes = text_len.checked_mul(4).unwrap_or(0);
    let total_bytes = total.checked_mul(hidden).and_then(|n| n.checked_mul(4)).unwrap_or(0);
    if hidden == 0
        || a.image_tokens == 0
        || total == 0
        || a.head_dim == 0
        || a.head_dim > 256
        || a.head_dim % 2 != 0
        || a.image.len() != image_len
        || a.text.len() != text_len
        || a.image_norm.len() != image_len
        || a.text_norm.len() != text_len
        || a.image_cos.len() != a.image_tokens.saturating_mul(pairs)
        || a.image_sin.len() != a.image_cos.len()
        || a.text_cos.len() != a.text_tokens.saturating_mul(pairs)
        || a.text_sin.len() != a.text_cos.len()
        || a.image_q_norm.len() != a.head_dim
        || a.image_k_norm.len() != a.head_dim
        || a.text_q_norm.len() != a.head_dim
        || a.text_k_norm.len() != a.head_dim
        || a.image_q_bias.len() != hidden
        || a.image_k_bias.len() != hidden
        || a.image_v_bias.len() != hidden
        || a.text_q_bias.len() != hidden
        || a.text_k_bias.len() != hidden
        || a.text_v_bias.len() != hidden
        || a.image_out_bias.len() != hidden
        || a.text_out_bias.len() != hidden
        || a.image_attn_gate.len() != hidden
        || a.text_attn_gate.len() != hidden
        || a.image_mlp_in_bias.len() == 0
        || a.text_mlp_in_bias.len() == 0
        || a.image_mlp_out_bias.len() != hidden
        || a.text_mlp_out_bias.len() != hidden
        || a.image_mlp_mod.len() != hidden.saturating_mul(2)
        || a.text_mlp_mod.len() != hidden.saturating_mul(2)
        || a.image_mlp_gate.len() != hidden
        || a.text_mlp_gate.len() != hidden
    {
        return false;
    }
    let Some(image_mlp_entry) = model.tensors.get(a.image_mlp_in) else {
        return false;
    };
    let inter = image_mlp_entry.shape.first().copied().unwrap_or(0);
    if inter == 0 || inter % 32 != 0 || hidden % 32 != 0 {
        return false;
    }
    if a.image_mlp_in_bias.len() != inter || a.text_mlp_in_bias.len() != inter {
        return false;
    }
    let dims_ok = |idx: usize, rows: usize, cols: usize| {
        model.tensors.get(idx).is_some_and(|e| {
            e.dtype == cortiq_core::TensorDtype::Q4TiledP
                && e.shape.as_slice() == [rows, cols]
        })
    };
    if !dims_ok(a.image_q, hidden, hidden)
        || !dims_ok(a.image_k, hidden, hidden)
        || !dims_ok(a.image_v, hidden, hidden)
        || !dims_ok(a.text_q, hidden, hidden)
        || !dims_ok(a.text_k, hidden, hidden)
        || !dims_ok(a.text_v, hidden, hidden)
        || !dims_ok(a.image_out, hidden, hidden)
        || !dims_ok(a.text_out, hidden, hidden)
        || !dims_ok(a.image_mlp_in, inter, hidden)
        || !dims_ok(a.image_mlp_out, hidden, inter)
        || !dims_ok(a.text_mlp_in, inter, hidden)
        || !dims_ok(a.text_mlp_out, hidden, inter)
    {
        return false;
    }
    if !qwen_storage_binding_fit(c, image_bytes, "chain image state")
        || (text_len > 0 && !qwen_storage_binding_fit(c, text_bytes, "chain text state"))
        || !qwen_storage_binding_fit(c, total_bytes, "chain joint state")
        || !qwen_storage_binding_fit(
            c,
            a.image_tokens.saturating_mul(inter).saturating_mul(4),
            "chain image MLP",
        )
        || (a.text_tokens > 0
            && !qwen_storage_binding_fit(
                c,
                a.text_tokens.saturating_mul(inter).saturating_mul(4),
                "chain text MLP",
            ))
    {
        return false;
    }

    // The existing matrix path serializes on this gate because its weight
    // and plane caches are shared.  The encoder-taking helper below never
    // takes it again, so all eight attention and four MLP projections can be
    // assembled under one lock without the historical recursive deadlock.
    let _gate = c.mm_gate.lock().unwrap();
    let pooled = |bytes: usize, usage: wgpu::BufferUsages, label: &'static str| {
        let want = bytes.max(4) as u64;
        let mut pool = c.dit_pool.lock().unwrap();
        if let Some((buf, cap)) = pool.get(label) {
            if *cap >= want {
                return buf.clone();
            }
        }
        let size = want.next_power_of_two();
        let buf = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size,
            usage,
            mapped_at_creation: false,
        });
        pool.insert(label, (buf.clone(), size));
        buf
    };
    let st = wgpu::BufferUsages::STORAGE;
    let img_state = pooled(
        image_bytes,
        st | wgpu::BufferUsages::COPY_DST,
        "qwen-chain-img-state",
    );
    let txt_state = pooled(
        text_bytes,
        st | wgpu::BufferUsages::COPY_DST,
        "qwen-chain-txt-state",
    );
    let img_norm = pooled(
        image_bytes,
        st | wgpu::BufferUsages::COPY_DST,
        "qwen-chain-img-norm",
    );
    let txt_norm = pooled(
        text_bytes,
        st | wgpu::BufferUsages::COPY_DST,
        "qwen-chain-txt-norm",
    );
    let img_q = pooled(image_bytes, st, "qwen-chain-img-q");
    let img_k = pooled(image_bytes, st, "qwen-chain-img-k");
    let img_v = pooled(image_bytes, st, "qwen-chain-img-v");
    let txt_q = pooled(text_bytes, st, "qwen-chain-txt-q");
    let txt_k = pooled(text_bytes, st, "qwen-chain-txt-k");
    let txt_v = pooled(text_bytes, st, "qwen-chain-txt-v");
    let joint_q = pooled(total_bytes, st, "qwen-chain-joint-q");
    let joint_k = pooled(total_bytes, st, "qwen-chain-joint-k");
    let joint_v = pooled(total_bytes, st, "qwen-chain-joint-v");
    let img_attn = pooled(
        image_bytes,
        st | wgpu::BufferUsages::COPY_SRC,
        "qwen-chain-img-attn",
    );
    let txt_attn = pooled(
        text_bytes,
        st | wgpu::BufferUsages::COPY_SRC,
        "qwen-chain-txt-attn",
    );
    let img_norm2 = pooled(image_bytes, st, "qwen-chain-img-norm2");
    let txt_norm2 = pooled(text_bytes, st, "qwen-chain-txt-norm2");
    let image_mid_bytes = a.image_tokens.saturating_mul(inter).saturating_mul(4);
    let text_mid_bytes = a.text_tokens.saturating_mul(inter).saturating_mul(4);
    let img_mid = pooled(image_mid_bytes, st, "qwen-chain-img-mid");
    let txt_mid = pooled(text_mid_bytes, st, "qwen-chain-txt-mid");
    let img_out = pooled(
        image_bytes,
        st | wgpu::BufferUsages::COPY_SRC,
        "qwen-chain-img-out",
    );
    let txt_out = pooled(
        text_bytes,
        st | wgpu::BufferUsages::COPY_SRC,
        "qwen-chain-txt-out",
    );
    let stage = pooled(
        image_bytes.saturating_add(text_bytes),
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "qwen-chain-stage",
    );
    c.queue
        .write_buffer(&img_state, 0, bytemuck::cast_slice(a.image));
    c.queue
        .write_buffer(&img_norm, 0, bytemuck::cast_slice(a.image_norm));
    if a.text_tokens > 0 {
        c.queue
            .write_buffer(&txt_state, 0, bytemuck::cast_slice(a.text));
        c.queue
            .write_buffer(&txt_norm, 0, bytemuck::cast_slice(a.text_norm));
    }

    let mut enc = c.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
        label: Some("qwen-chain-block"),
    });
    let _merge_guard = PassMergeGuard::new(&enc);
    if !qwen_q4tp_gemm_encode(
        c,
        model,
        a.image_q,
        &img_norm,
        0,
        &img_q,
        a.image_tokens,
        hidden,
        hidden,
        &mut enc,
    ) || !qwen_q4tp_gemm_encode(
        c,
        model,
        a.image_k,
        &img_norm,
        0,
        &img_k,
        a.image_tokens,
        hidden,
        hidden,
        &mut enc,
    ) || !qwen_q4tp_gemm_encode(
        c,
        model,
        a.image_v,
        &img_norm,
        0,
        &img_v,
        a.image_tokens,
        hidden,
        hidden,
        &mut enc,
    ) {
        return false;
    }
    if a.text_tokens > 0
        && (!qwen_q4tp_gemm_encode(
            c,
            model,
            a.text_q,
            &txt_norm,
            0,
            &txt_q,
            a.text_tokens,
            hidden,
            hidden,
            &mut enc,
        ) || !qwen_q4tp_gemm_encode(
            c,
            model,
            a.text_k,
            &txt_norm,
            0,
            &txt_k,
            a.text_tokens,
            hidden,
            hidden,
            &mut enc,
        ) || !qwen_q4tp_gemm_encode(
            c,
            model,
            a.text_v,
            &txt_norm,
            0,
            &txt_v,
            a.text_tokens,
            hidden,
            hidden,
            &mut enc,
        ))
    {
        return false;
    }

    let mut img_bias = Vec::with_capacity(3 * hidden);
    img_bias.extend_from_slice(a.image_q_bias);
    img_bias.extend_from_slice(a.image_k_bias);
    img_bias.extend_from_slice(a.image_v_bias);
    let mut txt_bias = Vec::with_capacity(3 * hidden);
    txt_bias.extend_from_slice(a.text_q_bias);
    txt_bias.extend_from_slice(a.text_k_bias);
    txt_bias.extend_from_slice(a.text_v_bias);
    let img_bias_b = qwen_f32_const(c, &img_bias, "qwen-chain-img-qkv-bias");
    let txt_bias_b = qwen_f32_const(c, &txt_bias, "qwen-chain-txt-qkv-bias");
    let img_qn_b = qwen_f32_const(c, a.image_q_norm, "qwen-chain-img-qnorm");
    let img_kn_b = qwen_f32_const(c, a.image_k_norm, "qwen-chain-img-knorm");
    let txt_qn_b = qwen_f32_const(c, a.text_q_norm, "qwen-chain-txt-qnorm");
    let txt_kn_b = qwen_f32_const(c, a.text_k_norm, "qwen-chain-txt-knorm");
    let img_cos_b = qwen_f32_const(c, a.image_cos, "qwen-chain-img-cos");
    let img_sin_b = qwen_f32_const(c, a.image_sin, "qwen-chain-img-sin");
    let txt_cos_b = qwen_f32_const(c, a.text_cos, "qwen-chain-txt-cos");
    let txt_sin_b = qwen_f32_const(c, a.text_sin, "qwen-chain-txt-sin");
    let rope_p = uniform_u32x8(
        c,
        [
            a.image_tokens as u32,
            a.text_tokens as u32,
            a.heads as u32,
            a.head_dim as u32,
            hidden as u32,
            total as u32,
            pairs as u32,
            0,
        ],
    );
    let rope_bg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("qwen-chain-rope-bg"),
        layout: &c.qwen_rope_pack.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, &img_q),
            bind_buf(1, &img_k),
            bind_buf(2, &img_v),
            bind_buf(3, &txt_q),
            bind_buf(4, &txt_k),
            bind_buf(5, &txt_v),
            bind_buf(6, &img_bias_b),
            bind_buf(7, &txt_bias_b),
            bind_buf(8, &img_qn_b),
            bind_buf(9, &img_kn_b),
            bind_buf(10, &txt_qn_b),
            bind_buf(11, &txt_kn_b),
            bind_buf(12, &img_cos_b),
            bind_buf(13, &img_sin_b),
            bind_buf(14, &txt_cos_b),
            bind_buf(15, &txt_sin_b),
            bind_buf(16, &joint_q),
            bind_buf(17, &joint_k),
            bind_buf(18, &joint_v),
            bind_buf(19, &rope_p),
        ],
    });
    let jobs = total.saturating_mul(a.heads) as u32;
    let mut pass = begin_pass_with(&mut enc, Some("qwen-chain-rope"), None);
    pass.set_pipeline(&c.qwen_rope_pack);
    pass.set_bind_group(0, &rope_bg, &[]);
    pass.dispatch_workgroups(jobs.min(65_535), jobs.div_ceil(65_535), 1);
    drop(pass);
    let Some(attn) = qwen_attention_encode(
        c,
        &mut enc,
        &joint_q,
        &joint_k,
        &joint_v,
        a.heads,
        total,
        a.head_dim,
    ) else {
        return false;
    };
    if !qwen_q4tp_gemm_encode(
        c,
        model,
        a.image_out,
        &attn,
        text_bytes as u64,
        &img_attn,
        a.image_tokens,
        hidden,
        hidden,
        &mut enc,
    ) || (a.text_tokens > 0
        && !qwen_q4tp_gemm_encode(
            c,
            model,
            a.text_out,
            &attn,
            0,
            &txt_attn,
            a.text_tokens,
            hidden,
            hidden,
            &mut enc,
        ))
    {
        return false;
    }
    let img_out_bias = qwen_f32_const(c, a.image_out_bias, "qwen-chain-img-out-bias");
    let txt_out_bias = qwen_f32_const(c, a.text_out_bias, "qwen-chain-txt-out-bias");
    encode_qwen_gelu_bias(
        c,
        &mut enc,
        &img_attn,
        &img_out_bias,
        image_len,
        hidden,
        false,
    );
    if a.text_tokens > 0 {
        encode_qwen_gelu_bias(
            c,
            &mut enc,
            &txt_attn,
            &txt_out_bias,
            text_len,
            hidden,
            false,
        );
    }
    let img_attn_gate = qwen_f32_const(c, a.image_attn_gate, "qwen-chain-img-attn-gate");
    encode_qwen_gated_residual(
        c,
        &mut enc,
        &img_state,
        &img_attn,
        &img_attn_gate,
        a.image_tokens,
        hidden,
    );
    if a.text_tokens > 0 {
        let txt_attn_gate = qwen_f32_const(c, a.text_attn_gate, "qwen-chain-txt-attn-gate");
        encode_qwen_gated_residual(
            c,
            &mut enc,
            &txt_state,
            &txt_attn,
            &txt_attn_gate,
            a.text_tokens,
            hidden,
        );
    }
    if !encode_qwen_layernorm_mod(
        c,
        &mut enc,
        &img_attn,
        a.image_mlp_mod,
        &img_norm2,
        a.image_tokens,
        hidden,
        1.0e-6,
    ) {
        return false;
    }
    if a.text_tokens > 0
        && !encode_qwen_layernorm_mod(
            c,
            &mut enc,
            &txt_attn,
            a.text_mlp_mod,
            &txt_norm2,
            a.text_tokens,
            hidden,
            1.0e-6,
        )
    {
        return false;
    }
    if !qwen_q4tp_gemm_encode(
        c,
        model,
        a.image_mlp_in,
        &img_norm2,
        0,
        &img_mid,
        a.image_tokens,
        inter,
        hidden,
        &mut enc,
    ) {
        return false;
    }
    let img_mlp_in_bias = qwen_f32_const(c, a.image_mlp_in_bias, "qwen-chain-img-mlp-in-bias");
    encode_qwen_gelu_bias(
        c,
        &mut enc,
        &img_mid,
        &img_mlp_in_bias,
        image_len / hidden * inter,
        inter,
        true,
    );
    if !qwen_q4tp_gemm_encode(
        c,
        model,
        a.image_mlp_out,
        &img_mid,
        0,
        &img_out,
        a.image_tokens,
        hidden,
        inter,
        &mut enc,
    ) {
        return false;
    }
    let img_mlp_out_bias = qwen_f32_const(c, a.image_mlp_out_bias, "qwen-chain-img-mlp-out-bias");
    encode_qwen_gelu_bias(
        c,
        &mut enc,
        &img_out,
        &img_mlp_out_bias,
        image_len,
        hidden,
        false,
    );
    let img_mlp_gate = qwen_f32_const(c, a.image_mlp_gate, "qwen-chain-img-mlp-gate");
    encode_qwen_gated_residual(
        c,
        &mut enc,
        &img_attn,
        &img_out,
        &img_mlp_gate,
        a.image_tokens,
        hidden,
    );
    if a.text_tokens > 0 {
        if !qwen_q4tp_gemm_encode(
            c,
            model,
            a.text_mlp_in,
            &txt_norm2,
            0,
            &txt_mid,
            a.text_tokens,
            inter,
            hidden,
            &mut enc,
        ) {
            return false;
        }
        let txt_mlp_in_bias = qwen_f32_const(c, a.text_mlp_in_bias, "qwen-chain-txt-mlp-in-bias");
        encode_qwen_gelu_bias(
            c,
            &mut enc,
            &txt_mid,
            &txt_mlp_in_bias,
            text_len / hidden * inter,
            inter,
            true,
        );
        if !qwen_q4tp_gemm_encode(
            c,
            model,
            a.text_mlp_out,
            &txt_mid,
            0,
            &txt_out,
            a.text_tokens,
            hidden,
            inter,
            &mut enc,
        ) {
            return false;
        }
        let txt_mlp_out_bias = qwen_f32_const(c, a.text_mlp_out_bias, "qwen-chain-txt-mlp-out-bias");
        encode_qwen_gelu_bias(
            c,
            &mut enc,
            &txt_out,
            &txt_mlp_out_bias,
            text_len,
            hidden,
            false,
        );
        let txt_mlp_gate = qwen_f32_const(c, a.text_mlp_gate, "qwen-chain-txt-mlp-gate");
        encode_qwen_gated_residual(
            c,
            &mut enc,
            &txt_attn,
            &txt_out,
            &txt_mlp_gate,
            a.text_tokens,
            hidden,
        );
    }
    drop(_merge_guard);
    if a.text_tokens == 0 {
        readback(
            c,
            enc,
            &img_out,
            &stage,
            image_bytes as u64,
            a.image,
        )
    } else {
        readback_two(
            c,
            enc,
            &img_out,
            image_bytes as u64,
            &txt_out,
            text_bytes as u64,
            &stage,
            a.image,
            a.text,
        )
    }
}

fn qwen_chain_chunk_size() -> usize {
    static SIZE: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
    *SIZE.get_or_init(|| {
        std::env::var("CMF_QWEN_IMAGE_CHAIN_CHUNK")
            .ok()
            .and_then(|v| v.parse::<usize>().ok())
            .unwrap_or(4)
            .max(1)
    })
}

fn qwen_chain_pooled(
    c: &Ctx,
    bytes: usize,
    usage: wgpu::BufferUsages,
    label: &'static str,
) -> wgpu::Buffer {
    let want = bytes.max(4) as u64;
    let mut pool = c.dit_pool.lock().unwrap();
    if let Some((buf, cap)) = pool.get(label) {
        if *cap >= want {
            return buf.clone();
        }
    }
    let size = want.next_power_of_two();
    let buf = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some(label),
        size,
        usage,
        mapped_at_creation: false,
    });
    pool.insert(label, (buf.clone(), size));
    buf
}

fn qwen_chain_validate_block(
    c: &Ctx,
    model: &Arc<CmfModel>,
    a: &crate::gpu::QwenImageChainArgs<'_>,
    b: &crate::gpu::QwenImageChainBlock<'_>,
    hidden: usize,
) -> Option<usize> {
    if b.image_mod.len() != hidden.saturating_mul(6)
        || b.text_mod.len() != hidden.saturating_mul(6)
        || b.image_q_norm.len() != a.head_dim
        || b.image_k_norm.len() != a.head_dim
        || b.text_q_norm.len() != a.head_dim
        || b.text_k_norm.len() != a.head_dim
        || b.image_q_bias.len() != hidden
        || b.image_k_bias.len() != hidden
        || b.image_v_bias.len() != hidden
        || b.text_q_bias.len() != hidden
        || b.text_k_bias.len() != hidden
        || b.text_v_bias.len() != hidden
        || b.image_out_bias.len() != hidden
        || b.text_out_bias.len() != hidden
        || b.image_attn_gate.len() != hidden
        || b.text_attn_gate.len() != hidden
        || b.image_mlp_out_bias.len() != hidden
        || b.text_mlp_out_bias.len() != hidden
        || b.image_mlp_in_bias.is_empty()
        || b.text_mlp_in_bias.is_empty()
    {
        return None;
    }
    let shape = |idx: usize, rows: usize, cols: usize| {
        model.tensors.get(idx).is_some_and(|e| {
            e.dtype == cortiq_core::TensorDtype::Q4TiledP
                && e.shape.as_slice() == [rows, cols]
        })
    };
    if !shape(b.image_q, hidden, hidden)
        || !shape(b.image_k, hidden, hidden)
        || !shape(b.image_v, hidden, hidden)
        || !shape(b.text_q, hidden, hidden)
        || !shape(b.text_k, hidden, hidden)
        || !shape(b.text_v, hidden, hidden)
        || !shape(b.image_out, hidden, hidden)
        || !shape(b.text_out, hidden, hidden)
    {
        return None;
    }
    let inter = model
        .tensors
        .get(b.image_mlp_in)
        .and_then(|e| e.shape.first().copied())?;
    if inter == 0
        || inter % 32 != 0
        || hidden % 32 != 0
        || b.image_mlp_in_bias.len() != inter
        || b.text_mlp_in_bias.len() != inter
        || !shape(b.image_mlp_in, inter, hidden)
        || !shape(b.image_mlp_out, hidden, inter)
        || !shape(b.text_mlp_in, inter, hidden)
        || !shape(b.text_mlp_out, hidden, inter)
    {
        return None;
    }
    let image_bytes = a.image_tokens.checked_mul(hidden)?.checked_mul(4)?;
    let text_bytes = a.text_tokens.checked_mul(hidden)?.checked_mul(4)?;
    let total_bytes = a
        .image_tokens
        .checked_add(a.text_tokens)?
        .checked_mul(hidden)?
        .checked_mul(4)?;
    if !qwen_storage_binding_fit(c, image_bytes, "qwen chain image")
        || (a.text_tokens > 0 && !qwen_storage_binding_fit(c, text_bytes, "qwen chain text"))
        || !qwen_storage_binding_fit(c, total_bytes, "qwen chain joint")
        || !qwen_storage_binding_fit(
            c,
            a.image_tokens.checked_mul(inter)?.checked_mul(4)?,
            "qwen chain image MLP",
        )
        || (a.text_tokens > 0
            && !qwen_storage_binding_fit(
                c,
                a.text_tokens.checked_mul(inter)?.checked_mul(4)?,
                "qwen chain text MLP",
            ))
    {
        return None;
    }
    Some(inter)
}

#[allow(clippy::too_many_arguments, clippy::too_many_lines)]
fn qwen_chain_encode_block(
    c: &Ctx,
    model: &Arc<CmfModel>,
    b: &crate::gpu::QwenImageChainBlock<'_>,
    image_tokens: usize,
    text_tokens: usize,
    heads: usize,
    head_dim: usize,
    hidden: usize,
    image_cos: &[f32],
    image_sin: &[f32],
    text_cos: &[f32],
    text_sin: &[f32],
    inter: usize,
    image_state: &wgpu::Buffer,
    text_state: &wgpu::Buffer,
    image_norm: &wgpu::Buffer,
    text_norm: &wgpu::Buffer,
    image_out: &wgpu::Buffer,
    text_out: &wgpu::Buffer,
    enc: &mut wgpu::CommandEncoder,
) -> bool {
    let image_len = image_tokens.saturating_mul(hidden);
    let text_len = text_tokens.saturating_mul(hidden);
    let image_bytes = image_len.saturating_mul(4);
    let text_bytes = text_len.saturating_mul(4);
    let total = image_tokens.saturating_add(text_tokens);
    let total_bytes = total.saturating_mul(hidden).saturating_mul(4);
    let st = wgpu::BufferUsages::STORAGE;
    let img_q = qwen_chain_pooled(c, image_bytes, st, "qwen-chain-img-q");
    let img_k = qwen_chain_pooled(c, image_bytes, st, "qwen-chain-img-k");
    let img_v = qwen_chain_pooled(c, image_bytes, st, "qwen-chain-img-v");
    let txt_q = qwen_chain_pooled(c, text_bytes, st, "qwen-chain-txt-q");
    let txt_k = qwen_chain_pooled(c, text_bytes, st, "qwen-chain-txt-k");
    let txt_v = qwen_chain_pooled(c, text_bytes, st, "qwen-chain-txt-v");
    let joint_q = qwen_chain_pooled(c, total_bytes, st, "qwen-chain-joint-q");
    let joint_k = qwen_chain_pooled(c, total_bytes, st, "qwen-chain-joint-k");
    let joint_v = qwen_chain_pooled(c, total_bytes, st, "qwen-chain-joint-v");
    let img_attn = qwen_chain_pooled(c, image_bytes, st, "qwen-chain-img-attn");
    let txt_attn = qwen_chain_pooled(c, text_bytes, st, "qwen-chain-txt-attn");
    let img_norm2 = qwen_chain_pooled(c, image_bytes, st, "qwen-chain-img-norm2");
    let txt_norm2 = qwen_chain_pooled(c, text_bytes, st, "qwen-chain-txt-norm2");
    let img_mid = qwen_chain_pooled(
        c,
        image_tokens.saturating_mul(inter).saturating_mul(4),
        st,
        "qwen-chain-img-mid",
    );
    let txt_mid = qwen_chain_pooled(
        c,
        text_tokens.saturating_mul(inter).saturating_mul(4),
        st,
        "qwen-chain-txt-mid",
    );

    if !qwen_q4tp_gemm_encode(
        c,
        model,
        b.image_q,
        image_norm,
        0,
        &img_q,
        image_tokens,
        hidden,
        hidden,
        enc,
    ) || !qwen_q4tp_gemm_encode(
        c,
        model,
        b.image_k,
        image_norm,
        0,
        &img_k,
        image_tokens,
        hidden,
        hidden,
        enc,
    ) || !qwen_q4tp_gemm_encode(
        c,
        model,
        b.image_v,
        image_norm,
        0,
        &img_v,
        image_tokens,
        hidden,
        hidden,
        enc,
    ) {
        return false;
    }
    if text_tokens > 0
        && (!qwen_q4tp_gemm_encode(
            c,
            model,
            b.text_q,
            text_norm,
            0,
            &txt_q,
            text_tokens,
            hidden,
            hidden,
            enc,
        ) || !qwen_q4tp_gemm_encode(
            c,
            model,
            b.text_k,
            text_norm,
            0,
            &txt_k,
            text_tokens,
            hidden,
            hidden,
            enc,
        ) || !qwen_q4tp_gemm_encode(
            c,
            model,
            b.text_v,
            text_norm,
            0,
            &txt_v,
            text_tokens,
            hidden,
            hidden,
            enc,
        ))
    {
        return false;
    }

    let mut image_bias = Vec::with_capacity(3 * hidden);
    image_bias.extend_from_slice(b.image_q_bias);
    image_bias.extend_from_slice(b.image_k_bias);
    image_bias.extend_from_slice(b.image_v_bias);
    let mut text_bias = Vec::with_capacity(3 * hidden);
    text_bias.extend_from_slice(b.text_q_bias);
    text_bias.extend_from_slice(b.text_k_bias);
    text_bias.extend_from_slice(b.text_v_bias);
    let image_bias_b = qwen_f32_const(c, &image_bias, "qwen-chain-img-qkv-bias");
    let text_bias_b = qwen_f32_const(c, &text_bias, "qwen-chain-txt-qkv-bias");
    let image_qn_b = qwen_f32_const(c, b.image_q_norm, "qwen-chain-img-qnorm");
    let image_kn_b = qwen_f32_const(c, b.image_k_norm, "qwen-chain-img-knorm");
    let text_qn_b = qwen_f32_const(c, b.text_q_norm, "qwen-chain-txt-qnorm");
    let text_kn_b = qwen_f32_const(c, b.text_k_norm, "qwen-chain-txt-knorm");
    let image_cos_b = qwen_f32_const(c, image_cos, "qwen-chain-img-cos");
    let image_sin_b = qwen_f32_const(c, image_sin, "qwen-chain-img-sin");
    let text_cos_b = qwen_f32_const(c, text_cos, "qwen-chain-txt-cos");
    let text_sin_b = qwen_f32_const(c, text_sin, "qwen-chain-txt-sin");
    let pairs = head_dim / 2;
    let rope_p = uniform_u32x8(
        c,
        [
            image_tokens as u32,
            text_tokens as u32,
            heads as u32,
            head_dim as u32,
            hidden as u32,
            total as u32,
            pairs as u32,
            0,
        ],
    );
    let rope_bg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("qwen-chain-rope-bg"),
        layout: &c.qwen_rope_pack.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, &img_q),
            bind_buf(1, &img_k),
            bind_buf(2, &img_v),
            bind_buf(3, &txt_q),
            bind_buf(4, &txt_k),
            bind_buf(5, &txt_v),
            bind_buf(6, &image_bias_b),
            bind_buf(7, &text_bias_b),
            bind_buf(8, &image_qn_b),
            bind_buf(9, &image_kn_b),
            bind_buf(10, &text_qn_b),
            bind_buf(11, &text_kn_b),
            bind_buf(12, &image_cos_b),
            bind_buf(13, &image_sin_b),
            bind_buf(14, &text_cos_b),
            bind_buf(15, &text_sin_b),
            bind_buf(16, &joint_q),
            bind_buf(17, &joint_k),
            bind_buf(18, &joint_v),
            bind_buf(19, &rope_p),
        ],
    });
    let jobs = total.saturating_mul(heads) as u32;
    let mut pass = begin_pass_with(enc, Some("qwen-chain-rope"), None);
    pass.set_pipeline(&c.qwen_rope_pack);
    pass.set_bind_group(0, &rope_bg, &[]);
    pass.dispatch_workgroups(jobs.min(65_535), jobs.div_ceil(65_535), 1);
    drop(pass);
    let Some(attn) = qwen_attention_encode(c, enc, &joint_q, &joint_k, &joint_v, heads, total, head_dim)
    else {
        return false;
    };
    if !qwen_q4tp_gemm_encode(
        c,
        model,
        b.image_out,
        &attn,
        text_bytes as u64,
        &img_attn,
        image_tokens,
        hidden,
        hidden,
        enc,
    ) || (text_tokens > 0
        && !qwen_q4tp_gemm_encode(
            c,
            model,
            b.text_out,
            &attn,
            0,
            &txt_attn,
            text_tokens,
            hidden,
            hidden,
            enc,
        ))
    {
        return false;
    }
    let image_out_bias = qwen_f32_const(c, b.image_out_bias, "qwen-chain-img-out-bias");
    let text_out_bias = qwen_f32_const(c, b.text_out_bias, "qwen-chain-txt-out-bias");
    encode_qwen_gelu_bias(c, enc, &img_attn, &image_out_bias, image_len, hidden, false);
    if text_tokens > 0 {
        encode_qwen_gelu_bias(c, enc, &txt_attn, &text_out_bias, text_len, hidden, false);
    }
    let image_attn_gate = qwen_f32_const(c, b.image_attn_gate, "qwen-chain-img-attn-gate");
    encode_qwen_gated_residual(
        c,
        enc,
        image_state,
        &img_attn,
        &image_attn_gate,
        image_tokens,
        hidden,
    );
    if text_tokens > 0 {
        let text_attn_gate = qwen_f32_const(c, b.text_attn_gate, "qwen-chain-txt-attn-gate");
        encode_qwen_gated_residual(
            c,
            enc,
            text_state,
            &txt_attn,
            &text_attn_gate,
            text_tokens,
            hidden,
        );
    }
    if !encode_qwen_layernorm_mod(
        c,
        enc,
        &img_attn,
        &b.image_mod[3 * hidden..5 * hidden],
        &img_norm2,
        image_tokens,
        hidden,
        1.0e-6,
    ) {
        return false;
    }
    if text_tokens > 0
        && !encode_qwen_layernorm_mod(
            c,
            enc,
            &txt_attn,
            &b.text_mod[3 * hidden..5 * hidden],
            &txt_norm2,
            text_tokens,
            hidden,
            1.0e-6,
        )
    {
        return false;
    }
    if !qwen_q4tp_gemm_encode(
        c,
        model,
        b.image_mlp_in,
        &img_norm2,
        0,
        &img_mid,
        image_tokens,
        inter,
        hidden,
        enc,
    ) {
        return false;
    }
    let image_mlp_in_bias = qwen_f32_const(c, b.image_mlp_in_bias, "qwen-chain-img-mlp-in-bias");
    encode_qwen_gelu_bias(
        c,
        enc,
        &img_mid,
        &image_mlp_in_bias,
        image_tokens.saturating_mul(inter),
        inter,
        true,
    );
    if !qwen_q4tp_gemm_encode(
        c,
        model,
        b.image_mlp_out,
        &img_mid,
        0,
        image_out,
        image_tokens,
        hidden,
        inter,
        enc,
    ) {
        return false;
    }
    let image_mlp_out_bias = qwen_f32_const(c, b.image_mlp_out_bias, "qwen-chain-img-mlp-out-bias");
    encode_qwen_gelu_bias(
        c,
        enc,
        image_out,
        &image_mlp_out_bias,
        image_len,
        hidden,
        false,
    );
    let image_mlp_gate = qwen_f32_const(
        c,
        &b.image_mod[5 * hidden..6 * hidden],
        "qwen-chain-img-mlp-gate",
    );
    encode_qwen_gated_residual(
        c,
        enc,
        &img_attn,
        image_out,
        &image_mlp_gate,
        image_tokens,
        hidden,
    );
    if text_tokens > 0 {
        if !qwen_q4tp_gemm_encode(
            c,
            model,
            b.text_mlp_in,
            &txt_norm2,
            0,
            &txt_mid,
            text_tokens,
            inter,
            hidden,
            enc,
        ) {
            return false;
        }
        let text_mlp_in_bias = qwen_f32_const(c, b.text_mlp_in_bias, "qwen-chain-txt-mlp-in-bias");
        encode_qwen_gelu_bias(
            c,
            enc,
            &txt_mid,
            &text_mlp_in_bias,
            text_tokens.saturating_mul(inter),
            inter,
            true,
        );
        if !qwen_q4tp_gemm_encode(
            c,
            model,
            b.text_mlp_out,
            &txt_mid,
            0,
            text_out,
            text_tokens,
            hidden,
            inter,
            enc,
        ) {
            return false;
        }
        let text_mlp_out_bias = qwen_f32_const(c, b.text_mlp_out_bias, "qwen-chain-txt-mlp-out-bias");
        encode_qwen_gelu_bias(
            c,
            enc,
            text_out,
            &text_mlp_out_bias,
            text_len,
            hidden,
            false,
        );
        let text_mlp_gate = qwen_f32_const(
            c,
            &b.text_mod[5 * hidden..6 * hidden],
            "qwen-chain-txt-mlp-gate",
        );
        encode_qwen_gated_residual(
            c,
            enc,
            &txt_attn,
            text_out,
            &text_mlp_gate,
            text_tokens,
            hidden,
        );
    }
    true
}

/// Keep the complete Qwen transformer forward resident across its layer
/// boundaries.  Each chunk submits ordered work without a hidden-state fence;
/// only the final chunk performs the one paired readback.  The explicit args
/// make this state lifetime visible to the caller and preserve fallback
/// behavior when a codec, shape, or backend contract declines.
pub fn qwen_image_chain(
    model: &Arc<CmfModel>,
    a: &mut crate::gpu::QwenImageChainArgs<'_>,
) -> bool {
    let Some(c) = ctx() else { return false };
    let hidden = a.heads.checked_mul(a.head_dim).unwrap_or(0);
    let total = a.image_tokens.checked_add(a.text_tokens).unwrap_or(0);
    let pairs = a.head_dim / 2;
    let image_len = a.image_tokens.checked_mul(hidden).unwrap_or(0);
    let text_len = a.text_tokens.checked_mul(hidden).unwrap_or(0);
    let image_bytes = image_len.checked_mul(4).unwrap_or(0);
    let text_bytes = text_len.checked_mul(4).unwrap_or(0);
    if hidden == 0
        || a.image_tokens == 0
        || total == 0
        || a.head_dim == 0
        || a.head_dim > 256
        || a.head_dim % 2 != 0
        || a.blocks.is_empty()
        || a.image.len() != image_len
        || a.text.len() != text_len
        || a.image_cos.len() != a.image_tokens.saturating_mul(pairs)
        || a.image_sin.len() != a.image_cos.len()
        || a.text_cos.len() != a.text_tokens.saturating_mul(pairs)
        || a.text_sin.len() != a.text_cos.len()
        || !qwen_storage_binding_fit(c, image_bytes, "qwen chain image state")
        || (text_len > 0 && !qwen_storage_binding_fit(c, text_bytes, "qwen chain text state"))
    {
        return false;
    }
    let mut inter = Vec::with_capacity(a.blocks.len());
    for b in a.blocks {
        let Some(width) = qwen_chain_validate_block(c, model, a, b, hidden) else {
            return false;
        };
        inter.push(width);
    }

    // The shared matrix cache is serialized once for the whole forward. No
    // helper below submits or maps an activation, so chunk boundaries only
    // order device work and never expose a partial hidden state to the host.
    let _gate = c.mm_gate.lock().unwrap();
    let st = wgpu::BufferUsages::STORAGE;
    let image_state = qwen_chain_pooled(
        c,
        image_bytes,
        st | wgpu::BufferUsages::COPY_DST,
        "qwen-chain-img-state",
    );
    let text_state = qwen_chain_pooled(
        c,
        text_bytes,
        st | wgpu::BufferUsages::COPY_DST,
        "qwen-chain-txt-state",
    );
    let image_norm = qwen_chain_pooled(
        c,
        image_bytes,
        st | wgpu::BufferUsages::COPY_DST,
        "qwen-chain-img-norm",
    );
    let text_norm = qwen_chain_pooled(
        c,
        text_bytes,
        st | wgpu::BufferUsages::COPY_DST,
        "qwen-chain-txt-norm",
    );
    let image_out = qwen_chain_pooled(
        c,
        image_bytes,
        st | wgpu::BufferUsages::COPY_SRC,
        "qwen-chain-img-out",
    );
    let text_out = qwen_chain_pooled(
        c,
        text_bytes,
        st | wgpu::BufferUsages::COPY_SRC,
        "qwen-chain-txt-out",
    );
    let stage = qwen_chain_pooled(
        c,
        image_bytes.saturating_add(text_bytes),
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "qwen-chain-stage",
    );
    c.queue
        .write_buffer(&image_state, 0, bytemuck::cast_slice(a.image));
    if a.text_tokens > 0 {
        c.queue
            .write_buffer(&text_state, 0, bytemuck::cast_slice(a.text));
    }

    let chunk_size = qwen_chain_chunk_size();
    let mut start = 0usize;
    let mut chunks = 0usize;
    while start < a.blocks.len() {
        let end = start.saturating_add(chunk_size).min(a.blocks.len());
        let mut enc = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("qwen-chain-forward"),
            });
        let merge_guard = PassMergeGuard::new(&enc);
        for index in start..end {
            let b = &a.blocks[index];
            let state_image = &image_out;
            let state_text = &text_out;
            // The first block reads the one initial upload. Every later block
            // reads the previous block's output in the same pooled buffer;
            // ordered command encoders make the in-place reuse explicit.
            let (state_image, state_text) = if index == 0 {
                (&image_state, &text_state)
            } else {
                (state_image, state_text)
            };
            if !encode_qwen_layernorm_mod(
                c,
                &mut enc,
                state_image,
                &b.image_mod[..2 * hidden],
                &image_norm,
                a.image_tokens,
                hidden,
                1.0e-6,
            ) || (a.text_tokens > 0
                && !encode_qwen_layernorm_mod(
                    c,
                    &mut enc,
                    state_text,
                    &b.text_mod[..2 * hidden],
                    &text_norm,
                    a.text_tokens,
                    hidden,
                    1.0e-6,
                ))
            {
                return false;
            }
            if !qwen_chain_encode_block(
                c,
                model,
                b,
                a.image_tokens,
                a.text_tokens,
                a.heads,
                a.head_dim,
                hidden,
                a.image_cos,
                a.image_sin,
                a.text_cos,
                a.text_sin,
                inter[index],
                state_image,
                state_text,
                &image_norm,
                &text_norm,
                &image_out,
                &text_out,
                &mut enc,
            ) {
                return false;
            }
        }
        drop(merge_guard);
        let final_chunk = end == a.blocks.len();
        if final_chunk {
            let ok = if a.text_tokens == 0 {
                readback(c, enc, &image_out, &stage, image_bytes as u64, a.image)
            } else {
                readback_two(
                    c,
                    enc,
                    &image_out,
                    image_bytes as u64,
                    &text_out,
                    text_bytes as u64,
                    &stage,
                    a.image,
                    a.text,
                )
            };
            if !ok {
                return false;
            }
        } else {
            submit(c, finish_enc(enc));
        }
        chunks += 1;
        start = end;
    }
    if std::env::var("CMF_QWEN_IMAGE_PROFILE").is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("on"))
        || std::env::var("CMF_GPU_DEBUG").is_ok()
    {
        eprintln!(
            "qwen_image chain: blocks={} chunks={} submissions={} readbacks=1 chunk_size={} hidden={} image_tokens={} text_tokens={}",
            a.blocks.len(),
            chunks,
            chunks,
            chunk_size,
            hidden,
            a.image_tokens,
            a.text_tokens,
        );
    }
    true
}

/// Qwen Image's exact double-stream attention half on WGPU. Every large
/// panel is device-resident from the normalized stream upload through the
/// two output projections. The only host boundary is the pair of projected
/// streams returned to the Qwen residual code. This is deliberately a
/// Qwen-specific path: its two independent streams and per-stream qk norms
/// do not fit the one-stream `dit_block` contract.
#[allow(clippy::too_many_lines)]
pub fn qwen_image_attention(
    model: &Arc<CmfModel>,
    a: &mut crate::gpu::QwenImageAttentionArgs<'_>,
) -> bool {
    let Some(c) = ctx() else { return false };
    let hidden = a.heads.checked_mul(a.head_dim).unwrap_or(0);
    let total = a.image_tokens.checked_add(a.text_tokens).unwrap_or(0);
    let pairs = a.head_dim / 2;
    if hidden == 0
        || a.image_tokens == 0
        || total == 0
        || a.head_dim == 0
        || a.head_dim > 256
        || a.head_dim % 2 != 0
        || a.image.len() != a.image_tokens.saturating_mul(hidden)
        || a.text.len() != a.text_tokens.saturating_mul(hidden)
        || a.image_proj.len() != a.image.len()
        || a.text_proj.len() != a.text.len()
        || a.image_cos.len() != a.image_tokens.saturating_mul(pairs)
        || a.image_sin.len() != a.image_cos.len()
        || a.text_cos.len() != a.text_tokens.saturating_mul(pairs)
        || a.text_sin.len() != a.text_cos.len()
        || a.image_q_norm.len() != a.head_dim
        || a.image_k_norm.len() != a.head_dim
        || a.text_q_norm.len() != a.head_dim
        || a.text_k_norm.len() != a.head_dim
        || a.image_q_bias.len() != hidden
        || a.image_k_bias.len() != hidden
        || a.image_v_bias.len() != hidden
        || a.text_q_bias.len() != hidden
        || a.text_k_bias.len() != hidden
        || a.text_v_bias.len() != hidden
        || a.image_out_bias.len() != hidden
        || a.text_out_bias.len() != hidden
    {
        return false;
    }
    let Some(total_bytes) = total.checked_mul(hidden).and_then(|n| n.checked_mul(4)) else {
        return false;
    };
    let Some(image_bytes) = a
        .image_tokens
        .checked_mul(hidden)
        .and_then(|n| n.checked_mul(4))
    else {
        return false;
    };
    let Some(text_bytes) = a
        .text_tokens
        .checked_mul(hidden)
        .and_then(|n| n.checked_mul(4))
    else {
        return false;
    };
    // Each head-major plane is still `[heads, total, head_dim]`, which is
    // exactly `total * hidden` elements; Q/K/V each use one such plane.
    let head_bytes = total_bytes;
    if !qwen_storage_binding_fit(c, image_bytes, "image stream")
        || !qwen_storage_binding_fit(c, text_bytes, "text stream")
        || !qwen_storage_binding_fit(c, total_bytes, "joint stream")
        || !qwen_storage_binding_fit(c, head_bytes, "head-major stream")
    {
        return false;
    }
    let dims_ok = |idx: usize, rows: usize, cols: usize| {
        model.tensors.get(idx).is_some_and(|e| {
            e.dtype == cortiq_core::TensorDtype::Q4TiledP
                && e.shape.as_slice() == [rows, cols]
        })
    };
    if !dims_ok(a.image_q, hidden, hidden)
        || !dims_ok(a.image_k, hidden, hidden)
        || !dims_ok(a.image_v, hidden, hidden)
        || !dims_ok(a.text_q, hidden, hidden)
        || !dims_ok(a.text_k, hidden, hidden)
        || !dims_ok(a.text_v, hidden, hidden)
        || !dims_ok(a.image_out, hidden, hidden)
        || !dims_ok(a.text_out, hidden, hidden)
    {
        return false;
    }
    let rope_pipe = &c.qwen_rope_pack;

    let st = wgpu::BufferUsages::STORAGE;
    // The shared attention scratch slots are also used by later pipeline
    // stages that may upload a host panel.  Keep COPY_DST on these pooled
    // buffers from their first Qwen allocation; WGPU usage flags cannot be
    // upgraded when Scratch::ensure reuses an existing buffer.
    let qkv_st = st | wgpu::BufferUsages::COPY_DST;
    // The block is called once per layer and its output readback completes
    // before the next layer starts.  Reuse one bounded buffer per role so a
    // 60-layer render does not turn eight harmless panels into hundreds of
    // driver allocations while still keeping the chain's live footprint
    // bounded by one block.
    let pooled = |bytes: usize, usage: wgpu::BufferUsages, label: &'static str| {
        let want = bytes.max(4) as u64;
        let mut pool = c.dit_pool.lock().unwrap();
        if let Some((buf, cap)) = pool.get(label) {
            if *cap >= want {
                return buf.clone();
            }
        }
        let buf = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size: want.next_power_of_two(),
            usage,
            mapped_at_creation: false,
        });
        pool.insert(label, (buf.clone(), want.next_power_of_two()));
        buf
    };
    let img_x = pooled(image_bytes, st | wgpu::BufferUsages::COPY_DST, "qwen-img-x");
    let txt_x = pooled(text_bytes, st | wgpu::BufferUsages::COPY_DST, "qwen-txt-x");
    c.queue
        .write_buffer(&img_x, 0, bytemuck::cast_slice(a.image));
    if !a.text.is_empty() {
        c.queue
            .write_buffer(&txt_x, 0, bytemuck::cast_slice(a.text));
    }

    // Each projection has its own destination: the historical q4tp helper
    // reuses one scratch result slot, which would make six independent
    // stream panels alias one another before the join pass.
    let img_q = pooled(image_bytes, st, "qwen-img-q");
    let img_k = pooled(image_bytes, st, "qwen-img-k");
    let img_v = pooled(image_bytes, st, "qwen-img-v");
    let txt_q = pooled(text_bytes, st, "qwen-txt-q");
    let txt_k = pooled(text_bytes, st, "qwen-txt-k");
    let txt_v = pooled(text_bytes, st, "qwen-txt-v");
    if !qwen_q4tp_gemm_keep(model, a.image_q, &img_x, &img_q, a.image_tokens, hidden, hidden)
        || !qwen_q4tp_gemm_keep(model, a.image_k, &img_x, &img_k, a.image_tokens, hidden, hidden)
        || !qwen_q4tp_gemm_keep(model, a.image_v, &img_x, &img_v, a.image_tokens, hidden, hidden)
        || (a.text_tokens > 0
            && (!qwen_q4tp_gemm_keep(
                model, a.text_q, &txt_x, &txt_q, a.text_tokens, hidden, hidden,
            ) || !qwen_q4tp_gemm_keep(
                model, a.text_k, &txt_x, &txt_k, a.text_tokens, hidden, hidden,
            ) || !qwen_q4tp_gemm_keep(
                model, a.text_v, &txt_x, &txt_v, a.text_tokens, hidden, hidden,
            )))
    {
        return false;
    }

    let mut img_bias = Vec::with_capacity(3 * hidden);
    img_bias.extend_from_slice(a.image_q_bias);
    img_bias.extend_from_slice(a.image_k_bias);
    img_bias.extend_from_slice(a.image_v_bias);
    let mut txt_bias = Vec::with_capacity(3 * hidden);
    txt_bias.extend_from_slice(a.text_q_bias);
    txt_bias.extend_from_slice(a.text_k_bias);
    txt_bias.extend_from_slice(a.text_v_bias);
    let img_bias_b = qwen_f32_const(c, &img_bias, "qwen-img-qkv-bias");
    let txt_bias_b = qwen_f32_const(c, &txt_bias, "qwen-txt-qkv-bias");
    let img_qn_b = qwen_f32_const(c, a.image_q_norm, "qwen-img-qnorm");
    let img_kn_b = qwen_f32_const(c, a.image_k_norm, "qwen-img-knorm");
    let txt_qn_b = qwen_f32_const(c, a.text_q_norm, "qwen-txt-qnorm");
    let txt_kn_b = qwen_f32_const(c, a.text_k_norm, "qwen-txt-knorm");
    let img_cos_b = qwen_f32_const(c, a.image_cos, "qwen-img-cos");
    let img_sin_b = qwen_f32_const(c, a.image_sin, "qwen-img-sin");
    let txt_cos_b = qwen_f32_const(c, a.text_cos, "qwen-txt-cos");
    let txt_sin_b = qwen_f32_const(c, a.text_sin, "qwen-txt-sin");
    let (qb, kb, vb) = {
        let mut sc = c.scratch.lock().unwrap();
        (
            Scratch::ensure(
                &c.device,
                &mut sc.dq,
                head_bytes as u64,
                qkv_st,
                "qwen-joint-q",
            ),
            Scratch::ensure(
                &c.device,
                &mut sc.dk,
                head_bytes as u64,
                qkv_st,
                "qwen-joint-k",
            ),
            Scratch::ensure(
                &c.device,
                &mut sc.dv,
                head_bytes as u64,
                qkv_st,
                "qwen-joint-v",
            ),
        )
    };
    let rope_p = c.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some("qwen-rope-p"),
        contents: bytemuck::cast_slice(&[
            a.image_tokens as u32,
            a.text_tokens as u32,
            a.heads as u32,
            a.head_dim as u32,
            hidden as u32,
            total as u32,
            pairs as u32,
            0u32,
        ]),
        usage: wgpu::BufferUsages::UNIFORM,
    });
    let rope_bg = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("qwen-rope-bg"),
        layout: &rope_pipe.get_bind_group_layout(0),
        entries: &[
            bind_buf(0, &img_q),
            bind_buf(1, &img_k),
            bind_buf(2, &img_v),
            bind_buf(3, &txt_q),
            bind_buf(4, &txt_k),
            bind_buf(5, &txt_v),
            bind_buf(6, &img_bias_b),
            bind_buf(7, &txt_bias_b),
            bind_buf(8, &img_qn_b),
            bind_buf(9, &img_kn_b),
            bind_buf(10, &txt_qn_b),
            bind_buf(11, &txt_kn_b),
            bind_buf(12, &img_cos_b),
            bind_buf(13, &img_sin_b),
            bind_buf(14, &txt_cos_b),
            bind_buf(15, &txt_sin_b),
            bind_buf(16, &qb),
            bind_buf(17, &kb),
            bind_buf(18, &vb),
            bind_buf(19, &rope_p),
        ],
    });
    let jobs = total.saturating_mul(a.heads) as u32;
    let mut rope_enc = c.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
        label: Some("qwen-rope-pack"),
    });
    {
        let mut pass = begin_pass(&mut rope_enc);
        pass.set_pipeline(rope_pipe);
        pass.set_bind_group(0, &rope_bg, &[]);
        pass.dispatch_workgroups(jobs.min(65_535), jobs.div_ceil(65_535), 1);
    }
    submit(c, finish_enc(rope_enc));

    // The existing joint attention path consumes the same scratch planes and
    // leaves its token-major result on the card when `keep` is requested.
    let mut attn_keep = None;
    if !dit_attention_inner(
        &[],
        &[],
        &[],
        a.heads,
        a.heads,
        total,
        a.head_dim,
        1.0 / (a.head_dim as f32).sqrt(),
        &mut [],
        true,
        Some(&mut attn_keep),
    ) {
        return false;
    }
    let Some(attn) = attn_keep else { return false };

    let img_out_b = pooled(image_bytes, st | wgpu::BufferUsages::COPY_SRC, "qwen-img-proj");
    let txt_out_b = pooled(text_bytes, st | wgpu::BufferUsages::COPY_SRC, "qwen-txt-proj");
    if !qwen_q4tp_gemm_keep_offset(
        model,
        a.image_out,
        &attn,
        text_bytes as u64,
        &img_out_b,
        a.image_tokens,
        hidden,
        hidden,
    ) || (a.text_tokens > 0
        && !qwen_q4tp_gemm_keep_offset(
            model,
            a.text_out,
            &attn,
            0,
            &txt_out_b,
            a.text_tokens,
            hidden,
            hidden,
        ))
    {
        return false;
    }
    let img_out_bias_b = qwen_f32_const(c, a.image_out_bias, "qwen-img-out-bias");
    let txt_out_bias_b = qwen_f32_const(c, a.text_out_bias, "qwen-txt-out-bias");
    let mut bias_enc = c.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
        label: Some("qwen-attn-output-bias"),
    });
    encode_qwen_gelu_bias(
        c,
        &mut bias_enc,
        &img_out_b,
        &img_out_bias_b,
        a.image_tokens * hidden,
        hidden,
        false,
    );
    if a.text_tokens > 0 {
        encode_qwen_gelu_bias(
            c,
            &mut bias_enc,
            &txt_out_b,
            &txt_out_bias_b,
            a.text_tokens * hidden,
            hidden,
            false,
        );
    }
    submit(c, finish_enc(bias_enc));

    let stage = {
        let mut sc = c.scratch.lock().unwrap();
        Scratch::ensure(
            &c.device,
            &mut sc.dstage,
            (image_bytes + text_bytes).max(4) as u64,
            wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            "qwen-attn-stage",
        )
    };
    let read_enc = c.device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
        label: Some("qwen-attn-readback"),
    });
    if a.text_tokens == 0 {
        return readback(c, read_enc, &img_out_b, &stage, image_bytes as u64, a.image_proj);
    }
    readback_two(
        c,
        read_enc,
        &img_out_b,
        image_bytes as u64,
        &txt_out_b,
        text_bytes as u64,
        &stage,
        a.image_proj,
        a.text_proj,
    )
}

/// One whole modulated DiT block on the card — the norms, the three
/// projections, qk-norm + RoPE, attention, both residuals and the
/// SwiGLU FFN in ONE command buffer. Only `x` crosses the boundary,
/// once in and once out, where the per-op path made six round trips a
/// block and the profile showed the transfers, not the math.
///
/// `segs` are token counts of independent sequences packed back to
/// back: attention runs per segment (so classifier-free guidance can
/// send both branches through as one batch and still never let them
/// meet in a score matrix), everything position-wise sees the whole
/// batch. A single-element slice is the plain one-sequence block.
pub fn dit_block_seg(
    model: &Arc<CmfModel>,
    a: &crate::gpu::DitBlockArgs,
    segs: &[usize],
    x: &mut [f32],
) -> bool {
    let bail = |why: &str| {
        if std::env::var("CMF_DIT_DBG").is_ok() {
            eprintln!("dit fused block declined: {why}");
        }
    };
    let Some(c) = ctx() else {
        bail("no ctx");
        return false;
    };
    let (n, hs, inter) = (a.n, a.hidden, a.inter);
    let (nh, nkv, hd) = (a.nh, a.nkv, a.hd);
    if hs % 32 != 0 || inter % 32 != 0 || hd % 2 != 0 || nkv == 0 {
        bail("dims");
        return false;
    }
    if nh % nkv != 0 || x.len() < n * hs || segs.iter().sum::<usize>() != n {
        bail("shape/segs");
        return false;
    }
    let pairs = hd / 2;
    if a.rope_cos.len() < n * pairs || a.rope_sin.len() < n * pairs {
        bail("rope len");
        return false;
    }
    let (Some(wq), Some(wk), Some(wv), Some(wo)) = (
        tensor_weight(c, model, a.wq, nh * hd, hs),
        tensor_weight(c, model, a.wk, nkv * hd, hs),
        tensor_weight(c, model, a.wv, nkv * hd, hs),
        tensor_weight(c, model, a.wo, hs, nh * hd),
    ) else {
        bail("weights");
        return false;
    };
    let (Some(w1), Some(w3), Some(w2)) = (
        tensor_weight(c, model, a.w1, inter, hs),
        tensor_weight(c, model, a.w3, inter, hs),
        tensor_weight(c, model, a.w2, hs, inter),
    ) else {
        bail("weights");
        return false;
    };
    let store = |data: &[f32], label: &'static str| -> wgpu::Buffer {
        let key = (data.as_ptr() as usize, data.len());
        let fp = fp_bytes(bytemuck::cast_slice(data));
        let mut cb = c.const_bufs.lock().unwrap();
        if let Some((b, f)) = cb.get_mut(&key) {
            if *f != fp {
                c.queue.write_buffer(b, 0, bytemuck::cast_slice(data));
                *f = fp;
            }
            return b.clone();
        }
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size: (data.len().max(1) * 4) as u64,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        c.queue.write_buffer(&b, 0, bytemuck::cast_slice(data));
        cb.insert(key, (b.clone(), fp));
        b
    };
    // Per-block vectors are content-stable across steps only for the
    // norms; the modulation changes with t, so it rides a fresh upload.
    let n1w = store(a.norm1, "dit-n1");
    let n2w = store(a.norm2, "dit-n2");
    let f1w = store(a.ffn_norm1, "dit-f1");
    let f2w = store(a.ffn_norm2, "dit-f2");
    let nqw = store(a.norm_q, "dit-nq");
    let nkw = store(a.norm_k, "dit-nk");
    let pooled = |els: usize, label: &'static str, usage: wgpu::BufferUsages| -> wgpu::Buffer {
        let want = (els.max(1) * 4) as u64;
        let mut pool = c.dit_pool.lock().unwrap();
        if let Some((b, have)) = pool.get(label) {
            if *have >= want {
                return b.clone();
            }
        }
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some(label),
            size: want,
            usage,
            mapped_at_creation: false,
        });
        pool.insert(label, (b.clone(), want));
        b
    };
    let up = |data: &[f32], label: &'static str| -> wgpu::Buffer {
        let b = pooled(
            data.len(),
            label,
            wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        );
        c.queue.write_buffer(&b, 0, bytemuck::cast_slice(data));
        b
    };
    let smsa = up(a.s_msa, "dit-smsa");
    let gmsa = up(a.gate_msa, "dit-gmsa");
    let smlp = up(a.s_mlp, "dit-smlp");
    let gmlp = up(a.gate_mlp, "dit-gmlp");
    let rcos = up(a.rope_cos, "dit-rcos");
    let rsin = up(a.rope_sin, "dit-rsin");
    let mk = |els: usize, label: &'static str, io: bool| {
        pooled(
            els,
            label,
            wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_DST
                | if io {
                    wgpu::BufferUsages::COPY_SRC
                } else {
                    wgpu::BufferUsages::empty()
                },
        )
    };
    let xb = mk(n * hs, "dit-x", true);
    if !a.resident_in {
        c.queue
            .write_buffer(&xb, 0, bytemuck::cast_slice(&x[..n * hs]));
    }
    let xn = mk(n * hs, "dit-xn", false);
    let qtok = mk(n * nh * hd, "dit-qt", false);
    let ktok = mk(n * nkv * hd, "dit-kt", false);
    let vtok = mk(n * nkv * hd, "dit-vt", false);
    let qhm = mk(nh * n * hd, "dit-qh", false);
    let khm = mk(nkv * n * hd, "dit-kh", false);
    let vhm = mk(nkv * n * hd, "dit-vh", false);
    let pan = mk(nh * n * hd, "dit-pan", false);
    let attn = mk(n * nh * hd, "dit-attn", false);
    let proj = mk(n * hs, "dit-proj", false);
    let gbuf = mk(n * inter, "dit-g", false);
    let ubuf = mk(n * inter, "dit-u", false);
    let abuf = mk(n * inter, "dit-a", false);
    let seg_max = segs.iter().copied().max().unwrap_or(n);
    let scb = mk(seg_max * seg_max * nh, "dit-sc", false);
    // The attention kernels' param block is EIGHT words; a 16-byte
    // uniform is rejected outright (min binding size 32).
    let p8 = |m: u32, k: u32, nn: u32, sc: f32| -> wgpu::Buffer {
        let raw = [m, k, nn, sc.to_bits(), 0u32, 0u32, 0u32, 0u32];
        let _ = &raw;
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("dit-blk-params"),
            size: 32,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        c.queue.write_buffer(&b, 0, bytemuck::cast_slice(&raw));
        b
    };
    let u4 = |v: [u32; 4]| uniform_u32x4(c, v);
    // THE IMAGE DiT'S BIGGEST REMAINING LEVER, and it is a known one:
    // this is the scalar arm — `q4tp_mm`, which unpacks the weight tile
    // inside the GEMM loop and repeats that for every 64-row tile of
    // activations. The matrix units are never touched. The same switch
    // on the video DiT (dequantize the plane ONCE into f16, then run the
    // pure f16 GEMM `q4tp_mm_coop_f16`) took its step 22.1 s → 14.7 s.
    //
    // Denoise is 13.0 s of a 20.4 s image and 88.5% of it is this block,
    // so the arm is worth ~4 s a render. The wiring is the work: each
    // weight needs its own plane from `dq_f16_plane` — which caches per
    // tensor, so 30 steps pay for it once — and every bind group here
    // grows a fifth entry for the activation scale (or the 0xFFFFFFFF
    // sentinel and a device-computed one, as `tp_matmat_impl` does).
    let mm = mm_pipeline(c, a.q4tp, false);
    // Matrix units when the device has them: unpack each weight plane
    // ONCE into f16 and run the pure f16 GEMM, instead of unpacking the
    // same tile inside the loop for every 64-row block of activations.
    // The activation scale is computed on the card — the host never
    // sees these panels — and without one the f16 operands overflow.
    // MEASURED AND REJECTED for this model, so opt-in
    // (`CMF_DIT_COOP16=1`): 11.23 s of block time against 11.02 for the
    // scalar arm. The plane unpack and the scale reduction are encoded
    // on EVERY call — every block, every step — and Lumina's weights
    // (2304×2304) are a quarter of the video DiT's (21504×5376), where
    // this same arm won 22.1 s → 14.7. The GEMM here does not get big
    // enough to pay for its own preparation. What would change the
    // verdict is caching the unpacked plane across steps rather than
    // re-encoding it: 30 steps pay 30 times for a weight that never
    // moves.
    //
    // The cache was then BUILT (`plane_cached`) and it did not change the
    // verdict: 11.36 s cached against 11.05 scalar. So the unpack was
    // never the cost, and the reason is bytes rather than work — an f16
    // plane is FOUR TIMES the q4 weight it came from (10.6 MB against
    // 2.65 here). A GEMM this narrow is bound by reading the weight, so
    // the in-loop unpack spends ALU that is free in that regime while
    // the plane spends bandwidth that is not.
    //
    // That is the rule, and it explains the video DiT's 22.1 → 14.7 too:
    // dequantize-once wins where the weight is re-read per activation
    // tile often enough to amortize four times the bytes. Judge it by
    // tokens per weight byte, not by "tensor cores are faster".
    let coop16 = std::env::var("CMF_DIT_COOP16").as_deref() == Ok("1")
        && a.q4tp
        && c.q4tp_mm_coop_f16.is_some()
        && c.q4tp_dq_f16.is_some()
        && ((c.act_amax_part.is_some() && c.act_amax_fold.is_some()) || c.act_absmax.is_some());
    // Say once whether this arm actually engaged. The flag needs three
    // pipelines and a `cols % 2` that a refusal does not report, so a
    // measurement that shows no win is otherwise indistinguishable from
    // one where the arm never ran — which is exactly how the plane cache
    // came to be judged on a model it never executed under.
    if std::env::var("CMF_DIT_COOP16").as_deref() == Ok("1") {
        use std::sync::atomic::{AtomicBool, Ordering};
        static SAID: AtomicBool = AtomicBool::new(false);
        if !SAID.swap(true, Ordering::Relaxed) {
            tracing::info!(
                "dit coop16 requested: q4tp {}, mm_coop_f16 {}, dq_f16 {}, amax {} -> {}",
                a.q4tp,
                c.q4tp_mm_coop_f16.is_some(),
                c.q4tp_dq_f16.is_some(),
                c.act_absmax.is_some() || c.act_amax_part.is_some(),
                if coop16 { "ACTIVE" } else { "declined" }
            );
        }
    }
    let uid = model.uid() as usize;
    let mm_enc = |enc: &mut wgpu::CommandEncoder,
                  idx: usize,
                  w: &wgpu::Buffer,
                  xs: &wgpu::Buffer,
                  y: &wgpu::Buffer,
                  rows: usize,
                  cols: usize,
                  k: usize| {
        if coop16 && cols % 2 == 0 {
            // Cached per weight: the unpack pass is encoded on the FIRST
            // step only, and the remaining twenty-nine read the plane.
            let dq = plane_cached(c, (uid, idx), w, rows, cols, 8192);
            if let (Some((plane, fresh)), Some(pdq), Some(pipe)) =
                (&dq, c.q4tp_dq_f16.as_ref(), c.q4tp_mm_coop_f16.as_ref())
            {
                if let Some(bind_dq) = fresh {
                    let mut pass = begin_pass(enc);
                    pass.set_pipeline(pdq);
                    pass.set_bind_group(0, bind_dq, &[]);
                    let wgs = ((rows * cols / 2) as u32).div_ceil(256);
                    pass.dispatch_workgroups(wgs.min(MAX_WG), wgs.div_ceil(MAX_WG), 1);
                }
                let asc = c.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("dit-ascale"),
                    size: 4,
                    usage: wgpu::BufferUsages::STORAGE,
                    mapped_at_creation: false,
                });
                if encode_act_absmax(c, enc, xs, k * cols, &asc) {
                    encode_q4_tile_mm_full(
                        c,
                        enc,
                        pipe,
                        plane,
                        xs,
                        y,
                        rows,
                        cols,
                        k,
                        0.0,
                        Some(&asc),
                    );
                    return;
                }
            }
        }
        encode_q4_tile_mm(c, enc, mm, w, xs, y, rows, cols, k);
    };
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dit-block"),
        });
    // E1 — attention pre-norm, modulated.
    let dm_p = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[n as u32, hs as u32, a.eps.to_bits(), 1u32]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let norm = |enc: &mut wgpu::CommandEncoder,
                src: &wgpu::Buffer,
                w: &wgpu::Buffer,
                sc: &wgpu::Buffer,
                dst: &wgpu::Buffer,
                p: &wgpu::Buffer| {
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.dit_rmsmod.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, src),
                bind_buf(1, w),
                bind_buf(2, sc),
                bind_buf(3, dst),
                bind_buf(4, p),
            ],
        });
        let mut pass = begin_pass(enc);
        pass.set_pipeline(&c.dit_rmsmod);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(n as u32, 1, 1);
    };
    norm(&mut enc, &xb, &n1w, &smsa, &xn, &dm_p);
    // E2 — the three projections.
    mm_enc(&mut enc, a.wq, &wq, &xn, &qtok, nh * hd, hs, n);
    mm_enc(&mut enc, a.wk, &wk, &xn, &ktok, nkv * hd, hs, n);
    mm_enc(&mut enc, a.wv, &wv, &xn, &vtok, nkv * hd, hs, n);
    // E3 — qk-norm + RoPE into head-major panels; v is packed as is.
    let scale = 1.0f32 / (hd as f32).sqrt();
    let rope_pass = |enc: &mut wgpu::CommandEncoder,
                     src: &wgpu::Buffer,
                     dst: &wgpu::Buffer,
                     w: &wgpu::Buffer,
                     heads: usize,
                     sc: f32| {
        let p = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&[
                    n as u32,
                    heads as u32,
                    hd as u32,
                    a.eps.to_bits(),
                    sc.to_bits(),
                    0u32,
                    0u32,
                    0u32,
                ]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.dit_ropepack.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, src),
                bind_buf(1, dst),
                bind_buf(2, w),
                bind_buf(3, &rcos),
                bind_buf(4, &rsin),
                bind_buf(5, &p),
            ],
        });
        let mut pass = begin_pass(enc);
        pass.set_pipeline(&c.dit_ropepack);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(n as u32, heads as u32, 1);
    };
    // The scale rides the qk kernel's params, as in dit_attention.
    rope_pass(&mut enc, &qtok, &qhm, &nqw, nh, 1.0);
    rope_pass(&mut enc, &ktok, &khm, &nkw, nkv, 1.0);
    // v needs the same token-major → head-major transpose, without the
    // norm or the rotation: a unit weight vector and identity rope would
    // do it, but the blit is cheaper and exact.
    {
        let p = c
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(&[
                    n as u32,
                    nkv as u32,
                    hd as u32,
                    0f32.to_bits(),
                    1.0f32.to_bits(),
                    1u32,
                    0u32,
                    0u32,
                ]),
                usage: wgpu::BufferUsages::UNIFORM,
            });
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.dit_ropepack.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &vtok),
                bind_buf(1, &vhm),
                bind_buf(2, &nqw),
                bind_buf(3, &rcos),
                bind_buf(4, &rsin),
                bind_buf(5, &p),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.dit_ropepack);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(n as u32, nkv as u32, 1);
    }
    // E4 — attention: one dispatch for ALL heads of a segment, the head
    // index riding wid.z. Per head it was 24 launches of three kernels
    // apiece, and the P·V grid — hd/64 by ns/64, 34 workgroups — left a
    // 150-SM card almost idle. The scores buffer is what forced it: shared
    // between heads, they had to run in turn. Giving each head its own
    // slice costs ns² · nh · 4 bytes, 104 MB at 512x512, and buys the
    // occupancy back.
    let hpk = nh / nkv;
    let p8h = |m: u32, k: u32, nn: u32, sc: f32, hpk: u32, ntok: u32| -> wgpu::Buffer {
        let raw = [m, k, nn, sc.to_bits(), 0u32, 0u32, hpk, ntok];
        let b = c.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("dit-attn-params"),
            size: 32,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        c.queue.write_buffer(&b, 0, bytemuck::cast_slice(&raw));
        b
    };
    let mut off = 0usize;
    for &ns in segs {
        let p_qk = p8h(ns as u32, hd as u32, ns as u32, scale, hpk as u32, n as u32);
        let p_sm = p8h(ns as u32, hd as u32, ns as u32, 1.0, hpk as u32, n as u32);
        let p_pv = p8h(ns as u32, ns as u32, hd as u32, 1.0, hpk as u32, n as u32);
        {
            let qoff = (off * hd * 4) as u64;
            let koff = (off * hd * 4) as u64;
            let hlen = (((nh - 1) * n + ns) * hd * 4) as u64;
            let klen = (((nkv - 1) * n + ns) * hd * 4) as u64;
            let slen = (nh * ns * ns * 4) as u64;
            let bind3 = |pipe: &wgpu::ComputePipeline,
                         a0: (&wgpu::Buffer, u64, u64),
                         a1: (&wgpu::Buffer, u64, u64),
                         a2: (&wgpu::Buffer, u64, u64),
                         pu: &wgpu::Buffer| {
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &pipe.get_bind_group_layout(0),
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                                buffer: a0.0,
                                offset: a0.1,
                                size: std::num::NonZeroU64::new(a0.2),
                            }),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                                buffer: a1.0,
                                offset: a1.1,
                                size: std::num::NonZeroU64::new(a1.2),
                            }),
                        },
                        wgpu::BindGroupEntry {
                            binding: 2,
                            resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                                buffer: a2.0,
                                offset: a2.1,
                                size: std::num::NonZeroU64::new(a2.2),
                            }),
                        },
                        wgpu::BindGroupEntry {
                            binding: 3,
                            resource: pu.as_entire_binding(),
                        },
                    ],
                })
            };
            let bg_qk = bind3(
                &c.dit_qk,
                (&qhm, qoff, hlen),
                (&khm, koff, klen),
                (&scb, 0, slen),
                &p_qk,
            );
            let bg_sm = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &c.dit_softmax.get_bind_group_layout(0),
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                            buffer: &scb,
                            offset: 0,
                            size: std::num::NonZeroU64::new(slen),
                        }),
                    },
                    wgpu::BindGroupEntry {
                        binding: 3,
                        resource: p_sm.as_entire_binding(),
                    },
                ],
            });
            let bg_pv = bind3(
                &c.dit_pv,
                (&scb, 0, slen),
                (&vhm, koff, klen),
                (&pan, qoff, hlen),
                &p_pv,
            );
            for (pipe, bg, gx, gy, gz) in [
                (
                    &c.dit_qk,
                    &bg_qk,
                    (ns as u32).div_ceil(64),
                    (ns as u32).div_ceil(64),
                    nh as u32,
                ),
                (&c.dit_softmax, &bg_sm, ns as u32, nh as u32, 1u32),
                (
                    &c.dit_pv,
                    &bg_pv,
                    (hd as u32).div_ceil(64),
                    (ns as u32).div_ceil(64),
                    nh as u32,
                ),
            ] {
                let mut pass = begin_pass(&mut enc);
                pass.set_pipeline(pipe);
                pass.set_bind_group(0, bg, &[]);
                pass.dispatch_workgroups(gx, gy, gz);
            }
        }
        off += ns;
    }
    // E5 — head-major panel back to token-major, then o and residual.
    {
        let p_un = p8(nh as u32, n as u32, hd as u32, 1.0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.dit_unstack.get_bind_group_layout(0),
            entries: &[bind_buf(0, &pan), bind_buf(2, &attn), bind_buf(3, &p_un)],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.dit_unstack);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups_flat(((nh * n * hd) as u32).div_ceil(256));
    }
    mm_enc(&mut enc, a.wo, &wo, &attn, &proj, hs, nh * hd, n);
    let gr_p = c
        .device
        .create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: None,
            contents: bytemuck::cast_slice(&[n as u32, hs as u32, a.eps.to_bits(), 0u32]),
            usage: wgpu::BufferUsages::UNIFORM,
        });
    let gated =
        |enc: &mut wgpu::CommandEncoder, d: &wgpu::Buffer, w: &wgpu::Buffer, g: &wgpu::Buffer| {
            let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &c.dit_gres.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, &xb),
                    bind_buf(1, d),
                    bind_buf(2, w),
                    bind_buf(3, g),
                    bind_buf(4, &gr_p),
                ],
            });
            let mut pass = begin_pass(enc);
            pass.set_pipeline(&c.dit_gres);
            pass.set_bind_group(0, &bind, &[]);
            pass.dispatch_workgroups(n as u32, 1, 1);
        };
    gated(&mut enc, &proj, &n2w, &gmsa);
    // E6 — FFN: modulated norm, SwiGLU, gated residual.
    norm(&mut enc, &xb, &f1w, &smlp, &xn, &dm_p);
    mm_enc(&mut enc, a.w1, &w1, &xn, &gbuf, inter, hs, n);
    mm_enc(&mut enc, a.w3, &w3, &xn, &ubuf, inter, hs, n);
    {
        let n1 = u4([(n * inter) as u32, 0, 0, 0]);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.silu.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &gbuf),
                bind_buf(1, &ubuf),
                bind_buf(2, &gbuf),
                bind_buf(3, &abuf),
                bind_buf(4, &n1),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.silu);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups_flat(((n * inter) as u32).div_ceil(256));
    }
    mm_enc(&mut enc, a.w2, &w2, &abuf, &proj, hs, inter, n);
    gated(&mut enc, &proj, &f2w, &gmlp);
    let bytes = (n * hs * 4) as u64;
    if a.resident_out {
        // Nothing to wait for: the next block reads `xb` where this one
        // left it, and the submission is the only thing that must happen.
        submit(c, finish_enc(enc));
        return true;
    }
    let stage = pooled(
        n * hs,
        "dit-block-stage",
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
    );
    readback(c, enc, &xb, &stage, bytes, &mut x[..n * hs])
}

/// Read the resident DiT hidden state back. Used when a chain of blocks
/// is interrupted — the state lives only on the device at that point.
pub fn dit_state_fetch(x: &mut [f32]) -> bool {
    let Some(c) = ctx() else { return false };
    let pool = c.dit_pool.lock().unwrap();
    let Some((xb, have)) = pool.get("dit-x").cloned() else {
        return false;
    };
    drop(pool);
    let bytes = (x.len() * 4) as u64;
    if have < bytes {
        return false;
    }
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dit-fetch-stage",
    );
    drop(sc);
    let enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dit-fetch"),
        });
    readback(c, enc, &xb, &stage, bytes, x)
}

/// The configured weight budget, initializing the device on first ask.
/// None when no GPU path is selected or init failed.
pub fn dsv4_vram_budget() -> Option<u64> {
    ctx().map(|c| c.vram_budget)
}

/// `dsv4_experts_fit` with the draft's own carve-out handed back: the draft
/// pack builder must see the room that was reserved FOR it, not the room
/// that remains after its own reservation.
pub fn dsv4_draft_fit(inter: usize, hidden: usize, gu_q2: bool, dn_q2: bool) -> usize {
    let base = dsv4_experts_fit(inter, hidden, gu_q2, dn_q2);
    if std::env::var("CMF_GPU_WORKSPACE_MB").is_ok() {
        // An explicit workspace is the operator's own split; the draft
        // competes inside it exactly as before the reservation existed.
        return base;
    }
    let gu = cortiq_core::quant::expected_nbytes(
        if gu_q2 {
            cortiq_core::TensorDtype::Q2TiledP
        } else {
            cortiq_core::TensorDtype::Q4TiledP
        },
        &[inter, hidden],
    )
    .unwrap_or(0);
    let dn = cortiq_core::quant::expected_nbytes(
        if dn_q2 {
            cortiq_core::TensorDtype::Q2TiledP
        } else {
            cortiq_core::TensorDtype::Q4TiledP
        },
        &[hidden, inter],
    )
    .unwrap_or(0);
    let per = (2 * gu + dn) as u64;
    if per == 0 {
        return base;
    }
    base + (DRAFT_PACK_RESERVE.load(std::sync::atomic::Ordering::Relaxed) / per) as usize
}

/// Can this layer's experts live on the card? Uploads them if they can, so a
/// caller that pre-flights every layer has also paid the upload before it
/// commits to the device path.
/// Ask for the attention weights BEFORE the experts, or the experts take the
/// card and the skeleton — two orders of magnitude smaller — has nowhere left.
/// Positional read that builds on every target: pread on unix,
/// seek_read on windows — the release CI's Windows and iOS lanes are
/// exactly the builds that do not have `std::os::unix`.
fn read_at(f: &std::fs::File, buf: &mut [u8], off: u64) -> std::io::Result<()> {
    #[cfg(unix)]
    {
        use std::os::unix::fs::FileExt;
        f.read_exact_at(buf, off)
    }
    #[cfg(windows)]
    {
        use std::os::windows::fs::FileExt;
        let mut done = 0usize;
        while done < buf.len() {
            let n = f.seek_read(&mut buf[done..], off + done as u64)?;
            if n == 0 {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::UnexpectedEof,
                    "seek_read hit EOF",
                ));
            }
            done += n;
        }
        Ok(())
    }
    #[cfg(not(any(unix, windows)))]
    {
        let _ = (f, buf, off);
        Err(std::io::Error::new(
            std::io::ErrorKind::Unsupported,
            "no positional read on this target",
        ))
    }
}

/// Contiguous host banks of expert payloads — the fill source FreeToken
/// calls `bank_sources`: one allocation per (layer, projection) holding
/// EVERY expert's bytes row-contiguous, so a slot fill is one memcpy from
/// a known offset instead of a page-cache walk that may reach disk.
/// Built once by a background sweep (sequential file order, page cache
/// dropped behind it — banks REPLACE the cache for these ranges, they do
/// not double it). `CMF_DSV4_HOST_BANKS=1` builds them; the fill path
/// falls back transparently while a bank is not ready yet.
struct HostBanks {
    g: Vec<u8>,
    u: Vec<u8>,
    d: Vec<u8>,
    gu_len: usize,
    d_len: usize,
}

fn host_banks() -> &'static Mutex<HashMap<(usize, usize), std::sync::Arc<HostBanks>>> {
    static M: std::sync::OnceLock<Mutex<HashMap<(usize, usize), std::sync::Arc<HostBanks>>>> =
        std::sync::OnceLock::new();
    M.get_or_init(|| Mutex::new(HashMap::new()))
}

pub fn host_banks_on() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| std::env::var("CMF_DSV4_HOST_BANKS").is_ok_and(|v| v != "0"))
}

/// Build one layer's banks from its expert tensor triples (w1, w3, w2 —
/// the same list a pack carries, shared expert last, which is skipped:
/// it is VRAM-resident for the life of the process). Reads by pread and
/// drops the cache behind itself.
pub fn dsv4_host_bank_build(
    model: &Arc<CmfModel>,
    pack_first: usize,
    experts: &[(usize, usize, usize)],
    inter: usize,
    hidden: usize,
    gu_q2: bool,
) {
    if !host_banks_on() {
        return;
    }
    let key = (model.uid() as usize, pack_first);
    if host_banks().lock().unwrap().contains_key(&key) {
        return;
    }
    let dt = |q2: bool| {
        if q2 {
            cortiq_core::TensorDtype::Q2TiledP
        } else {
            cortiq_core::TensorDtype::Q4TiledP
        }
    };
    let Some(gu_len) = cortiq_core::quant::expected_nbytes(dt(gu_q2), &[inter, hidden]) else {
        return;
    };
    let Some(d_len) = cortiq_core::quant::expected_nbytes(dt(false), &[hidden, inter]) else {
        return;
    };
    let n = experts.len().saturating_sub(1); // routed only
    let mut banks = HostBanks {
        g: vec![0u8; n * gu_len],
        u: vec![0u8; n * gu_len],
        d: vec![0u8; n * d_len],
        gu_len,
        d_len,
    };
    let Ok(f) = std::fs::File::open(&model.path) else {
        return;
    };
    let mut pull = |idx: usize, dst: &mut [u8]| -> bool {
        let Some(e) = model.tensors.get(idx) else {
            return false;
        };
        if e.nbytes as usize != dst.len() {
            return false;
        }
        let Some(abs) = model.entry_abs_offset(e) else {
            return false;
        };
        if read_at(&f, dst, abs as u64).is_err() {
            return false;
        }
        #[cfg(target_os = "linux")]
        unsafe {
            use std::os::unix::io::AsRawFd;
            libc::posix_fadvise(
                f.as_raw_fd(),
                abs as i64,
                dst.len() as i64,
                libc::POSIX_FADV_DONTNEED,
            );
        }
        true
    };
    for (gi, &(t1, t3, t2)) in experts[..n].iter().enumerate() {
        if !pull(t1, &mut banks.g[gi * gu_len..(gi + 1) * gu_len])
            || !pull(t3, &mut banks.u[gi * gu_len..(gi + 1) * gu_len])
            || !pull(t2, &mut banks.d[gi * d_len..(gi + 1) * d_len])
        {
            return; // partial banks are worse than none
        }
    }
    host_banks()
        .lock()
        .unwrap()
        .insert(key, std::sync::Arc::new(banks));
}

/// Fill counters for the dynamic expert slots (`dsv4_slot_fill`).
pub static DSV4_FILLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static DSV4_FILL_BYTES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static DSV4_FILL_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

fn dsv4_fill_profile_on() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| {
        std::env::var("CMF_DSV41_PROF").is_ok_and(|v| v != "0")
            || std::env::var("CMF_DSV4_PROF").is_ok_and(|v| v != "0")
    })
}

/// Overwrite ONE routed-expert slot of a layer's device bank buffers with
/// another expert's bytes — the FreeToken move: the packed subset follows
/// the router instead of staying whatever load-time frequency guessed.
///
/// Three `queue.write_buffer`s into slot offsets; the queue orders them
/// before any later submit, and the frame re-uploads its remap every call,
/// so the table and the bytes can never be seen out of step. Sources, in
/// order: the RAM tier (prefetched), pread (CMF_WEIGHT_PREAD), the mmap.
/// `pack_first` is the pack's first w1 tensor index — the bank cache key.
pub fn dsv4_slot_fill(
    model: &Arc<CmfModel>,
    pack_first: usize,
    slot: usize,
    gi: usize,
    t: (usize, usize, usize),
    inter: usize,
    hidden: usize,
    gu_q2: bool,
) -> bool {
    // The bank fast path: three memcpys from contiguous rows — no page
    // cache, no disk, no tier lookups. Falls through when the layer's
    // banks are not built (yet).
    if host_banks_on() {
        let bkey = (model.uid() as usize, pack_first);
        let banks = host_banks().lock().unwrap().get(&bkey).cloned();
        if let Some(bk) = banks {
            let (gl, dl) = (bk.gu_len, bk.d_len);
            if (gi + 1) * gl <= bk.g.len() && (gi + 1) * dl <= bk.d.len() {
                let key = (model.uid() as usize, pack_first);
                let Some((g, u, d)) =
                    ctx().and_then(|c| c.moe_expw.lock().unwrap().get(&key).cloned())
                else {
                    return false;
                };
                let c = ctx().unwrap();
                if (slot + 1) * gl as usize > g.size() as usize / 1 {}
                c.queue
                    .write_buffer(&g, (slot * gl) as u64, &bk.g[gi * gl..(gi + 1) * gl]);
                c.queue
                    .write_buffer(&u, (slot * gl) as u64, &bk.u[gi * gl..(gi + 1) * gl]);
                c.queue
                    .write_buffer(&d, (slot * dl) as u64, &bk.d[gi * dl..(gi + 1) * dl]);
                use std::sync::atomic::Ordering;
                DSV4_FILLS.fetch_add(1, Ordering::Relaxed);
                DSV4_FILL_BYTES.fetch_add((2 * gl + dl) as u64, Ordering::Relaxed);
                return true;
            }
        }
    }
    use std::sync::atomic::Ordering;
    let Some(c) = ctx() else { return false };
    let key = (model.uid() as usize, pack_first);
    let Some((g, u, d)) = c.moe_expw.lock().unwrap().get(&key).cloned() else {
        return false;
    };
    let plen4 = |rows: usize, cols: usize| -> Option<usize> {
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
    };
    let gu_len = if gu_q2 {
        cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q2TiledP, &[inter, hidden])
    } else {
        plen4(inter, hidden)
    };
    let (Some(gu_len), Some(d_len)) = (gu_len, plen4(hidden, inter)) else {
        return false;
    };
    let bytes = model.primary_bytes();
    let put = |buf: &wgpu::Buffer, idx: usize, plen: usize| -> bool {
        let Some(e) = model.tensors.get(idx) else {
            return false;
        };
        if e.nbytes as usize != plen {
            return false;
        }
        let Some(abs) = model.entry_abs_offset(e) else {
            return false;
        };
        if slot.checked_mul(plen).is_none() || (slot * plen + plen) as u64 > buf.size() {
            return false;
        }
        let key = (model.uid() as usize, idx);
        let tier = host_tier_get(key);
        let src: &[u8] = if let Some(v) = tier.as_deref() {
            if v.len() != plen {
                return false;
            }
            v
        } else if let Some(v) = pread_range(model, abs, plen) {
            // Fill-through: a fetched expert earned its trip by being
            // re-picked (the min-seen gate), so its NEXT eviction-refetch
            // is likelier than not — serve it from RAM then, not disk.
            c.queue.write_buffer(buf, (slot * plen) as u64, &v);
            host_tier_put(key, std::sync::Arc::new(v));
            return true;
        } else {
            let Some(sl) = bytes.get(abs..abs + plen) else {
                return false;
            };
            sl
        };
        c.queue.write_buffer(buf, (slot * plen) as u64, src);
        if tier.is_none() {
            host_tier_put(key, std::sync::Arc::new(src.to_vec()));
        }
        true
    };
    let ok = put(&g, t.0, gu_len) && put(&u, t.1, gu_len) && put(&d, t.2, d_len);
    if ok {
        DSV4_FILLS.fetch_add(1, Ordering::Relaxed);
        DSV4_FILL_BYTES.fetch_add((2 * gu_len + d_len) as u64, Ordering::Relaxed);
    }
    ok
}

/// Host→VRAM upload bandwidth over `rounds` writes of a `block`-sized
/// buffer, submits included — the fetch arm's real ceiling for the MoE
/// hybrid split (`cortiq bench --bw`). None when no device came up.
pub fn upload_bandwidth_probe(block: usize, rounds: usize) -> Option<f64> {
    let c = ctx()?;
    let src = vec![5u8; block];
    let buf = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("bw-probe"),
        size: block as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    // One warm round compiles nothing but pages the staging belt in.
    c.queue.write_buffer(&buf, 0, &src);
    c.queue.submit(None);
    let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
    let t0 = std::time::Instant::now();
    for _ in 0..rounds {
        c.queue.write_buffer(&buf, 0, &src);
        c.queue.submit(None);
    }
    let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
    Some((block * rounds) as f64 / t0.elapsed().as_secs_f64() / 1e9)
}

/// Automatic FreeToken-style refill policy for a partial DeepSeek-V4 pack.
/// Four small transfers distinguish a local high-bandwidth link from a
/// throttled/cloud staging path without adding a visible model-load tax.
/// Fast links refill on the first recurrence; slow links demand a second
/// observation so one-shot experts remain exact CPU cold picks.
pub fn dsv4_fetch_defaults() -> (usize, u16) {
    static POLICY: std::sync::OnceLock<(usize, u16)> = std::sync::OnceLock::new();
    *POLICY.get_or_init(|| {
        let gbs = upload_bandwidth_probe(8 * 1024 * 1024, 4).unwrap_or(0.0);
        let min_seen = if gbs >= 10.0 { 1 } else { 2 };
        tracing::info!(
            "dsv4 refill auto: upload {:.2} GB/s, quota 1, min_seen {}",
            gbs,
            min_seen
        );
        (1, min_seen)
    })
}

pub fn dsv4_experts_ready(
    model: &Arc<CmfModel>,
    experts: &[(usize, usize, usize)],
    inter: usize,
    hidden: usize,
    gu_q2: bool,
    dn_q2: bool,
) -> bool {
    let Some(c) = ctx() else { return false };
    moe_expert_bufs(c, model, experts, inter, hidden, true, gu_q2, dn_q2).is_some()
}

/// Seed the attention half's `post`/`comb` from the host. The frame's opening
/// expand reads what the PREVIOUS frame's tail left there; layer zero has no
/// previous frame, and neither does the layer after one that ran on the host.
/// Without this both read whatever was in the buffer — perplexity 1470.
/// Seed one token's hyper-connection mix for the run's FIRST layer.
///
/// Every later layer computes its own on the card, but the first takes the
/// host's. With one slot the batch's tokens would all get whichever seed was
/// written last — so this, like the state, is per token.
pub fn dsv4_hc_write_t(post: &[f32], comb: &[f32], tok: usize) -> bool {
    let Some(c) = ctx() else { return false };
    let pb = frame_buf_t(c, 43, tok, post.len() * 4, true);
    let cb = frame_buf_t(c, 44, tok, comb.len() * 4, true);
    c.queue.write_buffer(&pb, 0, bytemuck::cast_slice(post));
    c.queue.write_buffer(&cb, 0, bytemuck::cast_slice(comb));
    true
}

pub fn dsv4_hc_write(post: &[f32], comb: &[f32]) -> bool {
    let Some(c) = ctx() else { return false };
    let pb = frame_buf(c, 43, post.len() * 4, true);
    let cb = frame_buf(c, 44, comb.len() * 4, true);
    c.queue.write_buffer(&pb, 0, bytemuck::cast_slice(post));
    c.queue.write_buffer(&cb, 0, bytemuck::cast_slice(comb));
    true
}

/// Seed the layer-frame's hyper-connection state from the host (layer zero).
/// Seed one token's hyper-connection state on the card.
///
/// The batch needs one of these per token before the chain runs: each token
/// enters at its own embedding, and after the frame's buffers went per token
/// the single-slot writer can only reach token zero.
pub fn dsv4_state_write_t(state: &[f32], tok: usize) -> bool {
    let Some(c) = ctx() else { return false };
    let b = frame_buf_t(c, 40, tok, state.len() * 4, true);
    c.queue.write_buffer(&b, 0, bytemuck::cast_slice(state));
    true
}

pub fn dsv4_state_write(state: &[f32]) -> bool {
    let Some(c) = ctx() else { return false };
    let b = frame_buf(c, 40, state.len() * 4, true);
    c.queue.write_buffer(&b, 0, bytemuck::cast_slice(state));
    true
}

/// Read the hyper-connection state back (end of token, for the head).
pub fn dsv4_state_read(state: &mut [f32]) -> bool {
    let Some(c) = ctx() else { return false };
    let b = frame_buf(c, 40, state.len() * 4, true);
    let bytes = (state.len() * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dsv4-state-stage",
    );
    let enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("st") });
    let ok = readback(c, enc, &b, &stage, bytes, state);
    drop(sc);
    ok
}

/// Add the experts that did not fit in VRAM to the device-owned
/// hyper-connection state and bring the corrected state home.
///
/// The MoE frame has already expanded the resident experts as
/// `post[j] * resident + comb * residual`. Hyper-connections are linear in
/// the block output, so the exact correction is another expansion with an
/// identity `comb`: `state[j] += post[j] * cold`. Doing it here keeps the
/// policy independent of a layer number or a particular VRAM size.
/// `dsv4_state_add_cold`, but the post comes from the spare slot the layer
/// frame preserved before its next-layer fold rewrote the canonical one.
pub fn dsv4_state_add_cold_preserved(
    cold: &[f32],
    hc: usize,
    state_out: &mut [f32],
    kv_id: u64,
    li: usize,
) -> bool {
    dsv4_state_add_cold_inner(cold, hc, state_out, Some((kv_id, li)))
}

pub fn dsv4_state_add_cold(cold: &[f32], hc: usize, state_out: &mut [f32]) -> bool {
    dsv4_state_add_cold_inner(cold, hc, state_out, None)
}

fn dsv4_state_add_cold_inner(
    cold: &[f32],
    hc: usize,
    state_out: &mut [f32],
    preserved: Option<(u64, usize)>,
) -> bool {
    let Some(c) = ctx() else { return false };
    if hc == 0 || cold.is_empty() || state_out.len() != hc * cold.len() {
        return false;
    }
    let dim = cold.len();
    let state = frame_buf(c, 40, state_out.len() * 4, true);
    let corrected = frame_buf(c, 46, state_out.len() * 4, true);
    // `post` was produced by the FFN half's opening fold and is still live in
    // the canonical slot after the MoE frame returns its cold winners.
    let post = match preserved {
        Some((kv, li)) => store_slot(c, 48, kv, li, &[]),
        None => frame_buf(c, 43, hc * 4, true),
    };
    let cold_buf = frame_up(c, 121, bytemuck::cast_slice(cold));
    let mut identity = vec![0.0f32; hc * hc];
    for j in 0..hc {
        identity[j * hc + j] = 1.0;
    }
    let comb = frame_up(c, 122, bytemuck::cast_slice(&identity));
    let p = uniform_u32x8(c, [hc as u32, dim as u32, 0, 0, 0, 0, 0, 0]);
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-cold-state-fix"),
        });
    encode_hc_expand(
        c, &mut enc, &cold_buf, &state, &post, &comb, &corrected, &p, hc, dim,
    );
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&corrected, 0, &state, 0, (state_out.len() * 4) as u64);
    let bytes = (state_out.len() * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dsv4-cold-state-stage",
    );
    let ok = readback(c, enc, &state, &stage, bytes, state_out);
    drop(sc);
    ok
}

/// MoE routing on the device: scores in, chosen experts and their normalised
/// weights out, in selection order. `forced` is the hash layers' table row.
#[allow(clippy::too_many_arguments)]
/// The token-axis router against the same contract, `b` tokens at once:
/// `scores` is `[b, n]`, outputs are `[b, top_k+1]`. Exercises exactly the
/// kernel the batched frame dispatches.
#[allow(clippy::too_many_arguments)]
pub fn bt_moe_route_for_test(
    scores: &[f32],
    b: usize,
    bias: Option<&[f32]>,
    forced: Option<&[Vec<usize>]>,
    top_k: usize,
    route_scale: f32,
    idx_out: &mut Vec<usize>,
    w_out: &mut Vec<f32>,
) -> bool {
    let Some(c) = ctx() else { return false };
    let n = scores.len() / b.max(1);
    if n == 0 || n > 1024 || top_k == 0 || top_k > 64 || b == 0 {
        return false;
    }
    let slots = top_k + 1;
    let sb = storage_bytes(c, bytemuck::cast_slice(scores));
    let bb = match bias {
        Some(v) if v.len() >= n => storage_bytes(c, bytemuck::cast_slice(&v[..n])),
        _ => sb.clone(),
    };
    let mb = storage_bytes(c, bytemuck::cast_slice(&vec![1u32; n]));
    let fbuf = match forced {
        Some(rows) if rows.len() >= b => {
            let mut v = vec![0u32; b * top_k];
            for (t, r) in rows.iter().enumerate() {
                for (i, &e) in r.iter().take(top_k).enumerate() {
                    v[t * top_k + i] = e as u32;
                }
            }
            storage_bytes(c, bytemuck::cast_slice(&v))
        }
        _ => storage_bytes(c, bytemuck::cast_slice(&vec![0u32; b * top_k])),
    };
    let ib = rw_f32(c, b * slots, true);
    let wb = rw_f32(c, b * slots, true);
    let cnt = rw_f32(c, b, true);
    let rmb = storage_bytes(c, bytemuck::cast_slice(&vec![0u32; n]));
    let coldb = rw_f32(c, b * 4 * top_k, false);
    let flags = (bias.is_some_and(|v| v.len() >= n) as u32)
        | ((forced.is_some() as u32) << 2)
        | 8
        | ((n as u32) << 8);
    let p = uniform_mixed(c, [n as u32, top_k as u32, flags], route_scale);
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("bt-route-test"),
        });
    {
        let layout = c.bt_moe_route.get_bind_group_layout(0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layout,
            entries: &[
                bind_buf(0, &sb),
                bind_buf(1, &bb),
                bind_buf(2, &mb),
                bind_buf(3, &fbuf),
                bind_buf(4, &ib),
                bind_buf(5, &wb),
                bind_buf(6, &cnt),
                bind_buf(7, &p),
                bind_buf(8, &rmb),
                bind_buf(9, &coldb),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.bt_moe_route);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(b as u32, 1, 1);
    }
    let ib_bytes = (b * slots * 4) as u64;
    let mut iv = vec![0.0f32; b * slots];
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        ib_bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "bt-route-stage",
    );
    if !readback(c, enc, &ib, &stage, ib_bytes, &mut iv) {
        return false;
    }
    let iu: Vec<u32> = bytemuck::cast_slice(&iv).to_vec();
    idx_out.clear();
    idx_out.extend(iu.iter().map(|&x| x as usize));
    let mut enc2 = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("bt-route-test-w"),
        });
    let _ = &mut enc2;
    let mut wv = vec![0.0f32; b * slots];
    if !readback(c, enc2, &wb, &stage, ib_bytes, &mut wv) {
        return false;
    }
    w_out.clear();
    w_out.extend_from_slice(&wv);
    drop(sc);
    true
}

pub fn moe_route_for_test(
    scores: &[f32],
    bias: Option<&[f32]>,
    mask: Option<&[bool]>,
    forced: Option<&[usize]>,
    top_k: usize,
    route_scale: f32,
    // Pin the shared expert in slot `top_k` at weight 1, and write every slot
    // — the `msel`/`mwt` pair the batched expert kernels take. Off returns
    // just what the router chose.
    shared_slot: bool,
    idx_out: &mut Vec<usize>,
    w_out: &mut Vec<f32>,
) -> bool {
    moe_route_for_test_impl(
        scores,
        bias,
        mask,
        forced,
        top_k,
        route_scale,
        shared_slot,
        false,
        idx_out,
        w_out,
    )
}

/// Qwen MoE route parity hook. Qwen ranks the raw router logits and applies
/// a softmax over the selected experts, unlike the DeepSeek sqrt-softplus
/// route above. Keeping this as a separate entry point prevents accidental
/// changes to the established DSV4 test contract.
pub fn qwen_moe_route_for_test(
    scores: &[f32],
    top_k: usize,
    idx_out: &mut Vec<usize>,
    w_out: &mut Vec<f32>,
) -> bool {
    moe_route_for_test_impl(
        scores, None, None, None, top_k, 1.0, false, true, idx_out, w_out,
    )
}

#[allow(clippy::too_many_arguments)]
fn moe_route_for_test_impl(
    scores: &[f32],
    bias: Option<&[f32]>,
    mask: Option<&[bool]>,
    forced: Option<&[usize]>,
    top_k: usize,
    route_scale: f32,
    shared_slot: bool,
    qwen_softmax: bool,
    idx_out: &mut Vec<usize>,
    w_out: &mut Vec<f32>,
) -> bool {
    let Some(c) = ctx() else { return false };
    let n = scores.len();
    if n == 0 || n > 1024 || top_k == 0 || top_k > 64 {
        return false;
    }
    let sb = storage_bytes(c, bytemuck::cast_slice(scores));
    let bb = match bias {
        Some(b) if b.len() >= n => storage_bytes(c, bytemuck::cast_slice(&b[..n])),
        _ => sb.clone(),
    };
    let mb = match mask {
        Some(m) if m.len() >= n => {
            let v: Vec<u32> = m[..n].iter().map(|&x| x as u32).collect();
            storage_bytes(c, bytemuck::cast_slice(&v))
        }
        _ => storage_bytes(c, bytemuck::cast_slice(&vec![1u32; n])),
    };
    let fb = match forced {
        Some(f) if f.len() >= top_k => {
            let v: Vec<u32> = f[..top_k].iter().map(|&x| x as u32).collect();
            storage_bytes(c, bytemuck::cast_slice(&v))
        }
        _ => storage_bytes(c, bytemuck::cast_slice(&vec![0u32; top_k])),
    };
    let slots = top_k + shared_slot as usize;
    let ib = rw_f32(c, slots, true);
    let wb = rw_f32(c, slots, true);
    let cb = rw_f32(c, 1, true);
    let rmb = storage_bytes(c, bytemuck::cast_slice(&vec![0u32; n]));
    let coldb = rw_f32(c, 4 * top_k, false);
    let flags = (bias.is_some_and(|b| b.len() >= n) as u32)
        | ((mask.is_some_and(|m| m.len() >= n) as u32) << 1)
        | ((forced.is_some_and(|f| f.len() >= top_k) as u32) << 2)
        | ((shared_slot as u32) << 3)
        | ((qwen_softmax as u32) << 5)
        // Nothing is packed away here, so the shared expert sits at n — the
        // same slot the frame computes from its packing.
        | ((n as u32) << 8);
    let p = uniform_mixed(c, [n as u32, top_k as u32, flags], route_scale);
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("route"),
        });
    {
        let layout = c.moe_route.get_bind_group_layout(0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layout,
            entries: &[
                bind_buf(0, &sb),
                bind_buf(1, &bb),
                bind_buf(2, &mb),
                bind_buf(3, &fb),
                bind_buf(4, &ib),
                bind_buf(5, &wb),
                bind_buf(6, &cb),
                bind_buf(7, &p),
                // The router grew a remap and a winners buffer; a standalone
                // caller that skips them is a validation error, not a
                // silently different answer.
                bind_buf(8, &rmb),
                bind_buf(9, &coldb),
            ],
        });
        // The card's own clock around the whole MoE block. All three kernels
        // share one pass — splitting it to time them apart would add two
        // pass boundaries a layer and measure the split instead — so this is
        // the block's total, which is the number that says whether the card
        // is busy or waiting.
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.moe_route);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(1, 1, 1);
    }
    // idx | weights | count, one map.
    let bytes = ((2 * slots + 1) * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "route-stage",
    );
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&ib, 0, &stage, 0, (slots * 4) as u64);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&wb, 0, &stage, (slots * 4) as u64, (slots * 4) as u64);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&cb, 0, &stage, (2 * slots * 4) as u64, 4);
    submit(c, finish_enc(enc));
    let slice = stage.slice(..bytes);
    slice.map_async(wgpu::MapMode::Read, |_| {});
    if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        return false;
    }
    let mut ok = false;
    if let Ok(data) = slice.get_mapped_range() {
        let words: &[u32] = bytemuck::cast_slice(&data[..bytes as usize]);
        let ws: &[f32] = bytemuck::cast_slice(&data[slots * 4..2 * slots * 4]);
        // With a shared slot the caller wants every slot, filled or not; the
        // count is what the kernels use to skip nothing.
        let take = if shared_slot {
            slots
        } else {
            (words[2 * slots] as usize).min(top_k)
        };
        idx_out.clear();
        w_out.clear();
        idx_out.extend(words[..take].iter().map(|&x| x as usize));
        w_out.extend_from_slice(&ws[..take]);
        ok = true;
    }
    stage.unmap();
    drop(sc);
    ok
}

fn encode_rmsnorm(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    x: &wgpu::Buffer,
    w: &wgpu::Buffer,
    o: &wgpu::Buffer,
    n: usize,
    eps: f32,
    bkey: (u8, u64, usize),
) {
    let mut pass = begin_pass(enc);
    encode_rmsnorm_p(&mut pass, c, x, w, o, n, eps, bkey);
}

#[allow(clippy::too_many_arguments)]
fn encode_rmsnorm_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    x: &wgpu::Buffer,
    w: &wgpu::Buffer,
    o: &wgpu::Buffer,
    n: usize,
    eps: f32,
    bkey: (u8, u64, usize),
) {
    let bind = cached_bind(c, bkey, || {
        let p = uniform_u32x4(c, [n as u32, 0, eps.to_bits(), 0]);
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.rmsnorm.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, x),
                bind_buf(1, w),
                bind_buf(2, o),
                bind_buf(3, &p),
            ],
        })
    });
    pass.set_pipeline(&c.rmsnorm);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(1, 1, 1);
}

#[allow(clippy::too_many_arguments)]
fn encode_rope_heads(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    x: &wgpu::Buffer,
    freq: &wgpu::Buffer,
    posb: &wgpu::Buffer,
    nh: usize,
    hd: usize,
    rd: usize,
    rms: bool,
    inverse: bool,
    bkey: (u8, u64, usize),
) {
    let mut pass = begin_pass(enc);
    encode_rope_heads_p(&mut pass, c, x, freq, posb, nh, hd, rd, rms, inverse, bkey);
}

#[allow(clippy::too_many_arguments)]
fn encode_rope_heads_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    x: &wgpu::Buffer,
    freq: &wgpu::Buffer,
    posb: &wgpu::Buffer,
    nh: usize,
    hd: usize,
    rd: usize,
    rms: bool,
    inverse: bool,
    bkey: (u8, u64, usize),
) {
    let bind = cached_bind(c, bkey, || {
        let flags = (rms as u32) | ((inverse as u32) << 1);
        let p = uniform_u32x4(c, [nh as u32, hd as u32, rd as u32, flags]);
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.rope_heads.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, x),
                bind_buf(1, freq),
                bind_buf(2, &p),
                bind_buf(3, posb),
            ],
        })
    });
    pass.set_pipeline(&c.rope_heads);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(nh as u32, 1, 1);
}

/// Round an in-flight f32 activation buffer through the reference BF16
/// representation without taking it back to the host. The buffer is pooled
/// by frame role, so the extra dispatch adds no allocation or readback.
fn encode_bf16_round(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    x: &wgpu::Buffer,
    n: usize,
    bkey: (u8, u64, usize),
) {
    let bind = cached_bind(c, bkey, || {
        let p = uniform_u32x4(c, [n as u32, 0, 0, 0]);
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.bf16_round.get_bind_group_layout(0),
            entries: &[bind_buf(0, x), bind_buf(1, &p)],
        })
    });
    let mut pass = begin_pass(enc);
    pass.set_pipeline(&c.bf16_round);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((n as u32).div_ceil(256), 1, 1);
}

/// A constant vector (a norm weight, the sinks, the frequency table) parked
/// on the card and keyed on its host address — the same bytes arrive every
/// token, and re-uploading them 43 times a token is pure waste. The
/// fingerprint refreshes the SAME buffer in place when a reloaded model's
/// mapping reuses the address — never a new object, because the dsv4
/// bind-group cache holds the handle it was built from.
fn const_buf(c: &Ctx, data: &[u8]) -> wgpu::Buffer {
    let key = (data.as_ptr() as usize, data.len());
    let fp = fp_bytes(data);
    let mut m = c.const_bufs.lock().unwrap();
    if let Some((b, f)) = m.get_mut(&key) {
        if *f != fp {
            c.queue.write_buffer(b, 0, data);
            *f = fp;
        }
        return b.clone();
    }
    let b = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("dsv4-const"),
        size: data.len().max(4) as u64,
        // COPY_SRC too: the compressor frame slices one window's worth of
        // `ape` out of the parked table with a device-to-device copy,
        // which beats re-uploading that slice every token.
        usage: wgpu::BufferUsages::STORAGE
            | wgpu::BufferUsages::COPY_DST
            | wgpu::BufferUsages::COPY_SRC,
        mapped_at_creation: false,
    });
    c.queue.write_buffer(&b, 0, data);
    m.insert(key, (b.clone(), fp));
    b
}

/// A frame working buffer, created on first use and reused for good. `tag`
/// separates roles that happen to share a length — two buffers of the same
/// size are not interchangeable when both are live in one encoder.
/// A pooled scratch buffer, keyed by (role, size).
///
/// GROW-ONLY BY CONTRACT: an entry is never evicted or recreated, so the
/// handle a caller got last token is the same one it gets this token. The
/// bind-group cache depends on that — it holds groups built from these
/// buffers and only the GREW epoch invalidates them, so adding eviction
/// here would leave every cached group pointing at a dead buffer with
/// nothing to notice.
fn frame_buf(c: &Ctx, tag: u8, len_bytes: usize, upload: bool) -> wgpu::Buffer {
    let salt = dsv4_frame_salt();
    let tok = if salt == 0 {
        0
    } else {
        // Explicit token carriers occupy 0..FRAME_TOK_STRIDE. Keep implicit
        // scratch in a disjoint namespace so token 1's position buffer can
        // never alias token 0's temporary buffer with the same tag.
        FRAME_TOK_STRIDE + salt
    };
    frame_buf_t(c, tag, tok, len_bytes, upload)
}

/// The same pool, one buffer per token of the batch.
///
/// Scratch that is written and consumed inside a single layer can stay
/// shared: tokens are encoded one after another and a compute pass orders
/// its dispatches. What cannot be shared is anything that CARRIES a token
/// from one layer to the next — the hyper-connection state above all — since
/// the second token would overwrite the first one's before it is read.
fn frame_buf_t(c: &Ctx, tag: u8, tok: usize, len_bytes: usize, upload: bool) -> wgpu::Buffer {
    let mut m = c.dsv4_scratch.lock().unwrap();
    if let Some(b) = m.get(&(tag, tok, len_bytes)) {
        return b.clone();
    }
    let mut usage = wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC;
    if upload {
        usage |= wgpu::BufferUsages::COPY_DST;
    }
    let b = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("dsv4-frame"),
        size: len_bytes.max(4) as u64,
        usage,
        mapped_at_creation: false,
    });
    m.insert((tag, tok, len_bytes), b.clone());
    b
}

/// Upload into a reused buffer instead of minting one per call.
/// The position uniform, written once per TOKEN: its contents are the same
/// for every layer, and the per-layer `frame_up` was 43 redundant queue
/// writes a token.
/// The position uniform for one token of a batch.
///
/// It cannot share a slot: several `write_buffer` calls land before the
/// submission, so every bind group would read the LAST position written and
/// the whole batch would attend at one place. The write-skip cache is only
/// safe for the single-token path, so the batch always writes.
fn frame_up_pos_t(c: &Ctx, tag: u8, tok: usize, pos: usize, eps: f32) -> wgpu::Buffer {
    if tok == 0 {
        return frame_up_pos(c, tag, pos, eps);
    }
    let b = frame_buf_t(c, tag, tok, 8, true);
    c.queue
        .write_buffer(&b, 0, bytemuck::cast_slice(&[pos as f32, eps]));
    b
}

fn frame_up_pos(c: &Ctx, tag: u8, pos: usize, eps: f32) -> wgpu::Buffer {
    use std::sync::atomic::{AtomicU64, Ordering};
    static LAST: [AtomicU64; 256] = [const { AtomicU64::new(u64::MAX) }; 256];
    let stamp = ((pos as u64) << 32) | eps.to_bits() as u64;
    let b = frame_buf(c, tag, 8, true);
    if LAST[tag as usize].swap(stamp, Ordering::Relaxed) != stamp {
        c.queue
            .write_buffer(&b, 0, bytemuck::cast_slice(&[pos as f32, eps]));
    }
    b
}

fn frame_up(c: &Ctx, tag: u8, data: &[u8]) -> wgpu::Buffer {
    let b = frame_buf(c, tag, data.len(), true);
    c.queue.write_buffer(&b, 0, data);
    b
}

fn sa_split() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| std::env::var("CMF_SA_SPLIT").is_ok_and(|v| v != "0"))
}

/// The two-dispatch sparse attention: scores per head, then the weighted sum
/// over nh*hd independent outputs. Same numbers as the one-workgroup-per-head
/// kernel, spread across the card instead of 64 groups of it — which measured
/// 0.54 ms a layer and was the whole cost of the block once its encoding was
/// cached away.
#[allow(clippy::too_many_arguments)]
/// Take the next slot pair and record it under `which`, so the frame that
/// owns the encoder can resolve every pass it wrote.
fn ts_pair(c: &Ctx, which: usize) -> Option<wgpu::ComputePassTimestampWrites<'_>> {
    let (qs, _, _) = c.ts_query.as_ref()?;
    let slot = (TS_SLOT.fetch_add(2, std::sync::atomic::Ordering::Relaxed) % 254) as u32;
    TS_PAIRS.lock().unwrap().push((which, slot));
    Some(wgpu::ComputePassTimestampWrites {
        query_set: qs,
        beginning_of_pass_write_index: Some(slot),
        end_of_pass_write_index: Some(slot + 1),
    })
}

/// The one-kernel attention's bind group, shared by the timestamped path
/// (a pass of its own, so the query set has boundaries to write) and the
/// fused one (a dispatch inside the layer's pass).
#[allow(clippy::too_many_arguments)]
/// What a dispatch actually costs, and WHY.
///
/// Every fusion this session was bought or rejected on an assumed ~30 µs a
/// dispatch, and that number came from dividing a frame by its dispatch
/// count — which cannot tell a kernel LAUNCH from the memory barrier wgpu
/// inserts when two dispatches touch the same buffer. The two have opposite
/// remedies: launches want fewer dispatches, barriers want independent ones
/// grouped together. So measure both.
///
/// N trivial dispatches into ONE buffer (every pair a hazard) against N into
/// eight buffers round-robin (no hazard within a group of eight).
pub fn dispatch_bench() -> Vec<String> {
    let mut out = Vec::new();
    let Some(c) = ctx() else {
        out.push("нет устройства".into());
        return out;
    };
    const N: usize = 2000;
    let bufs: Vec<wgpu::Buffer> = (0..8)
        .map(|_| {
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("dbench"),
                size: 4096,
                usage: wgpu::BufferUsages::STORAGE,
                mapped_at_creation: false,
            })
        })
        .collect();
    let p = uniform_u32x4(c, [1024, 0, 0, 0]);
    let binds: Vec<wgpu::BindGroup> = bufs
        .iter()
        .map(|b| {
            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &c.zero.get_bind_group_layout(0),
                entries: &[bind_buf(0, b), bind_buf(1, &p)],
            })
        })
        .collect();
    for (name, stride) in [
        ("зависимые (один буфер)", 0usize),
        ("независимые (8 буферов)", 1),
    ] {
        // Warm, then time.
        for round in 0..2 {
            let t0 = std::time::Instant::now();
            let mut enc = c
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
            {
                let mut pass = begin_pass(&mut enc);
                pass.set_pipeline(&c.zero);
                for i in 0..N {
                    let b = if stride == 0 { 0 } else { i % 8 };
                    pass.set_bind_group(0, &binds[b], &[]);
                    pass.dispatch_workgroups(4, 1, 1);
                }
            }
            submit(c, finish_enc(enc));
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            if round == 1 {
                let us = t0.elapsed().as_secs_f64() * 1e6 / N as f64;
                out.push(format!("{name}: {us:.2} мкс на диспатч"));
            }
        }
    }
    out
}

/// How many position-chunks to cut a head into. One workgroup a chunk, so
/// this is the occupancy knob: 64 heads alone left the card at a few percent.
fn sa_chunks(m: usize) -> usize {
    if m == 0 {
        return 1;
    }
    // CMF_DSV4_SA_CHUNKS forces the count. The toys attend to fewer than 128
    // positions, so they would take one chunk and never exercise the merge
    // at all — the flag is what lets the known-good stands test it.
    static F: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
    let forced = *F.get_or_init(|| {
        std::env::var("CMF_DSV4_SA_CHUNKS")
            .ok()
            .and_then(|v| v.parse::<usize>().ok())
            .unwrap_or(0)
    });
    if forced > 0 {
        return forced.min(m).min(SA_MAX_CHUNKS);
    }
    // MEASURED, not reasoned: on the release, 2 chunks gave 22.9 tok/s, 4
    // gave 23.6 and 8 gave 24.5 — monotone, because each chunk is a
    // workgroup and 64 heads alone leave the card idle. ~16 positions a
    // chunk, so a short list still splits.
    m.div_ceil(16).clamp(1, SA_MAX_CHUNKS)
}

/// The scratch is nh·SA_MAX_CHUNKS·hd floats — 4 MB at the release's shape.
const SA_MAX_CHUNKS: usize = 32;

/// `CMF_DSV4_SA_SPLIT=0` reverts attention to the one-workgroup-per-head
/// kernel. The split changes the order the softmax sums in, so it is a
/// contract change and wants a flag.
fn sa_split_k() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| {
        std::env::var("CMF_DSV4_SA_SPLIT")
            .map(|v| v != "0")
            .unwrap_or(true)
    })
}

/// Sparse attention as two dispatches: a per-chunk pass and a merge.
#[allow(clippy::too_many_arguments)]
fn encode_sa_split_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    q: &wgpu::Buffer,
    kv: &wgpu::Buffer,
    ixb: &wgpu::Buffer,
    sink: &wgpu::Buffer,
    out: &wgpu::Buffer,
    nh: usize,
    hd: usize,
    m: usize,
    scale: f32,
    kv_id: u64,
    li: usize,
) {
    let nc = sa_chunks(m);
    let ck = m.div_ceil(nc).max(1);
    let acc = frame_buf(c, 112, nh * SA_MAX_CHUNKS * hd * 4, false);
    let mxb = frame_buf(c, 113, nh * SA_MAX_CHUNKS * 4, false);
    let lnb = frame_buf(c, 114, nh * SA_MAX_CHUNKS * 4, false);
    let pp = uni_slot8(
        c,
        176,
        kv_id,
        li,
        [
            nh as u32,
            hd as u32,
            m as u32,
            nc as u32,
            ck as u32,
            0,
            0,
            scale.to_bits(),
        ],
    );
    let bind = cached_bind(c, (180, kv_id, li), || {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.sa_part.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, q),
                bind_buf(1, kv),
                bind_buf(2, ixb),
                bind_buf(3, &acc),
                bind_buf(4, &mxb),
                bind_buf(5, &lnb),
                bind_buf(6, &pp),
            ],
        })
    });
    pass.set_pipeline(&c.sa_part);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(nh as u32, nc as u32, 1);

    let mp = uni_slot(c, 182, kv_id, li, [nh as u32, hd as u32, nc as u32, 0]);
    let bind2 = cached_bind(c, (184, kv_id, li), || {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.sa_merge.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &acc),
                bind_buf(1, &mxb),
                bind_buf(2, &lnb),
                bind_buf(3, sink),
                bind_buf(4, out),
                bind_buf(5, &mp),
            ],
        })
    });
    pass.set_pipeline(&c.sa_merge);
    pass.set_bind_group(0, &bind2, &[]);
    pass.dispatch_workgroups(nh as u32, 1, 1);
}

/// The fused hyper-connection join: expand, mix, fold, norm — one dispatch.
#[allow(clippy::too_many_arguments)]
fn encode_hc_block_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    x: &wgpu::Buffer,
    res: &wgpu::Buffer,
    post: &wgpu::Buffer,
    comb: &wgpu::Buffer,
    mixw: &wgpu::Buffer,
    sc: &wgpu::Buffer,
    base: &wgpu::Buffer,
    nw: &wgpu::Buffer,
    state: &wgpu::Buffer,
    fold: &wgpu::Buffer,
    norm: &wgpu::Buffer,
    hc: usize,
    dim: usize,
    mix_hc: usize,
    iters: usize,
    eps: f32,
    bkey: (u8, u64, usize),
) {
    let bind = cached_bind(c, bkey, || {
        let p = uniform_u32x8(
            c,
            [
                hc as u32,
                dim as u32,
                iters as u32,
                eps.to_bits(),
                mix_hc as u32,
                0,
                0,
                0,
            ],
        );
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.hc_block.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, x),
                bind_buf(1, res),
                bind_buf(2, post),
                bind_buf(3, comb),
                bind_buf(4, mixw),
                bind_buf(5, sc),
                bind_buf(6, base),
                bind_buf(7, nw),
                bind_buf(8, state),
                bind_buf(9, fold),
                bind_buf(10, norm),
                bind_buf(11, &p),
            ],
        })
    });
    pass.set_pipeline(&c.hc_block);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(1, 1, 1);
}

/// A device-to-device copy of `n` floats, as a dispatch — so it can sit
/// inside a compute pass instead of ending one.
#[allow(clippy::too_many_arguments)]
fn encode_blit_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    src: &wgpu::Buffer,
    dst: &wgpu::Buffer,
    n: usize,
    soff: usize,
    doff: usize,
    // Some: cache the group. The offsets move with the sequence, so they ride
    // a per-layer uniform slot that is rewritten rather than a content-keyed
    // buffer.
    bkey: Option<(u8, u64, usize)>,
) {
    let p = match bkey {
        Some((tag, kv, li)) => uni_slot(c, tag, kv, li, [n as u32, soff as u32, doff as u32, 0]),
        None => uniform_u32x4(c, [n as u32, soff as u32, doff as u32, 0]),
    };
    let mk = || {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.blit.get_bind_group_layout(0),
            entries: &[bind_buf(0, src), bind_buf(1, dst), bind_buf(2, &p)],
        })
    };
    let bind = match bkey {
        Some((tag, kv, li)) => cached_bind(c, (tag.wrapping_add(1), kv, li), mk),
        None => mk(),
    };
    pass.set_pipeline(&c.blit);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((n as u32).div_ceil(256), 1, 1);
}

fn sa_bind_single(
    c: &Ctx,
    q: &wgpu::Buffer,
    kv: &wgpu::Buffer,
    ixb: &wgpu::Buffer,
    sink: &wgpu::Buffer,
    out: &wgpu::Buffer,
    nh: usize,
    hd: usize,
    m: usize,
    scale: f32,
    bkey: Option<(u64, usize)>,
) -> wgpu::BindGroup {
    let p = match bkey {
        Some((kv_id, li)) => uni_slot(
            c,
            140,
            kv_id,
            li,
            [nh as u32, hd as u32, m as u32, scale.to_bits()],
        ),
        None => uniform_mixed(c, [nh as u32, hd as u32, m as u32], scale),
    };
    let mk = || {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.sparse_attend.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, q),
                bind_buf(1, kv),
                bind_buf(2, ixb),
                bind_buf(3, sink),
                bind_buf(4, out),
                bind_buf(5, &p),
            ],
        })
    };
    match bkey {
        Some((kv_id, li)) => cached_bind(c, (141, kv_id, li), mk),
        None => mk(),
    }
}

fn encode_sparse_attend2(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    q: &wgpu::Buffer,
    kv: &wgpu::Buffer,
    ixb: &wgpu::Buffer,
    sink: &wgpu::Buffer,
    out: &wgpu::Buffer,
    nh: usize,
    hd: usize,
    m: usize,
    scale: f32,
    // Some: cache the bind group under this (sequence, layer). The attended
    // count walks with the sequence, so the uniform is a per-layer slot that
    // is rewritten rather than a content-keyed buffer — same trade as the
    // prep encoders. None: build fresh (the standalone frames, whose buffers
    // are not the chain's).
    bkey: Option<(u64, usize)>,
) {
    // ONE workgroup per head after all. The split into scores + apply spread
    // the work across the card and bought 0.54 -> 0.49 ms a layer — nothing —
    // while moving the model's perplexity by 0.7% through a different
    // accumulation order. Faster would have justified that; a wash does not.
    // CMF_SA_SPLIT=1 runs the split pair for anyone who wants to retry it on
    // a part where occupancy actually bites.
    if !sa_split() {
        let bind = sa_bind_single(c, q, kv, ixb, sink, out, nh, hd, m, scale, bkey);
        let mut pass = begin_pass_with(enc, None, ts_pair(c, 0));
        pass.set_pipeline(&c.sparse_attend);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(nh as u32, 1, 1);
        return;
    }
    let wbuf = frame_buf(c, 9, nh * m.max(1) * 4, false);
    let p = uniform_mixed(c, [nh as u32, hd as u32, m as u32], scale);
    {
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.sa_scores.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, q),
                bind_buf(1, kv),
                bind_buf(2, ixb),
                bind_buf(3, sink),
                bind_buf(4, &wbuf),
                bind_buf(5, &p),
            ],
        });
        let mut pass = begin_pass_with(enc, None, ts_pair(c, 1));
        pass.set_pipeline(&c.sa_scores);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(nh as u32, 1, 1);
    }
    {
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.sa_apply.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, &wbuf),
                bind_buf(1, kv),
                bind_buf(2, ixb),
                bind_buf(3, out),
                bind_buf(4, &p),
            ],
        });
        let mut pass = begin_pass_with(enc, None, ts_pair(c, 2));
        pass.set_pipeline(&c.sa_apply);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(((nh * hd) as u32).div_ceil(256), 1, 1);
    }
}

/// `y += w·d`, both device-side, `n` floats.
fn encode_axpy(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    d: &wgpu::Buffer,
    y: &wgpu::Buffer,
    w: f32,
    n: usize,
    bkey: (u8, u64, usize),
) {
    let mut pass = begin_pass(enc);
    encode_axpy_p(&mut pass, c, d, y, w, n, bkey);
}

#[allow(clippy::too_many_arguments)]
fn encode_axpy_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    d: &wgpu::Buffer,
    y: &wgpu::Buffer,
    w: f32,
    n: usize,
    bkey: (u8, u64, usize),
) {
    encode_axpy_full_p(pass, c, d, y, w, n, false, 0, Some(bkey));
}

/// `y = w·x[soff..]` when `set`, else `y += w·x[soff..]`.
///
/// The uniform is written in the order the SHADER declares — `w` first, then
/// `n`. It used to be written `[n, 0, 0, w]` against a `{ w: f32, n: u32 }`
/// struct, so the kernel read `n` as zero and every invocation returned at
/// the bounds check: the whole op was a no-op wherever the device ran it.
#[allow(clippy::too_many_arguments)]
fn encode_axpy_full_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    d: &wgpu::Buffer,
    y: &wgpu::Buffer,
    w: f32,
    n: usize,
    set: bool,
    soff: usize,
    bkey: Option<(u8, u64, usize)>,
) {
    let vals = [w.to_bits(), n as u32, set as u32, soff as u32];
    let mk = || {
        let p = uniform_u32x4(c, vals);
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.axpy.get_bind_group_layout(0),
            entries: &[bind_buf(0, d), bind_buf(1, y), bind_buf(2, &p)],
        })
    };
    let bind = match bkey {
        // `soff` walks with the sequence on the compressor's bias add, so the
        // cached form takes a per-layer slot it can rewrite.
        Some((tag, kv, li)) => {
            let p = uni_slot(c, tag, kv, li, vals);
            cached_bind(c, (tag, kv, li), || {
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &c.axpy.get_bind_group_layout(0),
                    entries: &[bind_buf(0, d), bind_buf(1, y), bind_buf(2, &p)],
                })
            })
        }
        None => mk(),
    };
    pass.set_pipeline(&c.axpy);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((n as u32).div_ceil(256), 1, 1);
}

/// The compressor's fold: a softmax over the slot axis PER DIMENSION, then
/// the weighted sum. `width` here is the OUTPUT width — half the projection
/// when the windows overlap, which is also the stride the kernel assumes.
#[allow(clippy::too_many_arguments)]
fn encode_kv_pool(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    prev_kv: &wgpu::Buffer,
    prev_sc: &wgpu::Buffer,
    cur_kv: &wgpu::Buffer,
    cur_sc: &wgpu::Buffer,
    ape: &wgpu::Buffer,
    out: &wgpu::Buffer,
    g: Dsv4CompGeom,
    have_prev: bool,
    kind: u8,
    kvid: u64,
    li: usize,
) {
    let mut pass = begin_pass(enc);
    encode_kv_pool_p(
        &mut pass, c, prev_kv, prev_sc, cur_kv, cur_sc, ape, out, g, have_prev, kind, kvid, li,
    );
}

#[allow(clippy::too_many_arguments)]
fn encode_kv_pool_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    prev_kv: &wgpu::Buffer,
    prev_sc: &wgpu::Buffer,
    cur_kv: &wgpu::Buffer,
    cur_sc: &wgpu::Buffer,
    ape: &wgpu::Buffer,
    out: &wgpu::Buffer,
    g: Dsv4CompGeom,
    have_prev: bool,
    kind: u8,
    kvid: u64,
    li: usize,
) {
    let ew = if g.overlap { g.width / 2 } else { g.width };
    let slots = if g.overlap { 2 * g.ratio } else { g.ratio };
    // The bias is folded in on arrival when the windows overlap, so the
    // kernel must not add it a second time — it does so only for the flat
    // compressor, which has nowhere else to put it.
    let use_ape = !g.overlap;
    // `have_prev` flips once early in the sequence, so the uniform is a
    // rewritten slot rather than a content-keyed buffer.
    let flags = (g.overlap as u32) | ((have_prev as u32) << 1) | ((use_ape as u32) << 2);
    let p = uni_slot(
        c,
        130 + kind,
        kvid,
        li,
        [slots as u32, ew as u32, g.ratio as u32, flags],
    );
    let bind = cached_bind(c, (132 + kind, kvid, li), || {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.kv_pool.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, prev_kv),
                bind_buf(1, prev_sc),
                bind_buf(2, cur_kv),
                bind_buf(3, cur_sc),
                bind_buf(4, ape),
                bind_buf(5, out),
                bind_buf(6, &p),
            ],
        })
    });
    pass.set_pipeline(&c.kv_pool);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((ew as u32).div_ceil(256), 1, 1);
}

/// What one compressor needs from the model, by directory index (the two
/// projections are q4tp) and by value (the small f32 pieces).
#[derive(Clone)]
pub struct Dsv4CompW<'a> {
    pub wkv: usize,
    pub wgate: usize,
    pub norm: &'a [f32],
    /// `[ratio, width]` — the in-window position bias.
    pub ape: &'a [f32],
}

#[derive(Clone, Copy)]
pub struct Dsv4CompGeom {
    /// `wkv.rows()`; the folded entry is half this when the windows overlap.
    pub width: usize,
    pub hidden: usize,
    pub ratio: usize,
    pub overlap: bool,
    pub rope_dim: usize,
    pub eps: f32,
}

/// Advance ONE compressor by one token, entirely on the device, and fold the
/// window when it closes.
///
/// The host's only inputs are the position and the hidden state that is
/// already there; the pending and previous streams live in `dsv4_comp` and
/// are never read back. That is the whole point: the compressor was the
/// reason every layer had to come back to the CPU mid-token.
///
/// Returns `Some(offset)` — in floats, into the layer's cache buffer — when
/// this token closed a window and an entry was appended there, `Some` with
/// no write is impossible, and `None` when the window is still filling or
/// the frame declined. `n_comp` is how many entries the cache already holds.
#[allow(clippy::too_many_arguments)]
pub fn dsv4_compressor_frame(
    model: &Arc<CmfModel>,
    w: &Dsv4CompW,
    g: Dsv4CompGeom,
    kind: u8,
    kv_id: u64,
    li: usize,
    hidden: &wgpu::Buffer,
    pos: usize,
    inv_freq: &[f32],
    // Where a folded entry goes: the layer's cache and the float offset of
    // the first free compressed slot.
    dst: &wgpu::Buffer,
    dst_off: usize,
    enc: &mut wgpu::CommandEncoder,
) -> Option<usize> {
    let c = ctx()?;
    let _ew = if g.overlap { g.width / 2 } else { g.width };
    if g.width == 0 || g.ratio == 0 || g.hidden % 32 != 0 {
        return None;
    }
    let bytes = model.primary_bytes();
    let mut wb = Vec::with_capacity(2);
    for &idx in &[w.wkv, w.wgate] {
        let e = model.tensors.get(idx)?;
        if e.dtype != cortiq_core::TensorDtype::Q4TiledP || e.shape.len() != 2 {
            return None;
        }
        let abs = model.entry_abs_offset(e)?;
        let plen = e.nbytes as usize;
        bytes.get(abs..abs + plen)?;
        wb.push(weight_buffer(
            c,
            (model.uid() as usize, idx),
            &bytes[abs..abs + plen],
        )?);
    }

    let ckv = frame_buf(c, 70 + kind, g.width * 4, false);
    let csc = frame_buf(c, 72 + kind, g.width * 4, false);
    // The two projections and everything the state step does are one chain
    // of dependent dispatches: one pass for the whole compressor.
    comp_state_step(
        c,
        w,
        g,
        kind,
        kv_id,
        li,
        &ckv,
        &csc,
        pos,
        inv_freq,
        dst,
        dst_off,
        enc,
        Some((&wb[0], &wb[1], hidden)),
    )
}

/// Everything after the two projections: the slot bookkeeping, the fold when
/// the window closes, and the shuffle of pending into previous. Split out
/// because this — not the matvecs, which have their own parity tests — is
/// what is new here, and it can be driven from a test with the projections
/// handed in.
#[allow(clippy::too_many_arguments)]
fn comp_state_step(
    c: &Ctx,
    w: &Dsv4CompW,
    g: Dsv4CompGeom,
    kind: u8,
    kv_id: u64,
    li: usize,
    ckv: &wgpu::Buffer,
    csc: &wgpu::Buffer,
    pos: usize,
    inv_freq: &[f32],
    dst: &wgpu::Buffer,
    dst_off: usize,
    enc: &mut wgpu::CommandEncoder,
    // Some: the two projections that produce `ckv`/`csc`, encoded into this
    // step's own pass instead of two of their own. None: the caller already
    // filled them (the parity test hands them in).
    proj: Option<(&wgpu::Buffer, &wgpu::Buffer, &wgpu::Buffer)>,
) -> Option<usize> {
    let ew = if g.overlap { g.width / 2 } else { g.width };
    let span = g.ratio * g.width;
    let st = {
        let mut m = c.dsv4_comp.lock().unwrap();
        m.entry((kind, kv_id, li))
            .or_insert_with(|| {
                let mk = || {
                    c.device.create_buffer(&wgpu::BufferDescriptor {
                        label: Some("dsv4-comp-stream"),
                        size: (span * 4).max(4) as u64,
                        usage: wgpu::BufferUsages::STORAGE
                            | wgpu::BufferUsages::COPY_SRC
                            | wgpu::BufferUsages::COPY_DST,
                        mapped_at_creation: false,
                    })
                };
                [mk(), mk(), mk(), mk()]
            })
            .clone()
    };
    let (pend_kv, pend_sc, prev_kv, prev_sc) = (&st[0], &st[1], &st[2], &st[3]);
    // This token's slot in the window, and whether it closes it. Both are a
    // function of the position alone, so no counter has to live on the card.
    let slot = pos % g.ratio;
    let folds = slot + 1 == g.ratio;
    // A previous window exists once one has closed.
    let have_prev = pos + 1 > g.ratio;

    // One pass for the whole compressor step: the copies are dispatches now
    // (`blit`), which is what let them stop cutting the layer into pieces.
    let ape_all = const_buf(c, bytemuck::cast_slice(w.ape));
    let _ape_slot = frame_buf(c, 74 + kind, g.width * 4, true);
    let folded = frame_buf(c, 76 + kind, ew * 4, false);
    let normed = frame_buf(c, 78 + kind, ew * 4, false);
    let nw = const_buf(c, bytemuck::cast_slice(&w.norm[..ew]));
    {
        let mut pass = begin_pass(enc);
        if let Some((wkv, wgate, hidden)) = proj {
            encode_q4tp_mvw_p(
                &mut pass,
                c,
                wkv,
                hidden,
                ckv,
                g.width,
                g.hidden,
                (80 + kind, kv_id, li),
            );
            encode_q4tp_mvw_p(
                &mut pass,
                c,
                wgate,
                hidden,
                csc,
                g.width,
                g.hidden,
                (82 + kind, kv_id, li),
            );
        }
        if g.overlap {
            // The reference biases the score as the token ARRIVES and keeps
            // it biased across the shift, so `ape` is added once, here, and
            // the pooling kernel is told not to add it again.
            // The strided source reads this token's slot of `ape` directly:
            // the copy into a scratch buffer existed only because axpy could
            // not offset its input.
            encode_axpy_full_p(
                &mut pass,
                c,
                &ape_all,
                csc,
                1.0,
                g.width,
                false,
                slot * g.width,
                Some((84 + kind, kv_id, li)),
            );
        }
        encode_blit_p(
            &mut pass,
            c,
            ckv,
            pend_kv,
            g.width,
            0,
            slot * g.width,
            Some((146 + kind * 2, kv_id, li)),
        );
        encode_blit_p(
            &mut pass,
            c,
            csc,
            pend_sc,
            g.width,
            0,
            slot * g.width,
            Some((150 + kind * 2, kv_id, li)),
        );
        if !folds {
            return None;
        }
        encode_kv_pool_p(
            &mut pass, c, prev_kv, prev_sc, pend_kv, pend_sc, &ape_all, &folded, g, have_prev,
            kind, kv_id, li,
        );
        encode_rmsnorm_p(
            &mut pass,
            c,
            &folded,
            &nw,
            &normed,
            ew,
            g.eps,
            (86 + kind, kv_id, li),
        );
        // The entry carries a window key's rope tail, at the position of the
        // window's FIRST token — not this one.
        let freq = const_buf(c, bytemuck::cast_slice(&inv_freq[..g.rope_dim / 2]));
        let posb = frame_up_pos(c, 79 + kind, pos + 1 - g.ratio, g.eps);
        encode_rope_heads_p(
            &mut pass,
            c,
            &normed,
            &freq,
            &posb,
            1,
            ew,
            g.rope_dim,
            false,
            false,
            (88 + kind, kv_id, li),
        );
        encode_blit_p(
            &mut pass,
            c,
            &normed,
            dst,
            ew,
            0,
            dst_off,
            Some((154 + kind * 2, kv_id, li)),
        );
        if g.overlap {
            // The window that just closed becomes the previous one — the fold
            // reads half its dimensions from that stride.
            encode_blit_p(
                &mut pass,
                c,
                pend_kv,
                prev_kv,
                span,
                0,
                0,
                Some((158 + kind * 2, kv_id, li)),
            );
            encode_blit_p(
                &mut pass,
                c,
                pend_sc,
                prev_sc,
                span,
                0,
                0,
                Some((162 + kind * 2, kv_id, li)),
            );
        }
    }
    Some(dst_off)
}

/// Append this token's key to the sliding window, on the device: the KV
/// projection, its norm, its rope tail, and the shift that keeps the window
/// at capacity.
///
/// The last producer the host still owned. The reference keeps a ring; the
/// engine keeps the last N in order, which is the same SET but makes the
/// position list plain `0..win_len` — so the shift has to be a real shift.
/// At 128×512 floats that is a quarter of a megabyte of device-local copy a
/// layer, which is microseconds, and it buys the host's exit from the loop.
#[allow(clippy::too_many_arguments)]
pub fn dsv4_window_append(
    model: &Arc<CmfModel>,
    wkv: usize,
    kv_norm: &[f32],
    hidden: &wgpu::Buffer,
    cache: &wgpu::Buffer,
    // head_dim, the window's capacity in slots, how many are filled BEFORE
    // this token, the model's hidden width and the rope tail.
    hd: usize,
    window: usize,
    filled: usize,
    dim: usize,
    rope_dim: usize,
    eps: f32,
    pos: usize,
    inv_freq: &[f32],
    kv_id: u64,
    li: usize,
    enc: &mut wgpu::CommandEncoder,
) -> Option<usize> {
    let c = ctx()?;
    let bytes = model.primary_bytes();
    let e = model.tensors.get(wkv)?;
    if e.dtype != cortiq_core::TensorDtype::Q4TiledP {
        return None;
    }
    let abs = model.entry_abs_offset(e)?;
    let plen = e.nbytes as usize;
    bytes.get(abs..abs + plen)?;
    let wb = weight_buffer(c, (model.uid() as usize, wkv), &bytes[abs..abs + plen])?;

    let raw = frame_buf(c, 104, hd * 4, false);
    let _kv = frame_buf(c, 105, hd * 4, false);
    window_place(
        c,
        enc,
        &raw,
        kv_norm,
        hd,
        window,
        filled,
        rope_dim,
        eps,
        pos,
        inv_freq,
        cache,
        kv_id,
        li,
        Some((&wb, hidden, dim)),
    );
    Some((filled + 1).min(window))
}

/// The half of the window append that is NOT a matvec: the norm, the rope
/// tail and the slide that keeps the window at capacity. Split out so a test
/// can drive it with the projection handed in — the q4tp matvec has its own
/// parity test, these three did not have one at all.
#[allow(clippy::too_many_arguments)]
fn window_place(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    raw: &wgpu::Buffer,
    kv_norm: &[f32],
    hd: usize,
    window: usize,
    filled: usize,
    rope_dim: usize,
    eps: f32,
    pos: usize,
    inv_freq: &[f32],
    cache: &wgpu::Buffer,
    kv_id: u64,
    li: usize,
    // Some: the KV projection that fills `raw`, encoded into this step's own
    // pass. None: the caller filled it (the parity test hands it in).
    proj: Option<(&wgpu::Buffer, &wgpu::Buffer, usize)>,
) {
    let kv = frame_buf(c, 105, hd * 4, false);
    let nw = const_buf(c, bytemuck::cast_slice(&kv_norm[..hd]));
    let freq = const_buf(c, bytemuck::cast_slice(&inv_freq[..rope_dim / 2]));
    let posb = frame_up_pos(c, 108, pos, eps);
    let tmp = frame_buf(c, 110, (window.max(1) - 1).max(1) * hd * 4, true);
    let mut pass = begin_pass(enc);
    if let Some((wb, hidden, dim)) = proj {
        encode_q4tp_mvw_p(&mut pass, c, wb, hidden, raw, hd, dim, (106, kv_id, li));
    }
    encode_rmsnorm_p(&mut pass, c, raw, &nw, &kv, hd, eps, (107, kv_id, li));
    encode_rope_heads_p(
        &mut pass,
        c,
        &kv,
        &freq,
        &posb,
        1,
        hd,
        rope_dim,
        false,
        false,
        (109, kv_id, li),
    );

    // Where it lands, and whether the window has to slide first.
    let slot = if filled < window {
        filled
    } else {
        // Drop the oldest: everything moves down one slot. A copy whose
        // source and destination overlap is not allowed, so it goes through
        // a scratch buffer — still device-local, still no host.
        let n = (window - 1) * hd;
        encode_blit_p(&mut pass, c, cache, &tmp, n, hd, 0, Some((166, kv_id, li)));
        encode_blit_p(&mut pass, c, &tmp, cache, n, 0, 0, Some((168, kv_id, li)));
        window - 1
    };
    encode_blit_p(
        &mut pass,
        c,
        &kv,
        cache,
        hd,
        0,
        slot * hd,
        Some((170, kv_id, li)),
    );
}

/// Drive the norm, the rope tail and the slide from a test with the
/// projection handed in, and give back the whole window as attention would
/// read it.
#[allow(clippy::too_many_arguments)]
pub fn dsv4_window_place_for_test(
    raw: &[f32],
    kv_norm: &[f32],
    inv_freq: &[f32],
    seed: &[f32],
    hd: usize,
    window: usize,
    filled: usize,
    rope_dim: usize,
    eps: f32,
    pos: usize,
    out: &mut Vec<f32>,
) -> bool {
    let Some(c) = ctx() else { return false };
    // A POOLED buffer, not a fresh one: `encode_rmsnorm` caches its bind
    // group by key, so a new buffer every call would leave the cache
    // pointing at the first one. Production hands it `frame_buf` for the
    // same reason.
    let rb = frame_up(c, 112, bytemuck::cast_slice(&raw[..hd]));
    // Pooled for the SAME reason as `rb` above, and the reason is not
    // decorative: the final blit caches its bind group by key, so a fresh
    // buffer each call leaves that group pointing at the first one. The new
    // entry then lands in a buffer nobody reads and the slot stays zero —
    // which is what this test reported for two months as a kernel fault.
    let cache = frame_buf(c, 113, window * hd * 4, true);
    c.queue.write_buffer(&cache, 0, bytemuck::cast_slice(seed));
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
    window_place(
        c, &mut enc, &rb, kv_norm, hd, window, filled, rope_dim, eps, pos, inv_freq, &cache, 909,
        0, None,
    );
    out.clear();
    out.resize(window * hd, 0.0);
    let bytes = (window * hd * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "win-test-stage",
    );
    let ok = readback(c, enc, &cache, &stage, bytes, out);
    drop(sc);
    ok
}

/// What the indexer needs from the model: two q4tp projections by directory
/// index. Its own compressor goes through `dsv4_compressor_frame` as kind 1.
#[derive(Clone)]
pub struct Dsv4IxW {
    /// `[ih*idim, q_lora]` — reads the SHARED LoRA output, not attention's
    /// queries.
    pub wq_b: usize,
    /// `[ih, hidden]`.
    pub weights_proj: usize,
}

#[derive(Clone, Copy)]
pub struct Dsv4IxGeom {
    pub ih: usize,
    pub idim: usize,
    pub q_lora: usize,
    pub hidden: usize,
    pub rope_dim: usize,
    pub eps: f32,
    pub top_k: usize,
    /// The cache's window CAPACITY — where the compressed region starts.
    pub window: usize,
}

/// Score the compressed positions, take the top-k, and assemble the attended
/// list — all on the device, into `out_idx`.
///
/// Returns the list's length, which the host computes rather than reads back:
/// the window in use plus however many of the top-k there were positions for.
/// That is the whole trick — the CONTENTS are a device secret, the LENGTH
/// never was, so nothing has to come home mid-token.
#[allow(clippy::too_many_arguments)]
pub fn dsv4_indexer_frame(
    model: &Arc<CmfModel>,
    w: &Dsv4IxW,
    g: Dsv4IxGeom,
    kv_id: u64,
    li: usize,
    hidden: &wgpu::Buffer,
    qn: &wgpu::Buffer,
    index_kv: &wgpu::Buffer,
    // How many entries the indexer's own cache holds, and how many compressed
    // positions attention has: the reference scores the smaller of the two.
    n_ix: usize,
    n_comp: usize,
    win_len: usize,
    pos: usize,
    inv_freq: &[f32],
    out_idx: &wgpu::Buffer,
    enc: &mut wgpu::CommandEncoder,
) -> Option<usize> {
    let c = ctx()?;
    let limit = n_ix.min(n_comp);
    if g.ih == 0 || g.idim == 0 || limit == 0 || limit > 4096 {
        return None;
    }
    let bytes = model.primary_bytes();
    let mut wb = Vec::with_capacity(2);
    for &idx in &[w.wq_b, w.weights_proj] {
        let e = model.tensors.get(idx)?;
        if e.dtype != cortiq_core::TensorDtype::Q4TiledP || e.shape.len() != 2 {
            return None;
        }
        let abs = model.entry_abs_offset(e)?;
        let plen = e.nbytes as usize;
        bytes.get(abs..abs + plen)?;
        wb.push(weight_buffer(
            c,
            (model.uid() as usize, idx),
            &bytes[abs..abs + plen],
        )?);
    }

    let qi = frame_buf(c, 90, g.ih * g.idim * 4, false);
    let hw_raw = frame_buf(c, 91, g.ih * 4, false);
    let hw = frame_buf(c, 92, g.ih * 4, false);
    // Fixed capacities, not `limit`-sized: frame_buf keys on (tag, len), so
    // a length that walks with the sequence would change the buffer's
    // identity under the cached bind groups below. `limit` is capped at
    // 4096 on entry.
    let scores = frame_buf(c, 93, 4096 * 4, false);
    let pick = frame_buf(c, 94, g.top_k.max(1) * 4, false);
    let cnt = frame_buf(c, 95, 4, false);

    let freq = const_buf(c, bytemuck::cast_slice(&inv_freq[..g.rope_dim / 2]));
    let posb = frame_up_pos(c, 97, pos, g.eps);
    // The reference folds head_dim^-0.5 · n_heads^-0.5 into weights_proj's
    // output. A uniform positive factor cannot change which positions win,
    // but the scores are the kernel's contract, not just a ranking key.
    let sc_factor = (g.idim as f32).powf(-0.5) * (g.ih as f32).powf(-0.5);
    let k_actual = g.top_k.min(limit);
    // EIGHT dispatches, ONE pass. Each is a step of the one before it, and
    // dispatches inside a compute pass already run in order with the writes
    // of the previous one visible — so the eight passes this used to open
    // bought nothing but eight lots of driver bookkeeping, which on a small
    // layer is most of what the token costs.
    {
        let mut pass = begin_pass(enc);
        encode_q4tp_mvw_p(
            &mut pass,
            c,
            &wb[0],
            qn,
            &qi,
            g.ih * g.idim,
            g.q_lora,
            (96, kv_id, li),
        );
        encode_rope_heads_p(
            &mut pass,
            c,
            &qi,
            &freq,
            &posb,
            g.ih,
            g.idim,
            g.rope_dim,
            false,
            false,
            (98, kv_id, li),
        );
        encode_q4tp_mvw_p(
            &mut pass,
            c,
            &wb[1],
            hidden,
            &hw_raw,
            g.ih,
            g.hidden,
            (99, kv_id, li),
        );
        // `set`: hw = sc_factor·hw_raw. This was a whole-buffer zero-fill
        // followed by an accumulate — two dispatches to express an
        // assignment.
        encode_axpy_full_p(
            &mut pass,
            c,
            &hw_raw,
            &hw,
            sc_factor,
            g.ih,
            true,
            0,
            Some((101, kv_id, li)),
        );
        encode_index_scores_p(
            &mut pass,
            c,
            &qi,
            index_kv,
            &hw,
            &scores,
            g.ih,
            g.idim,
            limit,
            (kv_id, li),
        );
        encode_top_k_p(
            &mut pass,
            c,
            &scores,
            &pick,
            &cnt,
            limit,
            g.top_k,
            (kv_id, li),
        );
        encode_idx_build_p(
            &mut pass,
            c,
            &pick,
            out_idx,
            win_len,
            g.window,
            k_actual,
            Some((kv_id, li)),
        );
    }
    Some(win_len + k_actual)
}

fn encode_fill_zero(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    y: &wgpu::Buffer,
    n: usize,
    bkey: (u8, u64, usize),
) {
    let mut pass = begin_pass(enc);
    encode_fill_zero_p(&mut pass, c, y, n, bkey);
}

#[allow(clippy::too_many_arguments)]
fn encode_fill_zero_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    y: &wgpu::Buffer,
    n: usize,
    bkey: (u8, u64, usize),
) {
    let bind = cached_bind(c, bkey, || {
        let p = uniform_u32x4(c, [n as u32, 0, 0, 0]);
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.zero.get_bind_group_layout(0),
            entries: &[bind_buf(0, y), bind_buf(1, &p)],
        })
    });
    pass.set_pipeline(&c.zero);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((n as u32).div_ceil(256), 1, 1);
}

#[allow(clippy::too_many_arguments)]
fn encode_index_scores(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    q: &wgpu::Buffer,
    kv: &wgpu::Buffer,
    hw: &wgpu::Buffer,
    out: &wgpu::Buffer,
    nh: usize,
    hd: usize,
    n_pos: usize,
    bkey: (u64, usize),
) {
    let mut pass = begin_pass(enc);
    encode_index_scores_p(&mut pass, c, q, kv, hw, out, nh, hd, n_pos, bkey);
}

#[allow(clippy::too_many_arguments)]
fn encode_index_scores_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    q: &wgpu::Buffer,
    kv: &wgpu::Buffer,
    hw: &wgpu::Buffer,
    out: &wgpu::Buffer,
    nh: usize,
    hd: usize,
    n_pos: usize,
    bkey: (u64, usize),
) {
    // n_pos walks with the sequence, so the uniform lives in a per-layer
    // slot whose contents are rewritten each call; the bind group can then
    // survive across tokens.
    let p = uni_slot(
        c,
        134,
        bkey.0,
        bkey.1,
        [nh as u32, hd as u32, n_pos as u32, n_pos as u32],
    );
    let bind = cached_bind(c, (135, bkey.0, bkey.1), || {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.index_scores.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, q),
                bind_buf(1, kv),
                bind_buf(2, hw),
                bind_buf(3, out),
                bind_buf(4, &p),
            ],
        })
    });
    pass.set_pipeline(&c.index_scores);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((n_pos as u32).min(MAX_WG), 1, 1);
}

fn encode_top_k(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    scores: &wgpu::Buffer,
    pick: &wgpu::Buffer,
    cnt: &wgpu::Buffer,
    n: usize,
    k: usize,
    bkey: (u64, usize),
) {
    let mut pass = begin_pass(enc);
    encode_top_k_p(&mut pass, c, scores, pick, cnt, n, k, bkey);
}

#[allow(clippy::too_many_arguments)]
fn encode_top_k_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    scores: &wgpu::Buffer,
    pick: &wgpu::Buffer,
    cnt: &wgpu::Buffer,
    n: usize,
    k: usize,
    bkey: (u64, usize),
) {
    let p = uni_slot(c, 136, bkey.0, bkey.1, [n as u32, k as u32, 0, 0]);
    let bind = cached_bind(c, (137, bkey.0, bkey.1), || {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.top_k_index.get_bind_group_layout(0),
            entries: &[
                bind_buf(0, scores),
                bind_buf(1, pick),
                bind_buf(2, cnt),
                bind_buf(3, &p),
            ],
        })
    });
    pass.set_pipeline(&c.top_k_index);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups(1, 1, 1);
}

/// NOT cached, and this one bites: `win_len` grows with the sequence and `k`
/// changes with it, so the uniform — which `uniform_u32x4` keys on its
/// CONTENTS — becomes a different buffer while a cached bind group would
/// still point at the previous token's. The first call would be right and
/// every one after it would read a stale window length.
fn encode_idx_build(
    c: &Ctx,
    enc: &mut wgpu::CommandEncoder,
    pick: &wgpu::Buffer,
    out: &wgpu::Buffer,
    win_len: usize,
    window: usize,
    k: usize,
    // The indexer path binds stable buffers and can keep its group; the
    // prep fallback's pick buffer changes identity with its length, so it
    // passes None and builds a fresh group each call.
    bkey: Option<(u64, usize)>,
) {
    let mut pass = begin_pass(enc);
    encode_idx_build_p(&mut pass, c, pick, out, win_len, window, k, bkey);
}

#[allow(clippy::too_many_arguments)]
fn encode_idx_build_p(
    pass: &mut wgpu::ComputePass<'_>,
    c: &Ctx,
    pick: &wgpu::Buffer,
    out: &wgpu::Buffer,
    win_len: usize,
    window: usize,
    k: usize,
    // The indexer path binds stable buffers and can keep its group; the
    // prep fallback's pick buffer changes identity with its length, so it
    // passes None and builds a fresh group each call.
    bkey: Option<(u64, usize)>,
) {
    let n = win_len + k;
    let mk = |p: &wgpu::Buffer| {
        c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &c.idx_build.get_bind_group_layout(0),
            entries: &[bind_buf(0, pick), bind_buf(1, out), bind_buf(2, p)],
        })
    };
    let bind = match bkey {
        Some((kvid, li)) => {
            let p = uni_slot(
                c,
                138,
                kvid,
                li,
                [win_len as u32, window as u32, k as u32, 0],
            );
            cached_bind(c, (139, kvid, li), || mk(&p))
        }
        None => mk(&uniform_u32x4(
            c,
            [win_len as u32, window as u32, k as u32, 0],
        )),
    };
    pass.set_pipeline(&c.idx_build);
    pass.set_bind_group(0, &bind, &[]);
    pass.dispatch_workgroups((n as u32).div_ceil(256), 1, 1);
}

/// Drive the device-side compressor state with the projections handed in,
/// one token per call, and return the folded entry on the token that closes
/// a window. The frame does exactly this after its two matvecs; giving the
/// projections from the host is what lets a test compare the STATE — the
/// slots, the fold timing, the previous-window shuffle, the rope position —
/// against `dsv4::compressor_step` without a model file.
#[allow(clippy::too_many_arguments)]
pub fn dsv4_comp_state_for_test(
    ckv: &[f32],
    csc: &[f32],
    norm: &[f32],
    ape: &[f32],
    inv_freq: &[f32],
    width: usize,
    ratio: usize,
    overlap: bool,
    rope_dim: usize,
    eps: f32,
    pos: usize,
    kv_id: u64,
    out: &mut Vec<f32>,
) -> Option<bool> {
    let c = ctx()?;
    let g = Dsv4CompGeom {
        width,
        hidden: 0,
        ratio,
        overlap,
        rope_dim,
        eps,
    };
    let ew = if overlap { width / 2 } else { width };
    let w = Dsv4CompW {
        wkv: 0,
        wgate: 0,
        norm,
        ape,
    };
    // POOLED, all three. The steps inside cache their bind groups by key, so
    // a buffer freshly created per call leaves those groups bound to the
    // first call's memory: writes land where nothing reads them and the
    // result reads as a broken kernel. Production hands these in from the
    // frame pool for exactly this reason.
    let kb = frame_up(c, 114, bytemuck::cast_slice(&ckv[..width]));
    let sb = frame_up(c, 115, bytemuck::cast_slice(&csc[..width]));
    let dst = frame_buf(c, 116, ew * 4, true);
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
    let folded = comp_state_step(
        c, &w, g, 9, kv_id, 0, &kb, &sb, pos, inv_freq, &dst, 0, &mut enc, None,
    )
    .is_some();
    if !folded {
        submit(c, finish_enc(enc));
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        return Some(false);
    }
    out.clear();
    out.resize(ew, 0.0);
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        (ew * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "comp-test-stage",
    );
    let ok = readback(c, enc, &dst, &stage, (ew * 4) as u64, out);
    drop(sc);
    ok.then_some(true)
}

/// The position-list assembly alone, driven from a test: the picks come in
/// as they would from top-k, and what comes back is what `sparse_attend`
/// would read. The shift by the window's CAPACITY (not its fill) is the part
/// worth pinning — get it wrong and attention reads the wrong keys while
/// every shape still checks out.
pub fn dsv4_idx_build_for_test(
    pick: &[u32],
    win_len: usize,
    window: usize,
    out: &mut Vec<u32>,
) -> bool {
    let Some(c) = ctx() else { return false };
    let n = win_len + pick.len();
    // A zero-length binding is a validation error, and "the indexer picked
    // nothing" is a real state — pad rather than refuse.
    let padded: Vec<u32> = if pick.is_empty() {
        vec![0]
    } else {
        pick.to_vec()
    };
    let pb = storage_bytes(c, bytemuck::cast_slice(&padded));
    let ob = c.device.create_buffer(&wgpu::BufferDescriptor {
        label: None,
        size: (n.max(1) * 4) as u64,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        mapped_at_creation: false,
    });
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
    encode_idx_build(c, &mut enc, &pb, &ob, win_len, window, pick.len(), None);
    let bytes = (n * 4) as u64;
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        bytes,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "idx-build-stage",
    );
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&ob, 0, &stage, 0, bytes);
    submit(c, finish_enc(enc));
    let slice = stage.slice(..bytes);
    slice.map_async(wgpu::MapMode::Read, |_| {});
    if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        return false;
    }
    let mut ok = false;
    if let Ok(data) = slice.get_mapped_range() {
        out.clear();
        out.extend_from_slice(bytemuck::cast_slice(&data[..bytes as usize]));
        ok = true;
    }
    stage.unmap();
    drop(sc);
    ok
}

/// Drop a sequence's compressor streams (called with the rest of its state).
pub fn dsv4_compressor_forget(kv_id: u64) {
    if let Some(c) = ctx() {
        c.dsv4_comp.lock().unwrap().retain(|k, _| k.1 != kv_id);
    }
}

/// The one thing that has to survive between tokens. `off` and `data` are in
/// floats; the buffer is created on first use at `cap` and never shrinks.
pub fn dsv4_cache_write(kv_id: u64, li: usize, off: usize, data: &[f32], cap: usize) -> bool {
    let dbg = std::env::var("CMF_DSV4_FRAME_DEBUG").is_ok();
    let Some(c) = ctx() else {
        if dbg {
            eprintln!("кеш dsv4: нет контекста wgpu");
        }
        return false;
    };
    if off + data.len() > cap {
        if dbg {
            eprintln!("кеш dsv4: {off}+{} не влезает в {cap}", data.len());
        }
        return false;
    }
    // Storage buffers have a size ceiling of their own, well under VRAM, and
    // silently refusing at it reads as "no device" from the caller's side.
    if (cap * 4) as u64 > c.device.limits().max_storage_buffer_binding_size {
        tracing::warn!(
            "кеш dsv4: {} МБ превышает предел одного буфера {} МБ — слой остаётся на CPU",
            cap * 4 / (1 << 20),
            c.device.limits().max_storage_buffer_binding_size / (1 << 20)
        );
        return false;
    }
    let mut map = c.dsv4_kv.lock().unwrap();
    // Grow rather than refuse: the compressed axis lengthens as the sequence
    // does, and a buffer sized for token 100 is not a reason to fall off the
    // device at token 1000. The caller rewrites both regions each token, so
    // losing the old contents costs nothing.
    if map.get(&(kv_id, li)).is_some_and(|(_, have)| *have < cap) {
        map.remove(&(kv_id, li));
        GREW.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }
    let e = map.entry((kv_id, li)).or_insert_with(|| {
        (
            c.device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("dsv4-kv"),
                size: (cap * 4) as u64,
                // COPY_SRC as well: the window slide reads this buffer to
                // write it one slot down, so the cache is its own source.
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_DST
                    | wgpu::BufferUsages::COPY_SRC,
                mapped_at_creation: false,
            }),
            cap,
        )
    });
    if e.1 < off + data.len() {
        return false; // a longer context than the cache was built for
    }
    if !data.is_empty() {
        c.queue
            .write_buffer(&e.0, (off * 4) as u64, bytemuck::cast_slice(data));
    }
    true
}

/// ── the u8 tag registry ────────────────────────────────────────────────
///
/// Four separate maps key on a `u8` tag, and only the tag tells two call
/// sites apart. A collision does not crash: it hands one site the other's
/// bind group, pointing at the wrong buffers, and the model quietly gets
/// worse. The numbers in use, so the next one can be picked without reading
/// the file:
///
/// * `frame_buf` (tag, len)   — 1, 2, 9, 17, 26, 27, 40, 43, 44, 70–79,
///                              90–95, 104, 105, 108, 110, 111
/// * `frame_up_pos` (tag)     — 1, 79–80, 97, 108
/// * `uni_slot` (tag, kv, li) — 84–85, 101, 130–131, 134, 136, 138, 140,
///                              142–164 (blit, even), 166, 168, 170
/// * `cached_bind` (tag,kv,li)— 30–33, 36–42, 50–64, 80–101, 120–129,
///                              132–141, 143–165 (blit, odd), 167, 169, 171
///
/// `CMF_DSV4_SLOT_CHECK=1` turns a collision between two `uni_slot` or
/// `store_slot` sites into a panic instead of a silent wrong answer: a slot
/// written twice before one submission is either a collision or the
/// last-write-wins bug that made a run of layers route with the last
/// layer's router bias.
///
/// Bind groups for the dsv4 frames, keyed by role and layer. Their buffers
/// are pooled and stable between tokens, so building 11 of them per layer per
/// token — 473 a token on the release — was pure host overhead. The epoch
/// invalidates the lot whenever a pooled buffer is rebuilt underneath them.
fn cached_bind<F>(c: &Ctx, key: (u8, u64, usize), build: F) -> wgpu::BindGroup
where
    F: FnOnce() -> wgpu::BindGroup,
{
    use std::sync::atomic::Ordering;
    let key = (key.0, key.1, dsv4_salted_li(key.2));
    let epoch = GREW.load(Ordering::Relaxed);
    let mut m = c.dsv4_binds.lock().unwrap();
    if m.0 != epoch {
        m.0 = epoch;
        m.1.clear();
    }
    if let Some(b) = m.1.get(&key) {
        return b.clone();
    }
    let b = build();
    m.1.insert(key, b.clone());
    b
}

/// A per-(tag, sequence, layer) UNIFORM buffer, 16 bytes, written every
/// call: how a sequence-varying scalar coexists with a CACHED bind group.
/// The buffer's identity never changes, only its contents — the opposite
/// trade from the content-keyed uniform pool, whose identity IS its
/// contents.
/// `uni_slot`'s eight-word twin, for the kernels whose params do not fit in
/// four. Same contract: one write per key per submission.
fn uni_slot8(c: &Ctx, tag: u8, kv: u64, li: usize, vals: [u32; 8]) -> wgpu::Buffer {
    let li = dsv4_salted_li(li);
    let b = {
        let mut m = c.dsv4_uni.lock().unwrap();
        m.entry((tag, kv, li))
            .or_insert_with(|| {
                c.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("dsv4-uni-slot8"),
                    size: 32,
                    usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                })
            })
            .clone()
    };
    note_slot_write(c, "uni8", (tag, kv, li));
    c.queue.write_buffer(&b, 0, bytemuck::cast_slice(&vals));
    b
}

fn uni_slot(c: &Ctx, tag: u8, kv: u64, li: usize, vals: [u32; 4]) -> wgpu::Buffer {
    let li = dsv4_salted_li(li);
    let b = {
        let mut m = c.dsv4_uni.lock().unwrap();
        m.entry((tag, kv, li))
            .or_insert_with(|| {
                c.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("dsv4-uni-slot"),
                    size: 16,
                    usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                })
            })
            .clone()
    };
    note_slot_write(c, "uni", (tag, kv, li));
    c.queue.write_buffer(&b, 0, bytemuck::cast_slice(&vals));
    b
}

/// `uni_slot`'s storage twin: a per-(tag, sequence, layer) STORAGE buffer
/// rewritten every call. What it buys is the hash layers — their forced
/// expert list changes per token and used to go through the (tag, len) pool,
/// where every layer of a run shared one buffer and the last write won. With
/// a buffer per layer they route with their own table row and can share a
/// submission with everyone else.
fn store_slot(c: &Ctx, tag: u8, kv: u64, li: usize, data: &[u8]) -> wgpu::Buffer {
    let li = dsv4_salted_li(li);
    let size = (data.len().max(4) as u64).div_ceil(16) * 16;
    let b = {
        let mut m = c.dsv4_store.lock().unwrap();
        let e = m.entry((tag, kv, li));
        match e {
            std::collections::hash_map::Entry::Occupied(o) if o.get().size() >= size => {
                o.get().clone()
            }
            other => {
                let b = c.device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("dsv4-store-slot"),
                    size,
                    // COPY_SRC: the per-layer cold/post slots are read back
                    // after a run — a slot the host can never read is a slot
                    // that can only ever feed full-pack chains.
                    usage: wgpu::BufferUsages::STORAGE
                        | wgpu::BufferUsages::COPY_DST
                        | wgpu::BufferUsages::COPY_SRC,
                    mapped_at_creation: false,
                });
                match other {
                    std::collections::hash_map::Entry::Occupied(mut o) => {
                        o.insert(b.clone());
                        // A cached bind group would otherwise outlive the
                        // buffer it names.
                        GREW.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                    }
                    std::collections::hash_map::Entry::Vacant(v) => {
                        v.insert(b.clone());
                    }
                }
                b
            }
        }
    };
    if !data.is_empty() {
        note_slot_write(c, "store", (tag, kv, li));
        c.queue.write_buffer(&b, 0, data);
    }
    b
}

/// Bumped whenever a cache buffer is reallocated. A caller that writes only
/// the tail has to notice, because the new buffer holds nothing.
pub static GREW: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// Drop a conversation's caches (a new sequence, or the pipeline resetting).
pub fn dsv4_cache_clear(kv_id: u64) {
    if let Some(c) = ctx() {
        c.dsv4_kv.lock().unwrap().retain(|k, _| k.0 != kv_id);
    }
}

/// The quantized tensors one DeepSeek-V4 attention block reads, by directory
/// index, plus the two small f32 vectors it needs whole.
#[derive(Clone)]
pub struct Dsv4AttnW<'a> {
    pub wq_a: usize,
    pub wq_b: usize,
    pub wo_a: usize,
    pub wo_b: usize,
    pub q_norm: &'a [f32],
    pub sink: &'a [f32],
}

/// Shapes for one block. Separate from the weights so the caller can build it
/// once per layer and keep it.
#[derive(Clone, Copy)]
pub struct Dsv4AttnGeom {
    pub dim: usize,
    pub nh: usize,
    pub hd: usize,
    pub rd: usize,
    pub q_lora: usize,
    pub o_lora: usize,
    pub o_groups: usize,
    pub eps: f32,
    pub scale: f32,
    /// V4.1 materialises BF16 after q projection, attended output, grouped
    /// projection, and the final output. Generic DSV4 keeps its historical
    /// f32 frame when this is false.
    pub bf16: bool,
    /// Whether the query heads receive DSV4's second RMSNorm after `wq_b`.
    /// V4.1's official attention path has only the LoRA-rank `q_norm`.
    pub q_rms: bool,
}

/// DeepSeek-V4's attention block, start to finish, in ONE submission.
///
/// Eight operations that were eight round trips: the query LoRA and its two
/// norms, the rope tail, attention over the index list, the inverse rope, and
/// the grouped output projection. Nothing between them touches the host — the
/// intermediate vectors never leave the card, and the KV cache is already
/// there.
///
/// The kv vector itself stays on the CPU deliberately: the compressor's
/// pending windows are host state, and pulling one 512-wide vector back is
/// cheaper than moving that state too. That is the next frame to fuse, not a
/// thing forgotten here.
#[allow(clippy::too_many_arguments)]
/// The hyper-connection work that sits between a layer's two halves, so the
/// frame can do it instead of the host.
///
/// Measured: these are 19 ms of a 57 ms token on the CPU and hundredths of a
/// millisecond on the card. Nothing about them needs the host — they were
/// only there because the frames handed their output back.
pub struct Dsv4HcTail<'a> {
    /// The FFN half's projection, scales and base.
    pub fn_: &'a [f32],
    pub scale: &'a [f32; 3],
    pub base: &'a [f32],
    /// The norm applied to the fold that feeds the MoE half.
    pub norm: &'a [f32],
    pub hc: usize,
    pub sinkhorn_iters: usize,
    pub hc_eps: f32,
    /// RMS epsilon for the norm that follows the fold.
    pub eps: f32,
}

#[allow(clippy::too_many_arguments)]
pub fn dsv4_attn_frame(
    model: &Arc<CmfModel>,
    w: &Dsv4AttnW,
    g: Dsv4AttnGeom,
    hidden: &[f32],
    // The layer's own `q_norm(wq_a(x))`, when the caller already has it — the
    // indexer needs that vector on the host anyway, and computing it twice is
    // worse than uploading 1536 floats. `None` puts both ops in the frame.
    qn_in: Option<&[f32]>,
    // The final V4.1 query after wq_b, RoPE, and BF16 materialisation. When
    // present, the frame consumes this exact adapter output and skips the
    // generic DSV4 query projection/normalisation path.
    q_in: Option<&[f32]>,
    kv_id: u64,
    li: usize,
    idxs: &[u32],
    inv_freq: &[f32],
    pos: usize,
    // When present the frame also expands its output into the layer state,
    // folds the FFN half and norms it, and hands back THAT — the MoE half's
    // input — instead of the attention output.
    hc: Option<&Dsv4HcTail>,
    out: &mut [f32],
) -> bool {
    // A refusal used to be a silent `false`, and three of them in a row cost
    // an evening of guessing which guard had fired.
    macro_rules! no {
        ($($t:tt)*) => {{
            tracing::debug!("кадр dsv4 отклонён: {}", format_args!($($t)*));
            if std::env::var("CMF_DSV4_FRAME_DEBUG").is_ok() {
                eprintln!("кадр dsv4 отклонён: {}", format_args!($($t)*));
            }
            return false;
        }};
    }
    let Some(c) = ctx() else {
        no!("нет контекста wgpu")
    };
    // `hidden` is only read when the frame has to build the LoRA vector
    // itself; demanding it regardless refused every caller that had one.
    if (qn_in.is_none() && q_in.is_none() && hidden.len() < g.dim)
        // Empty is a CONTRACT, not a mistake: "leave the result on the
        // card". The layer-frame refactor hit this exact guard-versus-branch
        // ordering and documented it; this is the second instance.
        || (!out.is_empty() && out.len() < g.dim)
        || w.sink.len() < g.nh
        || w.q_norm.len() < g.q_lora
        || q_in.is_some_and(|v| v.len() < g.nh * g.hd)
        || idxs.is_empty()
        || idxs.len() > 1024
        || inv_freq.len() * 2 < g.rd
    {
        no!(
            "формы: hidden {} dim {} out {} sink {} nh {} q_norm {} q_lora {} idx {} freq {} rd {}",
            hidden.len(),
            g.dim,
            out.len(),
            w.sink.len(),
            g.nh,
            w.q_norm.len(),
            g.q_lora,
            idxs.len(),
            inv_freq.len(),
            g.rd
        );
    }
    let bytes = model.primary_bytes();
    // Every weight q4tp, or the frame declines: a mixed layer would need the
    // per-op branches back and this is not the place to guess a layout.
    let mut wb = Vec::with_capacity(4);
    for &idx in &[w.wq_a, w.wq_b, w.wo_a, w.wo_b] {
        let Some(e) = model.tensors.get(idx) else {
            no!("тензора {idx} нет в каталоге");
        };
        if e.dtype != cortiq_core::TensorDtype::Q4TiledP || e.shape.len() != 2 {
            no!("{} не q4tp ({:?}, {:?})", e.name, e.dtype, e.shape);
        }
        let Some(abs) = model.entry_abs_offset(e) else {
            no!("{} без абсолютного смещения", e.name);
        };
        let plen = e.nbytes as usize;
        if abs + plen > bytes.len() {
            no!("{} выходит за файл", e.name);
        }
        let Some(b) = weight_buffer_l(
            c,
            (model.uid() as usize, idx),
            &bytes[abs..abs + plen],
            layer_of_name(&model.tensors[idx].name),
        ) else {
            no!("{} не поместился в бюджет VRAM", e.name);
        };
        wb.push(b);
    }
    let cache = {
        let map = c.dsv4_kv.lock().unwrap();
        match map.get(&(kv_id, li)) {
            Some((b, _)) => b.clone(),
            None => no!("кеш ({kv_id}, {li}) не заведён"),
        }
    };

    if let Some(v) = qn_in {
        if v.len() < g.q_lora {
            no!("готовый qn короче q_lora: {} < {}", v.len(), g.q_lora);
        }
    }
    // Constants (q_norm, sink, inv_freq) go through the const cache keyed on
    // their address — they are the same bytes every token. Everything else is
    // a reused buffer written in place.
    let hb = match (qn_in, q_in) {
        (_, Some(_)) => frame_buf(c, 0, 4, true),
        (None, None) => frame_up(c, 0, bytemuck::cast_slice(&hidden[..g.dim])),
        (Some(_), None) => frame_buf(c, 0, 4, true),
    };
    // These three ARE model-owned and outlive the run, so address keying is
    // sound for them — unlike anything built per call.
    let qnw = const_buf(c, bytemuck::cast_slice(&w.q_norm[..g.q_lora]));
    let sink = const_buf(c, bytemuck::cast_slice(&w.sink[..g.nh]));
    let freq = const_buf(c, bytemuck::cast_slice(&inv_freq[..g.rd / 2]));
    let posb = frame_up_pos(c, 1, pos, g.eps);
    // The list length changes token to token; round the buffer up so it is
    // not reallocated on every step, and pass the true count in the uniform.
    let ixb = {
        let cap = idxs.len().next_power_of_two().max(64);
        let b = frame_buf(c, 2, cap * 4, true);
        c.queue.write_buffer(&b, 0, bytemuck::cast_slice(idxs));
        b
    };

    // Readable, all of them: `CMF_DSV4_FRAME_TAP` reads back an intermediate
    // instead of the output. Eight verified kernels can still be wired wrong,
    // and a single number at the end says only that they were.
    let qr = frame_buf(c, 3, g.q_lora * 4, false);
    let qn = match qn_in {
        Some(v) => frame_up(c, 4, bytemuck::cast_slice(&v[..g.q_lora])),
        None => frame_buf(c, 4, g.q_lora * 4, true),
    };
    let q = match q_in {
        // Tag 106 is dedicated to an uploaded final V4.1 query. Generic
        // DSV4 uses tag 5 as a non-upload scratch buffer.
        Some(v) => frame_up(c, 106, bytemuck::cast_slice(&v[..g.nh * g.hd])),
        None => frame_buf(c, 5, g.nh * g.hd * 4, false),
    };
    let attn = frame_buf(c, 6, g.nh * g.hd * 4, false);
    let mid = frame_buf(c, 7, g.o_groups * g.o_lora * 4, false);
    let yb = frame_buf(c, 8, g.dim * 4, false);

    let t_all = std::time::Instant::now();
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-attn"),
        });
    if qn_in.is_none() && q_in.is_none() {
        encode_q4tp_mv1(
            c,
            &mut enc,
            &wb[0],
            &hb,
            &qr,
            g.q_lora,
            g.dim,
            (30, kv_id, li),
        );
        encode_rmsnorm(
            c,
            &mut enc,
            &qr,
            &qnw,
            &qn,
            g.q_lora,
            g.eps,
            (31, kv_id, li),
        );
    }
    if q_in.is_none() {
        encode_q4tp_mv1(
            c,
            &mut enc,
            &wb[1],
            &qn,
            &q,
            g.nh * g.hd,
            g.q_lora,
            (32, kv_id, li),
        );
        if g.bf16 {
            encode_bf16_round(c, &mut enc, &q, g.nh * g.hd, (37, kv_id, li));
        }
        encode_rope_heads(
            c,
            &mut enc,
            &q,
            &freq,
            &posb,
            g.nh,
            g.hd,
            g.rd,
            g.q_rms,
            false,
            (33, kv_id, li),
        );
    }
    encode_sparse_attend2(
        c,
        &mut enc,
        &q,
        &cache,
        &ixb,
        &sink,
        &attn,
        g.nh,
        g.hd,
        idxs.len(),
        g.scale,
        None,
    );
    encode_rope_heads(
        c,
        &mut enc,
        &attn,
        &freq,
        &posb,
        g.nh,
        g.hd,
        g.rd,
        false,
        true,
        (34, kv_id, li),
    );
    if g.bf16 {
        encode_bf16_round(c, &mut enc, &attn, g.nh * g.hd, (40, kv_id, li));
    }
    {
        let rows = g.o_groups * g.o_lora;
        let cols = g.nh * g.hd / g.o_groups;
        let bind = cached_bind(c, (36, kv_id, li), || {
            let p = uniform_u32x4(c, [(cols / 32) as u32, rows as u32, g.o_lora as u32, 0]);
            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &c.o_lora_a.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, &wb[2]),
                    bind_buf(1, &attn),
                    bind_buf(2, &mid),
                    bind_buf(3, &p),
                ],
            })
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.o_lora_a);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups((rows as u32).min(MAX_WG), 1, 1);
    }
    if g.bf16 {
        encode_bf16_round(c, &mut enc, &mid, g.o_groups * g.o_lora, (38, kv_id, li));
    }
    encode_q4tp_mv1(
        c,
        &mut enc,
        &wb[3],
        &mid,
        &yb,
        g.dim,
        g.o_groups * g.o_lora,
        (35, kv_id, li),
    );
    if g.bf16 {
        encode_bf16_round(c, &mut enc, &yb, g.dim, (39, kv_id, li));
    }

    // ── the hyper-connections, when the caller handed them over ──
    // The same order the host's hc_block keeps: expand this half's output
    // into the layer state, mix, fold with the FFN half's parameters, norm.
    // What comes out is the MoE half's INPUT, so the host has nothing left
    // to do between the two frames — which is 19 ms of a 57 ms token.
    let hc_out = hc.map(|h| {
        let mix_hc = (2 + h.hc) * h.hc;
        let state = frame_buf(c, 40, h.hc * g.dim * 4, true);
        let state2 = frame_buf(c, 46, h.hc * g.dim * 4, true);
        let hpost = frame_buf(c, 43, h.hc * 4, true);
        let hcomb = frame_buf(c, 44, h.hc * h.hc * 4, true);
        let mixes = frame_buf(c, 41, mix_hc * 4, true);
        let folded = frame_buf(c, 42, g.dim * 4, true);
        let x2 = frame_buf(c, 45, g.dim * 4, true);
        let hcp = uniform_u32x8(
            c,
            [
                h.hc as u32,
                g.dim as u32,
                h.sinkhorn_iters as u32,
                h.hc_eps.to_bits(),
                0,
                0,
                0,
                0,
            ],
        );
        let ffn_fn = const_buf(c, bytemuck::cast_slice(h.fn_));
        let ffn_sc = const_buf(c, bytemuck::cast_slice(h.scale));
        let ffn_bs = const_buf(c, bytemuck::cast_slice(&h.base[..mix_hc]));
        let ffn_nw = const_buf(c, bytemuck::cast_slice(&h.norm[..g.dim]));
        encode_hc_expand(
            c, &mut enc, &yb, &state, &hpost, &hcomb, &state2, &hcp, h.hc, g.dim,
        );
        encode_f32matvec(c, &mut enc, &ffn_fn, &state2, &mixes, mix_hc, h.hc * g.dim);
        encode_hc_fold(
            c, &mut enc, &state2, &mixes, &ffn_sc, &ffn_bs, &folded, &hpost, &hcomb, &hcp,
        );
        encode_rmsnorm(
            c,
            &mut enc,
            &folded,
            &ffn_nw,
            &x2,
            g.dim,
            h.eps,
            (54, kv_id, li),
        );
        x2
    });

    let tap = std::env::var("CMF_DSV41_TAIL_TAP")
        .or_else(|_| std::env::var("CMF_DSV4_FRAME_TAP"))
        .unwrap_or_default();
    let (src, n) = match tap.as_str() {
        "qr" => (&qr, g.q_lora),
        "qn" => (&qn, g.q_lora),
        "q" => (&q, g.nh * g.hd),
        "attn" => (&attn, g.nh * g.hd),
        "mid" => (&mid, g.o_groups * g.o_lora),
        _ => (hc_out.as_ref().unwrap_or(&yb), g.dim),
    };
    // An EMPTY `out` means the caller wants the result left where it is: the
    // MoE frame reads it from the same buffer, so the token does not stop
    // here at all.
    if out.is_empty() {
        submit(c, finish_enc(enc));
        ATT_ENC_NS.fetch_add(
            t_all.elapsed().as_nanos() as u64,
            std::sync::atomic::Ordering::Relaxed,
        );
        return true;
    }
    if out.len() < n {
        no!("отвод {tap} нуждается в {n} значениях, дано {}", out.len());
    }
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        (n * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dsv4-attn-stage",
    );
    // Resolve every pair the passes above took, into the staging buffer
    // back to back — one copy per pass, all inside this frame's encoder.
    let pairs: Vec<(usize, u32)> = std::mem::take(&mut TS_PAIRS.lock().unwrap());
    if let Some((qs, resolve, tstage)) = &c.ts_query {
        for (n, (_, slot)) in pairs.iter().enumerate() {
            flush_pass(&enc);
            enc.resolve_query_set(qs, *slot..*slot + 2, resolve, 0);
            flush_pass(&enc);
            enc.copy_buffer_to_buffer(resolve, 0, tstage, (n as u64) * 16, 16);
        }
    }
    let t_enc = std::time::Instant::now();
    let ok = readback(c, enc, src, &stage, (n * 4) as u64, &mut out[..n]);
    // The card's clock, read after the frame's own fence.
    if ok && !pairs.is_empty() {
        if let Some((_, _, tstage)) = &c.ts_query {
            let bytes = (pairs.len() as u64) * 16;
            let (tx, rx) = std::sync::mpsc::channel();
            tstage.map_async(wgpu::MapMode::Read, ..bytes, move |r| {
                let _ = tx.send(r);
            });
            let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
            if rx.recv().map(|r| r.is_ok()).unwrap_or(false) {
                if let Ok(raw) = tstage.get_mapped_range(..bytes) {
                    let t: &[u64] = bytemuck::cast_slice(&raw);
                    for (n, (which, _)) in pairs.iter().enumerate() {
                        let d = t[2 * n + 1].saturating_sub(t[2 * n]);
                        let ns = (d as f64 * c.ts_period as f64) as u64;
                        ATT_GPU_NS[*which].fetch_add(ns, std::sync::atomic::Ordering::Relaxed);
                    }
                    drop(raw);
                    ATT_GPU_N.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                }
            }
            tstage.unmap();
        }
    }
    drop(sc);
    ATT_ENC_NS.fetch_add(
        t_enc.duration_since(t_all).as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    ATT_WAIT_NS.fetch_add(
        t_enc.elapsed().as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    ok
}

/// Time inside `dsv4_attn_frame`, split at the submit — the same question the
/// MoE frame already answers, asked of the block that now costs more.
pub static ATT_ENC_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static ATT_WAIT_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// One DeepSeek-V4 MoE block on the device: route, run the chosen experts and
/// the shared one, sum. One submission.
///
/// The experts arrive PACKED — a subset chosen by the host, shared expert
/// last — while routing remains in global numbering. The explicit task mask
/// closes experts before top-k; the remap then distinguishes an allowed cold
/// expert from an allowed resident slot. Keeping those fields separate lets a
/// mask-complete pack chain without changing the route by accident.
#[derive(Clone)]
pub struct Dsv4MoeW<'a> {
    /// The gate as dense f32 `[n_exp, hidden]`. Used when `logits` is empty,
    /// which is what a device-resident input forces: the host cannot score
    /// a vector it does not have.
    pub router: &'a [f32],
    /// `(gate, up, down)` directory indices per packed expert, shared LAST.
    pub experts: &'a [(usize, usize, usize)],
    /// Router logits over the packed routed experts (shared excluded).
    pub logits: &'a [f32],
    /// noaux_tc selection bias, same numbering. Absent on the hash layers.
    pub bias: Option<&'a [f32]>,
    /// Global 0/1 task mask. A closed expert is excluded before top-k; it is
    /// not the same thing as an allowed expert whose VRAM slot is cold.
    pub mask: Option<&'a [u32]>,
    /// Hash-layer row, already in packed numbering.
    pub forced: Option<&'a [usize]>,
    /// global expert id -> packed slot, `u32::MAX` where the expert did not
    /// fit. When present the router ranges over ALL experts and hands the
    /// cold picks back instead of avoiding them.
    pub remap: Option<&'a [u32]>,
    /// Model-wide segmented slot banks. When present `experts` is not a
    /// physical pack: every remap value is a flat slot in this pool and the
    /// shared expert has its own pinned flat slot.
    pub global: Option<Dsv4GlobalMoe>,
    /// DeepSeek carries a pinned shared expert after the routed top-k. Qwen
    /// does not. Keeping this explicit lets both architectures use the same
    /// segmented out-of-core bank without inventing a zero shared tensor.
    pub has_shared: bool,
    /// Host-computed sigmoid gate of the shared expert. DeepSeek's ungated
    /// shared branch leaves this at 1; Qwen evaluates its tiny 1×hidden gate
    /// on the host beside the router and passes the exact scalar through the
    /// route uniform's companion buffer.
    pub shared_weight: f32,
    /// Forced expert ids already carry their final CPU-normalized weights in
    /// `bias`. Used by Qwen's host predictor so GPU exp/reduction rounding
    /// cannot perturb an otherwise exact route. DeepSeek leaves this false.
    pub preweighted: bool,
    /// Qwen router semantics: raw-logit top-k and selected softmax. False is
    /// DeepSeek-V4's sqrt-softplus/noaux route.
    pub qwen_softmax: bool,
}

#[derive(Clone, Copy)]
pub struct Dsv4GlobalMoe {
    pub pool_uid: u64,
    pub shared_slot: u32,
    pub segment_slots: u32,
}

#[derive(Clone, Copy)]
pub struct Dsv4MoeGeom {
    pub hidden: usize,
    pub inter: usize,
    pub top_k: usize,
    pub route_scale: f32,
    pub swiglu_limit: f32,
    /// gate/up are q2tp against a q4tp down — the mixed 2-bit profile.
    pub gu_q2: bool,
    /// Preserve the reference BF16 tensor boundaries for DeepSeek-V4.  The
    /// generic MoE kernels intentionally stay f32 between quantized matvecs;
    /// V4.1 sets this only for its source-equivalent path.
    pub bf16: bool,
}

pub fn dsv4_moe_frame(
    model: &Arc<CmfModel>,
    w: &Dsv4MoeW,
    g: Dsv4MoeGeom,
    // EMPTY means the attention frame left this half's input on the card in
    // its own buffer, which is the whole point: with the hyper-connections
    // done there too, the host has nothing to carry between the halves.
    x: &[f32],
    // `(expert, weight)` pairs the device left for the host, empty when the
    // whole packing was resident.
    cold_out: &mut Vec<(usize, f32)>,
    // The normalized FFN input the cold experts consume. When `x` above is
    // empty that vector exists only in the device frame; return it beside the
    // cold IDs so disk-backed CPU completion has real activations.
    cold_x_out: &mut Vec<f32>,
    // The state handover, split the way the layer frame splits it and for
    // the same reason: the EXPANSION of this half's output into the state is
    // unconditional whenever the device owns the state — the last layer has
    // no next fold, but its MoE half still has to enter the state the head
    // reads, and making the whole tail conditional is exactly how it did
    // not. `hc_cur` drives the expand; `hc_next` the next layer's fold.
    hc_cur: Option<&Dsv4HcTail>,
    hc_next: Option<(&Dsv4HcTail, &[f32])>,
    out: &mut [f32],
) -> bool {
    macro_rules! no {
        ($($t:tt)*) => {{
            tracing::debug!("кадр MoE отклонён: {}", format_args!($($t)*));
            if std::env::var("CMF_DSV4_FRAME_DEBUG").is_ok() {
                eprintln!("кадр MoE отклонён: {}", format_args!($($t)*));
            }
            return false;
        }};
    }
    let t_all = std::time::Instant::now();
    let Some(c) = ctx() else {
        no!("нет контекста wgpu")
    };
    let global_bufs = w
        .global
        .and_then(|gl| c.dsv4_global_moe.lock().unwrap().get(&gl.pool_uid).cloned());
    if w.global.is_some() && global_bufs.is_none() {
        no!("глобальный пул модели не найден")
    }
    if global_bufs.as_ref().is_some_and(|b| b.gu_q2 != g.gu_q2) {
        no!("dtype глобального пула не совпадает с gate/up")
    }
    let n_pack = global_bufs
        .as_ref()
        .map_or_else(|| w.experts.len().saturating_sub(1), |b| b.capacity);
    let slots = g.top_k + usize::from(w.has_shared);
    let n_all = if w.logits.is_empty() {
        w.router.len() / g.hidden.max(1)
    } else {
        w.logits.len()
    };
    let subset = w.remap.is_some_and(|r| r.len() >= n_all) && n_all > 0;
    // ONE width for the whole routing side: the scores, the bias and the
    // uniform must agree, and they did not. The bias went in n_pack long
    // while the kernel ranked over n_all, so every index past the packing
    // boundary read the LAST bias entry — WGSL clamps an out-of-bounds read
    // rather than faulting, so it looked like a plausible number and the
    // router quietly preferred the packed experts.
    let n_route = if subset { n_all } else { n_pack };
    if n_pack == 0
        || (!w.logits.is_empty() && w.logits.len() < n_route)
        || (global_bufs.is_none() && n_pack > 1024)
        || g.top_k == 0
        || g.top_k > 63
        || (!x.is_empty() && x.len() < g.hidden)
        || out.len() < g.hidden
        || g.hidden % 32 != 0
        || g.inter % 32 != 0
    {
        no!(
            "формы: упаковано {n_pack} логитов {} top_k {} hidden {} inter {}",
            w.logits.len(),
            g.top_k,
            g.hidden,
            g.inter
        );
    }
    let t_bufs = std::time::Instant::now();
    let local_bufs = if global_bufs.is_none() {
        let Some(v) = moe_expert_bufs(c, model, w.experts, g.inter, g.hidden, true, g.gu_q2, false)
        else {
            no!("эксперты не поместились в бюджет VRAM");
        };
        Some(v)
    } else {
        None
    };
    MOE_BUFS_NS.fetch_add(
        t_bufs.elapsed().as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    let t_up = std::time::Instant::now();

    // ── routing, on the device, straight into the msel/mwt the kernels read ──
    let lg = if w.logits.is_empty() {
        // Score on the card, from the input that is already there.
        let rb = const_buf(c, bytemuck::cast_slice(&w.router[..n_route * g.hidden]));
        // upload=true even on the device-scored arm: the pool keys buffers
        // by (tag, tok, len) and the OTHER arm of this `if` fills the same
        // slot with `write_buffer`. Whichever arm ran first used to fix the
        // usage flags for the whole process — on a card small enough to
        // route some layers on the host and some on the device, the second
        // pattern died on a COPY_DST validation. An unused COPY_DST is free.
        let lb = frame_buf(c, 16, n_route * 4, true);
        let mut e0 = c
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        let xin = frame_buf(c, 45, g.hidden * 4, true);
        encode_f32matvec(c, &mut e0, &rb, &xin, &lb, n_route, g.hidden);
        submit(c, finish_enc(e0));
        lb
    } else {
        frame_up(c, 16, bytemuck::cast_slice(&w.logits[..n_route]))
    };
    // NOT const_buf: that cache is keyed on the host ADDRESS, which is only
    // meaningful for model weights that outlive the process. The bias arrives
    // in a Vec built per layer, and the allocator hands back the same address
    // layer after layer — so every layer was routed with layer zero's bias.
    // The toys never caught it because they carry no expert_bias at all.
    let bs = match w.bias {
        Some(b) if b.len() >= n_route => frame_up(c, 25, bytemuck::cast_slice(&b[..n_route])),
        _ => lg.clone(),
    };
    let rmb = match w.remap {
        Some(r) if subset => frame_up(c, 26, bytemuck::cast_slice(&r[..n_all])),
        _ => frame_buf(c, 26, n_all.max(1) * 4, true),
    };
    let coldb = frame_buf(c, 27, 4 * g.top_k * 4, false);
    let mk = match w.mask {
        Some(m) if m.len() >= n_route => const_buf(c, bytemuck::cast_slice(&m[..n_route])),
        _ => frame_buf(c, 17, n_route.max(1) * 4, true),
    };
    let mut forced_words = vec![0u32; slots];
    if let Some(f) = w.forced.filter(|f| f.len() >= g.top_k) {
        for (dst, &src) in forced_words[..g.top_k].iter_mut().zip(f) {
            *dst = src as u32;
        }
    }
    if w.has_shared {
        forced_words[g.top_k] = w.shared_weight.to_bits();
    }
    let fc = frame_up(c, 18, bytemuck::cast_slice(&forced_words));
    let msel = frame_buf(c, 19, slots * 4, false);
    let mwt = frame_buf(c, 20, slots * 4, false);
    let mcnt = frame_buf(c, 21, 4, false);
    let mact = frame_buf(c, 22, slots * g.inter * 4, false);
    let xb = if x.is_empty() {
        frame_buf(c, 45, g.hidden * 4, true)
    } else {
        frame_up(c, 23, bytemuck::cast_slice(&x[..g.hidden]))
    };
    let ob = frame_buf(c, 24, g.hidden * 4, false);

    let rflags = (w.bias.is_some_and(|b| b.len() >= n_route) as u32)
        | ((w.mask.is_some_and(|m| m.len() >= n_route) as u32) << 1)
        | ((w.forced.is_some_and(|f| f.len() >= g.top_k) as u32) << 2)
        | ((w.has_shared as u32) << 3)
        | ((subset as u32) << 4)
        | ((w.qwen_softmax as u32) << 5)
        | ((w.has_shared as u32) << 6)
        | ((w.preweighted as u32) << 7)
        | (w.global.map_or(n_pack as u32, |gl| gl.shared_slot) << 8);
    if std::env::var("CMF_DSV4_MOE_CHECK").is_ok() {
        eprintln!(
            "[маршрут] n_all={n_all} n_pack={n_pack} subset={subset} flags={rflags} \
             remap[0..12]={:?}",
            w.remap.map(|r| &r[..12.min(r.len())])
        );
    }
    // Ranking ranges over EVERY expert when the packing is a subset — that is
    // the whole point. Passing n_pack here silently turned it back into a
    // mask that also indexed the packed buffer with global ids.
    let rp = uniform_mixed(c, [n_route as u32, g.top_k as u32, rflags], g.route_scale);

    let stride16 = |rows: usize, cols: usize, q2: bool| -> u32 {
        let dt = if q2 {
            cortiq_core::TensorDtype::Q2TiledP
        } else {
            cortiq_core::TensorDtype::Q4TiledP
        };
        (cortiq_core::quant::expected_nbytes(dt, &[rows, cols]).unwrap_or(0) / 2) as u32
    };
    let gu_u = uniform_u32x8(
        c,
        [
            (g.hidden / 32) as u32,
            g.inter as u32,
            slots as u32,
            stride16(g.inter, g.hidden, g.gu_q2),
            g.swiglu_limit.to_bits(),
            w.global.map_or(0, |gl| gl.segment_slots),
            u32::from(g.bf16),
            0,
        ],
    );
    let dn_u4 = uniform_u32x4(
        c,
        [
            (g.inter / 32) as u32,
            g.hidden as u32,
            slots as u32,
            stride16(g.hidden, g.inter, false),
        ],
    );
    let dn_u8 = w.global.map(|gl| {
        uniform_u32x8(
            c,
            [
                (g.inter / 32) as u32,
                g.hidden as u32,
                slots as u32,
                stride16(g.hidden, g.inter, false),
                gl.segment_slots,
                u32::from(g.bf16),
                0,
                0,
            ],
        )
    });
    let gu_r4 = global_bufs.is_none() && g.inter % 4 == 0 && bt_gu_r4_on();
    let (p_gu, p_dn) = if let Some(gb) = global_bufs.as_ref() {
        let Some(pipes) = dsv4_global_moe_pipelines(c, g.gu_q2, gb.segments) else {
            no!("global MoE cache has no matching pipeline geometry");
        };
        pipes
    } else if g.gu_q2 {
        (
            if gu_r4 {
                &c.bt_moe_gate_up_q2tp_r4
            } else {
                &c.moe_gate_up_q2tp
            },
            &c.moe_down_q4tp,
        )
    } else {
        (
            if gu_r4 {
                &c.moe_gate_up_q4tp_b_r4
            } else {
                &c.moe_gate_up_q4tp_b
            },
            &c.moe_down_q4tp_b,
        )
    };
    let l_gu = p_gu.get_bind_group_layout(0);
    let l_dn = p_dn.get_bind_group_layout(0);

    MOE_UP_NS.fetch_add(
        t_up.elapsed().as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    let t_pass = std::time::Instant::now();
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("dsv4-moe"),
        });
    // The layer's identity for the bind cache: its first expert's directory
    // index, which is unique per layer and already at hand.
    let lkey = w.experts.first().map(|e| e.0).unwrap_or(0);
    if w.logits.is_empty() {
        let rb = const_buf(c, bytemuck::cast_slice(&w.router[..n_route * g.hidden]));
        let xin = frame_buf(c, 45, g.hidden * 4, true);
        encode_f32matvec(c, &mut enc, &rb, &xin, &lg, n_route, g.hidden);
    }
    {
        // NOT cached. This group holds `rp`, a CONTENT-keyed uniform: change a
        // flag and the uniform becomes a different buffer while the cached
        // group keeps pointing at the old one — the layer then routes with
        // yesterday's flags forever. Encoding it costs 0.01 ms a layer; being
        // wrong costs a model.
        let _ = lkey;
        let bind = {
            c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: None,
                layout: &c.moe_route.get_bind_group_layout(0),
                entries: &[
                    bind_buf(0, &lg),
                    bind_buf(1, &bs),
                    bind_buf(2, &mk),
                    bind_buf(3, &fc),
                    bind_buf(4, &msel),
                    bind_buf(5, &mwt),
                    bind_buf(6, &mcnt),
                    bind_buf(7, &rp),
                    bind_buf(8, &rmb),
                    bind_buf(9, &coldb),
                ],
            })
        };
        // The card's own clock around the whole MoE block. All three kernels
        // share this pass; splitting it to time them apart would add two
        // pass boundaries a layer and measure the split.
        //
        // A FRESH pair of slots: rewriting the same two on every frame of
        // every token leaves the queries unreset between submissions, and an
        // unreset timestamp query reads back as zero.
        let slot = (TS_SLOT.fetch_add(2, std::sync::atomic::Ordering::Relaxed) % 254) as u32;
        TS_LAST.store(slot, std::sync::atomic::Ordering::Relaxed);
        let tsw = c
            .ts_query
            .as_ref()
            .map(|(qs, _, _)| wgpu::ComputePassTimestampWrites {
                query_set: qs,
                beginning_of_pass_write_index: Some(slot),
                end_of_pass_write_index: Some(slot + 1),
            });
        let mut pass = begin_pass_with(&mut enc, None, tsw);
        pass.set_pipeline(&c.moe_route);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(1, 1, 1);

        if let Some(gb) = global_bufs.as_ref() {
            let gate_bindings: Vec<_> = gb
                .gate
                .iter()
                .map(wgpu::Buffer::as_entire_buffer_binding)
                .collect();
            let up_bindings: Vec<_> = gb
                .up
                .iter()
                .map(wgpu::Buffer::as_entire_buffer_binding)
                .collect();
            let down_bindings: Vec<_> = gb
                .down
                .iter()
                .map(wgpu::Buffer::as_entire_buffer_binding)
                .collect();
            let bg_gu = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("dsv4-global-gu"),
                layout: &l_gu,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: wgpu::BindingResource::BufferArray(&gate_bindings),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: wgpu::BindingResource::BufferArray(&up_bindings),
                    },
                    bind_buf(2, &xb),
                    bind_buf(3, &msel),
                    bind_buf(4, &mact),
                    bind_buf(5, &mwt),
                ],
            });
            let bg_gu_p = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("dsv4-global-gu-p"),
                layout: &p_gu.get_bind_group_layout(1),
                entries: &[bind_buf(0, &gu_u)],
            });
            pass.set_pipeline(p_gu);
            pass.set_bind_group(0, &bg_gu, &[]);
            pass.set_bind_group(1, &bg_gu_p, &[]);
            pass.dispatch_workgroups(g.inter as u32, slots as u32, 1);

            let bg_dn = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("dsv4-global-dn"),
                layout: &l_dn,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: wgpu::BindingResource::BufferArray(&down_bindings),
                    },
                    bind_buf(1, &mact),
                    bind_buf(2, &msel),
                    bind_buf(3, &mwt),
                    bind_buf(4, &ob),
                ],
            });
            let bg_dn_p = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("dsv4-global-dn-p"),
                layout: &p_dn.get_bind_group_layout(1),
                entries: &[bind_buf(0, dn_u8.as_ref().unwrap())],
            });
            pass.set_pipeline(p_dn);
            pass.set_bind_group(0, &bg_dn, &[]);
            pass.set_bind_group(1, &bg_dn_p, &[]);
            pass.dispatch_workgroups(g.hidden as u32, 1, 1);
        } else {
            let (gate_all, up_all, down_all) = local_bufs.as_ref().unwrap();
            let bg_gu = cached_bind(c, (41, 0, lkey), || {
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &l_gu,
                    entries: &[
                        bind_buf(0, gate_all),
                        bind_buf(1, up_all),
                        bind_buf(2, &xb),
                        bind_buf(3, &msel),
                        bind_buf(4, &mact),
                        bind_buf(5, &gu_u),
                    ],
                })
            });
            pass.set_pipeline(p_gu);
            pass.set_bind_group(0, &bg_gu, &[]);
            pass.dispatch_workgroups(
                (g.inter as u32).div_ceil(if gu_r4 { 4 } else { 1 }),
                slots as u32,
                1,
            );

            let bg_dn = cached_bind(c, (42, 0, lkey), || {
                c.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: None,
                    layout: &l_dn,
                    entries: &[
                        bind_buf(0, down_all),
                        bind_buf(1, &mact),
                        bind_buf(2, &msel),
                        bind_buf(3, &mwt),
                        bind_buf(4, &ob),
                        bind_buf(5, &dn_u4),
                    ],
                })
            });
            pass.set_pipeline(p_dn);
            pass.set_bind_group(0, &bg_dn, &[]);
            pass.dispatch_workgroups(g.hidden as u32, 1, 1);
        }
    }
    // ── the state handover, on the card ──
    if let Some(h) = hc_cur {
        let state = frame_buf(c, 40, h.hc * g.hidden * 4, true);
        let state2 = frame_buf(c, 46, h.hc * g.hidden * 4, true);
        let hpost = frame_buf(c, 43, h.hc * 4, true);
        let hcomb = frame_buf(c, 44, h.hc * h.hc * 4, true);
        let hcp = uniform_u32x8(
            c,
            [
                h.hc as u32,
                g.hidden as u32,
                h.sinkhorn_iters as u32,
                h.hc_eps.to_bits(),
                0,
                0,
                0,
                0,
            ],
        );
        encode_hc_expand(
            c, &mut enc, &ob, &state2, &hpost, &hcomb, &state, &hcp, h.hc, g.hidden,
        );
    }
    let hc_out = hc_next.map(|(h, next_norm)| {
        let mix_hc = (2 + h.hc) * h.hc;
        let state = frame_buf(c, 40, h.hc * g.hidden * 4, true);
        let _state2 = frame_buf(c, 46, h.hc * g.hidden * 4, true);
        let hpost = frame_buf(c, 43, h.hc * 4, true);
        let hcomb = frame_buf(c, 44, h.hc * h.hc * 4, true);
        let mixes = frame_buf(c, 41, mix_hc * 4, true);
        let folded = frame_buf(c, 42, g.hidden * 4, true);
        let x2 = frame_buf(c, 45, g.hidden * 4, true);
        let hcp = uniform_u32x8(
            c,
            [
                h.hc as u32,
                g.hidden as u32,
                h.sinkhorn_iters as u32,
                h.hc_eps.to_bits(),
                0,
                0,
                0,
                0,
            ],
        );
        let nfn = const_buf(c, bytemuck::cast_slice(h.fn_));
        let nsc = const_buf(c, bytemuck::cast_slice(h.scale));
        let nbs = const_buf(c, bytemuck::cast_slice(&h.base[..mix_hc]));
        let nnw = const_buf(c, bytemuck::cast_slice(&next_norm[..g.hidden]));
        encode_f32matvec(c, &mut enc, &nfn, &state, &mixes, mix_hc, h.hc * g.hidden);
        encode_hc_fold(
            c, &mut enc, &state, &mixes, &nsc, &nbs, &folded, &hpost, &hcomb, &hcp,
        );
        encode_rmsnorm(
            c,
            &mut enc,
            &folded,
            &nnw,
            &x2,
            g.hidden,
            h.eps,
            (55, 0, lkey),
        );
        x2
    });

    if let Some((qs, resolve, tstage)) = &c.ts_query {
        let slot = TS_LAST.load(std::sync::atomic::Ordering::Relaxed);
        // An UNCONDITIONAL value in the staging buffer before the copy. If it
        // comes back, the copy never landed; if it comes back zeroed, the
        // copy landed and the queries themselves are empty. Two answers, one
        // run — the same trick that separated a kernel from its readback in
        // the cold-expert path.
        if std::env::var("CMF_TS_DEBUG").is_ok() {
            let mark: [u64; 2] = [0xDEAD_BEEF_1111, 0xDEAD_BEEF_2222];
            c.queue.write_buffer(tstage, 0, bytemuck::cast_slice(&mark));
        }
        // Offset ZERO, always: a query resolve's destination offset has to be
        // 256-byte aligned, and slot*8 is not for any slot but the first
        // thirty-two. The pair still comes from the rotating slots; only
        // where it lands is fixed.
        flush_pass(&enc);
        enc.resolve_query_set(qs, slot..slot + 2, resolve, 0);
        flush_pass(&enc);
        enc.copy_buffer_to_buffer(resolve, 0, tstage, 0, 16);
    }
    let t_enc = std::time::Instant::now();
    MOE_PASS_NS.fetch_add(
        t_pass.elapsed().as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    let mut sc = c.scratch.lock().unwrap();
    // The cold list rides the SAME staging buffer and the SAME fence, so the
    // host learns which picks it owes without paying a second barrier.
    //
    // This block went missing once — `cold_out` was declared, threaded all
    // the way down and never filled — and the empty list read exactly like a
    // router that had chosen no cold experts. An unconditional probe written
    // from the kernel is what proved otherwise.
    let cold_bytes = (4 * g.top_k * 4) as u64;
    let x_off = (g.hidden * 4) as u64 + cold_bytes;
    let total = x_off + (g.hidden * 4) as u64;
    // ONE ensure for the whole readback. A first call sized to the hidden
    // state alone used to run before this one, on the same slot: it built a
    // buffer that the next line immediately outgrew and replaced.
    let stage2 = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        total,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "dsv4-moe-stage",
    );
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(
        hc_out.as_ref().unwrap_or(&ob),
        0,
        &stage2,
        0,
        (g.hidden * 4) as u64,
    );
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&coldb, 0, &stage2, (g.hidden * 4) as u64, cold_bytes);
    flush_pass(&enc);
    enc.copy_buffer_to_buffer(&xb, 0, &stage2, x_off, (g.hidden * 4) as u64);
    submit(c, finish_enc(enc));
    let slice = stage2.slice(..total);
    slice.map_async(wgpu::MapMode::Read, |_| {});
    if c.device.poll(wgpu::PollType::wait_indefinitely()).is_err() {
        return false;
    }
    let mut ok = false;
    if let Ok(data) = slice.get_mapped_range() {
        out[..g.hidden].copy_from_slice(bytemuck::cast_slice(&data[..g.hidden * 4]));
        let tail: &[u32] = bytemuck::cast_slice(&data[g.hidden * 4..total as usize]);
        cold_out.clear();
        for t in 0..g.top_k {
            if tail[2 * t] != u32::MAX {
                cold_out.push((tail[2 * t] as usize, f32::from_bits(tail[2 * t + 1])));
            }
        }
        cold_x_out.clear();
        cold_x_out.extend_from_slice(bytemuck::cast_slice(&data[x_off as usize..total as usize]));
        if std::env::var("CMF_DSV4_MOE_CHECK").is_ok() {
            let picks: Vec<(u32, f32)> = (0..g.top_k)
                .map(|t| {
                    (
                        tail[2 * g.top_k + 2 * t],
                        f32::from_bits(tail[2 * g.top_k + 2 * t + 1]),
                    )
                })
                .collect();
            eprintln!("[победители карты] {picks:?}");
        }
        ok = true;
    }
    stage2.unmap();
    // The card's clock, read after the frame's own wait — free, because the
    // fence has already been paid for.
    if let Some((_, _, tstage)) = &c.ts_query {
        let (tx, rx) = std::sync::mpsc::channel();
        tstage.map_async(wgpu::MapMode::Read, ..16, move |r| {
            let _ = tx.send(r);
        });
        let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
        if rx.recv().map(|r| r.is_ok()).unwrap_or(false) {
            if let Ok(raw) = tstage.get_mapped_range(..16) {
                let t: &[u64] = bytemuck::cast_slice(&raw);
                if std::env::var("CMF_TS_DEBUG").is_ok() {
                    eprintln!("[ts] t0={} t1={} period={}", t[0], t[1], c.ts_period);
                }
                let ns = (t[1].saturating_sub(t[0]) as f64 * c.ts_period as f64) as u64;
                drop(raw);
                MOE_GPU_NS[0].fetch_add(ns, std::sync::atomic::Ordering::Relaxed);
                MOE_GPU_N.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            }
        }
        tstage.unmap();
    }
    drop(sc);
    // Encoding and waiting are different problems with different fixes, and
    // the layer total cannot tell them apart. Costs one Instant per layer.
    MOE_ENC_NS.fetch_add(
        t_enc.duration_since(t_all).as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    MOE_WAIT_NS.fetch_add(
        t_enc.elapsed().as_nanos() as u64,
        std::sync::atomic::Ordering::Relaxed,
    );
    ok
}

/// Time inside `dsv4_moe_frame`, split at the submit. Read by the dsv4
/// profile so a slow block can be blamed on the right half.
pub static MOE_ENC_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static MOE_WAIT_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
/// The encode half again, split three ways — "encoding" turned out to be the
/// whole token's cost and "which part of it" is not guessable: the expert
/// buffer lookup, the per-call uploads, and the passes themselves.
/// GPU time inside the MoE frame, per kernel: routing, gate/up, down.
/// The host counters above measure the CPU's share of a frame; these are the
/// only thing that says what the CARD spends, and dsv4 had no equivalent —
/// `CMF_GPU_TS` instruments the general token graph, which this arch does
/// not use, so the profiler simply printed nothing.
pub static MOE_GPU_NS: [std::sync::atomic::AtomicU64; 3] = [
    std::sync::atomic::AtomicU64::new(0),
    std::sync::atomic::AtomicU64::new(0),
    std::sync::atomic::AtomicU64::new(0),
];
pub static MOE_GPU_N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
/// Card time of the attention frame's passes: 0 = the single-kernel sparse
/// attend, 1 = its scores half, 2 = its apply half.
pub static ATT_GPU_NS: [std::sync::atomic::AtomicU64; 3] = [
    std::sync::atomic::AtomicU64::new(0),
    std::sync::atomic::AtomicU64::new(0),
    std::sync::atomic::AtomicU64::new(0),
];
pub static ATT_GPU_N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
/// (which counter, slot) for every pass the current encoder stamped.
static TS_PAIRS: Mutex<Vec<(usize, u32)>> = Mutex::new(Vec::new());
/// Rotating timestamp slot, so no two frames in flight share a query.
static TS_SLOT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
/// The pair the last encoded pass took, for the frame that resolves it.
static TS_LAST: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(0);

/// Per-stage GPU time inside the BT chain frame: `CMF_GPU_TS=1` plus
/// `CMF_BT_TS=li[,li…]` pick the sampled layers; `begin_pass` stamps every
/// pass opened while a stage label is set, and the chain's own fence pays
/// for the readback. Labels index `BT_TS_NAMES`.
static BT_TS_STAGE: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
const BT_TS_NAMES: [&str; 16] = [
    "", "q", "comp", "окно", "ix-mv", "ix-score", "attend", "olora", "glue1", "nextq", "commit",
    "route", "gu", "dn", "glue2", "wo_b",
];

// `CMF_BATCH_KERNEL_TS=1` reserves query slots 1024..4095 for the actual
// resident batch kernels while the coarse batch-stage trace keeps 0..1023.
// This is a bounded microscope, not a production timing path.  The helper
// resets the records at every true batch entry and maps them only after the
// frame fence. Any overflow is counted per stage instead of silently reusing
// an older frame's records.
const BATCH_KERNEL_TS_BASE: u32 = 1024;
const BATCH_KERNEL_TS_LIMIT: u32 = 4096;
static BATCH_KERNEL_TS_SLOT: std::sync::atomic::AtomicU32 =
    std::sync::atomic::AtomicU32::new(BATCH_KERNEL_TS_BASE);
static BATCH_KERNEL_TS_PAIRS: Mutex<Vec<(u8, u32, u64)>> = Mutex::new(Vec::new());
static BATCH_KERNEL_TS_DROPPED: [std::sync::atomic::AtomicU32; 2] = [
    std::sync::atomic::AtomicU32::new(0),
    std::sync::atomic::AtomicU32::new(0),
];

fn batch_kernel_ts_on(c: &Ctx) -> bool {
    std::env::var("CMF_BATCH_KERNEL_TS").as_deref() == Ok("1") && c.ts_query.is_some()
}

fn batch_kernel_ts_begin(c: &Ctx) {
    if !batch_kernel_ts_on(c) {
        return;
    }
    BATCH_KERNEL_TS_SLOT.store(BATCH_KERNEL_TS_BASE, std::sync::atomic::Ordering::Relaxed);
    BATCH_KERNEL_TS_PAIRS.lock().unwrap().clear();
    for n in &BATCH_KERNEL_TS_DROPPED {
        n.store(0, std::sync::atomic::Ordering::Relaxed);
    }
}

fn batch_kernel_ts_pair(
    c: &Ctx,
    stage: u8,
    active_bytes: u64,
) -> Option<wgpu::ComputePassTimestampWrites<'_>> {
    if !batch_kernel_ts_on(c) {
        return None;
    }
    let slot = BATCH_KERNEL_TS_SLOT.fetch_add(2, std::sync::atomic::Ordering::Relaxed);
    if slot + 1 >= BATCH_KERNEL_TS_LIMIT {
        if let Some(n) = BATCH_KERNEL_TS_DROPPED.get(stage as usize) {
            n.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        }
        return None;
    }
    let (qs, _, _) = c.ts_query.as_ref()?;
    BATCH_KERNEL_TS_PAIRS
        .lock()
        .unwrap()
        .push((stage, slot, active_bytes));
    Some(wgpu::ComputePassTimestampWrites {
        query_set: qs,
        beginning_of_pass_write_index: Some(slot),
        end_of_pass_write_index: Some(slot + 1),
    })
}

fn batch_kernel_ts_resolve(c: &Ctx, enc: &mut wgpu::CommandEncoder) {
    if !batch_kernel_ts_on(c) {
        return;
    }
    let pairs = BATCH_KERNEL_TS_PAIRS.lock().unwrap().clone();
    if pairs.is_empty() {
        return;
    }
    let end = BATCH_KERNEL_TS_SLOT
        .load(std::sync::atomic::Ordering::Relaxed)
        .min(BATCH_KERNEL_TS_LIMIT);
    if end <= BATCH_KERNEL_TS_BASE {
        return;
    }
    if let Some((qs, resolve, stage)) = c.ts_query.as_ref() {
        flush_pass(enc);
        enc.resolve_query_set(
            qs,
            BATCH_KERNEL_TS_BASE..end,
            resolve,
            (BATCH_KERNEL_TS_BASE * 8) as u64,
        );
        flush_pass(enc);
        enc.copy_buffer_to_buffer(
            resolve,
            (BATCH_KERNEL_TS_BASE * 8) as u64,
            stage,
            (BATCH_KERNEL_TS_BASE * 8) as u64,
            ((end - BATCH_KERNEL_TS_BASE) * 8) as u64,
        );
    }
}

fn batch_kernel_ts_report(c: &Ctx, batch_tokens: usize) {
    if !batch_kernel_ts_on(c) {
        return;
    }
    let pairs = BATCH_KERNEL_TS_PAIRS.lock().unwrap().clone();
    if pairs.is_empty() {
        eprintln!("batch-kernel-ts: no timestamped kernels");
        return;
    }
    let Some((_, _, stage)) = c.ts_query.as_ref() else {
        return;
    };
    let end = BATCH_KERNEL_TS_SLOT
        .load(std::sync::atomic::Ordering::Relaxed)
        .min(BATCH_KERNEL_TS_LIMIT);
    let bytes = ((end - BATCH_KERNEL_TS_BASE) * 8) as u64;
    let (tx, rx) = std::sync::mpsc::channel();
    stage.map_async(
        wgpu::MapMode::Read,
        (BATCH_KERNEL_TS_BASE * 8) as u64..(BATCH_KERNEL_TS_BASE as u64 * 8 + bytes),
        move |r| {
            let _ = tx.send(r);
        },
    );
    let _ = c.device.poll(wgpu::PollType::wait_indefinitely());
    if !rx.recv().map(|r| r.is_ok()).unwrap_or(false) {
        stage.unmap();
        eprintln!("batch-kernel-ts: timestamp map failed");
        return;
    }
    let Ok(raw) = stage.get_mapped_range(
        (BATCH_KERNEL_TS_BASE * 8) as u64..(BATCH_KERNEL_TS_BASE as u64 * 8 + bytes),
    ) else {
        stage.unmap();
        eprintln!("batch-kernel-ts: timestamp range unavailable");
        return;
    };
    let ticks: &[u64] = bytemuck::cast_slice(&raw);
    let mut ms = [0.0f64; 2];
    let mut counts = [0u32; 2];
    let mut active = [0u64; 2];
    for &(which, slot, active_bytes) in &pairs {
        let off = ((slot - BATCH_KERNEL_TS_BASE) as usize).min(ticks.len().saturating_sub(2));
        let elapsed = ticks[off + 1].saturating_sub(ticks[off]) as f64 * c.ts_period as f64 / 1e6;
        let i = which as usize;
        if i < ms.len() {
            ms[i] += elapsed;
            counts[i] += 1;
            active[i] = active[i].saturating_add(active_bytes);
        }
    }
    drop(raw);
    stage.unmap();
    let q2_gbps = if ms[0] > 0.0 {
        active[0] as f64 / 1e9 / (ms[0] / 1e3)
    } else {
        0.0
    };
    let fwht_gbps = if ms[1] > 0.0 {
        active[1] as f64 / 1e9 / (ms[1] / 1e3)
    } else {
        0.0
    };
    let dropped_q2 = BATCH_KERNEL_TS_DROPPED[0].load(std::sync::atomic::Ordering::Relaxed);
    let dropped_fwht = BATCH_KERNEL_TS_DROPPED[1].load(std::sync::atomic::Ordering::Relaxed);
    eprintln!(
        "batch-kernel-ts: k={} q2-mm={:.2}ms/{}(+{} dropped) active={:.3}GB ({:.1}MB/token) logical_bw={:.2}GB/s fwht={:.2}ms/{}(+{} dropped) active={:.3}GB ({:.1}MB/token) logical_bw={:.2}GB/s slots={}/{}",
        batch_tokens,
        ms[0],
        counts[0],
        dropped_q2,
        active[0] as f64 / 1e9,
        active[0] as f64 / 1e6 / batch_tokens.max(1) as f64,
        q2_gbps,
        ms[1],
        counts[1],
        dropped_fwht,
        active[1] as f64 / 1e9,
        active[1] as f64 / 1e6 / batch_tokens.max(1) as f64,
        fwht_gbps,
        pairs.len(),
        BATCH_KERNEL_TS_LIMIT - BATCH_KERNEL_TS_BASE,
    );
}

fn bt_ts_lis() -> &'static [usize] {
    static L: std::sync::OnceLock<Vec<usize>> = std::sync::OnceLock::new();
    L.get_or_init(|| {
        std::env::var("CMF_BT_TS")
            .map(|v| v.split(',').filter_map(|s| s.trim().parse().ok()).collect())
            .unwrap_or_default()
    })
}

fn bt_ts(stage: usize) {
    BT_TS_STAGE.store(stage, std::sync::atomic::Ordering::Relaxed);
}
pub static MOE_BUFS_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static MOE_UP_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
pub static MOE_PASS_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// Attention over an index list with the learned sink, on the device.
///
/// Step two of the whole-token graph. Verified against `dsv4::sparse_attend`
/// before anything depends on it: a sink that contributes to the numerator,
/// or a denominator missing its share, changes every head's output by a
/// factor that no generated text would reveal.
#[allow(clippy::too_many_arguments)]
pub fn sparse_attend_for_test(
    q: &[f32],
    kv: &[f32],
    idxs: &[u32],
    sink: &[f32],
    scale: f32,
    nh: usize,
    hd: usize,
    out: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if q.len() != nh * hd || out.len() != nh * hd || sink.len() != nh || idxs.len() > 1024 {
        return false;
    }
    let qb = storage_bytes(c, bytemuck::cast_slice(q));
    let kvb = storage_bytes(c, bytemuck::cast_slice(kv));
    let ib = storage_bytes(c, bytemuck::cast_slice(idxs));
    let sb = storage_bytes(c, bytemuck::cast_slice(sink));
    let ob = rw_f32(c, nh * hd, true);
    let p = uniform_u32x4(
        c,
        [nh as u32, hd as u32, idxs.len() as u32, scale.to_bits()],
    );
    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("sa") });
    let _ = &p; // the split path builds its own params
    encode_sparse_attend2(
        c,
        &mut enc,
        &qb,
        &kvb,
        &ib,
        &sb,
        &ob,
        nh,
        hd,
        idxs.len(),
        scale,
        None,
    );
    let mut sc = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc.stage,
        (nh * hd * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "sa-stage",
    );
    let ok = readback(c, enc, &ob, &stage, (nh * hd * 4) as u64, out);
    drop(sc);
    ok
}

/// One hyper-connection join on the device: fold the copies (with the
/// Sinkhorn) and expand them back around a block output computed elsewhere.
///
/// Step one of the whole-token graph, and deliberately useless on its own —
/// it costs a submission to save none. It exists so the join can be checked
/// against the CPU before anything is built on top of it, because a
/// transposed mixing matrix or a Sinkhorn off by one iteration produces
/// output that looks entirely reasonable.
#[allow(clippy::too_many_arguments)]
pub fn hc_join_for_test(
    state: &[f32],
    mixes: &[f32],
    scale: &[f32; 3],
    base: &[f32],
    block_out: &[f32],
    hc: usize,
    dim: usize,
    iters: u32,
    eps: f32,
    folded: &mut [f32],
    expanded: &mut [f32],
) -> bool {
    let Some(c) = ctx() else { return false };
    if state.len() != hc * dim || folded.len() != dim || expanded.len() != hc * dim {
        return false;
    }
    let st = storage_bytes(c, bytemuck::cast_slice(state));
    let mx = storage_bytes(c, bytemuck::cast_slice(mixes));
    let sc = storage_bytes(c, bytemuck::cast_slice(&scale[..]));
    let bs = storage_bytes(c, bytemuck::cast_slice(base));
    let fo = rw_f32(c, dim, true);
    let po = rw_f32(c, hc, false);
    let cb = rw_f32(c, hc * hc, false);
    // HcP grew a fifth field, `nrm`, when the fold learned to emit the
    // normed vector — so the uniform is 32 bytes now, and a 16-byte one is
    // rejected outright. Zero: this hook wants the raw fold, not the norm.
    let params = uniform_u32x8(c, [hc as u32, dim as u32, iters, eps.to_bits(), 0, 0, 0, 0]);

    let mut enc = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("hc") });
    {
        let ones = storage_bytes(c, bytemuck::cast_slice(&vec![1.0f32; dim]));
        let normed_out = storage_bytes(c, bytemuck::cast_slice(&vec![0.0f32; dim]));
        let layout = c.hc_pre_fold.get_bind_group_layout(0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layout,
            entries: &[
                bind_buf(0, &st),
                bind_buf(1, &mx),
                bind_buf(2, &sc),
                bind_buf(3, &bs),
                bind_buf(4, &fo),
                bind_buf(5, &po),
                bind_buf(6, &cb),
                bind_buf(7, &params),
                // 8 and 9 arrived when the fold learned to emit the NORMED
                // vector alongside the raw one. The layout takes every
                // binding the entry point touches, so leaving them out is a
                // validation error, not a smaller bind group — this hook had
                // been failing on it. The test reads binding 4, the raw
                // fold; a unit norm and a scratch output satisfy the rest.
                bind_buf(8, &ones),
                bind_buf(9, &normed_out),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.hc_pre_fold);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(1, 1, 1);
    }
    let bo = storage_bytes(c, bytemuck::cast_slice(block_out));
    let ex = rw_f32(c, hc * dim, true);
    {
        let layout = c.hc_post_expand.get_bind_group_layout(0);
        let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &layout,
            entries: &[
                bind_buf(0, &bo),
                bind_buf(1, &st),
                bind_buf(2, &po),
                bind_buf(3, &cb),
                bind_buf(4, &ex),
                bind_buf(5, &params),
            ],
        });
        let mut pass = begin_pass(&mut enc);
        pass.set_pipeline(&c.hc_post_expand);
        pass.set_bind_group(0, &bind, &[]);
        pass.dispatch_workgroups(((hc * dim) as u32).div_ceil(256), 1, 1);
    }
    // Two readbacks because the two results have different lengths; this is
    // a check, not a hot path.
    let mut sc_lock = c.scratch.lock().unwrap();
    let stage = Scratch::ensure(
        &c.device,
        &mut sc_lock.stage,
        ((hc * dim).max(dim) * 4) as u64,
        wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        "hc-stage",
    );
    if !readback(c, enc, &fo, &stage, (dim * 4) as u64, folded) {
        return false;
    }
    let enc2 = c
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("hc2") });
    if !readback(c, enc2, &ex, &stage, ((hc * dim) * 4) as u64, expanded) {
        return false;
    }
    drop(sc_lock);
    true
}

pub fn adapter_report() -> Vec<String> {
    let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
        backends: wgpu::Backends::all(),
        flags: wgpu::InstanceFlags::default(),
        memory_budget_thresholds: Default::default(),
        backend_options: Default::default(),
        display: None,
    });
    // The [N] prefix is the CMF_GPU_ADAPTER index — the pin that lets a
    // second process take a second card (`run --gpus N`).
    let mut out: Vec<String> =
        pollster::block_on(instance.enumerate_adapters(wgpu::Backends::all()))
            .iter()
            .enumerate()
            .map(|(n, a)| {
                let i = a.get_info();
                let l = a.limits();
                format!(
                    "[{n}] {:?} | {} | {:?} | буфер до {:.1} ГБ | рабочая группа {}",
                    i.backend,
                    i.name,
                    i.device_type,
                    l.max_buffer_size as f64 / 1e9,
                    l.max_compute_workgroup_size_x
                )
            })
            .collect();
    if out.is_empty() {
        out.push("адаптеров не найдено".into());
    }
    match pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
        power_preference: wgpu::PowerPreference::HighPerformance,
        force_fallback_adapter: false,
        compatible_surface: None,
        apply_limit_buckets: false,
    })) {
        Ok(a) => out.push(format!("выбран: {}", a.get_info().name)),
        Err(e) => out.push(format!("выбрать не удалось: {e}")),
    }
    out
}

pub fn adapter_probe() -> bool {
    let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
        backends: wgpu::Backends::all(),
        flags: wgpu::InstanceFlags::default(),
        memory_budget_thresholds: Default::default(),
        backend_options: Default::default(),
        display: None,
    });
    pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
        power_preference: wgpu::PowerPreference::HighPerformance,
        force_fallback_adapter: false,
        compatible_surface: None,
        apply_limit_buckets: false,
    }))
    .is_ok()
}

#[cfg(test)]
mod buffer_ceiling_tests {
    use super::buffers_fit;

    /// The exact numbers that killed a 27B bake mid-phase-A: an A100
    /// caps one storage buffer at 4 292 870 144 bytes and the f32
    /// embedding of a 248 320-token vocabulary at hidden 5120 is
    /// 5 085 593 600. `create_buffer` calls that fatal, so the check has
    /// to happen before the call.
    #[test]
    fn a_large_vocab_head_is_declined_not_fatal() {
        const A100_MAX: u64 = 4_292_870_144;
        let embed_f32 = 248_320u64 * 5120 * 4;
        assert_eq!(embed_f32, 5_085_593_600);
        assert!(!buffers_fit(A100_MAX, &[embed_f32]));
        // The other operands of the same call are fine on their own —
        // one oversized buffer is enough to decline the whole call.
        assert!(!buffers_fit(A100_MAX, &[1024, embed_f32, 4096]));
        assert!(buffers_fit(A100_MAX, &[1024, 4096]));
    }

    /// Exactly at the limit is allowed; one byte over is not.
    #[test]
    fn the_ceiling_is_inclusive() {
        assert!(buffers_fit(100, &[100]));
        assert!(!buffers_fit(100, &[101]));
        assert!(buffers_fit(100, &[]));
    }
}