hf2q 0.1.1

//! mmproj GGUF weight loader (ADR-005 Phase 2c, Task #15 iter 31).
//!
//! Reads every required tensor from a parsed `GgufFile` onto the Metal
//! device as F32 buffers, dequantizing Q-type tensors on the CPU first
//! via `mlx_native::gguf::GgufFile::load_tensor_f32`. The produced
//! `LoadedMmprojWeights` holds one `MlxBuffer` per tensor keyed by its
//! GGUF name (e.g. `"v.patch_embd.weight"`).
//!
//! # Sequencing vs iter 30's validator
//!
//! Iter 30 (`validate_tensor_set`) proves the tensors EXIST at startup
//! before the expensive load runs. Iter 31 (this module) actually
//! reads them onto the GPU. Caller should invoke the validator first
//! and bail early on missing tensors — that keeps the operator's
//! error message specific (missing list) rather than a generic
//! "tensor not found" from mid-load.
//!
//! # GPU cost
//!
//! Gemma 4 vision tower ≈ 400 MB of F32 after dequant (221 tensors).
//! The load sequentially dispatches a small allocation per tensor;
//! total time on M5 Max ≈ 150-300ms for a cold-page-cache load of the
//! Gemma 4 mmproj. Deliberately NOT parallelized: mlx-native's
//! `load_tensor_f32` serializes through the GGUF `BufReader` mutex,
//! and the cost is already dominated by the page-cache fill rather
//! than CPU dequant.
//!
//! # Not in this iter
//!
//! - Handler wiring. The loader is usable in isolation; the
//!   `process_multimodal_content` short-circuit at 501 is unchanged.
//!   iter 32+ threads the loaded weights through `patch_embed_forward`
//!   etc. as the ViT forward pass ports block-by-block.
//! - Lazy tensor loading. Every required tensor is loaded eagerly at
//!   `load()` time. A future iter can add a per-layer lazy mode if a
//!   memory-constrained deployment needs it.

#![allow(dead_code)]

use std::collections::HashMap;
use std::path::Path;

use anyhow::{anyhow, Result};
use mlx_native::gguf::GgufFile;
use mlx_native::{MlxBuffer, MlxDevice};

// W59 ADR-005 Phase 2c iter-128: route F16-stored mmproj tensors through
// `gguf.load_tensor` (native dtype-preserving) instead of
// `gguf.load_tensor_f32` (CPU-dequant to F32). Every Gemma 4 ViT weight
// is GGML F16 in storage; pre-iter-128 the load path dequantized to F32
// at upload, then `vit_linear_gpu` re-cast F32→BF16 inside the matmul,
// costing 8x the per-element rounding budget vs peer's F16 staging
// (W58 iter-127 numerical bisect, ADR-005 iter-127 entry).
//
// This module uses `mlx_native::gguf::tensor_info(name).ggml_type` to
// gate the load branch — F16 tensors keep their native `DType::F16`
// MlxBuffer, every other type still dequants to F32 (norms, embeddings,
// scalars, and any non-F16 weight a future producer might emit).
//
// `vit_linear_gpu` reads the resulting buffer's `dtype()` and dispatches
// the matching tensor-core kernel:
//   - DType::F16  -> mlx_native::dense_matmul_f16_f32_tensor (NEW, 0.4.8)
//   - DType::BF16 -> existing dense_matmul_bf16_f32_tensor
//   - DType::F32  -> existing F32->BF16 cast + BF16 matmul (legacy path,
//                    kept for non-F16-stored weight types).
//
// Dispatch is determined by the buffer's storage dtype — natural,
// deterministic, no env-gated fallback (per the iter-128 prompt
// constraint and `feedback_no_shortcuts.md`).
use mlx_native::GgmlType;

use super::mmproj::{vit_layer_tensor, MmprojConfig};

/// Collection of mmproj tensors loaded onto a Metal device as F32.
///
/// Cheap to move; cloning requires the caller to pay the GPU-alloc
/// cost again (not implemented here — if a use case needs cheap
/// cloning, wrap in `Arc` at the call site).
pub struct LoadedMmprojWeights {
    /// Keyed by the tensor's GGUF name. Values are F32 `MlxBuffer`s
    /// with shape preserved from the source GGUF.
    tensors: HashMap<String, MlxBuffer>,
    /// Device handle kept alive for the lifetime of the buffers.
    /// Held for RAII even though public accessors go through `tensors`.
    _device: MlxDevice,
}

impl std::fmt::Debug for LoadedMmprojWeights {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LoadedMmprojWeights")
            .field("tensor_count", &self.tensors.len())
            .finish()
    }
}

impl LoadedMmprojWeights {
    /// Load every tensor from the GGUF file onto the supplied device
    /// as F32. Arch-agnostic — walks `gguf.tensor_names()` and doesn't
    /// assume a particular naming convention, so it transparently
    /// handles both Gemma 4's SigLIP-style tower AND classic CLIP
    /// producers. Callers should run `validate_tensor_set` + detect
    /// `ArchProfile` first to know what the forward-pass dispatch
    /// branch needs.
    ///
    /// `_cfg` is accepted but currently unused — kept in the signature
    /// because a future lazy/tiered loader will partition tensor loads
    /// by cfg (e.g., load only stem + first few blocks on cold start,
    /// lazy-load remaining blocks on first request).
    pub fn load(gguf: &GgufFile, cfg: &MmprojConfig, device: MlxDevice) -> Result<Self> {
        let names = gguf.tensor_names();
        let mut tensors = HashMap::with_capacity(names.len());
        for name in &names {
            // W59 ADR-005 Phase 2c iter-128: route F16-stored tensors
            // through `load_tensor` (native F16 MlxBuffer, no CPU
            // dequant) so the downstream matmul can dispatch the
            // matching mlx-native 0.4.8 F16 tensor-core kernel without
            // a lossy F16 -> F32 -> BF16 round-trip. Every other ggml
            // type (F32, Q4_0, Q8_0, Q4_K, Q5_K, Q6_K, I16) keeps the
            // legacy F32-dequant path — those are non-weight tensors
            // (norms, embeddings, scalars) where F32 is the right
            // intermediate, OR weight types this loader doesn't yet
            // wire to a non-BF16 kernel.
            //
            // GGUF tensor_info() returns None only when the lookup name
            // doesn't exist — but we're iterating gguf.tensor_names() so
            // every name is guaranteed present. The else-branch on
            // tensor_info absence is defensive only; falls back to the
            // F32 dequant path so a future GGUF format change can't
            // silently break the loader.
            let info = gguf.tensor_info(*name);
            let is_f16 = info.map(|i| i.ggml_type == GgmlType::F16).unwrap_or(false);
            let buf = if is_f16 {
                gguf.load_tensor(*name, &device)
                    .map_err(|e| anyhow!("mmproj load_tensor (F16-native) '{}': {e}", name))?
            } else {
                gguf.load_tensor_f32(*name, &device)
                    .map_err(|e| anyhow!("mmproj load_tensor_f32 '{}': {e}", name))?
            };
            tensors.insert((*name).to_string(), buf);
        }

        // -------------------------------------------------------------------
        // Wedge-4c.5: fused `attn_qkv` → split `attn_q/k/v` slice views.
        //
        // /opt/llama.cpp/convert_hf_to_gguf.py:4853-4972 emits Qwen3-VL's
        // ViT QKV as a single fused tensor named `v.blk.{N}.attn_qkv.weight`
        // (and optional `.bias`) per /opt/llama.cpp/tools/mtmd/clip-impl.h:78.
        // The runtime forward consumer at vit_gpu_qwen3vl.rs requests split
        // tensors by name (`attn_q.weight`, `attn_k.weight`, `attn_v.weight`)
        // — so when we detect a fused tensor, we install three slice-view
        // buffers under those split names. The slice views share the
        // fused tensor's underlying Metal buffer; no extra copy is paid.
        //
        // Layout (per /opt/llama.cpp/tools/mtmd/clip.cpp:339-352):
        //   fused weight `[3*hidden, hidden]` row-major (output dim first
        //   per hf2q's vit_linear_gpu convention) — Q rows are
        //   `[0..hidden][0..hidden]`, K rows `[hidden..2*hidden][0..hidden]`,
        //   V rows `[2*hidden..3*hidden][0..hidden]`. Each slice is
        //   exactly `hidden * hidden` contiguous floats.
        //
        //   fused bias `[3*hidden]` 1-D — split into three contiguous
        //   `[hidden]` slices at offsets 0 / hidden / 2*hidden.
        //
        // We reject MIXED state (a block has BOTH fused AND split tensors)
        // because the validator's mixed-state check would have already
        // rejected at startup; this is a defensive guard against a future
        // caller that bypasses the validator.
        Self::install_fused_attn_qkv_slice_views(&mut tensors, cfg)?;

        Ok(Self {
            tensors,
            _device: device,
        })
    }

