inferencelayer 0.2.4

//! CUDA tower — the GLM-OCR vision tower on the CUDA driver API + cuBLAS tensor cores.
//!
//! WHY THIS EXISTS: on Volta the tensor cores are unreachable from Vulkan (the driver only
//! exposes cooperative-matrix on Turing+ — probed), and the portable-WGSL GEMMs plateau at
//! ~11.5 TF/s vs ~100+ TF/s of sm70 HMMA. This arm reaches them through cuBLAS `GemmEx`
//! (f16 in, f32 accumulate — the same acceptance class as the wgpu f16a path).
//!
//! PORTABILITY CONTRACT: everything here is behind the `cudarc` cargo feature; cudarc loads
//! `libcuda`/`libnvrtc`/`libcublas` with dlopen at RUNTIME (zero link-time deps), and the
//! serving loader only tries this arm when `OSFKB_CUDA_TOWER=1` AND device init succeeds —
//! else the portable wgpu tower (default) or CPU tower serve exactly as today. The engine
//! (scheduler/KV/decode) never touches CUDA: this arm lives behind the tower's host-memory
//! boundary (patches in, embedding rows out), the same seam preemption replay relies on.
//!
//! Stage order, layouts and epsilon handling mirror `vision_glm_gpu::run_one` 1:1; the CPU
//! tower (transformer-oracle-gated) is the parity reference, same as for the wgpu arm.

use anyhow::{Context as _, Result};
use std::sync::Arc;

use cudarc::cublas::{result as blas, sys as blas_sys, CudaBlas};
use cudarc::driver::{
    CudaContext, CudaFunction, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig,
    PushKernelArg,
};
use cudarc::nvrtc::compile_ptx;
use half::f16;

use crate::vision::{rope_tables_2d, ImagePatches};
use crate::vision_glm::GlmVisionTower;

/// All device kernels. Compiled once with NVRTC at startup (PTX, JIT'd by the driver for the
/// present arch). Constants (erf approximation, eps placement) match the WGSL/CPU arms exactly.
const KERNELS: &str = r#"
// f32 -> f16 (round-nearest-even) without cuda_fp16.h — NVRTC ships no toolkit headers, and
// this is the single f16 touch-point in the kernels (cuBLAS handles the GEMM f16 side).
extern "C" __global__ void to_h(const float* x, unsigned short* y, int len) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < len) {
        unsigned short h;
        asm("cvt.rn.f16.f32 %0, %1;" : "=h"(h) : "f"(x[i]));
        y[i] = h;
    }
}

extern "C" __global__ void bias_add(float* y, const float* b, int m, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < m * n) y[i] += b[i % n];
}

extern "C" __global__ void accum(float* y, const float* b, int len) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < len) y[i] += b[i];
}

extern "C" __global__ void mul_silu(float* g, const float* u, int len) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < len) {
        float v = g[i];
        g[i] = (v / (1.0f + expf(-v))) * u[i];
    }
}

// exact gelu via the same Abramowitz-Stegun erf approximation as the WGSL/CPU arms
extern "C" __global__ void gelu_erf(float* y, int len) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < len) {
        float x = y[i];
        float s = x < 0.0f ? -1.0f : 1.0f;
        float ax = fabsf(x) * 0.7071067811865476f;
        float t = 1.0f / (1.0f + 0.3275911f * ax);
        float er = 1.0f - (((((1.061405429f*t - 1.453152027f)*t) + 1.421413741f)*t
                    - 0.284496736f)*t + 0.254829592f)*t*expf(-ax*ax);
        y[i] = 0.5f * x * (1.0f + s * er);
    }
}

