inferencelayer 0.2.1

//! The Qwen3.5-VL vision tower on wgpu.
//!
//! The CPU tower ([`crate::vision::VisionTower`]) is the correctness reference — bit-exact
//! preprocessing, cosine 1.00000000 against HF at every one of its 24 blocks — and 2× slower than
//! torch-CPU on a real page. This is the same tower on the GPU.
//!
//! ## Where it ended up
//!
//! On a 912-patch page, INTERLEAVED against torch in one window, min-of-N — the only comparison this
//! box supports. It matters WHICH window, and that is the finding, not a caveat:
//!
//! ```text
//!   quiet box (load ~5)          loaded box (load ~24)
//!     torch-CPU     552 ms         torch-CPU    2116 ms
//!     engine-wgpu   149 ms         engine-wgpu   188 ms      <- degrades 1.3x
//!     torch-MPS     145 ms         torch-MPS     370 ms      <- degrades 2.6x
//!     -> dead heat                 -> engine 0.51x, WINS by 2x
//! ```
//!
//! Under contention the engine holds (165-221 ms across every run) while torch-MPS scatters
//! (145-995 ms). One submit of one command buffer does not queue behind whatever else wants the GPU
//! the way ~300 MPSGraph dispatches do. On a serving fleet that is the number that matters.
//!
//! It started at 5.81x behind torch-MPS. Four things moved it, and every one reversed a
//! plausible-sounding guess:
//!
//! 1. **GEMM on the hardware matrix units** — [`crate::encoder::enc_gemm4_src`], 5.35 TF/s against
//!    the best portable scalar kernel's 3.3. The FIRST version was 2x SLOWER than scalar, entirely
//!    because of how it was written.
//!
//! 2. **Attention: no shared-memory staging** — [`vis_attn_src`]. It was 59% of the tower. TWO
//!    textbook fixes (stage K/V tiles, flash-attention style) LOST, because on Metal's 32 KB budget
//!    that leaves one workgroup resident per core. Reuse from REGISTERS instead: 189 ms -> 40.
//!
//! 3. **Attention on the matrix units too** — [`vis_attn_coop_src`], 40 ms -> ~17. Its Q.K^T and P.V
//!    are GEMMs; the only thing standing in the way was that online softmax must RESCALE the
//!    accumulator and a coop matrix cannot be scaled. Keep the running output in SHARED memory and
//!    load it back as the accumulator's initial value.
//!
//! 4. **Shallow k-tiles**: 128x128x8 = 5.35 TF/s, x16 = 4.89, x32 = 1.89. Deeper tiles amortize
//!    barriers and pay in occupancy.
//!
//! Three of those four are the same lesson: **shallow and resident beats deep and alone.** Arithmetic
//! intensity is not free; it is bought with occupancy, and on this hardware occupancy was worth more
//! every single time.
//!
//! ## Portability
//!
//! Cooperative matrix ACCELERATES, it never gates. An adapter without it (`ctx.coop_matrix == false`)
//! runs [`crate::encoder::enc_gemm3_src`], the portable WGSL kernel, and everything still works —
//! that path is exercised, not hypothetical (`OSFKB_ENC_COOP=0`).
//!
//! ## Why most of this is not new code
//!
//! Strip away the 2-D rope and the merger and this ViT is an ordinary pre-LN encoder, so it rides the
//! kernels the text encoders already tuned: [`crate::encoder`]'s GEMM, LayerNorm and residual add.
//! The position-embedding interpolation and the rope tables are computed on the HOST, by calling the
//! CPU tower's own `interpolate_pos` / `rope_tables` — O(tokens) work next to O(tokens × 300M) of
//! GEMM, and one copy of the block-major walk instead of two that have to agree.
//!
//! ## Precision
//!
//! The transposed weight (which v3 and v4 share) is f32: `coop_mat8x8<f32>` is f32-only today. Gated
//! against the CPU tower at cos 1.00000000 on both parity images.

use crate::forward::ShaderModuleTuned as _;
use anyhow::Result;
use wgpu::util::DeviceExt;

use crate::GpuCtx;
use crate::encoder::{EncKernels, act_code, gemm2_tier, gemm2_tile, v2_bgs};
use crate::encoder::{enc_gemm3_src, enc_gemm4_src};
use crate::encoder_weights::Act;
use crate::forward::{make_bg, uni};
use crate::vision::{ImagePatches, VisionConfig, VisionTower};
use crate::weights::f32_to_f16_bytes;

/// `[n, k]` (HF Linear layout) → `[k, n]`. Done ONCE, at upload: it is what lets GEMM v3 read a
/// contiguous run of columns at a fixed k, and so read shared memory as vec4 instead of scalars.
fn transpose(w: &[f32], n: usize, k: usize) -> Vec<f32> {
    let mut t = vec![0f32; n * k];
    for r in 0..n {
        for c in 0..k {
            t[c * n + r] = w[r * k + c];
        }
    }
    t
}

/// The v3 tile. 8 rows × 8 cols per thread (intensity 4), 16-deep k-tiles.
const V3: (usize, usize, usize) = (128, 128, 16);
/// Flash attention on the HARDWARE MATRIX UNITS. Both of its matmuls are GEMMs; this runs them as
/// GEMMs.
///
/// ## The obstacle, and the way round it
///
/// Online softmax must RESCALE the running output whenever a block raises the row maximum
/// (`O ← O·e^{m_old−m_new} + P·V`). A cooperative matrix has no elementwise operations — you cannot
/// scale one, let alone per-row. Every naive port dies here.
///
/// The way through: **the running output lives in SHARED memory, and is loaded back as the P·V
/// accumulator's INITIAL value.** Rescale `osh` with scalar FMAs, then
/// `coopLoad<C>(&osh)` — `coopMultiplyAdd` accumulates on top of it, and `coopStore` writes it back.
/// The rescale costs a few percent; no separate staging array exists at all.
///
/// ## What actually bounds this kernel
///
/// Not the matmuls. K and V are re-read once per QUERY BLOCK, so traffic goes as `T²·hd/RB` — with
/// `RB = 8` that is ~20 GB across the tower, i.e. ~40 ms of pure bandwidth, which is what the scalar
/// kernel was already paying. Widening `RB` is the only lever, and it is CHEAP here in a way it is
/// not for a scalar kernel: a coop accumulator is spread across the subgroup's lanes, so doubling the
/// rows costs ~2 f32/lane rather than 64 registers. `RB = 16` halves the traffic for 12 KB of shared
/// memory.
///
/// ## Shape
///
/// One workgroup = one subgroup (32 lanes) = `RB` query rows of one head. Per 64-key block:
///   * `S[RB][64] = Q[RB][64] · Kᵀ[64][64]`. Kᵀ needs no transpose in memory: a NON-transposed
///     (column-major) `coopLoad` of `K[key][dim]` IS `Kᵀ[dim][key]` — the layout falls out free.
///   * softmax on the RB×64 scores, spread across ALL 32 lanes (4 lanes per row). The first version
///     of this ran it on 8 lanes and left 24 idle.
///   * `O[RB][64] += P[RB][64] · V[64][64]`, accumulating from `osh`.
///
/// Keys past the span read clamped (finite) garbage and are masked to zero probability, so a ragged
/// tail needs no separate kernel.
fn vis_attn_coop_src(rb: usize) -> String {
    let rt = rb / 8; // row tiles
    let sh = rb * 64;
    let lanes_per_row = 32 / rb.min(32);
    let keys_per_lane = 64 / lanes_per_row.max(1);

    let mut s_decl = String::new();
    for i in 0..rt {
        for j in 0..8 {
            s_decl.push_str(&format!(
                "        var s{i}_{j} = coopLoadT<coop_mat8x8<f32, C>>(&zt[0], 8u);\n"
            ));
        }
    }
    let mut qk = String::new();
    for d in 0..8 {
        for i in 0..rt {
            qk.push_str(&format!(
                "        let a{i}_{d} = coopLoadT<coop_mat8x8<f32, A>>(&qsh[{}u * 64u + {}u], 64u);\n",
                i * 8, d * 8
            ));
        }
        for j in 0..8 {
            qk.push_str(&format!(
                "        {{ let b = coopLoad<coop_mat8x8<f32, B>>(&k[(kb + {}u) * hid + hoff + {}u], hid);\n",
                j * 8, d * 8
            ));
            for i in 0..rt {
                qk.push_str(&format!(
                    "          s{i}_{j} = coopMultiplyAdd(a{i}_{d}, b, s{i}_{j});\n"
                ));
            }
            qk.push_str("        }\n");
        }
    }
    let mut s_store = String::new();
    for i in 0..rt {
        for j in 0..8 {
            s_store.push_str(&format!(
                "        coopStoreT(s{i}_{j}, &ssh[{}u * 64u + {}u], 64u);\n",
                i * 8,
                j * 8
            ));
        }
    }
    // O accumulators initialized FROM the (already rescaled) running output.
    let mut o_init = String::new();
    for i in 0..rt {
        for e in 0..8 {
            o_init.push_str(&format!(
                "        var o{i}_{e} = coopLoadT<coop_mat8x8<f32, C>>(&osh[{}u * 64u + {}u], 64u);\n",
                i * 8, e * 8
            ));
        }
    }
    let mut pv = String::new();
    for j in 0..8 {
        for i in 0..rt {
            pv.push_str(&format!(
                "        let pa{i}_{j} = coopLoadT<coop_mat8x8<f32, A>>(&ssh[{}u * 64u + {}u], 64u);\n",
                i * 8, j * 8
            ));
        }
        for e in 0..8 {
            pv.push_str(&format!(
                "        {{ let vb = coopLoadT<coop_mat8x8<f32, B>>(&v[(kb + {}u) * hid + hoff + {}u], hid);\n",
                j * 8, e * 8
            ));
            for i in 0..rt {
                pv.push_str(&format!(
                    "          o{i}_{e} = coopMultiplyAdd(pa{i}_{j}, vb, o{i}_{e});\n"
                ));
            }
            pv.push_str("        }\n");
        }
    }
    let mut o_store = String::new();
    for i in 0..rt {
        for e in 0..8 {
            o_store.push_str(&format!(
                "        coopStoreT(o{i}_{e}, &osh[{}u * 64u + {}u], 64u);\n",
                i * 8,
                e * 8
            ));
        }
    }

    format!(
        r#"enable wgpu_cooperative_matrix;
struct Meta {{ rows: u32, nh: u32, hd: u32, pad: u32 }}
@group(0) @binding(0) var<storage, read>       q: array<f32>;
@group(0) @binding(1) var<storage, read>       k: array<f32>;
@group(0) @binding(2) var<storage, read>       v: array<f32>;
@group(0) @binding(3) var<storage, read_write> o: array<f32>;
@group(0) @binding(4) var<uniform>             mt: Meta;

const RB: u32 = {rb}u;
const LPR: u32 = {lanes_per_row}u;   // lanes per row in the softmax
const KPL: u32 = {keys_per_lane}u;   // keys each of those lanes owns

var<workgroup> qsh: array<f32, {sh}>;   // [RB][64] queries
var<workgroup> ssh: array<f32, {sh}>;   // [RB][64] scores, then probabilities
var<workgroup> osh: array<f32, {sh}>;   // [RB][64] running output — NOT a coop matrix, on purpose
var<workgroup> zt:  array<f32, 64>;
var<workgroup> red: array<f32, 32>;
var<workgroup> rmax: array<f32, {rb}>;
var<workgroup> rsum: array<f32, {rb}>;
var<workgroup> rsc:  array<f32, {rb}>;

@compute @workgroup_size(32)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) t: u32) {{
    let head = wg.y;
    let hid = mt.nh * 64u;
    let hoff = head * 64u;
    let r0 = wg.x * RB;
    let ri = t / LPR;          // this lane's row
    let li = t % LPR;          // ...and its slice of the 64 keys

    if (t < 64u) {{ zt[t] = 0.0; }}
    for (var i = t; i < RB * 64u; i += 32u) {{
        let rr = i / 64u;
        let dd = i % 64u;
        let gr = r0 + rr;
        var val = 0.0;
        if (gr < mt.rows) {{ val = q[gr * hid + hoff + dd]; }}
        qsh[i] = val;
        osh[i] = 0.0;
    }}
    for (var i = t; i < RB; i += 32u) {{ rmax[i] = -3.0e38; rsum[i] = 0.0; }}
    workgroupBarrier();

    let scale = 1.0 / sqrt(64.0);
    var kb = 0u;
    loop {{
        if (kb >= mt.rows) {{ break; }}
        // ---- S = Q · Kᵀ, on the matrix units ----
{s_decl}
{qk}
{s_store}
        workgroupBarrier();

        // ---- softmax, across ALL 32 lanes (LPR lanes per row, KPL keys each) ----
        {{
            var lmax = -3.0e38;
            for (var j = li * KPL; j < li * KPL + KPL; j++) {{
                if (kb + j < mt.rows) {{ lmax = max(lmax, ssh[ri * 64u + j] * scale); }}
            }}
            red[t] = lmax;
            workgroupBarrier();
            if (li == 0u) {{
                var mx = -3.0e38;
                for (var l = 0u; l < LPR; l++) {{ mx = max(mx, red[t + l]); }}
                let nm = max(rmax[ri], mx);
                rsc[ri] = exp(rmax[ri] - nm);
                rmax[ri] = nm;
            }}
            workgroupBarrier();
            var lsum = 0.0;
            for (var j = li * KPL; j < li * KPL + KPL; j++) {{
                var pr = 0.0;
                if (kb + j < mt.rows) {{ pr = exp(ssh[ri * 64u + j] * scale - rmax[ri]); }}
                ssh[ri * 64u + j] = pr;
                lsum += pr;
            }}
            red[t] = lsum;
            workgroupBarrier();
            if (li == 0u) {{
                var sm = 0.0;
                for (var l = 0u; l < LPR; l++) {{ sm += red[t + l]; }}
                rsum[ri] = rsum[ri] * rsc[ri] + sm;
            }}
            workgroupBarrier();
            // Rescale the running output IN SHARED — this is the step a coop matrix cannot do.
            for (var i = t; i < RB * 64u; i += 32u) {{ osh[i] *= rsc[i / 64u]; }}
            workgroupBarrier();
        }}

        // ---- O += P · V, accumulating FROM osh (no staging array) ----
{o_init}
{pv}
{o_store}
        workgroupBarrier();
        kb += 64u;
    }}

    for (var i = t; i < RB * 64u; i += 32u) {{
        let rr = i / 64u;
        let gr = r0 + rr;
        if (gr < mt.rows && rsum[rr] > 0.0) {{
            o[gr * hid + hoff + (i % 64u)] = osh[i] / rsum[rr];
        }}
    }}
}}
"#
    )
}