    /// Detect fused `v.blk.{N}.attn_qkv.{weight,bias}` per block and
    /// install three split-name slice views (`attn_q/k/v.{weight,bias}`)
    /// pointing at the fused tensor's underlying Metal storage. Idempotent
    /// for split-only mmprojs (no fused tensors to slice).
    ///
    /// See `LoadedMmprojWeights::load` for the layout reasoning. Reject
    /// MIXED state loud — a single block with both fused and split is
    /// ambiguous.
    fn install_fused_attn_qkv_slice_views(
        tensors: &mut HashMap<String, mlx_native::MlxBuffer>,
        cfg: &MmprojConfig,
    ) -> Result<()> {
        let hidden = cfg.hidden_size as usize;
        if hidden == 0 {
            // No vision tower → nothing to slice. Defensive: a zero hidden
            // size would also break every downstream consumer; the parser
            // already rejects this via from_gguf, but we don't re-check here.
            return Ok(());
        }
        for layer_idx in 0..cfg.num_hidden_layers as usize {
            let fused_w = vit_layer_tensor(layer_idx, "attn_qkv.weight");
            let fused_b = vit_layer_tensor(layer_idx, "attn_qkv.bias");
            let split_q_w = vit_layer_tensor(layer_idx, "attn_q.weight");
            let split_k_w = vit_layer_tensor(layer_idx, "attn_k.weight");
            let split_v_w = vit_layer_tensor(layer_idx, "attn_v.weight");

            let has_fused_w = tensors.contains_key(&fused_w);
            let has_split_w = tensors.contains_key(&split_q_w)
                || tensors.contains_key(&split_k_w)
                || tensors.contains_key(&split_v_w);

            if has_fused_w && has_split_w {
                return Err(anyhow!(
                    "mmproj loader: block {layer_idx} has BOTH fused '{}' AND \
                     split attn_q/k/v.weight tensors — refusing to mix conventions \
                     (validator should have caught this at startup; bypassing \
                     validator is unsupported)",
                    fused_w
                ));
            }
            if !has_fused_w {
                continue; // split-only block (or no QKV at all — that's a
                          // validator job). Nothing to slice.
            }

            // The fused-only path. Slice the weight.
            let fused_buf = tensors
                .get(&fused_w)
                .expect("has_fused_w=true contract violated by tensors.get");
            let elem_size = fused_buf.dtype().size_of();
            let chunk_elems = hidden * hidden;
            let chunk_bytes = chunk_elems * elem_size;
            // Validate the fused tensor's storage is exactly 3 * chunk_bytes.
            // Anything else means the converter wrote an off-spec shape; we
            // refuse rather than silently slice the wrong region.
            let expected_bytes = 3 * chunk_bytes;
            if fused_buf.byte_len() < expected_bytes {
                return Err(anyhow!(
                    "mmproj loader: fused '{fused_w}' byte_len {} < expected 3*hidden*hidden*{}={} \
                     (hidden={hidden}, dtype={:?}); converter likely wrote a wrong shape",
                    fused_buf.byte_len(),
                    elem_size,
                    expected_bytes,
                    fused_buf.dtype(),
                ));
            }
            let q_w = fused_buf.slice_view(0u64, chunk_elems);
            let k_w = fused_buf.slice_view(chunk_bytes as u64, chunk_elems);
            let v_w = fused_buf.slice_view((2 * chunk_bytes) as u64, chunk_elems);
            tensors.insert(split_q_w, q_w);
            tensors.insert(split_k_w, k_w);
            tensors.insert(split_v_w, v_w);

            // Slice the optional bias if present. Bias is 1-D `[3*hidden]`.
            if tensors.contains_key(&fused_b) {
                let split_q_b = vit_layer_tensor(layer_idx, "attn_q.bias");
                let split_k_b = vit_layer_tensor(layer_idx, "attn_k.bias");
                let split_v_b = vit_layer_tensor(layer_idx, "attn_v.bias");
                let fused_bias_buf = tensors
                    .get(&fused_b)
                    .expect("tensors.contains_key(&fused_b) contract violated by tensors.get");
                let bias_elem = fused_bias_buf.dtype().size_of();
                let bias_chunk_bytes = hidden * bias_elem;
                let bias_expected = 3 * bias_chunk_bytes;
                if fused_bias_buf.byte_len() < bias_expected {
                    return Err(anyhow!(
                        "mmproj loader: fused '{fused_b}' byte_len {} < expected 3*hidden*{}={} \
                         (hidden={hidden}, dtype={:?})",
                        fused_bias_buf.byte_len(),
                        bias_elem,
                        bias_expected,
                        fused_bias_buf.dtype(),
                    ));
                }
                let q_b = fused_bias_buf.slice_view(0u64, hidden);
                let k_b = fused_bias_buf.slice_view(bias_chunk_bytes as u64, hidden);
                let v_b = fused_bias_buf.slice_view((2 * bias_chunk_bytes) as u64, hidden);
                tensors.insert(split_q_b, q_b);
                tensors.insert(split_k_b, k_b);
                tensors.insert(split_v_b, v_b);
            }
        }
        Ok(())
    }

    /// Load from a GGUF file path. Opens the file, creates a default
    /// MlxDevice, and loads every tensor. Convenience wrapper for the
    /// common startup path.
    pub fn load_from_path(path: &Path, cfg: &MmprojConfig) -> Result<Self> {
        let gguf = GgufFile::open(path)
            .map_err(|e| anyhow!("open mmproj GGUF {}: {e}", path.display()))?;
        let device =
            MlxDevice::new().map_err(|e| anyhow!("create MlxDevice for mmproj load: {e}"))?;
        Self::load(&gguf, cfg, device)
    }

