inferencelayer 0.2.3

//! Qwen3.5-VL vision tower — native-resolution ViT, on the CPU, in pure Rust.
//!
//! This is the vision half of `Qwen3_5ForConditionalGeneration` (NuExtract-3 and friends). The TEXT
//! half already runs on the engine unchanged, because the decoder is the same dense Qwen3.5 the
//! claim-extractor serves; what was missing — and what this module adds — is everything that turns
//! image BYTES into the `[merged_tokens, out_hidden]` block the decoder consumes at its image
//! placeholder positions.
//!
//! It is NOT SigLIP. The engine's other tower ([`crate::encoder_cpu`]'s `CpuVision`) is a
//! fixed-resolution SigLIP ViT with a learned per-patch position table and a MAP attention-pooling
//! head. This one is native-resolution: the grid is whatever the image is, positions are BILINEARLY
//! INTERPOLATED out of a 48×48 learned table, attention carries a 2-D rope over (row, col), and the
//! head is a `merger` that concatenates 2×2 blocks of patches. Nothing but the linear algebra is
//! shared, which is why this is its own module rather than another `EncArch` arm.
//!
//! ## The three things that make a port of this silently wrong
//!
//! 1. **Patch order is BLOCK-MAJOR, not row-major.** The processor emits patches ordered
//!    `(block_row, block_col, row_in_block, col_in_block)` over `merge × merge` blocks — see
//!    [`patchify`]. This is the *whole* reason the merger may simply concatenate 4 CONSECUTIVE
//!    tokens. Feed it row-major patches and every layer still runs, every shape still matches, and
//!    the model reads a scrambled image.
//!
//! 2. **The resize is INTEGER arithmetic, and every detail of it matters.** The HF processor resizes
//!    the *uint8* tensor and only then rescales to `x/127.5 - 1` — verified against the oracle, where
//!    every value of `pixel_values · 127.5 + 127.5` is an exact integer (max |u - round(u)| = 0.0).
//!    But "resize in f32 and round at the end" still does not match: the reference's own float path
//!    differs from its uint8 path by up to 17 u8 steps. Bit-exactness needs all four of: bicubic
//!    **a = -0.5** (PIL's constant — torchvision's *antialias* path is a PIL port; the familiar -0.75
//!    belongs to the non-antialias kernel), **int16 weights** (i32 accumulate, round-and-shift), a
//!    **precision derived per axis** (15 downscaling, 14 upscaling — a fixed 15 overflows int16 when
//!    the centre tap nears 1.0), and a **u8 intermediate** between the two separable passes. Each was
//!    found by measurement, and each alone still leaves hundreds of pixels wrong. See
//!    [`resample_taps`].
//!
//!    Also: Python's `round` is HALF-TO-EVEN. A 400px side at factor 32 is exactly 12.5 factors, so
//!    Rust's `f64::round` yields a grid two patch rows too tall — see [`smart_resize`].
//!
//! 3. **Two different GELUs.** The blocks' MLP is `gelu_pytorch_tanh`; the merger's is `nn.GELU()`,
//!    the exact erf form. They differ by ~1e-3 at the tails — enough to move a logit, not enough to
//!    look like a bug.
//!
//! Gated end-to-end against `tests/fixtures/export_nuextract_vision.py`'s dump, layer by layer, so
//! the first mismatch names the layer that is wrong.

use anyhow::{Context, Result};
use rayon::prelude::*;
use std::path::Path;
use std::sync::OnceLock;

use crate::cpu_gemm::PackedWeight;
use crate::encoder_weights::Act;
use crate::weights::LazySt;

/// Geometry and hyper-parameters of the vision tower, from `config.json`'s `vision_config`.
#[derive(Clone, Debug)]
pub struct VisionConfig {
    pub depth: usize,
    pub hidden: usize,
    pub heads: usize,
    pub intermediate: usize,
    /// Width the merger projects INTO — the decoder's hidden size.
    pub out_hidden: usize,
    pub patch: usize,
    pub merge: usize,
    pub temporal: usize,
    pub in_channels: usize,
    /// Side of the learned position-embedding grid (`sqrt(num_position_embeddings)`).
    pub grid_side: usize,
    pub eps: f32,
    /// Blocks' MLP activation (`hidden_act`). The merger's is always exact GELU.
    pub act: Act,
    pub rope_theta: f32,
    /// Pixel-count bounds for `smart_resize` (the processor's `size.shortest_edge`/`longest_edge`,
    /// which are AREAS despite the names).
    pub min_pixels: usize,
    pub max_pixels: usize,
}

impl VisionConfig {
    pub fn head_dim(&self) -> usize {
        self.hidden / self.heads
    }

    /// Patches per merged token (`merge²`) — 4 here.
    pub fn merge_unit(&self) -> usize {
        self.merge * self.merge
    }

    /// Length of one flattened patch: `C · T · p · p`.
    pub fn patch_dim(&self) -> usize {
        self.in_channels * self.temporal * self.patch * self.patch
    }

    /// Parse from the FULL `config.json` (reads `vision_config`, and the processor bounds from
    /// `processor_config.json` if the caller supplies them via [`Self::with_pixel_bounds`]).
    pub fn from_config(v: &serde_json::Value) -> Result<Self> {
        let c = v
            .get("vision_config")
            .context("config.json vision_config")?;
        let usize_at = |k: &str| -> Result<usize> {
            c.get(k)
                .and_then(|x| x.as_u64())
                .map(|x| x as usize)
                .with_context(|| format!("vision_config.{k}"))
        };
        let num_pos = usize_at("num_position_embeddings")?;
        let grid_side = (num_pos as f64).sqrt().round() as usize;
        anyhow::ensure!(
            grid_side * grid_side == num_pos,
            "num_position_embeddings ({num_pos}) must be a perfect square; the pos table is a \
             square grid the interpolation walks"
        );
        let act = match c.get("hidden_act").and_then(|x| x.as_str()) {
            Some("gelu_pytorch_tanh") | Some("gelu_new") => Act::GeluTanh,
            Some("gelu") => Act::GeluErf,
            Some("silu") => Act::Silu,
            other => anyhow::bail!("unsupported vision hidden_act {other:?}"),
        };
        Ok(Self {
            depth: usize_at("depth")?,
            hidden: usize_at("hidden_size")?,
            heads: usize_at("num_heads")?,
            intermediate: usize_at("intermediate_size")?,
            out_hidden: usize_at("out_hidden_size")?,
            patch: usize_at("patch_size")?,
            merge: usize_at("spatial_merge_size")?,
            temporal: usize_at("temporal_patch_size")?,
            in_channels: c.get("in_channels").and_then(|x| x.as_u64()).unwrap_or(3) as usize,
            grid_side,
            eps: 1e-6,
            act,
            rope_theta: 10_000.0,
            // The processor's defaults; overridden from processor_config.json when present.
            min_pixels: 65536,
            max_pixels: 16_777_216,
        })
    }