/// Query rows per attention workgroup — the K/V traffic divisor. Swept: see the kernel doc.
/// Query rows per attention workgroup — the K/V traffic divisor, and the only knob that matters here.
///
/// Swept on the 912-patch page (interleaved, min of 7):
///
/// ```text
///   RB     shared    tower
///    2     1.5 KB    218.5 ms
///    4     3.0 KB    188.4
///    6     4.5 KB    178.0
///    8     6.0 KB    174.0   <- the knee
///   12     9.0 KB    187.4
///   16    12.0 KB    188.6
/// ```
///
/// It turns back UP past 8: more rows per workgroup means more live registers and more shared memory,
/// and occupancy starts paying for the traffic it saves. Same trade that made the shared-memory-staged
/// versions lose outright.
/// Query rows per COOP attention workgroup. Traffic goes as 1/RB and coop accumulators are
/// lane-distributed, so this can go wider than the scalar kernel's 8 for very little.
fn coop_attn_rb() -> usize {
    std::env::var("OSFKB_COOP_ATTN_RB")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(16)
}

fn attn_rb() -> usize {
    std::env::var("OSFKB_ATTN_RB")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(8)
}

/// The v4 (cooperative-matrix) tile — the SAME 128×128 as v3. Handing the matrix units a smaller
/// tile than the scalar kernel is exactly how the first version of this ended up 2× slower.
/// 128×128 output, **8-deep k-tiles**. Swept on the tower's fc1 (912×4096×1024):
///
/// ```text
///   128×128×8    5.35 TF/s   <- 8 KB shared, 4 workgroups resident
///   128×128×16   4.89        16 KB
///   128×128×32   1.89        32 KB — at Metal's limit, one workgroup resident
///   128×256×16   2.69
///   256×128×16   2.68
/// ```
///
/// Deeper k-tiles amortize barriers and lose to occupancy — the same trade that decides the attention
/// kernel. Shallow and resident beats deep and alone.
const V4: (usize, usize, usize) = (128, 128, 8);

/// Split the packed qkv projection into q/k/v, applying the 2-D rope to q and k.
///
/// The rope tables arrive precomputed (`cos`/`sin`, `[tokens, head_dim]`) from the CPU tower, so this
/// kernel is pure data movement plus one rotation — no trig, and no second implementation of the
/// block-major (row, col) walk that produced them.
///
/// `rotate_half`: element `d` pairs with `d ± hd/2`, negated on the low half. Writing it as a select
/// rather than a branch keeps the two halves' control flow uniform.
const VIS_QKV_ROPE: &str = r#"
struct Meta { rows: u32, nh: u32, hd: u32, pad: u32 }
@group(0) @binding(0) var<storage, read>       qkv: array<f32>;  // [rows, 3*hid]
@group(0) @binding(1) var<storage, read>       cs:  array<f32>;  // [rows, hd]
@group(0) @binding(2) var<storage, read>       sn:  array<f32>;  // [rows, hd]
@group(0) @binding(3) var<storage, read_write> q:   array<f32>;  // [rows, hid]
@group(0) @binding(4) var<storage, read_write> k:   array<f32>;  // [rows, hid]
@group(0) @binding(5) var<storage, read_write> v:   array<f32>;  // [rows, hid]
@group(0) @binding(6) var<uniform>             mt:  Meta;
@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let hid = mt.nh * mt.hd;
    let row = gid.y;
    let i = gid.x;
    if (row >= mt.rows || i >= hid) { return; }

    let head = i / mt.hd;
    let d = i % mt.hd;
    let half = mt.hd / 2u;
    let base = row * 3u * hid + head * mt.hd;

    let c = cs[row * mt.hd + d];
    let s = sn[row * mt.hd + d];

    // rotate_half(x)[d] = -x[d + half] if d < half else x[d - half]
    let partner = select(d - half, d + half, d < half);
    let sign = select(1.0, -1.0, d < half);

    let qv = qkv[base + d];
    let qr = qkv[base + partner] * sign;
    q[row * hid + i] = qv * c + qr * s;

    let kv = qkv[base + hid + d];
    let kr = qkv[base + hid + partner] * sign;
    k[row * hid + i] = kv * c + kr * s;

    v[row * hid + i] = qkv[base + 2u * hid + d];   // V is not roped
}
"#;