    /// Look up a tensor by its GGUF name. `None` when absent (optional
    /// tensors like biases — callers gate the forward-pass branch on
    /// `Some`).
    pub fn get(&self, name: &str) -> Option<&MlxBuffer> {
        self.tensors.get(name)
    }

    /// Read a tensor's contents as an owned `Vec<f32>`, regardless of
    /// the storage dtype. Use ONLY for CPU consumers — the production
    /// matmul path dispatches the dtype-matching tensor-core kernel
    /// directly via `vit_linear_gpu`.
    ///
    /// W59 ADR-005 Phase 2c iter-128: with `LoadedMmprojWeights::load`
    /// keeping F16 GGUF tensors as native F16 MlxBuffers, callers that
    /// need an `&[f32]` view (CPU patch_embed reference, test parity
    /// L2 distance, etc.) must explicitly convert. This helper performs
    /// the float-narrowing only when needed; for an F32 buffer it
    /// allocates and copies, for an F16 buffer it widens via `half::f16`.
    /// Returns `Err` for non-{F32,F16} dtypes (caller must handle
    /// quant/U8 storage, which currently doesn't appear on the mmproj
    /// load path).
    ///
    /// Cost: O(N) heap alloc + per-element widen. The Gemma 4 mmproj's
    /// largest single tensor is `v.patch_embd.weight` at 884,736 f16
    /// elements ≈ 3.5 MB allocation; the SigLIP CPU patch_embed call
    /// site does this once per image. Acceptable.
    ///
    /// # Errors
    ///
    /// Returns `Err` when (a) the tensor's `as_slice::<…>` fails for
    /// the underlying buffer's dtype, or (b) the dtype is anything
    /// other than F32 or F16.
    pub fn tensor_as_f32_owned(&self, buf: &MlxBuffer) -> Result<Vec<f32>> {
        use mlx_native::DType;
        match buf.dtype() {
            DType::F32 => buf
                .as_slice::<f32>()
                .map(|s| s.to_vec())
                .map_err(|e| anyhow!("tensor_as_f32_owned (F32): {e}")),
            DType::F16 => {
                let raw = buf
                    .as_slice::<u16>()
                    .map_err(|e| anyhow!("tensor_as_f32_owned (F16 u16 view): {e}"))?;
                Ok(raw
                    .iter()
                    .map(|&u| half::f16::from_bits(u).to_f32())
                    .collect())
            }
            other => Err(anyhow!(
                "tensor_as_f32_owned: unsupported dtype {other:?} (only F32/F16)"
            )),
        }
    }

    /// Build an empty `LoadedMmprojWeights` with no tensors. Useful for
    /// tests that need an `AppState.mmproj` shape but don't need to
    /// drive a forward pass. The shortcut accessors all return `Err`
    /// (as the real accessors would on a broken-producer file).
    pub fn empty(device: MlxDevice) -> Self {
        Self {
            tensors: HashMap::new(),
            _device: device,
        }
    }

    /// Test-only: build a `LoadedMmprojWeights` from a pre-populated
    /// tensor map. Used by parity tests that synthesize block weights
    /// in-process rather than load a real GGUF (which would require a
    /// fixture file on disk and the full 400 MB dequant cost).
    #[cfg(test)]
    pub fn from_tensors_for_test(tensors: HashMap<String, MlxBuffer>, device: MlxDevice) -> Self {
        Self {
            tensors,
            _device: device,
        }
    }

    /// Number of loaded tensors.
    pub fn len(&self) -> usize {
        self.tensors.len()
    }

    /// Empty when no tensors were loaded (only possible from an empty
    /// `expected_tensor_names`; normal load paths return ≥ 5 tensors).
    pub fn is_empty(&self) -> bool {
        self.tensors.is_empty()
    }

    // -----------------------------------------------------------------------
    // Stem shortcuts. Each returns the buffer when present OR errors with
    // a specific name — matches what a forward-pass call-site needs.
    // -----------------------------------------------------------------------

    pub fn patch_embd_weight(&self) -> Result<&MlxBuffer> {
        self.tensors
            .get(super::mmproj::TENSOR_PATCH_EMBD)
            .ok_or_else(|| anyhow!("mmproj missing '{}'", super::mmproj::TENSOR_PATCH_EMBD))
    }

    pub fn position_embd_weight(&self) -> Result<&MlxBuffer> {
        self.tensors
            .get(super::mmproj::TENSOR_POS_EMBD)
            .ok_or_else(|| anyhow!("mmproj missing '{}'", super::mmproj::TENSOR_POS_EMBD))
    }

    /// Read the gemma4v dual position-embed table as a typed
    /// `[2, pos_size, hidden]` 3-D view.
    ///
    /// Returns `(buf, pos_size, hidden)` where `buf` is the same backing
    /// `MlxBuffer` returned by `position_embd_weight()` (no copy). The
    /// gemma4v vision tower stores this as `model.embed_vision.
    /// position_embedding_table`, mapped to `v.position_embd.weight` by
    /// `src/backends/gguf.rs:1782-1786`. The first dim is fixed at 2
    /// (X-axis table at `[0, ..]`, Y-axis table at `[1, ..]`).
    ///
    /// # Errors
    ///
    /// - tensor missing
    /// - shape isn't 3-D, or first dim isn't 2
    /// - product of dims doesn't match the buffer's element count
    ///   (catches a stale GGUF write or a producer mismatch)
    ///
    /// # Why a sibling accessor instead of changing
    /// `position_embd_weight()`
    ///
    /// SigLIP-49's vision tower stores a 2-D
    /// `[num_patches (+1 cls), hidden]` table; gemma4v's is 3-D. The
    /// untyped accessor returns the raw buffer for both — callers that
    /// need typed shape information branch on `ArchProfile`. This
    /// addition lands the gemma4v branch without churning the SigLIP
    /// path.
    pub fn position_embd_table_3d(&self) -> Result<(&MlxBuffer, u32, u32)> {
        let buf = self.position_embd_weight()?;
        let shape = buf.shape();
        if shape.len() != 3 {
            return Err(anyhow!(
                "v.position_embd.weight: expected 3-D [2, pos_size, hidden], got shape {:?}",
                shape
            ));
        }
        if shape[0] != 2 {
            return Err(anyhow!(
                "v.position_embd.weight: expected first dim 2 (gemma4v dual table), got {}",
                shape[0]
            ));
        }
        let pos_size = shape[1] as u32;
        let hidden = shape[2] as u32;
        if pos_size == 0 || hidden == 0 {
            return Err(anyhow!(
                "v.position_embd.weight: pos_size ({pos_size}) and hidden ({hidden}) must be > 0"
            ));
        }
        // Buffer-element-count cross-check: 2 * pos_size * hidden f32s.
        let expected_bytes =
            2usize * (pos_size as usize) * (hidden as usize) * std::mem::size_of::<f32>();
        if buf.byte_len() < expected_bytes {
            return Err(anyhow!(
                "v.position_embd.weight: byte_len {} < expected {} (2 * {} * {} * 4)",
                buf.byte_len(),
                expected_bytes,
                pos_size,
                hidden
            ));
        }
        Ok((buf, pos_size, hidden))
    }