// weight-only RMSNorm, one block per row (256 threads, tree reduce)
extern "C" __global__ void rmsnorm(const float* x, const float* w, float* y, int h, float eps) {
    __shared__ float red[256];
    int row = blockIdx.x;
    const float* xr = x + (long)row * h;
    float s = 0.0f;
    for (int j = threadIdx.x; j < h; j += 256) { float v = xr[j]; s += v * v; }
    red[threadIdx.x] = s; __syncthreads();
    for (int st = 128; st > 0; st >>= 1) {
        if (threadIdx.x < st) red[threadIdx.x] += red[threadIdx.x + st];
        __syncthreads();
    }
    float inv = rsqrtf(red[0] / h + eps);
    for (int j = threadIdx.x; j < h; j += 256) y[(long)row * h + j] = xr[j] * inv * w[j];
}

// LayerNorm with weight+bias (merger), one block per row
extern "C" __global__ void layernorm(const float* x, const float* w, const float* b, float* y,
                                     int h, float eps) {
    __shared__ float red[256];
    int row = blockIdx.x;
    const float* xr = x + (long)row * h;
    float s = 0.0f;
    for (int j = threadIdx.x; j < h; j += 256) s += xr[j];
    red[threadIdx.x] = s; __syncthreads();
    for (int st = 128; st > 0; st >>= 1) {
        if (threadIdx.x < st) red[threadIdx.x] += red[threadIdx.x + st];
        __syncthreads();
    }
    float mean = red[0] / h; __syncthreads();
    float v = 0.0f;
    for (int j = threadIdx.x; j < h; j += 256) { float d = xr[j] - mean; v += d * d; }
    red[threadIdx.x] = v; __syncthreads();
    for (int st = 128; st > 0; st >>= 1) {
        if (threadIdx.x < st) red[threadIdx.x] += red[threadIdx.x + st];
        __syncthreads();
    }
    float inv = rsqrtf(red[0] / h + eps);
    for (int j = threadIdx.x; j < h; j += 256)
        y[(long)row * h + j] = (xr[j] - mean) * inv * w[j] + b[j];
}

// Per-(token,head) q/k RMSNorm + NEOX rope — 1:1 port of the WGSL QKNORM_ROPE.
extern "C" __global__ void qknorm_rope(float* qkv, const float* qw, const float* kw,
                                       const float* cs, const float* sn,
                                       int n, int heads, float eps, int hid) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n * heads) return;
    int tok = idx / heads, h = idx % heads;
    long qb = (long)tok * 3 * hid + h * 64;
    long kb = qb + hid;
    float q[64], k[64];
    float qs = 0.0f, ks = 0.0f;
    for (int j = 0; j < 64; j++) {
        q[j] = qkv[qb + j]; qs += q[j] * q[j];
        k[j] = qkv[kb + j]; ks += k[j] * k[j];
    }
    float qi = rsqrtf(qs / 64.0f + eps);
    float ki = rsqrtf(ks / 64.0f + eps);
    for (int j = 0; j < 64; j++) { q[j] *= qi * qw[j]; k[j] *= ki * kw[j]; }
    long cb = (long)tok * 64;
    for (int j = 0; j < 64; j++) {
        float rq = j < 32 ? -q[j + 32] : q[j - 32];
        float rk = j < 32 ? -k[j + 32] : k[j - 32];
        qkv[qb + j] = q[j] * cs[cb + j] + rq * sn[cb + j];
        qkv[kb + j] = k[j] * cs[cb + j] + rk * sn[cb + j];
    }
}