/// Tiled bidirectional attention — 8 query rows share each staged K/V chunk.
///
/// The encoder's `ENC_ATTN` runs one workgroup per (row, head) and each one re-streams the WHOLE K/V
/// span. At 912 patches × 16 heads that is 14592 workgroups every layer, each reading 912 keys — the
/// span is re-read 912 times per layer. It never mattered for text encoders (64-token spans) and it
/// is the single largest cost in this tower: with the GEMM at 3.3 TF/s the matmuls are only ~165 ms
/// of it, and attention is most of the rest.
///
/// This is the decoder's `ATTN_TILED` discipline applied to the encoder: a workgroup owns 8 query
/// rows of one head, stages 32 keys of K and V into shared memory ONCE, and all 8 rows consume them.
/// K/V traffic drops 8×. Flash-style online softmax (running max/denominator, rescaled accumulator),
/// so nothing is materialized at [T, T].
///
/// Threads: 256 = 8 rows × 32 lanes. During scoring a lane owns one KEY of the chunk; during the
/// value accumulation it owns two output DIMS (hd = 64). Shared: q 2 KB + K 8 KB + V 8 KB + 2 KB of
/// reduction scratch = 20 KB.
/// Multi-row bidirectional attention: one workgroup owns `RB` query rows of one head, and every
/// `k[j]`/`v[j]` it reads feeds all `RB` of them.
///
/// `ENC_ATTN` runs one workgroup per (row, head) and re-streams the WHOLE K/V span in each — at 912
/// patches × 16 heads that is 14592 workgroups per layer each reading 912 keys, and it was 59% of the
/// tower (189 ms of 323). Never mattered for text encoders with 64-token spans; dominates here.
///
/// Two earlier attempts staged K/V tiles in shared memory — the textbook flash-attention move — and
/// BOTH LOST, at 12 KB and 20 KB of shared memory respectively. On Metal's 32 KB budget that leaves
/// one or two workgroups resident per core, and the latency stops being hidden. The traffic was 8×
/// lower and it did not matter.
///
/// So: no staging. Keep the shared footprint at ~3 KB per RB=4 (only q, the probabilities, and the
/// reduction scratch) and get the reuse from REGISTERS instead — one `k[j]` read, `RB` dot products
/// accumulated. Traffic falls RB-fold with occupancy untouched. Measured: 1.70× on the whole tower at
/// RB=4.
///
/// All RB rows share ONE reduction tree (6 barriers per 64-key tile, not 6 per row).
fn vis_attn_src(rb: usize) -> String {
    let n = |i: usize, f: &str| format!("{f}{i}");
    let decl = |f: &str, init: &str| {
        (0..rb)
            .map(|i| format!("    var {} = {init};", n(i, f)))
            .collect::<Vec<_>>()
            .join("\n")
    };
    let dots = (0..rb)
        .map(|i| format!("                d{i} += qsh[{i}u * HD + d] * kv;"))
        .collect::<Vec<_>>()
        .join("\n");
    let scores = (0..rb)
        .map(|i| format!("            s{i} = d{i} * scale;"))
        .collect::<Vec<_>>()
        .join("\n");
    let wr = |src: &str, dst: &str| {
        (0..rb)
            .map(|i| format!("        {dst}[{i}u * 64u + t] = {src}{i};"))
            .collect::<Vec<_>>()
            .join("\n")
    };
    let red_max = (0..rb)
        .map(|i| format!("                red[{i}u * 64u + t] = max(red[{i}u * 64u + t], red[{i}u * 64u + t + s]);"))
        .collect::<Vec<_>>()
        .join("\n");
    let red_sum = (0..rb)
        .map(|i| format!("                red[{i}u * 64u + t] += red[{i}u * 64u + t + s];"))
        .collect::<Vec<_>>()
        .join("\n");
    let newm = (0..rb)
        .map(|i| format!("        let nm{i} = max(m{i}, red[{}u]);", i * 64))
        .collect::<Vec<_>>()
        .join("\n");
    let resc = (0..rb)
        .map(|i| format!("        let rs{i} = exp(m{i} - nm{i});"))
        .collect::<Vec<_>>()
        .join("\n");
    let probs = (0..rb)
        .map(|i| format!("        if (psh[{i}u * 64u + t] > -3.0e37) {{ p{i} = exp(psh[{i}u * 64u + t] - nm{i}); }}"))
        .collect::<Vec<_>>()
        .join("\n");
    let lupd = (0..rb)
        .map(|i| {
            format!(
                "        l{i} = l{i} * rs{i} + red[{}u];  a{i} *= rs{i};",
                i * 64
            )
        })
        .collect::<Vec<_>>()
        .join("\n");
    let vacc = (0..rb)
        .map(|i| format!("            a{i} += psh[{i}u * 64u + jj] * vv;"))
        .collect::<Vec<_>>()
        .join("\n");
    let mupd = (0..rb)
        .map(|i| format!("        m{i} = nm{i};"))
        .collect::<Vec<_>>()
        .join("\n");
    let store = (0..rb)
        .map(|i| format!("    if (r0 + {i}u < mt.rows && l{i} > 0.0) {{ o[(r0 + {i}u) * hid + ob] = a{i} / l{i}; }}"))
        .collect::<Vec<_>>()
        .join("\n");
    let sh = rb * 64;
    let qsh_len = rb * 64;
    let m_decl = decl("m", "-3.0e38");
    let l_decl = decl("l", "0.0");
    let a_decl = decl("a", "0.0");
    let s_decl = decl("s", "-3.0e38");
    let d_decl = (0..rb)
        .map(|i| format!("            var d{i} = 0.0;"))
        .collect::<Vec<_>>()
        .join("\n");
    let p_decl = (0..rb)
        .map(|i| format!("        var p{i} = 0.0;"))
        .collect::<Vec<_>>()
        .join("\n");
    let psh_s = wr("s", "psh");
    let red_s = wr("s", "red");
    let psh_p = wr("p", "psh");
    let red_p = wr("p", "red");

    format!(
        r#"
struct Meta {{ rows: u32, nh: u32, hd: u32, pad: u32 }}
@group(0) @binding(0) var<storage, read>       q: array<f32>;
@group(0) @binding(1) var<storage, read>       k: array<f32>;
@group(0) @binding(2) var<storage, read>       v: array<f32>;
@group(0) @binding(3) var<storage, read_write> o: array<f32>;
@group(0) @binding(4) var<uniform>             mt: Meta;
const RB: u32 = {rb}u;
const HD: u32 = 64u;
var<workgroup> qsh: array<f32, {qsh_len}>;
var<workgroup> psh: array<f32, {sh}>;
var<workgroup> red: array<f32, {sh}>;
@compute @workgroup_size(64)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) t: u32) {{
    let head = wg.y;
    let hid = mt.nh * HD;
    let r0 = wg.x * RB;
    for (var i = t; i < RB * HD; i += 64u) {{
        let rr = i / HD;
        let dd = i % HD;
        let gr = r0 + rr;
        var val = 0.0;
        if (gr < mt.rows) {{ val = q[gr * hid + head * HD + dd]; }}
        qsh[i] = val;
    }}
    workgroupBarrier();
    let scale = 1.0 / sqrt(f32(HD));
{m_decl}
{l_decl}
{a_decl}
    let ntiles = (mt.rows + 63u) / 64u;
    for (var tile = 0u; tile < ntiles; tile++) {{
        let j = tile * 64u + t;
{s_decl}
        if (j < mt.rows) {{
{d_decl}
            let kb = j * hid + head * HD;
            for (var d = 0u; d < HD; d++) {{
                let kv = k[kb + d];
{dots}
            }}
{scores}
        }}
{psh_s}
{red_s}
        workgroupBarrier();
        for (var s = 32u; s > 0u; s >>= 1u) {{
            if (t < s) {{
{red_max}
            }}
            workgroupBarrier();
        }}
{newm}
{resc}
        workgroupBarrier();
{p_decl}
{probs}
        workgroupBarrier();
{psh_p}
{red_p}
        workgroupBarrier();
        for (var s = 32u; s > 0u; s >>= 1u) {{
            if (t < s) {{
{red_sum}
            }}
            workgroupBarrier();
        }}
{lupd}
        let tlen = min(64u, mt.rows - tile * 64u);
        for (var jj = 0u; jj < tlen; jj++) {{
            let vv = v[(tile * 64u + jj) * hid + head * HD + t];
{vacc}
        }}
{mupd}
        workgroupBarrier();
    }}
    let ob = head * HD + t;
{store}
}}
"#
    )
}

/// Elementwise activation over a GEMM's output.
///
/// GEMM v4 cannot fuse an activation — a cooperative matrix has no elementwise operations worth the
/// name yet — so the activated GEMMs (the MLP's fc1, the merger's fc1) pay one extra pass over their
/// output. That is ~0.08 ms against a 1.8 ms GEMM: 4%, against v4's 1.4x. Worth it.
const ENC_ACT: &str = r#"
struct Meta { n: u32, act: u32, p0: u32, p1: u32 }
@group(0) @binding(0) var<storage, read_write> y: array<f32>;
@group(0) @binding(1) var<uniform> mt: Meta;
fn erf_as(x: f32) -> f32 {
    let s = sign(x);
    let a = abs(x);
    let t = 1.0 / (1.0 + 0.3275911 * a);
    let y2 = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t
        + 0.254829592) * t * exp(-a * a);
    return s * y2;
}
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    if (g.x >= mt.n) { return; }
    let v = y[g.x];
    switch mt.act {
        case 1u: { y[g.x] = 0.5 * v * (1.0 + erf_as(v * 0.70710678)); }
        // Clamped, exactly as ACT_FNS is: WGSL tanh overflows to NaN for |arg| >~ 88, and gelu's
        // argument reaches that at a pre-activation of only 13.8 — which a ViT hits on block one.
        case 2u: { let a = clamp(0.7978845608 * (v + 0.044715 * v * v * v), -20.0, 20.0);
                   y[g.x] = 0.5 * v * (1.0 + tanh(a)); }
        case 3u: { y[g.x] = v / (1.0 + exp(-v)); }
        case 4u: { y[g.x] = tanh(v); }
        case 5u: { y[g.x] = max(v, 0.0); }
        default: {}
    }
}
"#;

enum Mat {
    F16(wgpu::Buffer),
    F32(wgpu::Buffer),
}

impl Mat {
    fn buf(&self) -> &wgpu::Buffer {
        match self {
            Mat::F16(b) | Mat::F32(b) => b,
        }
    }
}

struct GpuLinear {
    w: Mat,
    /// The same weight as `[k, n]`, for GEMM v3 and v4.
    wt: Mat,
    b: wgpu::Buffer,
    /// The bias broadcast to `[8, n]`. GEMM v4 loads its accumulator's INITIAL value from this, so
    /// the bias is already in the result and the tile can be stored straight to `y` — no output
    /// staging array, which is what was crushing occupancy.
    b8: wgpu::Buffer,
    n: u32,
    k: u32,
}

struct GpuNorm {
    w: wgpu::Buffer,
    b: wgpu::Buffer,
}

struct GpuBlock {
    n1: GpuNorm,
    qkv: GpuLinear,
    proj: GpuLinear,
    n2: GpuNorm,
    fc1: GpuLinear,
    fc2: GpuLinear,
}