    pub fn post_ln_weight(&self) -> Result<&MlxBuffer> {
        self.tensors
            .get(super::mmproj::TENSOR_POST_LN_WEIGHT)
            .ok_or_else(|| anyhow!("mmproj missing '{}'", super::mmproj::TENSOR_POST_LN_WEIGHT))
    }

    /// Per-block tensor accessor.
    ///
    /// `suffix` is the block-relative name ("attn_q.weight",
    /// "ffn_down.weight", etc. — see `BLOCK_REQUIRED_SUFFIXES`).
    ///
    /// W41 iter-116i: vision-namespace tensor names migrated to
    /// llama.cpp's short-form convention in W34 iter-116e (writer
    /// side) but the runtime forward path still uses the
    /// pre-migration suffixes. `block_tensor` accepts both: if the
    /// caller asks for a legacy name we fall back to the canonical
    /// short form. The mapping is bidirectional so a producer
    /// emitting either convention loads cleanly.
    ///
    /// Mappings (legacy ↔ canonical short form, per
    /// `/opt/llama.cpp/tools/mtmd/clip-impl.h`):
    ///   - `attn_output.{w,b}`  ↔ `attn_out.{w,b}`     (TN_ATTN_OUTPUT, l.82)
    ///   - `post_ffw_norm.{w,b}` ↔ `ffn_post_norm.{w,b}` (TN_FFN_POST_NORM, l.95)
    pub fn block_tensor(&self, layer_idx: usize, suffix: &str) -> Result<&MlxBuffer> {
        let key = vit_layer_tensor(layer_idx, suffix);
        if let Some(b) = self.tensors.get(&key) {
            return Ok(b);
        }
        // Try the legacy/canonical alias.
        let alias_suffix: Option<&str> = match suffix {
            "attn_output.weight" => Some("attn_out.weight"),
            "attn_output.bias" => Some("attn_out.bias"),
            "attn_out.weight" => Some("attn_output.weight"),
            "attn_out.bias" => Some("attn_output.bias"),
            "post_ffw_norm.weight" => Some("ffn_post_norm.weight"),
            "post_ffw_norm.bias" => Some("ffn_post_norm.bias"),
            "ffn_post_norm.weight" => Some("post_ffw_norm.weight"),
            "ffn_post_norm.bias" => Some("post_ffw_norm.bias"),
            _ => None,
        };
        if let Some(alt) = alias_suffix {
            let alt_key = vit_layer_tensor(layer_idx, alt);
            if let Some(b) = self.tensors.get(&alt_key) {
                return Ok(b);
            }
        }
        Err(anyhow!("mmproj missing '{}'", key))
    }

    /// Projector head weight tensor.
    ///
    /// W41 iter-116i: looks up the CLIP-classic name `mm.0.weight` first,
    /// then falls back to gemma4v's `mm.input_projection.weight`
    /// (`TN_MM_INP_PROJ` at `/opt/llama.cpp/tools/mtmd/clip-impl.h:110`).
    /// Both name back the same logical tensor — the writer chose the
    /// projector-specific base per llama.cpp convention (clip.cpp:1937
    /// hard-requires `mm.input_projection` for `PROJECTOR_TYPE_GEMMA4V`),
    /// and the runtime forward path uses whichever is present.
    ///
    /// The accessor name is preserved (`mm_0_weight`) for source-compat
    /// across `vit_gpu.rs` callers; the fallback is invisible to them.
    pub fn mm_0_weight(&self) -> Result<&MlxBuffer> {
        if let Some(b) = self.tensors.get(super::mmproj::TENSOR_MM_0_WEIGHT) {
            return Ok(b);
        }
        if let Some(b) = self
            .tensors
            .get(super::mmproj::TENSOR_MM_INPUT_PROJECTION_WEIGHT)
        {
            return Ok(b);
        }
        Err(anyhow!(
            "mmproj missing '{}' (and gemma4v fallback '{}')",
            super::mmproj::TENSOR_MM_0_WEIGHT,
            super::mmproj::TENSOR_MM_INPUT_PROJECTION_WEIGHT,
        ))
    }

    pub fn mm_2_weight(&self) -> Result<&MlxBuffer> {
        self.tensors
            .get(super::mmproj::TENSOR_MM_2_WEIGHT)
            .ok_or_else(|| anyhow!("mmproj missing '{}'", super::mmproj::TENSOR_MM_2_WEIGHT))
    }

    // -----------------------------------------------------------------------
    // Gemma4ClippableLinear scalar bounds for `mm.0.weight`.
    //
    // Per `/opt/llama.cpp/tools/mtmd/clip.cpp:1935-1959`, gemma4v emits
    // four optional scalar f32 tensors as siblings of `mm.0.weight`:
    //   - `mm.0.input_min`, `mm.0.input_max` (clamps applied BEFORE matmul)
    //   - `mm.0.output_min`, `mm.0.output_max` (clamps applied AFTER matmul)
    //
    // Each is a 1-element f32 tensor (the converter `unsqueeze(0)`s the
    // 0-D scalar so GGUF round-trips it as a 1-D `[1]` tensor; see
    // `/opt/llama.cpp/convert_hf_to_gguf.py:7851-7853`).
    //
    // Returns `Some(value)` when the tensor is present and decodes to
    // exactly one f32, else `None`. Callers compose the four into a
    // `Gemma4ClippableLinearBounds` via `mm_0_bounds()`.
    // -----------------------------------------------------------------------

    fn read_scalar_f32(&self, name: &str) -> Option<f32> {
        let buf = self.tensors.get(name)?;
        let slice = buf.as_slice::<f32>().ok()?;
        // Defensive: clamp scalars are 1-element. If we ever load a
        // mis-shaped sibling (e.g. converter wrote a vector), reject
        // cleanly rather than silently picking element 0.
        if slice.len() != 1 {
            return None;
        }
        Some(slice[0])
    }

    /// Read a clamp-scalar bound under either the CLIP-classic
    /// `mm.0.<suffix>` or the gemma4v `mm.input_projection.<suffix>`
    /// base name. W41 iter-116i: gemma4v's projector head is named
    /// `mm.input_projection` (clip-impl.h:110 + clip.cpp:1937-1959),
    /// so the optional clamp scalars share that base. Returns the
    /// first match in (`mm.0`, `mm.input_projection`) order.
    fn read_projector_scalar(&self, suffix: &str) -> Option<f32> {
        let mm0 = format!("mm.0.{suffix}");
        if let Some(v) = self.read_scalar_f32(&mm0) {
            return Some(v);
        }
        let mm_inp = format!("mm.input_projection.{suffix}");
        self.read_scalar_f32(&mm_inp)
    }