// Full (non-causal) flash attention at hd=64: RB=16 query rows per block, 256 threads
// (4 sub-lanes x 64 lanes), online softmax, f32 K/V tiles PADDED +1 column — the pad shifts
// row stride to 65 words so fixed-column accesses across the 64 lanes are bank-conflict-free
// (the unpadded v2 hit 32-way conflicts in both the score and PV phases: 1135ms).
#define RB 16
#define TS 64
extern "C" __global__ void flash64(const float* qkv, float* out, int n, int heads,
                                   int hid, float scale) {
    __shared__ float qs[RB][65];
    __shared__ float ks[TS][65];
    __shared__ float vs[TS][65];
    __shared__ float sc[RB][65];
    __shared__ float m_s[RB], l_s[RB], corr_s[RB];
    int h = blockIdx.y;
    int q0 = blockIdx.x * RB;
    int tx = threadIdx.x;
    int t = tx & 63;   // dim / key lane
    int sub = tx >> 6; // 0..3
    for (int q = sub; q < RB; q += 4) {
        int tok = q0 + q;
        qs[q][t] = tok < n ? qkv[(long)tok * 3 * hid + h * 64 + t] : 0.0f;
    }
    if (tx < RB) { m_s[tx] = -1e30f; l_s[tx] = 0.0f; }
    float acc[4]; // queries sub, sub+4, sub+8, sub+12 at dim t
    for (int i = 0; i < 4; i++) acc[i] = 0.0f;
    __syncthreads();
    for (int c0 = 0; c0 < n; c0 += TS) {
        int lim = min(TS, n - c0);
        // K/V tile: 256 threads x 16 elems each, coalesced global reads
        for (int e = tx; e < lim * 64; e += 256) {
            int j = e >> 6, d = e & 63;
            long kb = (long)(c0 + j) * 3 * hid + hid + h * 64 + d;
            ks[j][d] = qkv[kb];
            vs[j][d] = qkv[kb + hid];
        }
        __syncthreads();
        // scores: thread (sub, t) -> key t, queries sub+4i. 4-query register blocking +
        // float4 smem reads: the naive per-query dot did 2 smem reads per FMA and V100 smem
        // bandwidth caps that at ~1/4 of ALU peak (measured 3.3 TF/s, 432ms/page).
        if (t < lim) {
            // kd is the only banked read (conflict-free: thread t -> row t, stride 65);
            // the four q-row reads are warp-uniform (sub is constant per warp) -> smem
            // BROADCASTS. Net: ~1 banked LDS per 4 FMAs vs 2 per FMA in v3.
            float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;
            for (int d = 0; d < 64; d++) {
                float kd = ks[t][d];
                s0 += qs[sub][d] * kd;
                s1 += qs[sub + 4][d] * kd;
                s2 += qs[sub + 8][d] * kd;
                s3 += qs[sub + 12][d] * kd;
            }
            sc[sub][t] = s0 * scale;
            sc[sub + 4][t] = s1 * scale;
            sc[sub + 8][t] = s2 * scale;
            sc[sub + 12][t] = s3 * scale;
        }
        __syncthreads();
        // online softmax state per query row
        if (tx < RB) {
            float mx = m_s[tx];
            for (int j = 0; j < lim; j++) mx = fmaxf(mx, sc[tx][j]);
            float corr = expf(m_s[tx] - mx);
            float sum = 0.0f;
            for (int j = 0; j < lim; j++) {
                float e = expf(sc[tx][j] - mx);
                sc[tx][j] = e; sum += e;
            }
            l_s[tx] = l_s[tx] * corr + sum;
            m_s[tx] = mx;
            corr_s[tx] = corr;
        }
        __syncthreads();
        // PV: 4 queries per thread at dim t, one vs read shared across the four
        {
            float a0 = acc[0] * corr_s[sub];
            float a1 = acc[1] * corr_s[sub + 4];
            float a2 = acc[2] * corr_s[sub + 8];
            float a3 = acc[3] * corr_s[sub + 12];
            for (int j = 0; j < lim; j++) {
                float vv = vs[j][t];
                a0 += sc[sub][j] * vv;
                a1 += sc[sub + 4][j] * vv;
                a2 += sc[sub + 8][j] * vv;
                a3 += sc[sub + 12][j] * vv;
            }
            acc[0] = a0; acc[1] = a1; acc[2] = a2; acc[3] = a3;
        }
        __syncthreads();
    }
    for (int i = 0; i < 4; i++) {
        int q = sub + 4 * i;
        int tok = q0 + q;
        if (tok < n) out[(long)tok * hid + h * 64 + t] = acc[i] / l_s[q];
    }
}
"#;

/// A device linear layer: f16 weights `[n, k]` row-major (+ optional f32 bias).
struct DLin {
    w: CudaSlice<f16>,
    b: Option<CudaSlice<f32>>,
    n: usize,
    k: usize,
}