/// The vision tower, resident on the GPU. Holds the CPU tower alongside it — for the host-side
/// position/rope tables (see the module note), and because it is the oracle the parity gate compares
/// against.
pub struct VisionGpu {
    cpu: VisionTower,
    kernels: EncKernels,
    rope_pl: wgpu::ComputePipeline,
    /// Tiled attention (8 query rows share each staged K/V chunk). `OSFKB_VIS_ATTN=0` runs the
    /// encoder's per-row `ENC_ATTN` instead — the A/B the measurement used.
    attn_pl: wgpu::ComputePipeline,
    /// Attention on the matrix units — only where the adapter has them.
    attn_coop_pl: Option<wgpu::ComputePipeline>,
    tiled_attn: bool,
    attn_rb: usize,
    gemm3_pl: wgpu::ComputePipeline,
    /// The cooperative-matrix kernel — present only when the adapter has hardware matrix units.
    gemm4_pl: Option<wgpu::ComputePipeline>,
    /// v4's un-fusable activation, as its own pass.
    act_pl: wgpu::ComputePipeline,
    /// `OSFKB_ENC_GEMM3=0` runs the shared v2 kernel instead — the A/B switch the measurement used.
    v3: bool,
    f16: bool,
    patch: GpuLinear,
    blocks: Vec<GpuBlock>,
    m_norm: GpuNorm,
    m_fc1: GpuLinear,
    m_fc2: GpuLinear,
}

impl VisionGpu {
    pub fn config(&self) -> &VisionConfig {
        self.cpu.config()
    }