    pub fn with_pixel_bounds(mut self, min_pixels: usize, max_pixels: usize) -> Self {
        self.min_pixels = min_pixels;
        self.max_pixels = max_pixels;
        self
    }
}

// ---------------------------------------------------------------------------------------------
// Preprocessing: bytes → normalized, block-major patches
// ---------------------------------------------------------------------------------------------

/// One preprocessed image: patches `[num_patches, patch_dim]` (row-major) plus its `(t, h, w)` grid
/// in PATCH units. `num_patches == t · h · w`.
#[derive(Clone, Debug)]
pub struct ImagePatches {
    pub patches: Vec<f32>,
    pub grid: [u32; 3],
}

impl ImagePatches {
    pub fn num_patches(&self) -> usize {
        (self.grid[0] * self.grid[1] * self.grid[2]) as usize
    }

    /// Tokens the decoder will see for this image (patches ÷ merge²).
    pub fn num_tokens(&self, cfg: &VisionConfig) -> usize {
        self.num_patches() / cfg.merge_unit()
    }
}

/// Qwen's `smart_resize`: round each side to a multiple of `factor` while holding the aspect ratio,
/// then scale the whole thing until the pixel COUNT lands inside `[min_pixels, max_pixels]`.
pub fn smart_resize(
    height: usize,
    width: usize,
    factor: usize,
    min_pixels: usize,
    max_pixels: usize,
) -> Result<(usize, usize)> {
    let (hf, wf) = (height as f64, width as f64);
    let ratio = hf.max(wf) / hf.min(wf);
    anyhow::ensure!(
        ratio <= 200.0,
        "absolute aspect ratio must be < 200, got {ratio:.1}"
    );
    let f = factor as f64;
    // Python's `round` is HALF-TO-EVEN; Rust's `f64::round` is half-away-from-zero. A 400px side at
    // factor 32 is exactly 12.5 factors — Python gives 12, Rust would give 13 — and the resulting
    // grid is two patch rows too tall, which silently invalidates every tensor downstream. (Found by
    // the parity gate on a 620×400 invoice, which is precisely the kind of ordinary size that hits
    // an exact half.)
    let round_half_even = |x: f64| -> f64 {
        let lo = x.floor();
        match (x - lo).partial_cmp(&0.5).expect("finite") {
            std::cmp::Ordering::Greater => lo + 1.0,
            std::cmp::Ordering::Less => lo,
            std::cmp::Ordering::Equal if (lo as i64) % 2 == 0 => lo,
            std::cmp::Ordering::Equal => lo + 1.0,
        }
    };
    let round_by = |x: f64| -> usize { (round_half_even(x / f) * f) as usize };
    let (mut h, mut w) = (round_by(hf), round_by(wf));
    // `round` can floor a small side to zero; the grid must have at least one patch on each axis.
    h = h.max(factor);
    w = w.max(factor);

    if h * w > max_pixels {
        let beta = ((hf * wf) / max_pixels as f64).sqrt();
        h = (((hf / beta) / f).floor() as usize * factor).max(factor);
        w = (((wf / beta) / f).floor() as usize * factor).max(factor);
    } else if h * w < min_pixels {
        let beta = (min_pixels as f64 / (hf * wf)).sqrt();
        h = ((hf * beta) / f).ceil() as usize * factor;
        w = ((wf * beta) / f).ceil() as usize * factor;
    }
    Ok((h, w))
}

/// Bicubic kernel, `a = -0.5` — **PIL's** convention, not torch's `-0.75`.
///
/// This is a trap worth stating plainly, because it is the opposite of the obvious guess. The HF
/// processor is a torchvision backend, so one reasonably assumes torch's bicubic constant. But
/// torchvision's ANTIALIAS resize path is a port of PIL's resampler and uses `-0.5`; the `-0.75` is
/// the *non*-antialias `upsample_bicubic2d`. Measured, not reasoned: `-0.75` left 5946 pixels wrong
/// (mean |Δ| 0.096 u8), `-0.5` leaves 1424 (mean 0.0073) — and zero once the integer arithmetic
/// below is right too.
fn bicubic(x: f64) -> f64 {
    const A: f64 = -0.5;
    let x = x.abs();
    if x < 1.0 {
        ((A + 2.0) * x - (A + 3.0)) * x * x + 1.0
    } else if x < 2.0 {
        (((x - 5.0) * x + 8.0) * x - 4.0) * A
    } else {
        0.0
    }
}

/// PIL's BILINEAR (triangle) kernel, support 1.0 — used by `RTDetrImageProcessor` (`resample=2`)
/// for the Docling layout model's 640×640 input. Same antialiased integer-tap machinery as the
/// bicubic; only the kernel and base support differ.
fn triangle(x: f64) -> f64 {
    let x = x.abs();
    if x < 1.0 { 1.0 - x } else { 0.0 }
}

/// One axis's resampling plan: the fixed-point precision, and per output pixel the first source
/// index and its integer taps.
struct Taps {
    /// Fractional bits in the weights. **Chosen per axis, not a constant** — see [`resample_taps`].
    precision: u32,
    rows: Vec<(usize, Vec<i32>)>,
}

/// Build one axis's taps.
///
/// Antialiased: when DOWNscaling, the filter's support widens with the scale factor so the kernel
/// averages over every source pixel in the output pixel's footprint. A naive bicubic that samples 4
/// taps regardless of scale aliases badly on a 4× downscale — the common case for documents.
///
/// The weights are then quantized to INTEGERS, because the reference's uint8 kernel is integer
/// arithmetic (its own float path differs from its uint8 path by up to 17 u8 steps, so "compute in
/// f32 and round at the end" cannot match it no matter how careful the float math is).
///
/// The precision is **derived, not fixed**: the taps must fit in an int16, so the rule is the largest
/// `p` with `round(w_max · 2^(p+1)) < 2^15`. That lands on 15 for a typical downscale — where the
/// weights spread thin — and drops to 14 when UPscaling, where the centre tap approaches 1.0 and
/// `1.0 · 2^15 = 32768` would overflow. Both were verified bit-exact against the reference
/// (0 of 700416 and 0 of 236544 channel values differ); a hardcoded 15 leaves the upscale 268 off.
fn resample_taps(in_size: usize, out_size: usize) -> Taps {
    resample_taps_with(in_size, out_size, bicubic, 2.0, None)
}