struct DBlock {
    norm1_w: CudaSlice<f32>,
    qkv: DLin,
    q_norm_w: CudaSlice<f32>,
    k_norm_w: CudaSlice<f32>,
    proj: DLin,
    norm2_w: CudaSlice<f32>,
    gate: DLin,
    up: DLin,
    down: DLin,
}

/// `OSFKB_CUDA_STAGE_TRACE=1`: sync after each stage group and print cumulative ms per
/// forward — the decomposition that decides which kernel gets tuned next. Costs sync stalls;
/// diagnostics only.
fn stage_trace() -> bool {
    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
    *ON.get_or_init(|| std::env::var("OSFKB_CUDA_STAGE_TRACE").as_deref() == Ok("1"))
}

pub struct CudaGlmTower {
    /// Cumulative ns in cuBLAS GEMMs / the flash kernel (populated only under
    /// OSFKB_CUDA_STAGE_TRACE=1; the syncs it inserts make it diagnostics-only).
    t_gemm: std::sync::atomic::AtomicU64,
    t_flash: std::sync::atomic::AtomicU64,
    _ctx: Arc<CudaContext>,
    stream: Arc<CudaStream>,
    blas: CudaBlas,
    f_to_h: CudaFunction,
    f_bias: CudaFunction,
    f_accum: CudaFunction,
    f_mulsilu: CudaFunction,
    f_gelu: CudaFunction,
    f_rms: CudaFunction,
    f_ln: CudaFunction,
    f_qkrope: CudaFunction,
    f_flash: CudaFunction,
    cfg: crate::vision::VisionConfig,
    patch: DLin,
    blocks: Vec<DBlock>,
    post_ln_w: CudaSlice<f32>,
    downsample: DLin,
    merger_proj: DLin,
    merger_norm_w: CudaSlice<f32>,
    merger_norm_b: CudaSlice<f32>,
    merger_gate: DLin,
    merger_up: DLin,
    merger_down: DLin,
}