    /// Move a loaded CPU tower onto the GPU.
    pub fn new(ctx: &GpuCtx, cpu: VisionTower) -> Result<Self> {
        let f16 = std::env::var("OSFKB_VISION_F32").is_err();
        let lin = |l: &crate::vision::Linear| -> GpuLinear {
            let w = if f16 {
                Mat::F16(ctx.storage_bytes(&f32_to_f16_bytes(&l.w)))
            } else {
                Mat::F32(ctx.storage(&l.w))
            };
            // The TRANSPOSED copy is always f32. The cooperative-matrix API is `coop_mat8x8<f32>` —
            // f32 only, today — so v4 must read f32, and v3 shares the same buffer rather than
            // holding a second f16 copy of a 330M-parameter tower.
            let t = transpose(&l.w, l.n, l.k);
            let wt = Mat::F32(ctx.storage(&t));
            let bias = l.b.as_deref().unwrap_or(&[]);
            let mut b8 = Vec::with_capacity(8 * l.n);
            for _ in 0..8 {
                b8.extend_from_slice(bias);
            }
            GpuLinear {
                w,
                wt,
                b: ctx.storage(bias),
                b8: ctx.storage(&b8),
                n: l.n as u32,
                k: l.k as u32,
            }
        };
        let norm = |n: &crate::vision::Norm| GpuNorm {
            w: ctx.storage(&n.w),
            b: ctx.storage(&n.b),
        };

        let blocks = cpu
            .blocks
            .iter()
            .map(|b| GpuBlock {
                n1: norm(&b.norm1),
                qkv: lin(&b.qkv),
                proj: lin(&b.proj),
                n2: norm(&b.norm2),
                fc1: lin(&b.fc1),
                fc2: lin(&b.fc2),
            })
            .collect();

        let rope_pl = ctx
            .device
            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("vis_qkv_rope"),
                layout: None,
                module: &ctx
                    .device
                    .shader_module_tuned(wgpu::ShaderModuleDescriptor {
                        label: Some("vis_qkv_rope"),
                        source: wgpu::ShaderSource::Wgsl(VIS_QKV_ROPE.into()),
                    }),
                entry_point: Some("main"),
                compilation_options: Default::default(),
                cache: None,
            });

        let gemm3_pl = ctx
            .device
            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("enc_gemm3"),
                layout: None,
                module: &ctx
                    .device
                    .shader_module_tuned(wgpu::ShaderModuleDescriptor {
                        label: Some("enc_gemm3"),
                        // f16 = false: v3 reads the same f32 transposed weight v4 needs.
                        source: wgpu::ShaderSource::Wgsl(
                            enc_gemm3_src(false, V3.0, V3.1, V3.2).into(),
                        ),
                    }),
                entry_point: Some("main"),
                compilation_options: Default::default(),
                cache: None,
            });
        let v3 = !matches!(std::env::var("OSFKB_ENC_GEMM3").ok().as_deref(), Some("0"));
        let attn_pl = ctx
            .device
            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("vis_attn_tiled"),
                layout: None,
                module: &ctx
                    .device
                    .shader_module_tuned(wgpu::ShaderModuleDescriptor {
                        label: Some("vis_attn_tiled"),
                        source: wgpu::ShaderSource::Wgsl(vis_attn_src(attn_rb()).into()),
                    }),
                entry_point: Some("main"),
                compilation_options: Default::default(),
                cache: None,
            });
        // Also OFF by default, also measured: 8 rows sharing a staged K/V chunk cuts K/V traffic 8x
        // and is still SLOWER (951 ms vs 871), because it pays ~14 workgroup barriers per 32-key
        // chunk x 29 chunks to get there. Attention here is not bandwidth-bound, it is barrier-bound,
        // and the fix is subgroup reductions (no barriers) — which cannot assume a 32-lane subgroup
        // on Metal (AdapterInfo reports 4..64). Kept, correct, and honest about not paying off yet.
        let tiled_attn = !matches!(std::env::var("OSFKB_VIS_ATTN").ok().as_deref(), Some("0"));
        // `OSFKB_VIS_ATTN_COOP=0` forces the scalar kernel — the A/B the measurement uses, and the
        // path every adapter without matrix units takes anyway.
        let attn_coop_pl = (ctx.coop_matrix
            && !matches!(
                std::env::var("OSFKB_VIS_ATTN_COOP").ok().as_deref(),
                Some("0")
            ))
        .then(|| {
            ctx.device
                .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                    label: Some("vis_attn_coop"),
                    layout: None,
                    module: &ctx
                        .device
                        .shader_module_tuned(wgpu::ShaderModuleDescriptor {
                            label: Some("vis_attn_coop"),
                            source: wgpu::ShaderSource::Wgsl(
                                vis_attn_coop_src(coop_attn_rb()).into(),
                            ),
                        }),
                    entry_point: Some("main"),
                    compilation_options: Default::default(),
                    cache: None,
                })
        });
        let act_pl = ctx
            .device
            .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("enc_act"),
                layout: None,
                module: &ctx
                    .device
                    .shader_module_tuned(wgpu::ShaderModuleDescriptor {
                        label: Some("enc_act"),
                        source: wgpu::ShaderSource::Wgsl(ENC_ACT.into()),
                    }),
                entry_point: Some("main"),
                compilation_options: Default::default(),
                cache: None,
            });

        // Hardware matrix units, when the adapter has them. `OSFKB_ENC_COOP=0` forces the WGSL path —
        // which is also what every adapter WITHOUT them runs, so the fallback is exercised, not
        // hypothetical.
        // ON wherever the hardware has it. It runs the tower's GEMMs at 4.35 TF/s against v3's
        // 3.04 — and the first version of this kernel was 2x SLOWER than v3, which was entirely the
        // kernel's fault and worth recording: it staged its output tile in 16 KB of shared memory
        // (one workgroup resident per core, no latency hiding) and used a 64x64 tile with HALF the
        // arithmetic intensity of v3's 128x128. Give the matrix units the same tile as the scalar
        // kernel and fold the bias into the accumulator's initial value so nothing has to be staged,
        // and they win. `OSFKB_ENC_COOP=0` forces the WGSL path — which is also what every adapter
        // WITHOUT matrix units runs, so the fallback is exercised, not hypothetical.
        let use_coop = ctx.coop_matrix
            && !matches!(std::env::var("OSFKB_ENC_COOP").ok().as_deref(), Some("0"));
        let gemm4_pl = use_coop.then(|| {
            ctx.device
                .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                    label: Some("enc_gemm4_coop"),
                    layout: None,
                    module: &ctx
                        .device
                        .shader_module_tuned(wgpu::ShaderModuleDescriptor {
                            label: Some("enc_gemm4_coop"),
                            source: wgpu::ShaderSource::Wgsl(
                                enc_gemm4_src(V4.0, V4.1, V4.2).into(),
                            ),
                        }),
                    entry_point: Some("main"),
                    compilation_options: Default::default(),
                    cache: None,
                })
        });

        Ok(Self {
            act_pl,
            attn_pl,
            attn_coop_pl,
            attn_rb: attn_rb(),
            tiled_attn,
            gemm4_pl,
            gemm3_pl,
            v3,
            patch: lin(&cpu.patch),
            m_norm: norm(&cpu.merger_norm),
            m_fc1: lin(&cpu.merger_fc1),
            m_fc2: lin(&cpu.merger_fc2),
            blocks,
            kernels: EncKernels::new(ctx),
            rope_pl,
            f16,
            cpu,
        })
    }

    /// Run the tower over one or more preprocessed images. Returns `[total_merged_tokens,
    /// out_hidden]`, images concatenated in order — the same contract as the CPU tower.
    pub fn forward(&self, ctx: &GpuCtx, images: &[ImagePatches]) -> Result<Vec<f32>> {
        let mut out = Vec::new();
        for img in images {
            out.extend(self.forward_one(ctx, img)?);
        }
        Ok(out)
    }

    fn forward_one(&self, ctx: &GpuCtx, img: &ImagePatches) -> Result<Vec<f32>> {
        let cfg = self.config().clone();
        let (hid, heads, hd) = (cfg.hidden, cfg.heads, cfg.head_dim());
        let n = img.num_patches();
        let unit = cfg.merge_unit();
        anyhow::ensure!(
            n.is_multiple_of(unit),
            "patch count {n} is not a multiple of merge² ({unit})"
        );
        let tokens = n / unit;

        // ---- host tables (shared with the CPU tower — see the module note) ----
        let pos = self.cpu.interpolate_pos(img.grid);
        let (cos, sin) = self.cpu.rope_tables(img.grid);

        // ---- buffers ----
        let b_patch = ctx.storage(&img.patches);
        let b_pos = ctx.storage(&pos);
        let b_cos = ctx.storage(&cos);
        let b_sin = ctx.storage(&sin);
        let x = ctx.empty(n * hid);
        let normed = ctx.empty(n * hid);
        let qkv = ctx.empty(n * 3 * hid);
        let bq = ctx.empty(n * hid);
        let bk = ctx.empty(n * hid);
        let bv = ctx.empty(n * hid);
        let attn = ctx.empty(n * hid);
        let mid = ctx.empty(n * cfg.intermediate);
        let m1 = ctx.empty(tokens * hid * unit);
        let m2 = ctx.empty(tokens * cfg.out_hidden);

        // Attention over ONE image: a single ragged segment covering every patch. This is exactly
        // what `ENC_ATTN` was built for, so the ViT's "attend within an image, never across" needs no
        // new kernel — and batching several images later is just more segments.
        let b_starts = ctx.storage_u32(&[0u32, n as u32]);
        let b_seq_of = ctx.storage_u32(&vec![0u32; n]);
        let b_valid = ctx.storage_u32(&vec![1u32; n]);

        // ---- bind groups ----
        // Built per call because the buffers are sized by this image. That costs microseconds against
        // a forward measured in tens of milliseconds; caching them per patch-count would be an
        // optimization looking for a workload where image sizes repeat.
        let gemm = |x: &wgpu::Buffer,
                    l: &GpuLinear,
                    y: &wgpu::Buffer,
                    m: usize,
                    act: Option<Act>|
         -> Result<GemmStep> {
            let flags = 1u32 | (act_code(act) << 8); // bit 0 = bias present
            let meta = uni(ctx, bytemuck::cast_slice(&[m as u32, l.n, l.k, flags]));
            let bufs = [x, l.w.buf(), &l.b, y];
            // v3 needs the TRANSPOSED weight and its own (auto-layout) bind group — wgpu bind groups
            // are exclusive to the pipeline that created them.
            let bufs3 = [x, l.wt.buf(), &l.b, y];
            let bg3 = make_bg(ctx, &self.gemm3_pl, &bufs3, &meta);
            // v4 takes the SAME [k, n] weight v3 does (the transpose is already paid for) but the
            // BROADCAST bias, which it folds into the accumulator's initial value. It also needs
            // m and n to be multiples of 8: `coopStore` does not bounds-check, so a partial 8×8 block
            // would write outside `y`. Every shape in this tower satisfies it (912 = 114·8); anything
            // that does not silently falls back to v3.
            let coop_ok = m.is_multiple_of(8) && (l.n as usize).is_multiple_of(8);
            let bufs4 = [x, l.wt.buf(), &l.b8, y];
            let bg4 = self
                .gemm4_pl
                .as_ref()
                .filter(|_| coop_ok)
                .map(|pl| make_bg(ctx, pl, &bufs4, &meta));
            // v4 cannot fuse the activation; give it its own pass when there is one.
            let act = act.map(|a| {
                let am = uni(
                    ctx,
                    bytemuck::cast_slice(&[
                        (m * l.n as usize) as u32,
                        act_code(Some(a)),
                        0u32,
                        0u32,
                    ]),
                );
                let abg = make_bg(ctx, &self.act_pl, &[y], &am);
                ActStep {
                    bg: abg,
                    groups: ((m * l.n as usize) as u32).div_ceil(256),
                    _meta: am,
                }
            });
            Ok(GemmStep {
                bgs: v2_bgs(ctx, &self.kernels, self.f16, &bufs, &meta)?,
                bg3,
                bg4,
                act,
                n: l.n as usize,
                m,
                _meta: meta,
            })
        };

        let ln = |src: &wgpu::Buffer, nm: &GpuNorm, dst: &wgpu::Buffer| -> LnStep {
            // flags: bit1 = bias present. No residual fold (bit0): a pre-LN block needs the residual
            // to LAND in `x`, not merely to be added on the way into the norm.
            let meta = uni(
                ctx,
                bytemuck::cast_slice(&[hid as u32, 2u32, cfg.eps.to_bits(), 0u32]),
            );
            let bg = make_bg(
                ctx,
                self.kernels.ln_pl(),
                &[src, src, &nm.w, &nm.b, dst],
                &meta,
            );
            LnStep { bg, _meta: meta }
        };

        let add = |dst: &wgpu::Buffer, src: &wgpu::Buffer, len: usize| -> AddStep {
            let meta = uni(ctx, bytemuck::cast_slice(&[len as u32, 0u32, 0u32, 0u32]));
            let bg = make_bg(ctx, self.kernels.add_pl(), &[dst, src], &meta);
            AddStep {
                bg,
                groups: (len as u32).div_ceil(256),
                _meta: meta,
            }
        };

        // patch embed + interpolated positions
        let s_patch = gemm(&b_patch, &self.patch, &x, n, None)?;
        let s_pos = add(&x, &b_pos, n * hid);

        struct Blk {
            ln1: LnStep,
            qkv: GemmStep,
            rope: (wgpu::BindGroup, wgpu::Buffer),
            attn: (wgpu::BindGroup, wgpu::Buffer),
            tattn: (wgpu::BindGroup, wgpu::Buffer),
            cattn: Option<wgpu::BindGroup>,
            proj: GemmStep,
            add1: AddStep,
            ln2: LnStep,
            fc1: GemmStep,
            fc2: GemmStep,
            add2: AddStep,
        }
        let mut steps: Vec<Blk> = Vec::with_capacity(cfg.depth);
        for b in &self.blocks {
            let rope_meta = uni(
                ctx,
                bytemuck::cast_slice(&[n as u32, heads as u32, hd as u32, 0u32]),
            );
            let rope_bg = make_bg(
                ctx,
                &self.rope_pl,
                &[&qkv, &b_cos, &b_sin, &bq, &bk, &bv],
                &rope_meta,
            );
            // mode 0 = bidirectional, window 0 = unrestricted, n_kv_heads = n_heads (no GQA).
            let attn_meta = uni(
                ctx,
                bytemuck::cast_slice(&[
                    n as u32,
                    heads as u32,
                    hd as u32,
                    0u32,
                    0u32,
                    heads as u32,
                    0u32,
                    0u32,
                ]),
            );
            let attn_bg = make_bg(
                ctx,
                self.kernels.attn_pl(),
                &[&bq, &bk, &bv, &attn, &b_starts, &b_seq_of, &b_valid],
                &attn_meta,
            );
            // The tiled kernel needs no seq tables: one image is one span.
            let tmeta = uni(
                ctx,
                bytemuck::cast_slice(&[n as u32, heads as u32, hd as u32, 0u32]),
            );
            let tattn_bg = make_bg(ctx, &self.attn_pl, &[&bq, &bk, &bv, &attn], &tmeta);
            // wgpu auto-layout bind groups are EXCLUSIVE to the pipeline that created them, so the
            // coop kernel needs its own even though the bindings are identical.
            let cattn_bg = self
                .attn_coop_pl
                .as_ref()
                .map(|pl| make_bg(ctx, pl, &[&bq, &bk, &bv, &attn], &tmeta));
            steps.push(Blk {
                ln1: ln(&x, &b.n1, &normed),
                qkv: gemm(&normed, &b.qkv, &qkv, n, None)?,
                rope: (rope_bg, rope_meta),
                attn: (attn_bg, attn_meta),
                tattn: (tattn_bg, tmeta),
                cattn: cattn_bg,
                proj: gemm(&attn, &b.proj, &normed, n, None)?,
                add1: add(&x, &normed, n * hid),
                ln2: ln(&x, &b.n2, &normed),
                // The blocks' MLP activation folds into the GEMM epilogue — gelu_pytorch_tanh here,
                // and note the merger below uses the OTHER gelu.
                fc1: gemm(&normed, &b.fc1, &mid, n, Some(cfg.act))?,
                fc2: gemm(&mid, &b.fc2, &normed, n, None)?,
                add2: add(&x, &normed, n * hid),
            });
        }

        // Merger. The LN is per PATCH (width `hidden`); the concat of 4 consecutive patches is a
        // no-op on the buffer, because `[n, hidden]` row-major IS `[n/4, hidden·4]` row-major — the
        // whole reason the patches are laid out block-major in the first place. So the merger's first
        // GEMM simply reads the same bytes with m = n/4 and k = hidden·4.
        let s_mln = ln(&x, &self.m_norm, &normed);
        let s_m1 = gemm(&normed, &self.m_fc1, &m1, tokens, Some(Act::GeluErf))?;
        let s_m2 = gemm(&m1, &self.m_fc2, &m2, tokens, None)?;

        // ---- one encoder, one pass, one submit ----
        let mut enc = ctx
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
        {
            let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
            let run_gemm = |p: &mut wgpu::ComputePass, g: &GemmStep| {
                // Hardware matrix units first, when they exist.
                if let (Some(pl), Some(bg)) = (&self.gemm4_pl, &g.bg4) {
                    p.set_pipeline(pl);
                    p.set_bind_group(0, bg, &[]);
                    p.dispatch_workgroups(
                        (g.n as u32).div_ceil(V4.1 as u32),
                        (g.m as u32).div_ceil(V4.0 as u32),
                        1,
                    );
                    if let Some(a) = &g.act {
                        p.set_pipeline(&self.act_pl);
                        p.set_bind_group(0, &a.bg, &[]);
                        p.dispatch_workgroups(a.groups, 1, 1);
                    }
                    return;
                }
                // v3 only when its 128×128 tile still fills the GPU; otherwise fall back to v2's
                // occupancy-aware tile table (a 64-workgroup grid on a 40-core GPU is mostly idle,
                // and no amount of arithmetic intensity fixes an idle machine).
                let wgs3 = g.n.div_ceil(V3.1) * g.m.div_ceil(V3.0);
                if self.v3 && wgs3 >= 96 {
                    p.set_pipeline(&self.gemm3_pl);
                    p.set_bind_group(0, &g.bg3, &[]);
                    p.dispatch_workgroups(
                        (g.n as u32).div_ceil(V3.1 as u32),
                        (g.m as u32).div_ceil(V3.0 as u32),
                        1,
                    );
                    return;
                }
                let tile = gemm2_tile(g.m, g.n);
                let pl = self
                    .kernels
                    .gemm2_pipeline_tile(self.f16, tile)
                    .expect("gemm2 pipeline");
                p.set_pipeline(pl);
                p.set_bind_group(0, &g.bgs[gemm2_tier(tile)], &[]);
                p.dispatch_workgroups(
                    (g.n as u32).div_ceil(tile.1 as u32),
                    (g.m as u32).div_ceil(tile.0 as u32),
                    1,
                );
            };

            run_gemm(&mut p, &s_patch);
            p.set_pipeline(self.kernels.add_pl());
            p.set_bind_group(0, &s_pos.bg, &[]);
            p.dispatch_workgroups(s_pos.groups, 1, 1);

            for b in &steps {
                p.set_pipeline(self.kernels.ln_pl());
                p.set_bind_group(0, &b.ln1.bg, &[]);
                p.dispatch_workgroups(n as u32, 1, 1);

                run_gemm(&mut p, &b.qkv);

                p.set_pipeline(&self.rope_pl);
                p.set_bind_group(0, &b.rope.0, &[]);
                p.dispatch_workgroups((hid as u32).div_ceil(64), n as u32, 1);

                if let (Some(pl), Some(bg)) = (&self.attn_coop_pl, &b.cattn) {
                    // Attention on the matrix units: 8 query rows per subgroup.
                    p.set_pipeline(pl);
                    p.set_bind_group(0, bg, &[]);
                    p.dispatch_workgroups(
                        (n as u32).div_ceil(coop_attn_rb() as u32),
                        heads as u32,
                        1,
                    );
                } else if self.tiled_attn {
                    p.set_pipeline(&self.attn_pl);
                    p.set_bind_group(0, &b.tattn.0, &[]);
                    p.dispatch_workgroups(
                        (n as u32).div_ceil(self.attn_rb as u32),
                        heads as u32,
                        1,
                    );
                } else {
                    p.set_pipeline(self.kernels.attn_pl());
                    p.set_bind_group(0, &b.attn.0, &[]);
                    p.dispatch_workgroups(n as u32, heads as u32, 1);
                }

                run_gemm(&mut p, &b.proj);
                p.set_pipeline(self.kernels.add_pl());
                p.set_bind_group(0, &b.add1.bg, &[]);
                p.dispatch_workgroups(b.add1.groups, 1, 1);

                p.set_pipeline(self.kernels.ln_pl());
                p.set_bind_group(0, &b.ln2.bg, &[]);
                p.dispatch_workgroups(n as u32, 1, 1);

                run_gemm(&mut p, &b.fc1);
                run_gemm(&mut p, &b.fc2);
                p.set_pipeline(self.kernels.add_pl());
                p.set_bind_group(0, &b.add2.bg, &[]);
                p.dispatch_workgroups(b.add2.groups, 1, 1);
            }

            p.set_pipeline(self.kernels.ln_pl());
            p.set_bind_group(0, &s_mln.bg, &[]);
            p.dispatch_workgroups(n as u32, 1, 1);
            run_gemm(&mut p, &s_m1);
            run_gemm(&mut p, &s_m2);
        }
        ctx.queue.submit([enc.finish()]);
        ctx.read(&m2, tokens * cfg.out_hidden)
    }
}