/// `fixed_precision`: `None` derives the largest precision whose taps fit an **int16** — that is
/// torchvision's uint8 resampler, which the Qwen-VL/GLM processors ride and the bicubic gates
/// proved bit-exact. `Some(22)` is **PIL proper** (`Resample.c PRECISION_BITS = 32-8-2`, INT32
/// coefficients) — what `RTDetrImageProcessor` uses via `to_pil_image().resize()`. The two grids
/// disagree by ±1 u8 step on ~0.06% of pixels, so picking the right one per caller is what
/// "bit-exact" means here; both share every other convention (center = (i+0.5)·scale, clipped
/// support, normalize-then-quantize, round-half-away, u8 intermediate between passes).
fn resample_taps_with(
    in_size: usize,
    out_size: usize,
    kernel: fn(f64) -> f64,
    base_support: f64,
    fixed_precision: Option<u32>,
) -> Taps {
    let scale = in_size as f64 / out_size as f64;
    let filter_scale = scale.max(1.0);
    let support = base_support * filter_scale;
    let inv = 1.0 / filter_scale;

    // Pass 1: float weights, normalized. At the borders the support is clipped, and the taps must
    // still sum to 1 or the edge pixels darken.
    let float_rows: Vec<(usize, Vec<f64>)> = (0..out_size)
        .map(|i| {
            let center = (i as f64 + 0.5) * scale;
            let lo = ((center - support + 0.5).floor() as isize).max(0) as usize;
            let hi = ((center + support + 0.5).floor() as isize).min(in_size as isize) as usize;
            let w: Vec<f64> = (lo..hi)
                .map(|j| kernel((j as f64 + 0.5 - center) * inv))
                .collect();
            let sum: f64 = w.iter().sum();
            let w = if sum != 0.0 {
                w.into_iter().map(|v| v / sum).collect()
            } else {
                w
            };
            (lo, w)
        })
        .collect();

    // Pass 2: the precision — fixed (PIL) or the largest the int16 taps allow (torchvision).
    let precision = fixed_precision.unwrap_or_else(|| {
        let w_max = float_rows
            .iter()
            .flat_map(|(_, w)| w.iter())
            .fold(0f64, |a, &b| a.max(b));
        let mut p = 0u32;
        while p < 22 && ((0.5 + w_max * (1i64 << (p + 1)) as f64) as i64) < (1 << 15) {
            p += 1;
        }
        p
    });
    let one = (1i64 << precision) as f64;

    let rows = float_rows
        .into_iter()
        .map(|(lo, w)| {
            let q = w
                .iter()
                .map(|v| {
                    let x = v * one;
                    // round half AWAY FROM ZERO, matching the reference's coefficient quantization.
                    (if x < 0.0 { x - 0.5 } else { x + 0.5 }) as i32
                })
                .collect();
            (lo, q)
        })
        .collect();
    Taps { precision, rows }
}

/// Finish one fixed-point tap: `clamp((Σ src·w + 2^(p-1)) >> p, 0, 255)`.
#[inline]
fn tap_u8(acc: i32, precision: u32) -> u8 {
    let v = (acc + (1 << (precision - 1))) >> precision;
    v.clamp(0, 255) as u8
}

/// Separable bicubic resize of an interleaved RGB8 image, in integer arithmetic, with a **uint8
/// intermediate** between the horizontal and vertical passes — all of which are load-bearing for
/// bit-exactness with the reference processor.
pub(crate) fn resize_rgb8(src: &[u8], sw: usize, sh: usize, dw: usize, dh: usize) -> Vec<u8> {
    resize_rgb8_taps(
        src,
        sw,
        sh,
        dw,
        resample_taps(sw, dw),
        resample_taps(sh, dh),
    )
}

/// PIL-BILINEAR twin of [`resize_rgb8`] — the `RTDetrImageProcessor` resize (`resample=2`).
/// Same integer arithmetic and u8 intermediate; only the kernel (triangle, support 1) differs.
pub fn resize_rgb8_bilinear(src: &[u8], sw: usize, sh: usize, dw: usize, dh: usize) -> Vec<u8> {
    resize_rgb8_taps(
        src,
        sw,
        sh,
        dw,
        resample_taps_with(sw, dw, triangle, 1.0, Some(22)),
        resample_taps_with(sh, dh, triangle, 1.0, Some(22)),
    )
}

/// PIL-BICUBIC (a = -0.5, support 2) twin with PIL's INT32/PRECISION_BITS=22 taps —
/// `Image.resize`'s DEFAULT filter. The docling-parse v4 backend rides it for page images
/// ("render at 1.5× the requested scale, then `.resize()` down"), so the docling pipeline's
/// raster chain needs this exact kernel to match the reference's layout-model input.
pub fn resize_rgb8_bicubic(src: &[u8], sw: usize, sh: usize, dw: usize, dh: usize) -> Vec<u8> {
    resize_rgb8_taps(
        src,
        sw,
        sh,
        dw,
        resample_taps_with(sw, dw, bicubic, 2.0, Some(22)),
        resample_taps_with(sh, dh, bicubic, 2.0, Some(22)),
    )
}

/// Decode PNG/JPEG/WebP bytes to interleaved RGB8. Public because the RT-DETR/document callers
/// (and their parity tests) need the raw pixels, not a model-specific patch pipeline.
pub fn decode_rgb8(bytes: &[u8]) -> Result<(Vec<u8>, usize, usize)> {
    let img = image::load_from_memory(bytes).context("decode image")?;
    let rgb = img.to_rgb8();
    let (w, h) = (rgb.width() as usize, rgb.height() as usize);
    Ok((rgb.into_raw(), w, h))
}

fn resize_rgb8_taps(src: &[u8], sw: usize, sh: usize, dw: usize, hx: Taps, hy: Taps) -> Vec<u8> {
    const C: usize = 3;
    let dh = hy.rows.len();

    // Horizontal pass → [sh, dw, C] as u8 (the intermediate really is requantized).
    let mut tmp = vec![0u8; sh * dw * C];
    tmp.par_chunks_mut(dw * C).enumerate().for_each(|(y, row)| {
        let srow = &src[y * sw * C..(y + 1) * sw * C];
        for (x, (lo, w)) in hx.rows.iter().enumerate() {
            let mut acc = [0i32; C];
            for (t, &wt) in w.iter().enumerate() {
                let p = &srow[(lo + t) * C..(lo + t) * C + C];
                for c in 0..C {
                    acc[c] += wt * p[c] as i32;
                }
            }
            for c in 0..C {
                row[x * C + c] = tap_u8(acc[c], hx.precision);
            }
        }
    });

    // Vertical pass → [dh, dw, C] u8.
    let mut out = vec![0u8; dh * dw * C];
    out.par_chunks_mut(dw * C).enumerate().for_each(|(y, row)| {
        let (lo, w) = &hy.rows[y];
        for x in 0..dw {
            let mut acc = [0i32; C];
            for (t, &wt) in w.iter().enumerate() {
                let p = &tmp[((lo + t) * dw + x) * C..((lo + t) * dw + x) * C + C];
                for c in 0..C {
                    acc[c] += wt * p[c] as i32;
                }
            }
            for c in 0..C {
                row[x * C + c] = tap_u8(acc[c], hy.precision);
            }
        }
    });
    out
}