    /// Read the `mm.0.input_min` (or gemma4v's `mm.input_projection.input_min`)
    /// scalar bound (clamp BEFORE matmul). `None` when absent OR
    /// mis-shaped — caller treats absence as `f32::NEG_INFINITY` (no-op)
    /// per llama.cpp's default.
    pub fn mm_0_input_min(&self) -> Option<f32> {
        self.read_projector_scalar("input_min")
    }
    /// See `mm_0_input_min`.
    pub fn mm_0_input_max(&self) -> Option<f32> {
        self.read_projector_scalar("input_max")
    }
    /// Read the output_min scalar bound (clamp AFTER matmul).
    pub fn mm_0_output_min(&self) -> Option<f32> {
        self.read_projector_scalar("output_min")
    }
    /// See `mm_0_output_min`.
    pub fn mm_0_output_max(&self) -> Option<f32> {
        self.read_projector_scalar("output_max")
    }

    /// Compose the four clamp scalars into a single
    /// `Gemma4ClippableLinearBounds`. All-`None` result means the
    /// projector is byte-equivalent to a plain Linear (no clamps).
    pub fn mm_0_bounds(&self) -> super::vit::Gemma4ClippableLinearBounds {
        super::vit::Gemma4ClippableLinearBounds {
            input_min: self.mm_0_input_min(),
            input_max: self.mm_0_input_max(),
            output_min: self.mm_0_output_min(),
            output_max: self.mm_0_output_max(),
        }
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    /// Gemma 4 26B mmproj — present on this dev machine. Tests gate on
    /// existence so CI without the fixture skips them cleanly.
    const GEMMA4_MMPROJ_PATH: &str =
        "/opt/hf2q/models/gemma-4-26B-A4B-it-ara-abliterated-dwq/gemma-4-26B-A4B-it-ara-abliterated-dwq-mmproj.gguf";

    #[test]
    fn load_gemma4_mmproj_populates_arch_tensors() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // Real Gemma 4 mmproj (SigLIP variant): 356 tensors total —
        //   5 non-block (patch_embd, pos_embd, std_bias, std_scale, mm.0.weight)
        //   13/block × 27 blocks = 351
        //   No v.post_ln.weight, no mm.2.weight.
        // See /opt/hf2q/docs/ADR-005 iter 31 for the real tensor manifest.
        let path = Path::new(GEMMA4_MMPROJ_PATH);
        if !path.exists() {
            eprintln!(
                "skipping: mmproj fixture not found at {}",
                GEMMA4_MMPROJ_PATH
            );
            return;
        }
        let gguf = GgufFile::open(path).expect("open gemma4 mmproj");
        let cfg = MmprojConfig::from_gguf(&gguf).expect("parse mmproj config");
        // Sanity: Gemma 4's 27-layer SigLIP at 224×224 with 16×16 patches.
        assert_eq!(cfg.num_hidden_layers, 27);
        assert_eq!(cfg.image_size, 224);
        assert_eq!(cfg.patch_size, 16);
        assert_eq!(cfg.hidden_size, 1152);
        // W41 iter-116i: hf2q-emitted gemma4 mmproj writes
        // `clip.projector_type = "gemma4v"` (matches llama.cpp's
        // `PROJECTOR_TYPE_GEMMA4V` literal at clip-impl.h:323).
        // Pre-iter-116i the loader parsed this to `Other("gemma4v")`
        // and `is_supported()` returned false, blocking serve startup.
        assert_eq!(cfg.projector, ProjectorType::Gemma4v);

        let device = MlxDevice::new().expect("create device");
        let weights = LoadedMmprojWeights::load(&gguf, &cfg, device).expect("load weights");
        // Gemma 4 mmproj has 356 tensors total.
        assert_eq!(weights.len(), 356);
        // Arch-agnostic shortcuts present.
        weights.patch_embd_weight().expect("patch_embd_weight");
        weights
            .position_embd_weight()
            .expect("position_embd_weight");
        weights.mm_0_weight().expect("mm_0_weight");
        // post_ln + mm.2 do NOT exist in Gemma 4 mmproj.
        assert!(weights.post_ln_weight().is_err());
        assert!(weights.mm_2_weight().is_err());
        // Every layer's arch-agnostic QKV+output suffixes present.
        // W41/W42 iter-116i: vision-namespace short-form `attn_out` per
        // `TN_ATTN_OUTPUT = "%s.blk.%d.attn_out.%s"`
        // (`/opt/llama.cpp/tools/mtmd/clip-impl.h:82`); W34 iter-116e
        // fixed the writer to emit this short form and `validate_tensor_set`
        // requires the same. The pre-iter-116e long-form
        // `attn_output.weight` is no longer present.
        for layer_idx in 0..27 {
            for suffix in [
                "attn_q.weight",
                "attn_k.weight",
                "attn_v.weight",
                "attn_out.weight",
            ] {
                weights
                    .block_tensor(layer_idx, suffix)
                    .unwrap_or_else(|_| panic!("layer {} {}", layer_idx, suffix));
            }
        }
    }