fn ecfg(len: usize) -> LaunchConfig {
    LaunchConfig {
        grid_dim: (len.div_ceil(256) as u32, 1, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 0,
    }
}

impl CudaGlmTower {
    pub fn new(cpu: &GlmVisionTower) -> Result<Self> {
        let ctx = CudaContext::new(0).context("cuda device 0")?;
        let stream = ctx.default_stream();
        let blas = CudaBlas::new(stream.clone()).context("cublas")?;
        let ptx = compile_ptx(KERNELS).map_err(|e| anyhow::anyhow!("nvrtc: {e:?}"))?;
        let module = ctx.load_module(ptx).context("load ptx")?;
        let f = |name: &str| module.load_function(name).context(name.to_string());

        let up_lin = |l: &crate::vision::Linear| -> Result<DLin> {
            let wh: Vec<f16> = l.w.iter().map(|&v| f16::from_f32(v)).collect();
            Ok(DLin {
                w: stream.memcpy_stod(&wh)?,
                b: match &l.b {
                    Some(b) => Some(stream.memcpy_stod(b)?),
                    None => None,
                },
                n: l.n,
                k: l.k,
            })
        };
        let up_f32 = |v: &[f32]| -> Result<CudaSlice<f32>> { Ok(stream.memcpy_stod(v)?) };

        let mut blocks = Vec::with_capacity(cpu.blocks.len());
        for b in &cpu.blocks {
            blocks.push(DBlock {
                norm1_w: up_f32(&b.norm1_w)?,
                qkv: up_lin(&b.qkv)?,
                q_norm_w: up_f32(&b.q_norm_w)?,
                k_norm_w: up_f32(&b.k_norm_w)?,
                proj: up_lin(&b.proj)?,
                norm2_w: up_f32(&b.norm2_w)?,
                gate: up_lin(&b.gate)?,
                up: up_lin(&b.up)?,
                down: up_lin(&b.down)?,
            });
        }
        Ok(Self {
            t_gemm: std::sync::atomic::AtomicU64::new(0),
            t_flash: std::sync::atomic::AtomicU64::new(0),
            f_to_h: f("to_h")?,
            f_bias: f("bias_add")?,
            f_accum: f("accum")?,
            f_mulsilu: f("mul_silu")?,
            f_gelu: f("gelu_erf")?,
            f_rms: f("rmsnorm")?,
            f_ln: f("layernorm")?,
            f_qkrope: f("qknorm_rope")?,
            f_flash: f("flash64")?,
            cfg: cpu.cfg.clone(),
            patch: up_lin(&cpu.patch)?,
            blocks,
            post_ln_w: up_f32(&cpu.post_ln_w)?,
            downsample: up_lin(&cpu.downsample)?,
            merger_proj: up_lin(&cpu.merger_proj)?,
            merger_norm_w: up_f32(&cpu.merger_post_norm.w)?,
            merger_norm_b: up_f32(&cpu.merger_post_norm.b)?,
            merger_gate: up_lin(&cpu.merger_gate)?,
            merger_up: up_lin(&cpu.merger_up)?,
            merger_down: up_lin(&cpu.merger_down)?,
            _ctx: ctx,
            stream,
            blas,
        })
    }

    /// `y_f32[m, l.n] = x_f32[m, l.k] · Wᵀ (+ b)` — converts x to f16 into `xh`, then cuBLAS
    /// GemmEx (f16 inputs, CUBLAS_COMPUTE_32F ⇒ tensor cores with fp32 accumulate).
    fn gemm(
        &self,
        x: &CudaSlice<f32>,
        xh: &mut CudaSlice<f16>,
        m: usize,
        l: &DLin,
        y: &mut CudaSlice<f32>,
    ) -> Result<()> {
        let _t0 = if stage_trace() {
            self.stream.synchronize()?;
            Some(std::time::Instant::now())
        } else {
            None
        };
        let len = (m * l.k) as i32;
        unsafe {
            self.stream
                .launch_builder(&self.f_to_h)
                .arg(x)
                .arg(&mut *xh)
                .arg(&len)
                .launch(ecfg(m * l.k))?;
        }
        let alpha: f32 = 1.0;
        let beta: f32 = 0.0;
        {
        let (pw, _gw) = l.w.device_ptr(&self.stream);
        let (px, _gx) = xh.device_ptr(&self.stream);
        let (py, _gy) = y.device_ptr_mut(&self.stream);
        unsafe {
            blas::gemm_ex(
                *self.blas.handle(),
                blas_sys::cublasOperation_t::CUBLAS_OP_T,
                blas_sys::cublasOperation_t::CUBLAS_OP_N,
                l.n as i32,
                m as i32,
                l.k as i32,
                (&alpha) as *const f32 as *const _,
                pw as *const std::ffi::c_void,
                blas_sys::cudaDataType_t::CUDA_R_16F,
                l.k as i32,
                px as *const std::ffi::c_void,
                blas_sys::cudaDataType_t::CUDA_R_16F,
                l.k as i32,
                (&beta) as *const f32 as *const _,
                py as *mut std::ffi::c_void,
                blas_sys::cudaDataType_t::CUDA_R_32F,
                l.n as i32,
                blas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
                blas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT,
            )
            .map_err(|e| anyhow::anyhow!("gemm_ex: {e:?}"))?;
        }
        }
        if let Some(b) = &l.b {
            let (mi, ni) = (m as i32, l.n as i32);
            unsafe {
                self.stream
                    .launch_builder(&self.f_bias)
                    .arg(&mut *y)
                    .arg(b)
                    .arg(&mi)
                    .arg(&ni)
                    .launch(ecfg(m * l.n))?;
            }
        }
        if let Some(t0) = _t0 {
            self.stream.synchronize()?;
            self.t_gemm.fetch_add(
                t0.elapsed().as_nanos() as u64,
                std::sync::atomic::Ordering::Relaxed,
            );
        }
        Ok(())
    }

    fn run_one(&self, img: &ImagePatches) -> Result<Vec<f32>> {
        let cfg = &self.cfg;
        let (hid, heads, inter) = (cfg.hidden, cfg.heads, cfg.intermediate);
        let n = img.num_patches();
        anyhow::ensure!(img.patches.len() == n * cfg.patch_dim(), "patch buffer mismatch");
        let st = &self.stream;

        // scratch (per image; freed on drop)
        let patches: CudaSlice<f32> = st.memcpy_stod(&img.patches)?;
        let mut xh: CudaSlice<f16> = st.alloc_zeros(n * cfg.patch_dim().max(3 * hid).max(inter))?;
        let mut x: CudaSlice<f32> = st.alloc_zeros(n * hid)?;
        let mut normed: CudaSlice<f32> = st.alloc_zeros(n * hid)?;
        let mut qkv: CudaSlice<f32> = st.alloc_zeros(n * 3 * hid)?;
        let mut merged: CudaSlice<f32> = st.alloc_zeros(n * hid)?;
        let mut tmp: CudaSlice<f32> = st.alloc_zeros(n * hid)?;
        let mut g: CudaSlice<f32> = st.alloc_zeros(n * inter)?;
        let mut u: CudaSlice<f32> = st.alloc_zeros(n * inter)?;

        self.gemm(&patches, &mut xh, n, &self.patch, &mut x)?;

        let (cos, sin) = rope_tables_2d(cfg, img.grid);
        let cos_b: CudaSlice<f32> = st.memcpy_stod(&cos)?;
        let sin_b: CudaSlice<f32> = st.memcpy_stod(&sin)?;

        let scale = 1.0f32 / 8.0; // 1/sqrt(64)
        let eps = cfg.eps;
        let (ni, hi, hidi) = (n as i32, heads as i32, hid as i32);

        for blk in &self.blocks {
            // ---- attention ----
            let hf = hid as i32;
            unsafe {
                st.launch_builder(&self.f_rms)
                    .arg(&x).arg(&blk.norm1_w).arg(&mut normed).arg(&hf).arg(&eps)
                    .launch(LaunchConfig { grid_dim: (n as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
            }
            self.gemm(&normed, &mut xh, n, &blk.qkv, &mut qkv)?;
            unsafe {
                st.launch_builder(&self.f_qkrope)
                    .arg(&mut qkv).arg(&blk.q_norm_w).arg(&blk.k_norm_w)
                    .arg(&cos_b).arg(&sin_b).arg(&ni).arg(&hi).arg(&eps).arg(&hidi)
                    .launch(ecfg(n * heads))?;
            }
            let _tf = if stage_trace() {
                st.synchronize()?;
                Some(std::time::Instant::now())
            } else {
                None
            };
            unsafe {
                st.launch_builder(&self.f_flash)
                    .arg(&qkv).arg(&mut merged).arg(&ni).arg(&hi).arg(&hidi).arg(&scale)
                    .launch(LaunchConfig {
                        grid_dim: (n.div_ceil(16) as u32, heads as u32, 1),
                        block_dim: (256, 1, 1),
                        shared_mem_bytes: 0,
                    })?;
            }
            if let Some(t0) = _tf {
                st.synchronize()?;
                self.t_flash.fetch_add(
                    t0.elapsed().as_nanos() as u64,
                    std::sync::atomic::Ordering::Relaxed,
                );
            }
            self.gemm(&merged, &mut xh, n, &blk.proj, &mut tmp)?;
            let len = (n * hid) as i32;
            unsafe {
                st.launch_builder(&self.f_accum).arg(&mut x).arg(&tmp).arg(&len)
                    .launch(ecfg(n * hid))?;
            }
            // ---- gated MLP ----
            unsafe {
                st.launch_builder(&self.f_rms)
                    .arg(&x).arg(&blk.norm2_w).arg(&mut normed).arg(&hf).arg(&eps)
                    .launch(LaunchConfig { grid_dim: (n as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
            }
            self.gemm(&normed, &mut xh, n, &blk.gate, &mut g)?;
            self.gemm(&normed, &mut xh, n, &blk.up, &mut u)?;
            let li = (n * inter) as i32;
            unsafe {
                st.launch_builder(&self.f_mulsilu).arg(&mut g).arg(&u).arg(&li)
                    .launch(ecfg(n * inter))?;
            }
            self.gemm(&g, &mut xh, n, &blk.down, &mut tmp)?;
            unsafe {
                st.launch_builder(&self.f_accum).arg(&mut x).arg(&tmp).arg(&len)
                    .launch(ecfg(n * hid))?;
            }
        }

        // post RMSNorm → downsample (conv-as-linear over 4-consecutive tokens) → merger
        let hf = hid as i32;
        unsafe {
            st.launch_builder(&self.f_rms)
                .arg(&x).arg(&self.post_ln_w).arg(&mut normed).arg(&hf).arg(&eps)
                .launch(LaunchConfig { grid_dim: (n as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
        }
        let unit = cfg.merge_unit();
        anyhow::ensure!(n % unit == 0, "patch count {n} not a multiple of merge²");
        let tokens = n / unit;
        let oh = cfg.out_hidden;
        let inner = self.merger_gate.n;
        let mut ds: CudaSlice<f32> = st.alloc_zeros(tokens * oh)?;
        // [n, hid] IS [tokens, unit·hid]
        self.gemm(&normed, &mut xh, tokens, &self.downsample, &mut ds)?;
        let mut mp: CudaSlice<f32> = st.alloc_zeros(tokens * oh)?;
        self.gemm(&ds, &mut xh, tokens, &self.merger_proj, &mut mp)?;
        let mut ln: CudaSlice<f32> = st.alloc_zeros(tokens * oh)?;
        let ohi = oh as i32;
        let ln_eps = 1e-5f32; // the CPU path hard-codes it
        unsafe {
            st.launch_builder(&self.f_ln)
                .arg(&mp).arg(&self.merger_norm_w).arg(&self.merger_norm_b).arg(&mut ln)
                .arg(&ohi).arg(&ln_eps)
                .launch(LaunchConfig { grid_dim: (tokens as u32, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 })?;
        }
        let gl = (tokens * oh) as i32;
        unsafe {
            st.launch_builder(&self.f_gelu).arg(&mut ln).arg(&gl).launch(ecfg(tokens * oh))?;
        }
        let mut mg: CudaSlice<f32> = st.alloc_zeros(tokens * inner)?;
        let mut mu: CudaSlice<f32> = st.alloc_zeros(tokens * inner)?;
        self.gemm(&ln, &mut xh, tokens, &self.merger_gate, &mut mg)?;
        self.gemm(&ln, &mut xh, tokens, &self.merger_up, &mut mu)?;
        let mi = (tokens * inner) as i32;
        unsafe {
            st.launch_builder(&self.f_mulsilu).arg(&mut mg).arg(&mu).arg(&mi)
                .launch(ecfg(tokens * inner))?;
        }
        let mut out: CudaSlice<f32> = st.alloc_zeros(tokens * oh)?;
        self.gemm(&mg, &mut xh, tokens, &self.merger_down, &mut out)?;
        st.synchronize()?;
        if stage_trace() {
            use std::sync::atomic::Ordering::Relaxed;
            eprintln!(
                "[cuda-tower] gemm {:.0} ms | flash {:.0} ms (cumulative)",
                self.t_gemm.load(Relaxed) as f64 / 1e6,
                self.t_flash.load(Relaxed) as f64 / 1e6,
            );
        }
        Ok(st.memcpy_dtov(&out)?)
    }

    /// Run the tower over preprocessed images → `[total_merged_tokens, out_hidden]` — the same
    /// rows [`GlmVisionTower::forward`] returns.
    pub fn forward(&self, images: &[ImagePatches]) -> Result<Vec<f32>> {
        let mut out = Vec::new();
        for img in images {
            out.extend(self.run_one(img)?);
        }
        Ok(out)
    }
}