struct GemmStep {
    bgs: Vec<wgpu::BindGroup>,
    bg3: wgpu::BindGroup,
    bg4: Option<wgpu::BindGroup>,
    /// Only when this GEMM has an activation AND v4 runs it (v2/v3 fuse it into the epilogue).
    act: Option<ActStep>,
    n: usize,
    m: usize,
    _meta: wgpu::Buffer,
}
struct ActStep {
    bg: wgpu::BindGroup,
    groups: u32,
    _meta: wgpu::Buffer,
}
struct LnStep {
    bg: wgpu::BindGroup,
    _meta: wgpu::Buffer,
}
struct AddStep {
    bg: wgpu::BindGroup,
    groups: u32,
    _meta: wgpu::Buffer,
}

/// `DeviceExt` is only pulled in for `storage_u32`; keep the import honest.
trait StorageU32 {
    fn storage_u32(&self, v: &[u32]) -> wgpu::Buffer;
}

impl StorageU32 for GpuCtx {
    fn storage_u32(&self, v: &[u32]) -> wgpu::Buffer {
        self.device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: None,
                contents: bytemuck::cast_slice(v),
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_SRC
                    | wgpu::BufferUsages::COPY_DST,
            })
    }
}

// =====================================================================================================
// GLM-OCR vision tower on wgpu.
//
// The Qwen tower above and this one share every GPU kernel — GEMM v3/v4, the tiled/coop attention,
// LayerNorm, the residual add — because a ViT block is a ViT block. What differs is GLM's block
// internals, and each maps onto an existing kernel bar two:
//   * weight-only **RMSNorm** → the encoder's LayerNorm in RMS mode (flag bit 2, no bias);
//   * **per-head qk-RMSNorm BEFORE the rope** → folded into the split/rope kernel below;
//   * **SwiGLU** MLP → two GEMMs + one silu·mul pass (below) instead of fc1→act→fc2;
//   * **no learned position embedding** → the pos-add is simply gone;
//   * the merger's 2×2 conv-downsample is a Linear over the block-major 4-token concat, so it, too,
//     is just a GEMM (m = tokens, k = merge²·hidden).
// Gated cos 1.0 against the CPU [`crate::vision_glm::GlmVisionTower`], itself gated against HF.
// =====================================================================================================

/// Split the packed qkv, apply per-head RMS q/k-norm (dim `hd`, BEFORE the rope — GLM's order), then
/// the 2-D rope. One workgroup per (head, row), `hd` lanes (GLM's head_dim is 64, as the attention
/// kernels assume). This is the one kernel GLM needs that [`VIS_QKV_ROPE`] does not provide.
const GLM_QKV_QKNORM_ROPE: &str = r#"
struct Meta { rows: u32, nh: u32, hd: u32, eps: f32 }
@group(0) @binding(0) var<storage, read>       qkv: array<f32>;  // [rows, 3*hid]
@group(0) @binding(1) var<storage, read>       cs:  array<f32>;  // [rows, hd]
@group(0) @binding(2) var<storage, read>       sn:  array<f32>;  // [rows, hd]
@group(0) @binding(3) var<storage, read>       qw:  array<f32>;  // [hd]
@group(0) @binding(4) var<storage, read>       kw:  array<f32>;  // [hd]
@group(0) @binding(5) var<storage, read_write> q:   array<f32>;  // [rows, hid]
@group(0) @binding(6) var<storage, read_write> k:   array<f32>;
@group(0) @binding(7) var<storage, read_write> v:   array<f32>;
@group(0) @binding(8) var<uniform>             mt:  Meta;
var<workgroup> qs:  array<f32, 64>;
var<workgroup> ks:  array<f32, 64>;
var<workgroup> red: array<f32, 64>;
@compute @workgroup_size(64)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) t: u32) {
    let hd = mt.hd;
    let hid = mt.nh * hd;
    let head = wg.x;
    let row = wg.y;
    if (row >= mt.rows || head >= mt.nh) { return; }
    let half = hd / 2u;
    let base = row * 3u * hid + head * hd;

    let qv = qkv[base + t];
    let kv = qkv[base + hid + t];

    // per-head RMS over the 64 dims — one reduction for q, one for k
    red[t] = qv * qv;
    workgroupBarrier();
    for (var s = 32u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
    let qinv = 1.0 / sqrt(red[0] / f32(hd) + mt.eps);
    workgroupBarrier();
    red[t] = kv * kv;
    workgroupBarrier();
    for (var s = 32u; s > 0u; s >>= 1u) { if (t < s) { red[t] += red[t + s]; } workgroupBarrier(); }
    let kinv = 1.0 / sqrt(red[0] / f32(hd) + mt.eps);
    workgroupBarrier();

    qs[t] = qv * qinv * qw[t];
    ks[t] = kv * kinv * kw[t];
    workgroupBarrier();

    // rope (contiguous halves): out[d] = x[d]·c[d] + rotate_half(x)[d]·s[d]
    let c = cs[row * hd + t];
    let s2 = sn[row * hd + t];
    let partner = select(t - half, t + half, t < half);
    let sgn = select(1.0, -1.0, t < half);
    q[row * hid + head * hd + t] = qs[t] * c + sgn * qs[partner] * s2;
    k[row * hid + head * hd + t] = ks[t] * c + sgn * ks[partner] * s2;
    v[row * hid + head * hd + t] = qkv[base + 2u * hid + t];  // V is not roped
}
"#;

/// `gate = silu(gate) · up`, elementwise over two `[rows, intermediate]` buffers — GLM's SwiGLU
/// combine, which has SEPARATE gate/up projections (so [`GLU_SPLIT`]'s single interleaved buffer does
/// not apply).
const GLM_SILU_MUL: &str = r#"
struct Meta { n: u32, p0: u32, p1: u32, p2: u32 }
@group(0) @binding(0) var<storage, read_write> gate: array<f32>;
@group(0) @binding(1) var<storage, read>       up:   array<f32>;
@group(0) @binding(2) var<uniform>             mt:   Meta;
@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    if (g.x >= mt.n) { return; }
    let x = gate[g.x];
    gate[g.x] = (x / (1.0 + exp(-x))) * up[g.x];
}
"#;

struct GlmGpuBlock {
    n1_w: wgpu::Buffer,
    qkv: GpuLinear,
    q_norm_w: wgpu::Buffer,
    k_norm_w: wgpu::Buffer,
    proj: GpuLinear,
    n2_w: wgpu::Buffer,
    gate: GpuLinear,
    up: GpuLinear,
    down: GpuLinear,
}

/// The GLM-OCR vision tower, resident on the GPU. Holds the CPU tower alongside — for the host-side
/// rope tables and as the parity oracle.
pub struct GlmVisionGpu {
    cpu: crate::vision_glm::GlmVisionTower,
    kernels: EncKernels,
    qkrope_pl: wgpu::ComputePipeline,
    silumul_pl: wgpu::ComputePipeline,
    attn_pl: wgpu::ComputePipeline,
    attn_coop_pl: Option<wgpu::ComputePipeline>,
    gemm3_pl: wgpu::ComputePipeline,
    gemm4_pl: Option<wgpu::ComputePipeline>,
    act_pl: wgpu::ComputePipeline,
    v3: bool,
    f16: bool,
    attn_rb: usize,
    patch: GpuLinear,
    blocks: Vec<GlmGpuBlock>,
    post_ln_w: wgpu::Buffer,
    downsample: GpuLinear,
    merger_proj: GpuLinear,
    merger_norm: GpuNorm,
    merger_gate: GpuLinear,
    merger_up: GpuLinear,
    merger_down: GpuLinear,
}