/// Cut the resized image into flattened patches in **block-major** order.
///
/// Mirrors the processor's
/// `reshape(C, gh/m, m, p, gw/m, m, p).permute(gh/m, gw/m, m, m, C, p, p)`: tokens walk merge-blocks
/// first, then the `m × m` patches inside each block. Each patch vector is laid out `[C][T][row][col]`
/// with the temporal axis a DUPLICATE of the still frame (the processor `expand`s it — a still image
/// has no second frame, and the Conv3d simply sums two identical contributions).
pub(crate) fn patchify(img: &[u8], h: usize, w: usize, cfg: &VisionConfig) -> ImagePatches {
    // Qwen3.5-VL's rescale-then-normalize with mean .5 / std .5 ⇒ v/127.5 - 1.
    patchify_with(img, h, w, cfg, [0.5; 3], [0.5; 3])
}

/// [`patchify`] with explicit per-channel normalization (`(v/255 - mean)/std`) — the GLM tower uses
/// CLIP statistics, not Qwen's ±1 rescale.
pub(crate) fn patchify_with(
    img: &[u8],
    h: usize,
    w: usize,
    cfg: &VisionConfig,
    mean: [f32; 3],
    std: [f32; 3],
) -> ImagePatches {
    let (p, m, c, tp) = (cfg.patch, cfg.merge, cfg.in_channels, cfg.temporal);
    let (gh, gw) = (h / p, w / p);
    debug_assert_eq!(gh % m, 0, "grid height must be a multiple of merge");
    debug_assert_eq!(gw % m, 0, "grid width must be a multiple of merge");
    let pd = cfg.patch_dim();
    let mut out = vec![0f32; gh * gw * pd];

    // One task per output patch: index arithmetic only, no shared state.
    out.par_chunks_mut(pd).enumerate().for_each(|(tok, dst)| {
        // Undo the block-major walk: which patch of the grid is token `tok`?
        let per_block = m * m;
        let block = tok / per_block;
        let within = tok % per_block;
        let blocks_w = gw / m;
        let (br, bc) = (block / blocks_w, block % blocks_w);
        let (ir, ic) = (within / m, within % m);
        let (prow, pcol) = (br * m + ir, bc * m + ic);

        for ch in 0..c {
            for t in 0..tp {
                for py in 0..p {
                    let sy = prow * p + py;
                    for px in 0..p {
                        let sx = pcol * p + px;
                        let v = img[(sy * w + sx) * c + ch] as f32;
                        // rescale 1/255 then per-channel normalize. The ±1 case keeps the original
                        // fused expression: `(v/255-.5)/.5` differs from `v/127.5-1` in final ulps,
                        // and the Qwen path is gated near-bit-exact.
                        let idx = ((ch * tp + t) * p + py) * p + px;
                        dst[idx] = if mean[ch] == 0.5 && std[ch] == 0.5 {
                            v / 127.5 - 1.0
                        } else {
                            (v / 255.0 - mean[ch]) / std[ch]
                        };
                    }
                }
            }
        }
    });

    ImagePatches {
        patches: out,
        grid: [1, gh as u32, gw as u32],
    }
}

/// Full preprocessing of one interleaved RGB8 image: smart-resize → bicubic (u8) → normalize →
/// block-major patchify.
pub fn preprocess_rgb8(
    rgb: &[u8],
    width: usize,
    height: usize,
    cfg: &VisionConfig,
) -> Result<ImagePatches> {
    anyhow::ensure!(
        rgb.len() == width * height * cfg.in_channels,
        "rgb buffer is {} bytes, expected {}×{}×{}",
        rgb.len(),
        width,
        height,
        cfg.in_channels
    );
    let factor = cfg.patch * cfg.merge;
    let (rh, rw) = smart_resize(height, width, factor, cfg.min_pixels, cfg.max_pixels)?;
    let resized = if (rh, rw) == (height, width) {
        rgb.to_vec()
    } else {
        resize_rgb8(rgb, width, height, rw, rh)
    };
    Ok(patchify(&resized, rh, rw, cfg))
}

/// Decode encoded image bytes (PNG/JPEG/WebP/…) and preprocess. Pure Rust — the `image` crate, not
/// PIL, which is the entire point.
pub fn preprocess_bytes(bytes: &[u8], cfg: &VisionConfig) -> Result<ImagePatches> {
    let img = image::load_from_memory(bytes).context("decode image")?;
    let rgb = img.to_rgb8();
    let (w, h) = (rgb.width() as usize, rgb.height() as usize);
    preprocess_rgb8(rgb.as_raw(), w, h, cfg)
}

// ---------------------------------------------------------------------------------------------
// M-RoPE positions
// ---------------------------------------------------------------------------------------------