    #[test]
    fn load_gemma4_mmproj_patch_embd_has_expected_shape_and_values() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // `v.patch_embd.weight` in Gemma 4 is a 2D tensor [hidden,
        // 3*patch*patch] = [1152, 768] = 884,736 elements. The GGUF
        // stores it as F16 (per W58 iter-127 audit + gguf_dump.py
        // verification); pre-iter-128 the loader dequantized to F32 at
        // upload, post-iter-128 the loader keeps F16 native so the
        // matmul can dispatch the mlx-native 0.4.8 F16 tensor-core
        // kernel (8x tighter per-element rounding than BF16 staging,
        // closes the 1.16x/block ViT cascade compound).
        //
        // This test asserts shape (element count) AND non-trivial
        // content (non-zero patch weights) without depending on the
        // storage dtype — works whether the loader keeps F16 native or
        // dequantizes to F32.
        let path = Path::new(GEMMA4_MMPROJ_PATH);
        if !path.exists() {
            eprintln!(
                "skipping: mmproj fixture not found at {}",
                GEMMA4_MMPROJ_PATH
            );
            return;
        }
        let gguf = GgufFile::open(path).expect("open gemma4 mmproj");
        let cfg = MmprojConfig::from_gguf(&gguf).expect("parse mmproj config");
        let device = MlxDevice::new().expect("create device");
        let weights = LoadedMmprojWeights::load(&gguf, &cfg, device).expect("load weights");
        let patch = weights.patch_embd_weight().expect("patch_embd");
        let expected_elems =
            (cfg.hidden_size as usize) * 3 * (cfg.patch_size as usize) * (cfg.patch_size as usize);
        // Element-count check is dtype-agnostic via element_count(); it
        // matches expected_elems regardless of F16/F32 storage.
        assert_eq!(
            patch.element_count(),
            expected_elems,
            "patch_embd element count"
        );
        // Non-zero sanity: read the underlying bytes (works for both
        // F16 and F32). For F16 storage, every f16 has ≥1 nonzero bit
        // when its value is nonzero; for F32 the same holds. A patch
        // weight tensor with the first 1024 elements all zero would be
        // a load bug — assert that fewer than 95% of the first 2048
        // bytes (= 1024 f16 OR 512 f32) are zero.
        let raw: &[u8] = patch.as_slice().expect("as_slice raw bytes");
        let scan = raw.len().min(2048);
        let nonzero = raw[..scan].iter().filter(|&&b| b != 0).count();
        assert!(
            nonzero > scan * 5 / 100,
            "patch_embd loads to mostly-zero bytes (probable load bug): \
             {nonzero}/{scan} nonzero in first {scan} bytes"
        );
        // Dtype-specific spot-check: the gemma4v patch_embd is F16 in
        // storage; if the iter-128 path is wired correctly the buffer
        // dtype reflects that. (For SigLIP / classic CLIP producers
        // the tensor may be F32 instead — both paths are valid here;
        // we just print so the test record shows which one was loaded.)
        eprintln!(
            "patch_embd dtype: {:?}, element_count: {}",
            patch.dtype(),
            patch.element_count()
        );
    }

    #[test]
    fn load_from_path_wraps_gguf_open_and_device_create() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let path = Path::new(GEMMA4_MMPROJ_PATH);
        if !path.exists() {
            eprintln!(
                "skipping: mmproj fixture not found at {}",
                GEMMA4_MMPROJ_PATH
            );
            return;
        }
        let gguf = GgufFile::open(path).expect("open for cfg");
        let cfg = MmprojConfig::from_gguf(&gguf).expect("cfg");
        let weights = LoadedMmprojWeights::load_from_path(path, &cfg).expect("load_from_path");
        assert_eq!(weights.len(), 356);
    }

    #[test]
    fn accessors_return_err_with_specific_name_when_missing() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // Synthetic LoadedMmprojWeights with an empty tensor map — every
        // accessor should return Err naming the missing tensor.
        let weights = LoadedMmprojWeights {
            tensors: HashMap::new(),
            _device: MlxDevice::new().expect("device"),
        };
        let err = weights.patch_embd_weight().unwrap_err();
        assert!(format!("{err}").contains("v.patch_embd.weight"));

        let err = weights.position_embd_weight().unwrap_err();
        assert!(format!("{err}").contains("v.position_embd.weight"));

        let err = weights.block_tensor(5, "attn_q.weight").unwrap_err();
        assert!(format!("{err}").contains("v.blk.5.attn_q.weight"));

        let err = weights.mm_0_weight().unwrap_err();
        assert!(format!("{err}").contains("mm.0.weight"));
    }

    #[test]
    fn empty_weights_report_len_zero_and_is_empty_true() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let weights = LoadedMmprojWeights {
            tensors: HashMap::new(),
            _device: MlxDevice::new().expect("device"),
        };
        assert_eq!(weights.len(), 0);
        assert!(weights.is_empty());
    }

    #[test]
    fn get_returns_none_for_absent_tensor() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let weights = LoadedMmprojWeights {
            tensors: HashMap::new(),
            _device: MlxDevice::new().expect("device"),
        };
        assert!(weights.get("v.patch_embd.weight").is_none());
    }

    #[test]
    fn position_embd_table_3d_rejects_non_3d_shape() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // Synthesize a LoadedMmprojWeights with a 2-D position-embd
        // (the SigLIP shape). The 3-D accessor must reject it cleanly.
        let device = MlxDevice::new().expect("device");
        let buf = device
            .alloc_buffer(64 * 4, mlx_native::DType::F32, vec![8, 8])
            .expect("alloc");
        let mut tensors = HashMap::new();
        tensors.insert(super::super::mmproj::TENSOR_POS_EMBD.to_string(), buf);
        let weights = LoadedMmprojWeights {
            tensors,
            _device: device,
        };
        let err = weights.position_embd_table_3d().unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("expected 3-D"), "wrong error msg: {msg}");
    }

    #[test]
    fn position_embd_table_3d_rejects_first_dim_not_two() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let device = MlxDevice::new().expect("device");
        let buf = device
            .alloc_buffer(96 * 4, mlx_native::DType::F32, vec![3, 4, 8])
            .expect("alloc");
        let mut tensors = HashMap::new();
        tensors.insert(super::super::mmproj::TENSOR_POS_EMBD.to_string(), buf);
        let weights = LoadedMmprojWeights {
            tensors,
            _device: device,
        };
        let err = weights.position_embd_table_3d().unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("first dim 2"), "wrong error msg: {msg}");
    }

    #[test]
    fn position_embd_table_3d_returns_dims_for_valid_shape() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let device = MlxDevice::new().expect("device");
        let pos_size = 27usize;
        let hidden = 1152usize;
        let buf = device
            .alloc_buffer(
                2 * pos_size * hidden * 4,
                mlx_native::DType::F32,
                vec![2, pos_size, hidden],
            )
            .expect("alloc");
        let mut tensors = HashMap::new();
        tensors.insert(super::super::mmproj::TENSOR_POS_EMBD.to_string(), buf);
        let weights = LoadedMmprojWeights {
            tensors,
            _device: device,
        };
        let (buf, ps, h) = weights.position_embd_table_3d().expect("3d ok");
        assert_eq!(ps, pos_size as u32);
        assert_eq!(h, hidden as u32);
        assert_eq!(buf.shape(), &[2, pos_size, hidden]);
    }

    #[test]
    fn position_embd_table_3d_propagates_missing_tensor_error() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let device = MlxDevice::new().expect("device");
        let weights = LoadedMmprojWeights::empty(device);
        let err = weights.position_embd_table_3d().unwrap_err();
        assert!(format!("{err}").contains("v.position_embd.weight"));
    }

    #[test]
    fn empty_constructor_produces_zero_tensor_weights() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // `empty(device)` is the pub constructor for test/scaffolding
        // call sites that need a LoadedMmprojWeights shape without a
        // real 400MB load. Should len == 0, is_empty == true, and
        // every shortcut accessor should return Err.
        let device = MlxDevice::new().expect("device");
        let weights = LoadedMmprojWeights::empty(device);
        assert_eq!(weights.len(), 0);
        assert!(weights.is_empty());
        assert!(weights.patch_embd_weight().is_err());
        assert!(weights.position_embd_weight().is_err());
        assert!(weights.post_ln_weight().is_err());
        assert!(weights.mm_0_weight().is_err());
        assert!(weights.mm_2_weight().is_err());
        assert!(weights.block_tensor(0, "attn_q.weight").is_err());
    }

    /// iter-115: mm.0 Gemma4ClippableLinear scalar bounds accessors.
    /// Build a synthetic `LoadedMmprojWeights` carrying only the four
    /// 1-element clamp-scalar tensors and assert each accessor returns
    /// the expected scalar.
    #[test]
    fn mm_0_clamp_scalar_accessors_round_trip_single_element_tensors() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        use mlx_native::DType;
        let device = MlxDevice::new().expect("device");
        let put =
            |tensors: &mut HashMap<String, MlxBuffer>, dev: &MlxDevice, name: &str, value: f32| {
                // 1-element f32 tensor with shape [1] — matches what
                // convert_hf_to_gguf.py emits for the unsqueeze(0)'d scalar.
                let buf = dev
                    .alloc_buffer(4, DType::F32, vec![1])
                    .expect("alloc scalar");
                let s: &mut [f32] =
                    unsafe { std::slice::from_raw_parts_mut(buf.contents_ptr() as *mut f32, 1) };
                s[0] = value;
                tensors.insert(name.to_string(), buf);
            };
        let mut tensors: HashMap<String, MlxBuffer> = HashMap::new();
        put(&mut tensors, &device, "mm.0.input_min", -2.5);
        put(&mut tensors, &device, "mm.0.input_max", 2.5);
        put(&mut tensors, &device, "mm.0.output_min", -10.0);
        put(&mut tensors, &device, "mm.0.output_max", 10.0);
        let weights = LoadedMmprojWeights::from_tensors_for_test(tensors, device);

        assert_eq!(weights.mm_0_input_min(), Some(-2.5));
        assert_eq!(weights.mm_0_input_max(), Some(2.5));
        assert_eq!(weights.mm_0_output_min(), Some(-10.0));
        assert_eq!(weights.mm_0_output_max(), Some(10.0));

        let bounds = weights.mm_0_bounds();
        assert!(bounds.any());
        assert_eq!(bounds.input_min, Some(-2.5));
        assert_eq!(bounds.output_max, Some(10.0));
    }

    /// Absence of any clamp-scalar tensor → all accessors return None,
    /// `mm_0_bounds().any()` is false (the projector degrades to a
    /// plain Linear, byte-equivalent to the no-clamp path).
    #[test]
    fn mm_0_clamp_scalar_accessors_return_none_when_absent() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let device = MlxDevice::new().expect("device");
        let weights = LoadedMmprojWeights::empty(device);
        assert_eq!(weights.mm_0_input_min(), None);
        assert_eq!(weights.mm_0_input_max(), None);
        assert_eq!(weights.mm_0_output_min(), None);
        assert_eq!(weights.mm_0_output_max(), None);
        let bounds = weights.mm_0_bounds();
        assert!(!bounds.any());
    }

    // -----------------------------------------------------------------------
    // Wedge-4c.5: install_fused_attn_qkv_slice_views.
    // -----------------------------------------------------------------------

    /// Build a small synthetic Qwen3-VL-style MmprojConfig that
    /// `install_fused_attn_qkv_slice_views` consumes for hidden_size +
    /// num_hidden_layers. Other fields don't affect the slice-view path.
    fn synth_qwen3vl_loader_cfg(hidden: u32, num_layers: u32) -> MmprojConfig {
        MmprojConfig {
            image_size: 32,
            patch_size: 16,
            num_patches_side: 2,
            hidden_size: hidden,
            intermediate_size: hidden * 4,
            num_attention_heads: 4,
            num_hidden_layers: num_layers,
            layer_norm_eps: 1e-6,
            projector: super::super::mmproj::ProjectorType::Qwen3VlMerger,
            image_mean: [0.5, 0.5, 0.5],
            image_std: [0.5, 0.5, 0.5],
            spatial_merge_size: Some(2),
            projection_dim: Some(hidden),
            deepstack_indexes: Some(vec![]),
        }
    }

    /// Allocate an MlxBuffer of `n_elements` F32 values, populated by
    /// `f(i) -> f32` for i in 0..n_elements.
    fn alloc_f32_with<F: FnMut(usize) -> f32>(
        device: &MlxDevice,
        n_elements: usize,
        shape: Vec<usize>,
        mut f: F,
    ) -> MlxBuffer {
        use mlx_native::DType;
        let mut buf = device
            .alloc_buffer(n_elements * 4, DType::F32, shape)
            .expect("alloc f32 buffer");
        {
            let dst: &mut [f32] = buf.as_mut_slice::<f32>().expect("as_mut_slice f32");
            for (i, slot) in dst.iter_mut().enumerate().take(n_elements) {
                *slot = f(i);
            }
        }
        buf
    }

    #[test]
    fn install_fused_attn_qkv_splits_weight_into_three_slice_views() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // Single block, hidden=4 → fused weight is 3*4*4 = 48 floats.
        // Q chunk = floats [0..16] @ byte_offset 0,
        // K = [16..32] @ byte_offset 64 (16 floats × 4 bytes),
        // V = [32..48] @ byte_offset 128 (32 floats × 4 bytes).
        //
        // We verify the slice views by inspecting (byte_offset, shape,
        // dtype, element_count). Direct CPU readback via `as_slice`
        // does NOT honor `byte_offset` — `MlxBuffer::contents_ptr()`
        // returns the start of the whole storage — so we read the
        // backing storage manually using `contents_ptr` + the recorded
        // `byte_offset` to verify the kernel-dispatch contract.
        // This pattern matches what the encoder does at
        // /opt/mlx-native/src/encoder.rs:218-220
        // (`set_buffer(index, metal_buffer(), buf.byte_offset())`).
        let device = MlxDevice::new().expect("device");
        let cfg = synth_qwen3vl_loader_cfg(4, 1);
        let mut tensors: HashMap<String, MlxBuffer> = HashMap::new();
        let fused_buf = alloc_f32_with(&device, 48, vec![12, 4], |i| i as f32);
        tensors.insert("v.blk.0.attn_qkv.weight".to_string(), fused_buf);
        LoadedMmprojWeights::install_fused_attn_qkv_slice_views(&mut tensors, &cfg)
            .expect("install on fused-only tensor map must succeed");
        let q = tensors.get("v.blk.0.attn_q.weight").expect("Q view");
        let k = tensors.get("v.blk.0.attn_k.weight").expect("K view");
        let v = tensors.get("v.blk.0.attn_v.weight").expect("V view");
        // byte_offset matches the per-chunk offset.
        assert_eq!(q.byte_offset(), 0);
        assert_eq!(k.byte_offset(), 64); // 16 floats × 4 bytes.
        assert_eq!(v.byte_offset(), 128); // 32 floats × 4 bytes.
                                          // shape was flattened to a 1-D view of n_elements = hidden*hidden.
        assert_eq!(q.element_count(), 16);
        assert_eq!(k.element_count(), 16);
        assert_eq!(v.element_count(), 16);
        // Verify the underlying bytes are the right region by reading
        // contents_ptr + byte_offset directly. This is what the encoder
        // does on dispatch.
        let read_slice = |buf: &MlxBuffer, n: usize| -> Vec<f32> {
            let ptr = buf.contents_ptr() as *const u8;
            let off = buf.byte_offset() as usize;
            // SAFETY: synthetic test buffer alloc'd above; we hold a
            // shared ref via `tensors.get` and no GPU work is in flight.
            unsafe { std::slice::from_raw_parts((ptr.add(off)) as *const f32, n).to_vec() }
        };
        let q_s = read_slice(q, 16);
        let k_s = read_slice(k, 16);
        let v_s = read_slice(v, 16);
        for i in 0..16 {
            assert_eq!(q_s[i], i as f32, "Q[{i}] must equal fused[{i}]");
            assert_eq!(
                k_s[i],
                (16 + i) as f32,
                "K[{i}] must equal fused[{}]",
                16 + i
            );
            assert_eq!(
                v_s[i],
                (32 + i) as f32,
                "V[{i}] must equal fused[{}]",
                32 + i
            );
        }
    }

    #[test]
    fn install_fused_attn_qkv_handles_optional_bias() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // hidden=8, single block. Fused bias is 3*8 = 24 f32 values.
        // Q bias = [0..8] @ off 0; K bias = [8..16] @ off 32;
        // V bias = [16..24] @ off 64.
        let device = MlxDevice::new().expect("device");
        let cfg = synth_qwen3vl_loader_cfg(8, 1);
        let mut tensors: HashMap<String, MlxBuffer> = HashMap::new();
        let weight_buf = alloc_f32_with(&device, 3 * 8 * 8, vec![24, 8], |i| i as f32);
        let bias_buf = alloc_f32_with(&device, 24, vec![24], |i| -(i as f32));
        tensors.insert("v.blk.0.attn_qkv.weight".to_string(), weight_buf);
        tensors.insert("v.blk.0.attn_qkv.bias".to_string(), bias_buf);
        LoadedMmprojWeights::install_fused_attn_qkv_slice_views(&mut tensors, &cfg)
            .expect("install with optional bias must succeed");
        let q_b = tensors.get("v.blk.0.attn_q.bias").expect("Q bias view");
        let k_b = tensors.get("v.blk.0.attn_k.bias").expect("K bias view");
        let v_b = tensors.get("v.blk.0.attn_v.bias").expect("V bias view");
        assert_eq!(q_b.byte_offset(), 0);
        assert_eq!(k_b.byte_offset(), 32); // 8 f32 = 32 bytes.
        assert_eq!(v_b.byte_offset(), 64);
        assert_eq!(q_b.element_count(), 8);
        assert_eq!(k_b.element_count(), 8);
        assert_eq!(v_b.element_count(), 8);
        // Read the underlying region via contents_ptr + byte_offset.
        let read_slice = |buf: &MlxBuffer, n: usize| -> Vec<f32> {
            let ptr = buf.contents_ptr() as *const u8;
            let off = buf.byte_offset() as usize;
            unsafe { std::slice::from_raw_parts(ptr.add(off) as *const f32, n).to_vec() }
        };
        for (label, view, n_off) in [("Q", q_b, 0usize), ("K", k_b, 8), ("V", v_b, 16)] {
            let s = read_slice(view, 8);
            for i in 0..8 {
                assert_eq!(s[i], -((n_off + i) as f32), "{label} bias [{i}] mismatch");
            }
        }
    }

    #[test]
    fn install_fused_attn_qkv_split_only_is_a_noop() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // Pre-existing split tensors must pass through untouched —
        // the helper is a no-op on split-only inputs.
        let device = MlxDevice::new().expect("device");
        let cfg = synth_qwen3vl_loader_cfg(4, 1);
        let mut tensors: HashMap<String, MlxBuffer> = HashMap::new();
        for suffix in ["attn_q.weight", "attn_k.weight", "attn_v.weight"] {
            let key = format!("v.blk.0.{suffix}");
            let buf = alloc_f32_with(&device, 16, vec![4, 4], |i| i as f32);
            tensors.insert(key, buf);
        }
        let n_before = tensors.len();
        LoadedMmprojWeights::install_fused_attn_qkv_slice_views(&mut tensors, &cfg)
            .expect("split-only must be a no-op");
        assert_eq!(
            tensors.len(),
            n_before,
            "split-only tensor map must be unchanged"
        );
    }

    #[test]
    fn install_fused_attn_qkv_rejects_mixed_block() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // A block with BOTH fused and split is a producer bug — the
        // validator catches it normally; this test guards the
        // defense-in-depth check inside the loader.
        let device = MlxDevice::new().expect("device");
        let cfg = synth_qwen3vl_loader_cfg(4, 1);
        let mut tensors: HashMap<String, MlxBuffer> = HashMap::new();
        tensors.insert(
            "v.blk.0.attn_qkv.weight".to_string(),
            alloc_f32_with(&device, 48, vec![12, 4], |i| i as f32),
        );
        // Add a stray split tensor as well — this should trigger the
        // mixed-state error path.
        tensors.insert(
            "v.blk.0.attn_q.weight".to_string(),
            alloc_f32_with(&device, 16, vec![4, 4], |_| 0.0),
        );
        let err = LoadedMmprojWeights::install_fused_attn_qkv_slice_views(&mut tensors, &cfg)
            .expect_err("mixed fused+split per block must fail loud");
        let msg = format!("{err}");
        assert!(
            msg.contains("BOTH fused") && msg.contains("attn_qkv"),
            "loader error must call out the mixed-state case; got: {msg}"
        );
    }

    #[test]
    fn install_fused_attn_qkv_rejects_undersized_fused_weight() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // Fused weight that's smaller than 3*hidden*hidden floats is a
        // converter bug — slicing would silently return wrong data.
        // Reject loud.
        let device = MlxDevice::new().expect("device");
        let cfg = synth_qwen3vl_loader_cfg(4, 1);
        let mut tensors: HashMap<String, MlxBuffer> = HashMap::new();
        // Allocate only 2*4*4 = 32 floats instead of 48.
        let undersized = alloc_f32_with(&device, 32, vec![8, 4], |i| i as f32);
        tensors.insert("v.blk.0.attn_qkv.weight".to_string(), undersized);
        let err = LoadedMmprojWeights::install_fused_attn_qkv_slice_views(&mut tensors, &cfg)
            .expect_err("undersized fused weight must fail loud");
        let msg = format!("{err}");
        assert!(
            msg.contains("byte_len") && msg.contains("expected"),
            "loader error must name byte_len + expected size; got: {msg}"
        );
    }

    #[test]
    fn install_fused_attn_qkv_multi_block_batches_correctly() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // Multi-block: every block independently slices its fused
        // tensor. Block N's Q view points at block N's storage at
        // byte_offset 0 — distinct backing storage from block M (M≠N).
        let device = MlxDevice::new().expect("device");
        let cfg = synth_qwen3vl_loader_cfg(4, 3);
        let mut tensors: HashMap<String, MlxBuffer> = HashMap::new();
        for layer_idx in 0..3 {
            let fused = alloc_f32_with(&device, 48, vec![12, 4], |i| (layer_idx * 100 + i) as f32);
            tensors.insert(format!("v.blk.{layer_idx}.attn_qkv.weight"), fused);
        }
        LoadedMmprojWeights::install_fused_attn_qkv_slice_views(&mut tensors, &cfg)
            .expect("multi-block install must succeed");
        let read_slice = |buf: &MlxBuffer, n: usize| -> Vec<f32> {
            let ptr = buf.contents_ptr() as *const u8;
            let off = buf.byte_offset() as usize;
            unsafe { std::slice::from_raw_parts(ptr.add(off) as *const f32, n).to_vec() }
        };
        for layer_idx in 0..3 {
            let q = tensors
                .get(&format!("v.blk.{layer_idx}.attn_q.weight"))
                .unwrap_or_else(|| panic!("Q for block {layer_idx}"));
            assert_eq!(q.byte_offset(), 0, "Q view always at fused-tensor start");
            let q_s = read_slice(q, 16);
            for i in 0..16 {
                assert_eq!(
                    q_s[i],
                    (layer_idx * 100 + i) as f32,
                    "block {layer_idx} Q[{i}] mismatch"
                );
            }
        }
    }
}