impl GlmVisionGpu {
    pub fn config(&self) -> &VisionConfig {
        &self.cpu.cfg
    }

    /// Move a loaded CPU GLM tower onto the GPU.
    pub fn new(ctx: &GpuCtx, cpu: crate::vision_glm::GlmVisionTower) -> Result<Self> {
        anyhow::ensure!(
            cpu.cfg.head_dim() == 64,
            "GLM vision GPU assumes head_dim 64, got {}",
            cpu.cfg.head_dim()
        );
        let f16 = std::env::var("OSFKB_VISION_F32").is_err();
        let lin = |l: &crate::vision::Linear| -> GpuLinear {
            let w = if f16 {
                Mat::F16(ctx.storage_bytes(&f32_to_f16_bytes(&l.w)))
            } else {
                Mat::F32(ctx.storage(&l.w))
            };
            let t = transpose(&l.w, l.n, l.k);
            let wt = Mat::F32(ctx.storage(&t));
            // Bias-less linears (the merger projections) get an explicit zero bias, so bit-0 stays set
            // and the accumulator starts at zero — no per-linear flag branching.
            let bias: Vec<f32> = l.b.clone().unwrap_or_else(|| vec![0.0; l.n]);
            let mut b8 = Vec::with_capacity(8 * l.n);
            for _ in 0..8 {
                b8.extend_from_slice(&bias);
            }
            GpuLinear {
                w,
                wt,
                b: ctx.storage(&bias),
                b8: ctx.storage(&b8),
                n: l.n as u32,
                k: l.k as u32,
            }
        };
        let norm = |n: &crate::vision::Norm| GpuNorm {
            w: ctx.storage(&n.w),
            b: ctx.storage(&n.b),
        };
        let pipe = |label: &str, src: String| -> wgpu::ComputePipeline {
            ctx.device
                .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                    label: Some(label),
                    layout: None,
                    module: &ctx
                        .device
                        .shader_module_tuned(wgpu::ShaderModuleDescriptor {
                            label: Some(label),
                            source: wgpu::ShaderSource::Wgsl(src.into()),
                        }),
                    entry_point: Some("main"),
                    compilation_options: Default::default(),
                    cache: None,
                })
        };

        let blocks = cpu
            .blocks
            .iter()
            .map(|b| GlmGpuBlock {
                n1_w: ctx.storage(&b.norm1_w),
                qkv: lin(&b.qkv),
                q_norm_w: ctx.storage(&b.q_norm_w),
                k_norm_w: ctx.storage(&b.k_norm_w),
                proj: lin(&b.proj),
                n2_w: ctx.storage(&b.norm2_w),
                gate: lin(&b.gate),
                up: lin(&b.up),
                down: lin(&b.down),
            })
            .collect();

        let use_coop = ctx.coop_matrix
            && !matches!(std::env::var("OSFKB_ENC_COOP").ok().as_deref(), Some("0"));
        Ok(Self {
            patch: lin(&cpu.patch),
            post_ln_w: ctx.storage(&cpu.post_ln_w),
            downsample: lin(&cpu.downsample),
            merger_proj: lin(&cpu.merger_proj),
            merger_norm: norm(&cpu.merger_post_norm),
            merger_gate: lin(&cpu.merger_gate),
            merger_up: lin(&cpu.merger_up),
            merger_down: lin(&cpu.merger_down),
            blocks,
            qkrope_pl: pipe("glm_qkv_qknorm_rope", GLM_QKV_QKNORM_ROPE.to_string()),
            silumul_pl: pipe("glm_silu_mul", GLM_SILU_MUL.to_string()),
            attn_pl: pipe("vis_attn_tiled", vis_attn_src(attn_rb())),
            attn_coop_pl: (ctx.coop_matrix
                && !matches!(
                    std::env::var("OSFKB_VIS_ATTN_COOP").ok().as_deref(),
                    Some("0")
                ))
            .then(|| pipe("vis_attn_coop", vis_attn_coop_src(coop_attn_rb()))),
            gemm3_pl: pipe("enc_gemm3", enc_gemm3_src(false, V3.0, V3.1, V3.2)),
            gemm4_pl: use_coop.then(|| pipe("enc_gemm4_coop", enc_gemm4_src(V4.0, V4.1, V4.2))),
            act_pl: pipe("enc_act", ENC_ACT.to_string()),
            v3: !matches!(std::env::var("OSFKB_ENC_GEMM3").ok().as_deref(), Some("0")),
            attn_rb: attn_rb(),
            kernels: EncKernels::new(ctx),
            f16,
            cpu,
        })
    }

    /// Run the tower over one or more preprocessed images → `[total_merged_tokens, out_hidden]`.
    pub fn forward(&self, ctx: &GpuCtx, images: &[ImagePatches]) -> Result<Vec<f32>> {
        let mut out = Vec::new();
        for img in images {
            out.extend(self.forward_one(ctx, img)?);
        }
        Ok(out)
    }

    fn forward_one(&self, ctx: &GpuCtx, img: &ImagePatches) -> Result<Vec<f32>> {
        let cfg = self.cpu.cfg.clone();
        let (hid, heads, hd) = (cfg.hidden, cfg.heads, cfg.head_dim());
        let n = img.num_patches();
        let unit = cfg.merge_unit();
        anyhow::ensure!(
            n.is_multiple_of(unit),
            "patch count {n} not a multiple of merge² ({unit})"
        );
        let tokens = n / unit;
        let intermediate = cfg.intermediate;

        // Host-side rope tables (shared with the CPU tower — see the module note on the Qwen twin).
        let (cos, sin) = crate::vision::rope_tables_2d(&cfg, img.grid);

        let b_patch = ctx.storage(&img.patches);
        let b_cos = ctx.storage(&cos);
        let b_sin = ctx.storage(&sin);
        let x = ctx.empty(n * hid);
        let normed = ctx.empty(n * hid);
        let qkv = ctx.empty(n * 3 * hid);
        let bq = ctx.empty(n * hid);
        let bk = ctx.empty(n * hid);
        let bv = ctx.empty(n * hid);
        let attn = ctx.empty(n * hid);
        let gate = ctx.empty(n * intermediate);
        let up = ctx.empty(n * intermediate);
        let oh_ds = self.downsample.n as usize;
        let oh = self.merger_proj.n as usize;
        let inner = self.merger_gate.n as usize;
        let ds = ctx.empty(tokens * oh_ds);
        let mproj = ctx.empty(tokens * oh);
        let mln = ctx.empty(tokens * oh);
        let mg = ctx.empty(tokens * inner);
        let mu = ctx.empty(tokens * inner);
        let mout = ctx.empty(tokens * oh);

        // ---- step builders (mirror the Qwen tower's, with GLM's norms/MLP) ----
        let gemm = |x: &wgpu::Buffer,
                    l: &GpuLinear,
                    y: &wgpu::Buffer,
                    m: usize,
                    act: Option<Act>|
         -> Result<GemmStep> {
            let flags = 1u32 | (act_code(act) << 8);
            let meta = uni(ctx, bytemuck::cast_slice(&[m as u32, l.n, l.k, flags]));
            let bufs = [x, l.w.buf(), &l.b, y];
            let bufs3 = [x, l.wt.buf(), &l.b, y];
            let bg3 = make_bg(ctx, &self.gemm3_pl, &bufs3, &meta);
            let coop_ok = m.is_multiple_of(8) && (l.n as usize).is_multiple_of(8);
            let bufs4 = [x, l.wt.buf(), &l.b8, y];
            let bg4 = self
                .gemm4_pl
                .as_ref()
                .filter(|_| coop_ok)
                .map(|pl| make_bg(ctx, pl, &bufs4, &meta));
            let act = act.map(|a| {
                let am = uni(
                    ctx,
                    bytemuck::cast_slice(&[
                        (m * l.n as usize) as u32,
                        act_code(Some(a)),
                        0u32,
                        0u32,
                    ]),
                );
                let abg = make_bg(ctx, &self.act_pl, &[y], &am);
                ActStep {
                    bg: abg,
                    groups: ((m * l.n as usize) as u32).div_ceil(256),
                    _meta: am,
                }
            });
            Ok(GemmStep {
                bgs: v2_bgs(ctx, &self.kernels, self.f16, &bufs, &meta)?,
                bg3,
                bg4,
                act,
                n: l.n as usize,
                m,
                _meta: meta,
            })
        };
        // Weight-only RMSNorm: LayerNorm kernel, RMS mode (flag 4), no bias — bind the weight as the
        // (unread) bias slot.
        let ln_rms = |src: &wgpu::Buffer, w: &wgpu::Buffer, dst: &wgpu::Buffer| -> LnStep {
            let meta = uni(
                ctx,
                bytemuck::cast_slice(&[hid as u32, 4u32, cfg.eps.to_bits(), 0u32]),
            );
            let bg = make_bg(ctx, self.kernels.ln_pl(), &[src, src, w, w, dst], &meta);
            LnStep { bg, _meta: meta }
        };
        // LayerNorm with bias (the merger's post-projection norm), width `h`.
        let ln_bias =
            |src: &wgpu::Buffer, nm: &GpuNorm, dst: &wgpu::Buffer, h: usize, eps: f32| -> LnStep {
                let meta = uni(
                    ctx,
                    bytemuck::cast_slice(&[h as u32, 2u32, eps.to_bits(), 0u32]),
                );
                let bg = make_bg(
                    ctx,
                    self.kernels.ln_pl(),
                    &[src, src, &nm.w, &nm.b, dst],
                    &meta,
                );
                LnStep { bg, _meta: meta }
            };
        let add = |dst: &wgpu::Buffer, src: &wgpu::Buffer, len: usize| -> AddStep {
            let meta = uni(ctx, bytemuck::cast_slice(&[len as u32, 0u32, 0u32, 0u32]));
            let bg = make_bg(ctx, self.kernels.add_pl(), &[dst, src], &meta);
            AddStep {
                bg,
                groups: (len as u32).div_ceil(256),
                _meta: meta,
            }
        };
        // silu(gate)·up, in place on `gate` — reuses AddStep as a (bg, groups, meta) holder.
        let silu = |g: &wgpu::Buffer, u: &wgpu::Buffer, len: usize| -> AddStep {
            let meta = uni(ctx, bytemuck::cast_slice(&[len as u32, 0u32, 0u32, 0u32]));
            let bg = make_bg(ctx, &self.silumul_pl, &[g, u], &meta);
            AddStep {
                bg,
                groups: (len as u32).div_ceil(256),
                _meta: meta,
            }
        };
        // Standalone activation over a buffer (the merger's GELU).
        let actf = |y: &wgpu::Buffer, len: usize, a: Act| -> ActStep {
            let am = uni(
                ctx,
                bytemuck::cast_slice(&[len as u32, act_code(Some(a)), 0u32, 0u32]),
            );
            let bg = make_bg(ctx, &self.act_pl, &[y], &am);
            ActStep {
                bg,
                groups: (len as u32).div_ceil(256),
                _meta: am,
            }
        };

        let s_patch = gemm(&b_patch, &self.patch, &x, n, None)?; // NO position embedding

        struct GBlk {
            ln1: LnStep,
            qkv: GemmStep,
            qkrope_bg: wgpu::BindGroup,
            _qkrope_meta: wgpu::Buffer,
            tattn_bg: wgpu::BindGroup,
            cattn_bg: Option<wgpu::BindGroup>,
            _attn_meta: wgpu::Buffer,
            proj: GemmStep,
            add1: AddStep,
            ln2: LnStep,
            gate: GemmStep,
            up: GemmStep,
            silu: AddStep,
            down: GemmStep,
            add2: AddStep,
        }
        let mut steps: Vec<GBlk> = Vec::with_capacity(cfg.depth);
        for b in &self.blocks {
            let qkrope_meta = uni(
                ctx,
                bytemuck::cast_slice(&[n as u32, heads as u32, hd as u32, cfg.eps.to_bits()]),
            );
            let qkrope_bg = make_bg(
                ctx,
                &self.qkrope_pl,
                &[
                    &qkv,
                    &b_cos,
                    &b_sin,
                    &b.q_norm_w,
                    &b.k_norm_w,
                    &bq,
                    &bk,
                    &bv,
                ],
                &qkrope_meta,
            );
            let attn_meta = uni(
                ctx,
                bytemuck::cast_slice(&[n as u32, heads as u32, hd as u32, 0u32]),
            );
            let tattn_bg = make_bg(ctx, &self.attn_pl, &[&bq, &bk, &bv, &attn], &attn_meta);
            let cattn_bg = self
                .attn_coop_pl
                .as_ref()
                .map(|pl| make_bg(ctx, pl, &[&bq, &bk, &bv, &attn], &attn_meta));
            steps.push(GBlk {
                ln1: ln_rms(&x, &b.n1_w, &normed),
                qkv: gemm(&normed, &b.qkv, &qkv, n, None)?,
                qkrope_bg,
                _qkrope_meta: qkrope_meta,
                tattn_bg,
                cattn_bg,
                _attn_meta: attn_meta,
                proj: gemm(&attn, &b.proj, &normed, n, None)?,
                add1: add(&x, &normed, n * hid),
                ln2: ln_rms(&x, &b.n2_w, &normed),
                gate: gemm(&normed, &b.gate, &gate, n, None)?,
                up: gemm(&normed, &b.up, &up, n, None)?,
                silu: silu(&gate, &up, n * intermediate),
                down: gemm(&gate, &b.down, &normed, n, None)?,
                add2: add(&x, &normed, n * hid),
            });
        }

        // Merger: post-RMSNorm → conv-downsample (as a GEMM over the 4-token concat) → proj →
        // LayerNorm(+bias) → GELU → SwiGLU.
        let s_postln = ln_rms(&x, &self.post_ln_w, &normed);
        let s_ds = gemm(&normed, &self.downsample, &ds, tokens, None)?;
        let s_proj = gemm(&ds, &self.merger_proj, &mproj, tokens, None)?;
        let s_mln = ln_bias(&mproj, &self.merger_norm, &mln, oh, 1e-5);
        let s_gelu = actf(&mln, tokens * oh, Act::GeluErf);
        let s_mg = gemm(&mln, &self.merger_gate, &mg, tokens, None)?;
        let s_mu = gemm(&mln, &self.merger_up, &mu, tokens, None)?;
        let s_msilu = silu(&mg, &mu, tokens * inner);
        let s_mdown = gemm(&mg, &self.merger_down, &mout, tokens, None)?;

        // ---- one submit ----
        let mut enc = ctx
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("glm_vision"),
            });
        {
            let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
            let run_gemm = |p: &mut wgpu::ComputePass, g: &GemmStep| {
                if let (Some(pl), Some(bg)) = (&self.gemm4_pl, &g.bg4) {
                    p.set_pipeline(pl);
                    p.set_bind_group(0, bg, &[]);
                    p.dispatch_workgroups(
                        (g.n as u32).div_ceil(V4.1 as u32),
                        (g.m as u32).div_ceil(V4.0 as u32),
                        1,
                    );
                    if let Some(a) = &g.act {
                        p.set_pipeline(&self.act_pl);
                        p.set_bind_group(0, &a.bg, &[]);
                        p.dispatch_workgroups(a.groups, 1, 1);
                    }
                    return;
                }
                let wgs3 = g.n.div_ceil(V3.1) * g.m.div_ceil(V3.0);
                if self.v3 && wgs3 >= 96 {
                    p.set_pipeline(&self.gemm3_pl);
                    p.set_bind_group(0, &g.bg3, &[]);
                    p.dispatch_workgroups(
                        (g.n as u32).div_ceil(V3.1 as u32),
                        (g.m as u32).div_ceil(V3.0 as u32),
                        1,
                    );
                    return;
                }
                let tile = gemm2_tile(g.m, g.n);
                let pl = self
                    .kernels
                    .gemm2_pipeline_tile(self.f16, tile)
                    .expect("gemm2 pipeline");
                p.set_pipeline(pl);
                p.set_bind_group(0, &g.bgs[gemm2_tier(tile)], &[]);
                p.dispatch_workgroups(
                    (g.n as u32).div_ceil(tile.1 as u32),
                    (g.m as u32).div_ceil(tile.0 as u32),
                    1,
                );
            };
            let run_act = |p: &mut wgpu::ComputePass, a: &ActStep| {
                p.set_pipeline(&self.act_pl);
                p.set_bind_group(0, &a.bg, &[]);
                p.dispatch_workgroups(a.groups, 1, 1);
            };

            run_gemm(&mut p, &s_patch);

            for b in &steps {
                p.set_pipeline(self.kernels.ln_pl());
                p.set_bind_group(0, &b.ln1.bg, &[]);
                p.dispatch_workgroups(n as u32, 1, 1);

                run_gemm(&mut p, &b.qkv);

                p.set_pipeline(&self.qkrope_pl);
                p.set_bind_group(0, &b.qkrope_bg, &[]);
                p.dispatch_workgroups(heads as u32, n as u32, 1);

                // Matrix-unit attention when the adapter has it, else the tiled kernel. Both are the
                // shared vision kernels — the GLM attention is an ordinary softmax(QKᵀ/√d)·V.
                if let (Some(pl), Some(bg)) = (&self.attn_coop_pl, &b.cattn_bg) {
                    p.set_pipeline(pl);
                    p.set_bind_group(0, bg, &[]);
                    p.dispatch_workgroups(
                        (n as u32).div_ceil(coop_attn_rb() as u32),
                        heads as u32,
                        1,
                    );
                } else {
                    p.set_pipeline(&self.attn_pl);
                    p.set_bind_group(0, &b.tattn_bg, &[]);
                    p.dispatch_workgroups(
                        (n as u32).div_ceil(self.attn_rb as u32),
                        heads as u32,
                        1,
                    );
                }

                run_gemm(&mut p, &b.proj);
                p.set_pipeline(self.kernels.add_pl());
                p.set_bind_group(0, &b.add1.bg, &[]);
                p.dispatch_workgroups(b.add1.groups, 1, 1);

                p.set_pipeline(self.kernels.ln_pl());
                p.set_bind_group(0, &b.ln2.bg, &[]);
                p.dispatch_workgroups(n as u32, 1, 1);

                run_gemm(&mut p, &b.gate);
                run_gemm(&mut p, &b.up);
                p.set_pipeline(&self.silumul_pl);
                p.set_bind_group(0, &b.silu.bg, &[]);
                p.dispatch_workgroups(b.silu.groups, 1, 1);
                run_gemm(&mut p, &b.down);
                p.set_pipeline(self.kernels.add_pl());
                p.set_bind_group(0, &b.add2.bg, &[]);
                p.dispatch_workgroups(b.add2.groups, 1, 1);
            }

            // merger
            p.set_pipeline(self.kernels.ln_pl());
            p.set_bind_group(0, &s_postln.bg, &[]);
            p.dispatch_workgroups(n as u32, 1, 1);
            run_gemm(&mut p, &s_ds);
            run_gemm(&mut p, &s_proj);
            p.set_pipeline(self.kernels.ln_pl());
            p.set_bind_group(0, &s_mln.bg, &[]);
            p.dispatch_workgroups(tokens as u32, 1, 1);
            run_act(&mut p, &s_gelu);
            run_gemm(&mut p, &s_mg);
            run_gemm(&mut p, &s_mu);
            p.set_pipeline(&self.silumul_pl);
            p.set_bind_group(0, &s_msilu.bg, &[]);
            p.dispatch_workgroups(s_msilu.groups, 1, 1);
            run_gemm(&mut p, &s_mdown);
        }
        ctx.queue.submit([enc.finish()]);
        ctx.read(&mout, tokens * oh)
    }
}