/// Per-token 3-D positions `(t, h, w)` for a prompt whose image placeholders have been expanded,
/// plus the position the first GENERATED token takes.
///
/// This is why a VLM decoder cannot just reuse the 1-D position counter:
///
/// * **Text** advances all three axes together (`t = h = w = i`), which is precisely why a
///   text-only prompt is bit-identical under M-RoPE and ordinary RoPE — and why this engine has been
///   serving Qwen3.5-VL text correctly all along without knowing M-RoPE existed.
/// * **An image** is a 2-D object: every one of its tokens shares one `t`, while `h` and `w` walk the
///   merged grid. The image occupies `llm_h · llm_w` token slots but advances the position counter by
///   only `max(llm_h, llm_w)` — so text AFTER an image does not resume where the token index says it
///   should. Get this wrong and the model still runs, still emits fluent JSON, and quietly attends
///   with a skewed geometry.
///
/// `grids` are the images' `(t, h, w)` PATCH grids, in the order their placeholders appear.
pub fn mrope_positions(
    tokens: &[u32],
    image_token_id: u32,
    grids: &[[u32; 3]],
    merge: usize,
) -> Result<(Vec<[u32; 3]>, u32)> {
    let m = merge as u32;
    let mut pos = Vec::with_capacity(tokens.len());
    let mut cur = 0u32;
    let mut next_grid = 0usize;
    let mut i = 0usize;

    while i < tokens.len() {
        if tokens[i] != image_token_id {
            pos.push([cur, cur, cur]);
            cur += 1;
            i += 1;
            continue;
        }
        // A run of image placeholders: one image.
        let grid = *grids.get(next_grid).with_context(|| {
            format!(
                "prompt has more image placeholder runs than images ({} given)",
                grids.len()
            )
        })?;
        next_grid += 1;
        let (gt, lh, lw) = (grid[0], grid[1] / m, grid[2] / m);
        let want = (gt * lh * lw) as usize;
        let run = tokens[i..]
            .iter()
            .take_while(|&&t| t == image_token_id)
            .count();
        anyhow::ensure!(
            run == want,
            "image placeholder run is {run} tokens but grid {grid:?} (merge {merge}) needs {want}"
        );

        for t in 0..gt {
            for r in 0..lh {
                for c in 0..lw {
                    // meshgrid(t, h, w) in 'ij' order — row-major over the merged grid, which is
                    // exactly the order the merger emits its tokens.
                    pos.push([cur + t, cur + r, cur + c]);
                }
            }
        }
        // NOT `+= want`: an image advances the counter by its LONGEST side, not its token count.
        cur += lh.max(lw);
        i += run;
    }
    anyhow::ensure!(
        next_grid == grids.len(),
        "{} images supplied but the prompt has {next_grid} placeholder runs",
        grids.len()
    );

    // Generation resumes one past the furthest position any axis reached.
    let next = pos
        .iter()
        .flat_map(|p| p.iter())
        .copied()
        .max()
        .map_or(0, |m| m + 1);
    Ok((pos, next))
}

/// Everything the decoder needs for one multimodal prompt: run the tower over the images, place its
/// output against the prompt's placeholder tokens, and compute the 3-D positions.
///
/// `tokens` must already have each image's placeholder expanded to `llm_h · llm_w` copies of
/// `image_token_id` — that is the tokenizer/chat-template's job, and [`ImagePatches::num_tokens`]
/// says how many each image needs.
/// [`prepare_prompt`] on the GPU tower. Same contract, same placement, ~4x faster — this is the one
/// the server runs.
pub fn prepare_prompt_gpu(
    ctx: &crate::GpuCtx,
    tower: &crate::vision_gpu::VisionGpu,
    tokens: &[u32],
    image_token_id: u32,
    images: &[ImagePatches],
) -> Result<crate::server::VisionPrompt> {
    let cfg = tower.config();
    let grids: Vec<[u32; 3]> = images.iter().map(|i| i.grid).collect();
    let (mpos, next_pos) = mrope_positions(tokens, image_token_id, &grids, cfg.merge)?;
    let embeds = tower.forward(ctx, images)?;
    place_image_rows(
        tokens,
        image_token_id,
        embeds,
        cfg.out_hidden,
        mpos,
        next_pos,
    )
}

/// Shared by both towers: check the row count and map each placeholder to its row.
fn place_image_rows(
    tokens: &[u32],
    image_token_id: u32,
    embeds: Vec<f32>,
    out_hidden: usize,
    mpos: Vec<[u32; 3]>,
    next_pos: u32,
) -> Result<crate::server::VisionPrompt> {
    let rows = embeds.len() / out_hidden;
    let mut embed_index = vec![-1i32; tokens.len()];
    let mut next_row = 0i32;
    for (i, &t) in tokens.iter().enumerate() {
        if t == image_token_id {
            embed_index[i] = next_row;
            next_row += 1;
        }
    }
    anyhow::ensure!(
        next_row as usize == rows,
        "the prompt has {next_row} image placeholders but the tower produced {rows} rows"
    );
    Ok(crate::server::VisionPrompt {
        embeds,
        mpos,
        embed_index,
        next_pos,
    })
}

pub fn prepare_prompt(
    tower: &VisionTower,
    tokens: &[u32],
    image_token_id: u32,
    images: &[ImagePatches],
) -> Result<crate::server::VisionPrompt> {
    let cfg = tower.config();
    let grids: Vec<[u32; 3]> = images.iter().map(|i| i.grid).collect();
    let (mpos, next_pos) = mrope_positions(tokens, image_token_id, &grids, cfg.merge)?;

    let embeds = tower.forward(images)?;
    let rows = embeds.len() / cfg.out_hidden;
    let want: usize = images.iter().map(|i| i.num_tokens(cfg)).sum();
    anyhow::ensure!(
        rows == want,
        "the tower produced {rows} tokens but the images need {want}"
    );

    // Image placeholders consume the tower's rows in prompt order — the same order the images were
    // supplied and the tower concatenated them.
    let mut embed_index = vec![-1i32; tokens.len()];
    let mut next_row = 0i32;
    for (i, &t) in tokens.iter().enumerate() {
        if t == image_token_id {
            embed_index[i] = next_row;
            next_row += 1;
        }
    }
    anyhow::ensure!(
        next_row as usize == rows,
        "the prompt has {next_row} image placeholders but the tower produced {rows} rows"
    );

    Ok(crate::server::VisionPrompt {
        embeds,
        mpos,
        embed_index,
        next_pos,
    })
}

// ---------------------------------------------------------------------------------------------
// The tower
// ---------------------------------------------------------------------------------------------

pub(crate) struct Linear {
    pub(crate) w: Vec<f32>,
    pub(crate) b: Option<Vec<f32>>,
    pub(crate) n: usize,
    pub(crate) k: usize,
    packed: OnceLock<PackedWeight>,
}

impl Linear {
    pub(crate) fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        let w = st.tensor_f32(&format!("{prefix}.weight"))?;
        let b = st.tensor_f32(&format!("{prefix}.bias")).ok();
        let n = b.as_ref().map(Vec::len).unwrap_or(0);
        anyhow::ensure!(n > 0, "{prefix}.bias is required by this tower");
        let k = w.len() / n;
        Ok(Self {
            w,
            b,
            n,
            k,
            packed: OnceLock::new(),
        })
    }

    /// Load with an EXPLICIT output width (for bias-less weights, where `n` cannot be inferred from
    /// the bias). Used by the GLM tower, whose merger projections carry no biases.
    pub(crate) fn load_shaped(st: &LazySt, prefix: &str, n: usize) -> Result<Self> {
        let w = st.tensor_f32(&format!("{prefix}.weight"))?;
        let b = st.tensor_f32(&format!("{prefix}.bias")).ok();
        anyhow::ensure!(w.len() % n == 0, "{prefix}.weight not divisible by n={n}");
        let k = w.len() / n;
        Ok(Self {
            w,
            b,
            n,
            k,
            packed: OnceLock::new(),
        })
    }

    /// Build from an already-materialized weight (the GLM downsample conv, permuted at load).
    pub(crate) fn from_parts(w: Vec<f32>, b: Option<Vec<f32>>, n: usize, k: usize) -> Self {
        Self {
            w,
            b,
            n,
            k,
            packed: OnceLock::new(),
        }
    }

    /// `x` is `[m, k]` → `y` is `[m, n]`, through the prepacked NEON GEMM (packed once, reused for
    /// the life of the model — see [`crate::cpu_gemm`]).
    pub(crate) fn forward(&self, y: &mut [f32], x: &[f32]) {
        let m = x.len() / self.k;
        debug_assert_eq!(x.len(), m * self.k);
        let packed = self
            .packed
            .get_or_init(|| PackedWeight::new(&self.w, self.n, self.k));
        crate::cpu_gemm::gemm_packed(&mut y[..m * self.n], x, packed, m, self.b.as_deref());
    }
}

pub(crate) struct Norm {
    pub(crate) w: Vec<f32>,
    pub(crate) b: Vec<f32>,
}

impl Norm {
    pub(crate) fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        Ok(Self {
            w: st.tensor_f32(&format!("{prefix}.weight"))?,
            b: st.tensor_f32(&format!("{prefix}.bias"))?,
        })
    }
}

/// LayerNorm over rows of width `h`, in place.
pub(crate) fn layer_norm(x: &mut [f32], h: usize, n: &Norm, eps: f32) {
    x.par_chunks_mut(h).for_each(|row| {
        let mean = row.iter().sum::<f32>() / h as f32;
        let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / h as f32;
        let inv = 1.0 / (var + eps).sqrt();
        for (j, v) in row.iter_mut().enumerate() {
            *v = (*v - mean) * inv * n.w[j] + n.b[j];
        }
    });
}

pub(crate) fn act_inplace(x: &mut [f32], act: Act) {
    let c = (2.0 / std::f32::consts::PI).sqrt();
    x.par_chunks_mut(4096).for_each(|chunk| {
        for v in chunk {
            *v = match act {
                Act::GeluErf => 0.5 * *v * (1.0 + libm::erff(*v * std::f32::consts::FRAC_1_SQRT_2)),
                Act::GeluTanh => 0.5 * *v * (1.0 + (c * (*v + 0.044715 * *v * *v * *v)).tanh()),
                Act::Silu => *v / (1.0 + (-*v).exp()),
                Act::Tanh => v.tanh(),
                Act::Relu => v.max(0.0),
            };
        }
    });
}

pub(crate) struct Block {
    pub(crate) norm1: Norm,
    pub(crate) qkv: Linear,
    pub(crate) proj: Linear,
    pub(crate) norm2: Norm,
    pub(crate) fc1: Linear,
    pub(crate) fc2: Linear,
}

/// The vision tower: patch embedding, interpolated position embeddings, `depth` pre-LN blocks with a
/// 2-D rope, and the merger that produces the decoder's image tokens.
pub struct VisionTower {
    pub(crate) cfg: VisionConfig,
    pub(crate) patch: Linear,
    /// Learned position table `[grid_side², hidden]`, bilinearly resampled onto each image's grid.
    pub(crate) pos_embed: Vec<f32>,
    pub(crate) blocks: Vec<Block>,
    pub(crate) merger_norm: Norm,
    pub(crate) merger_fc1: Linear,
    pub(crate) merger_fc2: Linear,
}

impl VisionTower {
    /// Load `model.visual.*` from a Qwen3.5-VL checkpoint directory.
    pub fn load(dir: &Path) -> Result<Self> {
        let cfg_json: serde_json::Value = serde_json::from_slice(
            &std::fs::read(dir.join("config.json"))
                .with_context(|| format!("read {}", dir.join("config.json").display()))?,
        )?;
        let mut cfg = VisionConfig::from_config(&cfg_json)?;

        // The processor's pixel bounds live in a separate file and materially change the grid (and
        // so the token count). Read them when present rather than silently using the defaults.
        let pcfg = dir.join("processor_config.json");
        if pcfg.exists() {
            let p: serde_json::Value = serde_json::from_slice(&std::fs::read(&pcfg)?)?;
            if let Some(size) = p.pointer("/image_processor/size") {
                let get = |k: &str| size.get(k).and_then(|x| x.as_u64()).map(|x| x as usize);
                if let (Some(lo), Some(hi)) = (get("shortest_edge"), get("longest_edge")) {
                    cfg = cfg.with_pixel_bounds(lo, hi);
                }
            }
        }

        let st = LazySt::open(dir)?;
        Self::from_st(&st, cfg)
    }

    fn from_st(st: &LazySt, cfg: VisionConfig) -> Result<Self> {
        // The checkpoint may or may not carry the `model.` prefix depending on how it was exported.
        let p = if st.has("model.visual.pos_embed.weight") {
            "model.visual"
        } else {
            "visual"
        };
        anyhow::ensure!(
            st.has(&format!("{p}.pos_embed.weight")),
            "no vision tower in this checkpoint (looked for {p}.pos_embed.weight)"
        );

        let blocks = (0..cfg.depth)
            .map(|i| {
                let b = format!("{p}.blocks.{i}");
                Ok(Block {
                    norm1: Norm::load(st, &format!("{b}.norm1"))?,
                    qkv: Linear::load(st, &format!("{b}.attn.qkv"))?,
                    proj: Linear::load(st, &format!("{b}.attn.proj"))?,
                    norm2: Norm::load(st, &format!("{b}.norm2"))?,
                    fc1: Linear::load(st, &format!("{b}.mlp.linear_fc1"))?,
                    fc2: Linear::load(st, &format!("{b}.mlp.linear_fc2"))?,
                })
            })
            .collect::<Result<Vec<_>>>()?;

        // Conv3d `[hidden, C, T, p, p]` flattens to exactly the linear `[hidden, C·T·p·p]` the
        // patchify order feeds it — the conv IS a linear over one flattened patch.
        let patch = Linear::load(st, &format!("{p}.patch_embed.proj"))?;
        anyhow::ensure!(
            patch.k == cfg.patch_dim() && patch.n == cfg.hidden,
            "patch_embed is [{}, {}], expected [{}, {}]",
            patch.n,
            patch.k,
            cfg.hidden,
            cfg.patch_dim()
        );

        Ok(Self {
            pos_embed: st.tensor_f32(&format!("{p}.pos_embed.weight"))?,
            patch,
            merger_norm: Norm::load(st, &format!("{p}.merger.norm"))?,
            merger_fc1: Linear::load(st, &format!("{p}.merger.linear_fc1"))?,
            merger_fc2: Linear::load(st, &format!("{p}.merger.linear_fc2"))?,
            blocks,
            cfg,
        })
    }

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

    /// Bilinearly resample the `side × side` learned position table onto this image's `h × w` patch
    /// grid, emitting rows in the same block-major order as [`patchify`].
    pub(crate) fn interpolate_pos(&self, grid: [u32; 3]) -> Vec<f32> {
        let (t, h, w) = (grid[0] as usize, grid[1] as usize, grid[2] as usize);
        let (side, hid, m) = (self.cfg.grid_side, self.cfg.hidden, self.cfg.merge);

        // linspace(0, side-1, n) — the sample coordinates of this grid inside the learned table.
        let coord = |n: usize| -> Vec<f32> {
            if n == 1 {
                vec![0.0]
            } else {
                let step = (side - 1) as f32 / (n - 1) as f32;
                (0..n).map(|i| i as f32 * step).collect()
            }
        };
        let (hc, wc) = (coord(h), coord(w));

        let mut out = vec![0f32; t * h * w * hid];
        out.par_chunks_mut(hid).enumerate().for_each(|(tok, dst)| {
            // Block-major → (row, col), identical to `patchify`'s walk.
            let hw = h * w;
            let within_frame = tok % hw;
            let per_block = m * m;
            let block = within_frame / per_block;
            let within = within_frame % per_block;
            let blocks_w = w / m;
            let (br, bc) = (block / blocks_w, block % blocks_w);
            let (r, c) = (br * m + within / m, bc * m + within % m);

            let (hy, wx) = (hc[r], wc[c]);
            let (h0, w0) = (hy as usize, wx as usize);
            let (h1, w1) = ((h0 + 1).min(side - 1), (w0 + 1).min(side - 1));
            let (hf, wf) = (hy - h0 as f32, wx - w0 as f32);

            let corners = [
                ((h0 * side + w0), (1.0 - hf) * (1.0 - wf)),
                ((h0 * side + w1), (1.0 - hf) * wf),
                ((h1 * side + w0), hf * (1.0 - wf)),
                ((h1 * side + w1), hf * wf),
            ];
            for (idx, wt) in corners {
                let row = &self.pos_embed[idx * hid..(idx + 1) * hid];
                for (d, v) in dst.iter_mut().zip(row) {
                    *d += wt * v;
                }
            }
        });
        out
    }

    /// 2-D rope tables for this grid: `cos`/`sin`, each `[tokens, head_dim]`.
    ///
    /// The rotary half-dim is `head_dim/2`, split evenly between the ROW and COLUMN axes: the first
    /// half of the frequencies rotate by the patch's row, the second by its column. `emb` is then
    /// `cat(freqs, freqs)` so the usual `rotate_half` application covers the full head.
    pub(crate) fn rope_tables(&self, grid: [u32; 3]) -> (Vec<f32>, Vec<f32>) {
        rope_tables_2d(&self.cfg, grid)
    }
}

/// 2-D rope tables over a block-major patch grid: freqs `[row·inv | col·inv]`, `emb = cat(freqs,
/// freqs)`, one `[hd]` cos/sin row per token. Shared by the Qwen and GLM towers (identical math).
pub(crate) fn rope_tables_2d(cfg: &VisionConfig, grid: [u32; 3]) -> (Vec<f32>, Vec<f32>) {
    {
        let (t, h, w) = (grid[0] as usize, grid[1] as usize, grid[2] as usize);
        let (hd, m) = (cfg.head_dim(), cfg.merge);
        let half = hd / 2; // rotary dim
        let nf = half / 2; // frequencies per axis
        let inv: Vec<f32> = (0..nf)
            .map(|i| 1.0 / cfg.rope_theta.powf((2 * i) as f32 / half as f32))
            .collect();

        let n = t * h * w;
        let mut cos = vec![0f32; n * hd];
        let mut sin = vec![0f32; n * hd];
        cos.par_chunks_mut(hd)
            .zip(sin.par_chunks_mut(hd))
            .enumerate()
            .for_each(|(tok, (co, si))| {
                let hw = h * w;
                let within_frame = tok % hw;
                let per_block = m * m;
                let block = within_frame / per_block;
                let within = within_frame % per_block;
                let blocks_w = w / m;
                let (br, bc) = (block / blocks_w, block % blocks_w);
                let (r, c) = (br * m + within / m, bc * m + within % m);

                for i in 0..nf {
                    let (fr, fc) = (r as f32 * inv[i], c as f32 * inv[i]);
                    // freqs = [row·inv | col·inv]; emb = cat(freqs, freqs).
                    for (off, f) in [(i, fr), (nf + i, fc)] {
                        let (cv, sv) = (f.cos(), f.sin());
                        co[off] = cv;
                        co[off + half] = cv;
                        si[off] = sv;
                        si[off + half] = sv;
                    }
                }
            });
        (cos, sin)
    }
}

impl VisionTower {
    /// Run the tower over one or more preprocessed images.
    ///
    /// Attention is FULL within an image and never crosses image boundaries (the reference packs the
    /// batch and passes `cu_seqlens`; we simply run each image's span independently, which is the
    /// same computation). Returns `[total_merged_tokens, out_hidden]`, images concatenated in order.
    pub fn forward(&self, images: &[ImagePatches]) -> Result<Vec<f32>> {
        let mut out = Vec::new();
        for img in images {
            out.extend(self.run(img, false)?.1);
        }
        Ok(out)
    }

    /// [`Self::forward`] for one image, also returning every intermediate: `[0]` is post-(patch
    /// embed + position embed), `[i]` is the output of block `i`. Used ONLY by the parity gate — so
    /// that a mismatch names the layer that is wrong instead of leaving a 24-block bisect by hand.
    /// It runs the same body as production, so the two cannot drift.
    pub fn debug_forward(&self, img: &ImagePatches) -> Result<(Vec<Vec<f32>>, Vec<f32>)> {
        self.run(img, true)
    }

    fn run(&self, img: &ImagePatches, capture: bool) -> Result<(Vec<Vec<f32>>, Vec<f32>)> {
        let cfg = &self.cfg;
        let (hid, heads, hd) = (cfg.hidden, cfg.heads, cfg.head_dim());
        let n = img.num_patches();
        anyhow::ensure!(
            img.patches.len() == n * cfg.patch_dim(),
            "patch buffer {} does not match grid {:?}",
            img.patches.len(),
            img.grid
        );

        // Patch embed + interpolated position embeddings.
        let mut x = vec![0f32; n * hid];
        self.patch.forward(&mut x, &img.patches);
        let pos = self.interpolate_pos(img.grid);
        x.par_iter_mut()
            .zip(pos.par_iter())
            .for_each(|(a, b)| *a += b);

        let (cos, sin) = self.rope_tables(img.grid);

        let mut qkv = vec![0f32; n * 3 * hid];
        let mut attn = vec![0f32; n * hid];
        let mut mid = vec![0f32; n * cfg.intermediate];
        let mut normed = vec![0f32; n * hid];
        let mut states: Vec<Vec<f32>> = Vec::new();
        if capture {
            states.push(x.clone());
        }

        for blk in &self.blocks {
            // ---- attention ----
            normed.copy_from_slice(&x);
            layer_norm(&mut normed, hid, &blk.norm1, cfg.eps);
            blk.qkv.forward(&mut qkv, &normed);
            self.attention(&mut attn, &qkv, &cos, &sin, n, heads, hd);
            blk.proj.forward(&mut normed, &attn);
            x.par_iter_mut()
                .zip(normed.par_iter())
                .for_each(|(a, b)| *a += b);

            // ---- mlp ----
            normed.copy_from_slice(&x);
            layer_norm(&mut normed, hid, &blk.norm2, cfg.eps);
            blk.fc1.forward(&mut mid, &normed);
            act_inplace(&mut mid, cfg.act);
            blk.fc2.forward(&mut normed, &mid);
            x.par_iter_mut()
                .zip(normed.par_iter())
                .for_each(|(a, b)| *a += b);

            if capture {
                states.push(x.clone());
            }
        }

        // ---- merger ----
        // Pre-shuffle LayerNorm (per PATCH, width `hidden`), then 4 consecutive patches — one 2×2
        // block, which is exactly what the block-major order made adjacent — concatenate into one
        // row of width `hidden · merge²`.
        layer_norm(&mut x, hid, &self.merger_norm, cfg.eps);
        let unit = cfg.merge_unit();
        let tokens = n / unit;
        anyhow::ensure!(
            n.is_multiple_of(unit),
            "patch count {n} is not a multiple of merge² ({unit})"
        );
        // The concat is a no-op on the buffer: [n, hid] row-major IS [n/unit, hid·unit] row-major.
        let mut m1 = vec![0f32; tokens * hid * unit];
        self.merger_fc1.forward(&mut m1, &x);
        act_inplace(&mut m1, Act::GeluErf); // merger is nn.GELU() — exact, NOT the tanh form
        let mut m2 = vec![0f32; tokens * cfg.out_hidden];
        self.merger_fc2.forward(&mut m2, &m1);
        Ok((states, m2))
    }

    /// Full self-attention over one image's patches, with the 2-D rope applied to q and k.
    fn attention(
        &self,
        out: &mut [f32],
        qkv: &[f32],
        cos: &[f32],
        sin: &[f32],
        n: usize,
        heads: usize,
        hd: usize,
    ) {
        let hid = heads * hd;
        let scale = 1.0 / (hd as f32).sqrt();
        let half = hd / 2;

        // qkv rows are [q(hid) | k(hid) | v(hid)], each [heads][hd].
        let rope = |vec: &mut [f32], tok: usize| {
            let (c, s) = (
                &cos[tok * hd..(tok + 1) * hd],
                &sin[tok * hd..(tok + 1) * hd],
            );
            for h in 0..heads {
                let v = &mut vec[h * hd..(h + 1) * hd];
                let orig: Vec<f32> = v.to_vec();
                for j in 0..hd {
                    // rotate_half: [-x[half..], x[..half]]
                    let rot = if j < half {
                        -orig[j + half]
                    } else {
                        orig[j - half]
                    };
                    v[j] = orig[j] * c[j] + rot * s[j];
                }
            }
        };

        // Materialize rope'd q/k and plain v, per head-major layout [heads][n][hd] so each head's
        // GEMMs are contiguous.
        let mut q = vec![0f32; n * hid];
        let mut k = vec![0f32; n * hid];
        let mut v = vec![0f32; n * hid];
        q.par_chunks_mut(hid)
            .zip(k.par_chunks_mut(hid))
            .zip(v.par_chunks_mut(hid))
            .enumerate()
            .for_each(|(tok, ((qr, kr), vr))| {
                let row = &qkv[tok * 3 * hid..(tok + 1) * 3 * hid];
                qr.copy_from_slice(&row[..hid]);
                kr.copy_from_slice(&row[hid..2 * hid]);
                vr.copy_from_slice(&row[2 * hid..]);
                rope(qr, tok);
                rope(kr, tok);
            });

        // One task per head: scores [n, n] → softmax → [n, hd].
        let mut heads_out: Vec<Vec<f32>> = vec![Vec::new(); heads];
        heads_out.par_iter_mut().enumerate().for_each(|(h, slot)| {
            let mut scores = vec![0f32; n * n];
            // SAFETY: disjoint per-head output; q/k are read-only. Strides view the [n, hid]
            // tables as this head's [n, hd] block (element[r][c] = r·hid + c).
            unsafe {
                gemm::gemm(
                    n,
                    n,
                    hd,
                    scores.as_mut_ptr(),
                    1,
                    n as isize,
                    false,
                    q.as_ptr().add(h * hd),
                    1,
                    hid as isize,
                    k.as_ptr().add(h * hd),
                    hid as isize,
                    1,
                    0.0,
                    scale,
                    false,
                    false,
                    false,
                    gemm::Parallelism::None,
                );
            }
            for row in scores.chunks_mut(n) {
                let max = row.iter().copied().fold(f32::NEG_INFINITY, f32::max);
                let mut sum = 0.0;
                for s in row.iter_mut() {
                    *s = (*s - max).exp();
                    sum += *s;
                }
                let inv = 1.0 / sum;
                for s in row.iter_mut() {
                    *s *= inv;
                }
            }
            let mut ctx = vec![0f32; n * hd];
            // SAFETY: as above — `scores` is [n, n] dense, `v` viewed as this head's [n, hd].
            unsafe {
                gemm::gemm(
                    n,
                    hd,
                    n,
                    ctx.as_mut_ptr(),
                    1,
                    hd as isize,
                    false,
                    scores.as_ptr(),
                    1,
                    n as isize,
                    v.as_ptr().add(h * hd),
                    1,
                    hid as isize,
                    0.0,
                    1.0,
                    false,
                    false,
                    false,
                    gemm::Parallelism::None,
                );
            }
            *slot = ctx;
        });

        for (h, ctx) in heads_out.iter().enumerate() {
            for tok in 0..n {
                out[tok * hid + h * hd..tok * hid + (h + 1) * hd]
                    .copy_from_slice(&ctx[tok * hd..(tok + 1) * hd]);
            }
        }
    }
}