cera 0.5.4

//! GGUF quantization formats: block layouts, dequantization, and the
//! `vec_dot` kernels that consume a block without materializing it.
//!
//! Each format has a `#[repr(C)]` block struct mirroring its on-disk layout, a
//! `dequantize_*_row`/`_matrix` pair, and a `vec_dot_*` that fuses the dequant
//! into the dot product. The f16 converters live here too, open-coded rather
//! than taken from the `half` crate, which is a trap in a hot loop.
#![warn(missing_docs, clippy::missing_docs_in_private_items)]
//
// Both halves: `missing_docs` for the public items, the clippy one for the
// private ones. Scoped here rather than crate-wide (646 and 887 items
// elsewhere) because these files
// have repeatedly lost a doc comment to an insertion or deletion above an item:
// a doc binds to the next item below it, so both operations silently reassign
// it, and the item left bare is otherwise silent. There is no `missing_docs`
// for private items by default, and rustdoc stays green because the intra-doc
// links still resolve.
//
// Known limit: clippy skips `#[cfg(test)]`, so this does not cover test
// modules. A `#[test]` that loses its attribute is caught by `dead_code`
// instead (it becomes an uncalled private fn), but a doc that merely moves
// between two live test functions is caught by neither.

#[cfg_attr(not(feature = "parallel"), allow(unused_imports))]
use crate::par::{IndexedParallelIterator, ParallelIterator, ParallelSlice, ParallelSliceMut};

/// Convert one IEEE-754 half (raw bits) to f32. Accumulation downstream stays
/// f32 for softmax stability.
///
/// Open-coded rather than `half::f16::to_f32`, which is a trap in a loop: its
/// aarch64 path runs `is_aarch64_feature_detected!("fp16")` per call and lowers
/// to two out-of-line `bl`s (`detect_and_initialize`, then `f16_to_f32_fp16`).
/// Neither inlines, so every call clobbers the caller-saved registers and spills
/// whatever the caller had live. The disassembly of `attn_scores_f16` showed
/// exactly that, and this function is called from the *tail* loops of the NEON
/// kernels, not only from the scalar fallback. An earlier revision of this doc
/// claimed the wrapper "lowers to a single `fcvt`"; it does not, and that claim
/// went unchecked.
///
/// Integer-only on purpose. The well-known branchless form multiplies by
/// `2^112` to renormalize subnormals, which makes the result depend on
/// `FPCR.FZ`: under flush-to-zero every subnormal half would widen to the wrong
/// value. This version cannot be perturbed by the floating-point environment.
/// The branches are cheap because normals dominate overwhelmingly, and
/// `f16_widen_matches_half_crate_exhaustively` pins all 65536 patterns against
/// the `half` crate.
#[inline(always)]
pub fn f16_to_f32(bits: u16) -> f32 {
    #[cfg(target_arch = "aarch64")]
    unsafe {
        let out: f32;
        std::arch::asm!(
            "fcvt {0:s}, {1:h}",
            out(vreg) out,
            in(vreg) bits,
            options(pure, nomem, nostack)
        );
        out
    }

    #[cfg(not(target_arch = "aarch64"))]
    {
        let sign = (u32::from(bits) & 0x8000) << 16;
        let exp = (u32::from(bits) >> 10) & 0x1F;
        let mant = u32::from(bits) & 0x03FF;
        let rest = match exp {
            // Zero, or a subnormal that has to be renormalized into a normal f32.
            0 if mant == 0 => 0,
            0 => {
                let shift = mant.leading_zeros() - 21;
                ((113 - shift) << 23) | ((mant << (shift + 13)) & 0x007F_FFFF)
            }
            0x1F if mant == 0 => 0x7F80_0000,
            0x1F => 0x7FC0_0000 | (mant << 13),
            _ => ((exp + 112) << 23) | (mant << 13),
        };
        f32::from_bits(sign | rest)
    }
}

/// Convert one bfloat16 (raw bits) to f32.
///
/// bf16 is the top 16 bits of an f32: same 8-bit exponent, same bias, so every
/// finite value, zero and infinity widens exactly by a 16-bit shift, with no
/// rounding and no subnormal renormalization. Only NaN needs a case, and only
/// for its quiet bit: IEEE 754 says widening a signaling NaN quiets it, which is
/// what `half::bf16::to_f32` does and what [`f16_to_f32`] does on the
/// same input. A bare shift would leave sNaN signaling and disagree with
/// `MmapWeight::dequantize_row` on those patterns.
///
/// Exists so the widen has one spelling. It was open-coded in three places (the
/// bf16 GEMV, the embedding lookup, and a test reference) while
/// `model/weights.rs` used the `half` crate, which is exactly the drift that
/// lets two copies disagree; that site now calls this too.
/// `bf16_widen_matches_half_crate_exhaustively` pins all 65536 patterns.
#[inline(always)]
pub fn bf16_to_f32(bits: u16) -> f32 {
    // Exponent all ones with a non-zero significand is NaN; 0x7F80 itself is
    // infinity and must not gain a mantissa bit.
    let quiet = if (bits & 0x7FFF) > 0x7F80 {
        0x0040_0000
    } else {
        0
    };
    f32::from_bits((u32::from(bits) << 16) | quiet)
}

/// Convert one f32 to IEEE-754 half (raw bits), round-to-nearest-even.
///
/// The narrowing twin of [`f16_to_f32`], and open-coded for the same reason:
/// `half::f16::from_f32` is an out-of-line call with a runtime feature probe,
/// and the disassembly put 17 of them in hot paths (TurboQuant key/value
/// append, the LFM2 and Llama prefill and attention blocks, the KV rope shift,
/// and the Q8_0 activation quantizer).
///
/// Correctness here is load-bearing: this writes the f16 KV cache and the
/// TurboQuant norms, so a rounding slip corrupts inference rather than merely
/// slowing it. The full 2^32 input space was swept against the `half` crate
/// once, off-CI, with zero mismatches; `f32_to_f16_matches_half_crate` is what
/// remains of that in the suite, and it enumerates the classes a hand-rolled
/// rounder actually gets wrong rather than re-running the sweep.
#[inline(always)]
pub(crate) fn f32_to_f16(v: f32) -> u16 {
    let b = v.to_bits();
    let sign = ((b >> 16) & 0x8000) as u16;
    let exp = ((b >> 23) & 0xFF) as i32;
    let mant = b & 0x007F_FFFF;

    if exp == 0xFF {
        // Infinity keeps a zero significand. A NaN must stay a NaN even when
        // its payload truncates to zero, so force the quiet bit rather than
        // shifting alone.
        return if mant == 0 {
            sign | 0x7C00
        } else {
            sign | 0x7E00 | ((mant >> 13) as u16 & 0x03FF)
        };
    }

    // Rebias 127 -> 15.
    let e = exp - 127 + 15;
    if e >= 0x1F {
        return sign | 0x7C00; // overflow saturates to infinity
    }
    if e <= 0 {
        // Below the subnormal range entirely: round toward the nearer of zero
        // and the smallest subnormal, which for e < -10 is always zero.
        if e < -10 {
            return sign;
        }
        // Restore the implicit leading 1, then shift into subnormal position
        // with round-to-nearest-even.
        let m = mant | 0x0080_0000;
        let shift = (14 - e) as u32; // 14..=24
        let round = ((m >> (shift - 1)) & 1)
            & (((m & ((1 << (shift - 1)) - 1)) != 0) as u32 | ((m >> shift) & 1));
        return sign | ((m >> shift) + round) as u16;
    }

    // Normal. Round the 23-bit significand to 10 bits, ties to even; a carry
    // out of the significand bumps the exponent, and may reach infinity.
    let lsb = (mant >> 13) & 1;
    let rounded = mant + 0x0FFF + lsb;
    let carry = rounded >> 23;
    let e = e as u32 + carry;
    if e >= 0x1F {
        return sign | 0x7C00;
    }
    sign | ((e << 10) as u16) | (((rounded >> 13) & 0x03FF) as u16)
}

// ── Block layouts ────────────────────────────────────────────────────────────

/// Q4_0 quantization block: 32 values in 18 bytes.
///
/// Layout:
///   d: f16 (2 bytes) — scale factor
///   qs: [u8; 16] (16 bytes) — 32 4-bit unsigned quantized values (offset by 8)
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct BlockQ4_0 {
    /// Scale, f16 as raw bits.
    pub d: u16,
    /// 32 4-bit quants, two per byte, biased by 8.
    pub qs: [u8; 16],
}

const _: () = assert!(size_of::<BlockQ4_0>() == 18);

/// Q4_1 quantization block: 32 values in 20 bytes.
///
/// Layout:
///   d:  f16 (2 bytes) — scale
///   m:  f16 (2 bytes) — minimum
///   qs: [u8; 16]      — two 4-bit quants per byte
///
/// Differs from Q4_0 in more than the extra field: Q4_0 recenters its nibble
/// around zero (`(q - 8) * d`), while Q4_1 carries an explicit minimum and does
/// not recenter (`q * d + m`). The nibble *packing* is identical — element `i`
/// is the low nibble of `qs[i]` and element `i + 16` the high nibble — so only
/// the arithmetic changes, not the unpacking.
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct BlockQ4_1 {
    /// Scale, f16 as raw bits.
    pub d: u16,
    /// Minimum, f16 as raw bits.
    pub m: u16,
    /// 32 4-bit quants, two per byte, unbiased.
    pub qs: [u8; 16],
}

const _: () = assert!(size_of::<BlockQ4_1>() == 20);

/// Q8_0 quantization block: 32 values in 34 bytes.
///
/// Layout:
///   delta: f16 (2 bytes) — scale factor
///   quants: [i8; 32] (32 bytes) — quantized values
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct BlockQ8_0 {
    /// Scale, f16 as raw bits.
    pub delta: u16,
    /// 32 signed 8-bit quants.
    pub quants: [i8; 32],
}

const _: () = assert!(size_of::<BlockQ8_0>() == 34);

/// Q4_K_M quantization block: 256 values in 144 bytes.
///
/// Layout:
///   d: f16 (2 bytes) — super-block scale
///   dmin: f16 (2 bytes) — super-block minimum
///   scales: [u8; 12] (12 bytes) — packed sub-block scales and mins
///   qs: [u8; 128] (128 bytes) — 256 4-bit quantized values
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct BlockQ4KM {
    /// Super-scale applied to the 6-bit block scales, f16 as raw bits.
    pub d: u16,
    /// Super-scale applied to the 6-bit block minima, f16 as raw bits.
    pub dmin: u16,
    /// Eight 6-bit scale/min pairs, packed into 12 bytes.
    pub scales: [u8; 12],
    /// 256 4-bit quants, two per byte.
    pub qs: [u8; 128],
}

const _: () = assert!(size_of::<BlockQ4KM>() == 144);

/// Q6_K quantization block: 256 values in 210 bytes.
///
/// Layout (from ggml-common.h):
///   ql: [u8; 128] — lower 4 bits of 6-bit quants
///   qh: [u8; 64]  — upper 2 bits of 6-bit quants
///   scales: [i8; 16] — per-16-element sub-block scales (8-bit signed)
///   d: f16 (2 bytes) — super-block scale
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct BlockQ6K {
    /// Low 4 bits of each of the 256 quants.
    pub ql: [u8; 128],
    /// High 2 bits of each of the 256 quants.
    pub qh: [u8; 64],
    /// Sixteen signed 8-bit block scales.
    pub scales: [i8; 16],
    /// Super-scale, f16 as raw bits.
    pub d: u16,
}

const _: () = assert!(size_of::<BlockQ6K>() == 210);

/// Q5_K quantization block: 256 values in 176 bytes.
///
/// Layout (from ggml-common.h `block_q5_K`):
///   d: f16 (2 bytes) — super-block scale for the 6-bit sub-block scales
///   dmin: f16 (2 bytes) — super-block scale for the 6-bit sub-block mins
///   scales: [u8; 12] — 8 sub-block scales + 8 mins, 6-bit packed (identical
///     layout to Q4_K, decoded via `decode_q4km_scales`)
///   qh: [u8; 32] — the 5th (high) bit of each of the 256 quants
///   qs: [u8; 128] — the low 4 bits of each of the 256 quants
#[repr(C, packed)]
#[derive(Debug, Clone, Copy)]
pub struct BlockQ5K {
    /// Super-scale applied to the 6-bit block scales, f16 as raw bits.
    pub d: u16,
    /// Super-scale applied to the 6-bit block minima, f16 as raw bits.
    pub dmin: u16,
    /// Eight 6-bit scale/min pairs, packed into 12 bytes.
    pub scales: [u8; 12],
    /// The fifth bit of each of the 256 quants.
    pub qh: [u8; 32],
    /// Low 4 bits of each of the 256 quants.
    pub qs: [u8; 128],
}

const _: () = assert!(size_of::<BlockQ5K>() == 176);

/// Define a `dequantize_*_matrix` from its block width, block struct and per-row
/// helper.
///
/// Six of these existed verbatim, differing only in those three things and in
/// the function name repeated inside their own assertion messages, a block of
/// `debug_assert`s per quant kept in agreement by hand.
///
/// Those asserts are a debug-build check on the caller, not a release guard:
/// they compile out, and `par_chunks(row_bytes)` on a `k` that is not a whole
/// number of blocks then splits the rows at the wrong offsets and dequantizes
/// garbage without failing. What actually upholds the invariant in release is
/// the caller. Every one of these runs behind `batched_gemm_supports`, which
/// requires `k.is_multiple_of(256)` for the K-quants, and the 32-wide formats
/// get it from GGUF itself, which cannot store a quantized row that is not a
/// whole number of blocks. A new caller outside that gate would need its own
/// check.
///
/// What this does *not* do is make a wrong pairing impossible. `$elems`,
/// `$block` and `$row` are three independent arguments, and nothing ties them
/// to each other: a `$block` from one quant with a `$row` from another still
/// compiles, and shows up as a tripped `debug_assert` rather than an error.
/// What it removes is the hand-copied byte arithmetic and the six chances to
/// mistype a name inside an assertion message.
///
/// Invoked in place in each format's section rather than all together, so the
/// file stays organized by quant format.
macro_rules! dequantize_matrix {
    ($name:ident, $elems:expr, $block:ty, $row:path, $($summary:expr),+ $(,)?) => {
        $(#[doc = $summary])+
        ///
        /// `src` is the raw packed block bytes; `out` must have space for
        /// `m * k` f32s. Rows are dequantized in parallel with rayon, whose
        /// split-on-demand runs tiny inputs on a single worker, so there is no
        /// manual cutoff.
        pub fn $name(src: &[u8], m: usize, k: usize, out: &mut [f32]) {
            debug_assert_eq!(
                k % $elems,
                0,
                concat!(
                    stringify!($name),
                    ": k must be a multiple of ",
                    stringify!($elems)
                )
            );
            let row_bytes = (k / $elems) * size_of::<$block>();
            debug_assert_eq!(
                src.len(),
                m * row_bytes,
                concat!(stringify!($name), ": src length mismatch")
            );
            debug_assert_eq!(
                out.len(),
                m * k,
                concat!(stringify!($name), ": out length mismatch")
            );

            let num_threads = crate::par::current_num_threads().max(1);
            let rows_per_chunk = (m / num_threads).max(1);
            let dst_chunk_len = rows_per_chunk * k;
            let src_chunk_len = rows_per_chunk * row_bytes;

            out.par_chunks_mut(dst_chunk_len)
                .zip(src.par_chunks(src_chunk_len))
                .for_each(|(dst_chunk, src_chunk)| {
                    for (dst_row, src_row) in dst_chunk.chunks_mut(k).zip(src_chunk.chunks(row_bytes)) {
                        $row(src_row, dst_row);
                    }
                });
        }
    };
}

// ── Q4_0 dequantization ─────────────────────────────────────────────────────

/// Dequantize a single Q4_0 block to 32 f32 values.
///
/// Each byte in qs holds two 4-bit unsigned values (low nibble, high nibble).
/// Values are offset by -8 to center around zero: value = (nibble - 8) * d.
pub fn dequantize_q4_0_block(block: &BlockQ4_0) -> [f32; 32] {
    let d = f16_to_f32(block.d);
    let mut out = [0.0f32; 32];

    for i in 0..16 {
        let byte = block.qs[i];
        let lo = (byte & 0xF) as i32 - 8;
        let hi = (byte >> 4) as i32 - 8;
        out[i] = lo as f32 * d;
        out[i + 16] = hi as f32 * d;
    }
    out
}

/// Dequantize a row of Q4_0 blocks. `src` is raw bytes, `dst` is f32 output.
#[inline]
pub fn dequantize_q4_0_row(src: &[u8], dst: &mut [f32]) {
    let block_size = size_of::<BlockQ4_0>();
    let n_blocks = src.len() / block_size;
    debug_assert_eq!(src.len() % block_size, 0);
    debug_assert_eq!(dst.len(), n_blocks * 32);

    #[cfg(target_arch = "aarch64")]
    unsafe {
        use std::arch::aarch64::*;
        let mask_lo = vdupq_n_u8(0x0F);
        let offset_8 = vdupq_n_s8(0x8);

        for i in 0..n_blocks {
            let block_ptr = src.as_ptr().add(i * block_size) as *const BlockQ4_0;
            let block = &*block_ptr;
            let d_val = f16_to_f32(block.d);
            let d_vec = vdupq_n_f32(d_val);

            let v = vld1q_u8(block.qs.as_ptr());
            let v_lo = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(v, mask_lo)), offset_8);
            let v_hi = vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8::<4>(v)), offset_8);

            let v_lo_s16_lo = vmovl_s8(vget_low_s8(v_lo));
            let v_lo_s16_hi = vmovl_high_s8(v_lo);
            let v_lo_0 = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(v_lo_s16_lo))), d_vec);
            let v_lo_1 = vmulq_f32(vcvtq_f32_s32(vmovl_high_s16(v_lo_s16_lo)), d_vec);
            let v_lo_2 = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(v_lo_s16_hi))), d_vec);
            let v_lo_3 = vmulq_f32(vcvtq_f32_s32(vmovl_high_s16(v_lo_s16_hi)), d_vec);

            let v_hi_s16_lo = vmovl_s8(vget_low_s8(v_hi));
            let v_hi_s16_hi = vmovl_high_s8(v_hi);
            let v_hi_0 = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(v_hi_s16_lo))), d_vec);
            let v_hi_1 = vmulq_f32(vcvtq_f32_s32(vmovl_high_s16(v_hi_s16_lo)), d_vec);
            let v_hi_2 = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(v_hi_s16_hi))), d_vec);
            let v_hi_3 = vmulq_f32(vcvtq_f32_s32(vmovl_high_s16(v_hi_s16_hi)), d_vec);

            let out_ptr = dst.as_mut_ptr().add(i * 32);
            vst1q_f32(out_ptr, v_lo_0);
            vst1q_f32(out_ptr.add(4), v_lo_1);
            vst1q_f32(out_ptr.add(8), v_lo_2);
            vst1q_f32(out_ptr.add(12), v_lo_3);

            vst1q_f32(out_ptr.add(16), v_hi_0);
            vst1q_f32(out_ptr.add(20), v_hi_1);
            vst1q_f32(out_ptr.add(24), v_hi_2);
            vst1q_f32(out_ptr.add(28), v_hi_3);
        }
    }

    #[cfg(not(target_arch = "aarch64"))]
    for i in 0..n_blocks {
        let block_bytes = &src[i * block_size..(i + 1) * block_size];
        let block = unsafe { &*(block_bytes.as_ptr() as *const BlockQ4_0) };
        let values = dequantize_q4_0_block(block);
        dst[i * 32..(i + 1) * 32].copy_from_slice(&values);
    }
}

dequantize_matrix!(
    dequantize_q4_0_matrix,
    32,
    BlockQ4_0,
    dequantize_q4_0_row,
    "Dequantize a Q4_0 matrix of shape `[m, k]` (row-major) to `out`."
);

/// Dot product of a Q4_0 block with an f32 vector of length 32. Scalar version.
pub fn vec_dot_q4_0_f32_scalar(block: &BlockQ4_0, y: &[f32]) -> f32 {
    debug_assert_eq!(y.len(), 32);
    let d = f16_to_f32(block.d);
    let mut sum = 0.0f32;

    for i in 0..16 {
        let byte = block.qs[i];
        let lo = (byte & 0xF) as i32 - 8;
        let hi = (byte >> 4) as i32 - 8;
        sum += lo as f32 * y[i];
        sum += hi as f32 * y[i + 16];
    }
    sum * d
}

// ── Q4_1 dequantization ─────────────────────────────────────────────────────

/// Dequantize a single Q4_1 block to 32 f32 values.
///
/// `q * d + m`, with `q` the raw nibble in `[0, 15]` — no `- 8` recentering.
pub fn dequantize_q4_1_block(block: &BlockQ4_1) -> [f32; 32] {
    let d = f16_to_f32(block.d);
    let m = f16_to_f32(block.m);
    let mut out = [0.0f32; 32];

    for i in 0..16 {
        let byte = block.qs[i];
        let lo = (byte & 0xF) as i32;
        let hi = (byte >> 4) as i32;
        out[i] = lo as f32 * d + m;
        out[i + 16] = hi as f32 * d + m;
    }
    out
}

/// Dequantize a row of Q4_1 blocks. `src` is raw bytes, `dst` is f32 output.
pub fn dequantize_q4_1_row(src: &[u8], dst: &mut [f32]) {
    let block_size = size_of::<BlockQ4_1>();
    let n_blocks = src.len() / block_size;
    debug_assert_eq!(src.len() % block_size, 0);
    debug_assert_eq!(dst.len(), n_blocks * 32);

    for i in 0..n_blocks {
        let block_bytes = &src[i * block_size..(i + 1) * block_size];
        // SAFETY: `BlockQ4_1` is `repr(C, packed)` over plain integers, and the
        // slice above is exactly `size_of::<BlockQ4_1>()` bytes.
        let block = unsafe { &*(block_bytes.as_ptr() as *const BlockQ4_1) };
        let values = dequantize_q4_1_block(block);
        dst[i * 32..(i + 1) * 32].copy_from_slice(&values);
    }
}

dequantize_matrix!(
    dequantize_q4_1_matrix,
    32,
    BlockQ4_1,
    dequantize_q4_1_row,
    "Dequantize a Q4_1 matrix of shape `[m, k]` (row-major) to `out`."
);

/// Dot product of a Q4_1 block with an f32 vector of length 32.
///
/// The scalar reference for Q4_1: the batched int8 GEMMs (`gemm_q4_1_q8_0_neon`
/// on aarch64, `gemm_q4_1_q8_0` on x86) and the Metal kernels are checked against
/// it. Those int8 kernels cannot reuse the Q4_0 `dpbusd` path unchanged — that
/// sign trick assumes a zero-centred quant — so they carry a separate `m·Σ(x)`
/// correction term, reusing the K-quant activation col-sum machinery (added here,
/// where the K-quant `dmin` term subtracts).
///
/// `sum(q_i * y_i) * d + m * sum(y_i)`: the minimum is a per-block constant, so
/// it factors out of the dot product rather than being added per element.
pub fn vec_dot_q4_1_f32(block: &BlockQ4_1, y: &[f32]) -> f32 {
    debug_assert_eq!(y.len(), 32);
    let d = f16_to_f32(block.d);
    let m = f16_to_f32(block.m);
    let mut qsum = 0.0f32;
    let mut ysum = 0.0f32;

    for i in 0..16 {
        let byte = block.qs[i];
        let lo = (byte & 0xF) as i32;
        let hi = (byte >> 4) as i32;
        qsum += lo as f32 * y[i];
        qsum += hi as f32 * y[i + 16];
        ysum += y[i] + y[i + 16];
    }
    qsum * d + m * ysum
}

// ── Q8_0 dequantization ─────────────────────────────────────────────────────

/// Dequantize a single Q8_0 block to 32 f32 values.
pub fn dequantize_q8_0_block(block: &BlockQ8_0) -> [f32; 32] {
    let d = f16_to_f32(block.delta);
    let mut out = [0.0f32; 32];
    for (o, &q) in out.iter_mut().zip(block.quants.iter()) {
        *o = q as f32 * d;
    }
    out
}

/// Dequantize a row of Q8_0 blocks. `src` is raw bytes, `dst` is f32 output.
#[inline]
pub fn dequantize_q8_0_row(src: &[u8], dst: &mut [f32]) {
    let block_size = size_of::<BlockQ8_0>();
    let n_blocks = src.len() / block_size;
    debug_assert_eq!(src.len() % block_size, 0);
    debug_assert_eq!(dst.len(), n_blocks * 32);

    #[cfg(target_arch = "aarch64")]
    unsafe {
        use std::arch::aarch64::*;
        for i in 0..n_blocks {
            let block = &*(src.as_ptr().add(i * block_size) as *const BlockQ8_0);
            let d_val = f16_to_f32(block.delta);
            let d_vec = vdupq_n_f32(d_val);

            let q0 = vld1q_s8(block.quants.as_ptr());
            let q1 = vld1q_s8(block.quants.as_ptr().add(16));

            let q0_s16_lo = vmovl_s8(vget_low_s8(q0));
            let q0_s16_hi = vmovl_high_s8(q0);
            let q1_s16_lo = vmovl_s8(vget_low_s8(q1));
            let q1_s16_hi = vmovl_high_s8(q1);

            let f0 = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(q0_s16_lo))), d_vec);
            let f1 = vmulq_f32(vcvtq_f32_s32(vmovl_high_s16(q0_s16_lo)), d_vec);
            let f2 = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(q0_s16_hi))), d_vec);
            let f3 = vmulq_f32(vcvtq_f32_s32(vmovl_high_s16(q0_s16_hi)), d_vec);

            let f4 = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(q1_s16_lo))), d_vec);
            let f5 = vmulq_f32(vcvtq_f32_s32(vmovl_high_s16(q1_s16_lo)), d_vec);
            let f6 = vmulq_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(q1_s16_hi))), d_vec);
            let f7 = vmulq_f32(vcvtq_f32_s32(vmovl_high_s16(q1_s16_hi)), d_vec);

            let out_ptr = dst.as_mut_ptr().add(i * 32);
            vst1q_f32(out_ptr, f0);
            vst1q_f32(out_ptr.add(4), f1);
            vst1q_f32(out_ptr.add(8), f2);
            vst1q_f32(out_ptr.add(12), f3);
            vst1q_f32(out_ptr.add(16), f4);
            vst1q_f32(out_ptr.add(20), f5);
            vst1q_f32(out_ptr.add(24), f6);
            vst1q_f32(out_ptr.add(28), f7);
        }
    }

    #[cfg(not(target_arch = "aarch64"))]
    for i in 0..n_blocks {
        let block_bytes = &src[i * block_size..(i + 1) * block_size];
        let block = unsafe { &*(block_bytes.as_ptr() as *const BlockQ8_0) };
        let values = dequantize_q8_0_block(block);
        dst[i * 32..(i + 1) * 32].copy_from_slice(&values);
    }
}

dequantize_matrix!(
    dequantize_q8_0_matrix,
    32,
    BlockQ8_0,
    dequantize_q8_0_row,
    "Dequantize a Q8_0 matrix of shape `[m, k]` (row-major) to `out`."
);

dequantize_matrix!(
    dequantize_q4_k_m_matrix,
    256,
    BlockQ4KM,
    dequantize_q4_k_m_row,
    "Dequantize a Q4_K matrix of shape `[m, k]` (row-major) to `out`.",
    "",
    "Superblocks are 256 wide, so `k` must be a multiple of 256 (not 32).",
);

dequantize_matrix!(
    dequantize_q5_k_matrix,
    256,
    BlockQ5K,
    dequantize_q5_k_row,
    "Dequantize a Q5_K matrix of shape `[m, k]` (row-major) to `out`.",
    "",
    "Superblocks are 256 wide, so `k` must be a multiple of 256 (not 32).",
    "",
    "Exists for the BLAS prefill route, which dequantizes the weight and SGEMMs",
    "rather than running an int8 kernel: Q5_K is the one shipped K-quant with no",
    "int8 GEMM, so before this it fell through to the per-token GEMV and prefill",
    "collapsed (measured 4-5 tok/s against Q4_K_M's 228 on the same model and host).",
);

dequantize_matrix!(
    dequantize_q6_k_matrix,
    256,
    BlockQ6K,
    dequantize_q6_k_row,
    "Dequantize a Q6_K matrix of shape `[m, k]` (row-major) to `out`.",
    "",
    "Superblocks are 256 wide, so `k` must be a multiple of 256 (not 32).",
);

/// Dot product of a Q8_0 block with an f32 vector of length 32. Scalar version.
pub fn vec_dot_q8_0_f32_scalar(block: &BlockQ8_0, y: &[f32]) -> f32 {
    debug_assert_eq!(y.len(), 32);
    let d = f16_to_f32(block.delta);
    let sum: f32 = block
        .quants
        .iter()
        .zip(y.iter())
        .map(|(&q, &y)| q as f32 * y)
        .sum();
    sum * d
}

// ── Q4_K_M dequantization ───────────────────────────────────────────────────

/// Decode the packed sub-block scales and minimums from Q4_K_M's 12-byte scales array.
///
/// Q4_K_M has 8 sub-blocks of 32 values each. The 12 bytes encode:
/// - 8 6-bit scales and 8 6-bit minimums
///
/// Bytes 0-3: low 4 bits of scales[0..3] and mins[0..3]  (packed as scale|min per byte)
///   Wait — actually llama.cpp packs them differently.
///
/// From ggml-quants.c (get_scale_min_k4):
///   j < 4:  sc = scales[j] & 63,      m = scales[j+4] & 63
///   j >= 4: sc = (scales[j+4] & 0xF) | ((scales[j-4] >> 6) << 4),
///           m  = (scales[j+4] >> 4)   | ((scales[j-0] >> 6) << 4)
///
/// Returns (scales[8], mins[8]).
pub(crate) fn decode_q4km_scales(scales: &[u8; 12]) -> ([u8; 8], [u8; 8]) {
    let mut sc = [0u8; 8];
    let mut mn = [0u8; 8];

    for j in 0..4 {
        sc[j] = scales[j] & 63;
        mn[j] = scales[j + 4] & 63;
    }
    for j in 4..8 {
        sc[j] = (scales[j + 4] & 0xF) | ((scales[j - 4] >> 6) << 4);
        mn[j] = (scales[j + 4] >> 4) | ((scales[j] >> 6) << 4);
    }

    (sc, mn)
}

/// Dequantize a single Q4_K_M block to 256 f32 values.
///
/// Ported from llama.cpp's dequantize_row_q4_K.
pub fn dequantize_q4_k_m_block(block: &BlockQ4KM) -> [f32; 256] {
    let d = f16_to_f32(block.d);
    let dmin = f16_to_f32(block.dmin);
    let (sc, mn) = decode_q4km_scales(&block.scales);

    let mut out = [0.0f32; 256];
    let qs = &block.qs;

    for j in 0..8 {
        // Each sub-block has 32 values
        let sc_val = d * sc[j] as f32;
        let mn_val = dmin * mn[j] as f32;

        // First 16 values: low nibble of qs[j*16..j*16+16]
        // Second 16 values: high nibble of qs[j*16..j*16+16]
        // But the layout is actually:
        //   sub-blocks 0-3 use qs[0..64], lower nibble for 0-1, upper for 2-3
        //   sub-blocks 4-7 use qs[64..128], lower nibble for 4-5, upper for 6-7
        //
        // Actually from llama.cpp:
        //   for (int l = 0; l < 32; ++l) {
        //     *y++ = d * sc[is] * ((q[l] & 0xF) - (m ? dmin * mn[is] : 0))
        //   but that's not right either.
        //
        // Let me re-read the llama.cpp source carefully.
        // The actual layout from dequantize_row_q4_K:
        //
        //   q = qs (pointer to start of qs array)
        //   for j in 0..QK_K/64:     (QK_K=256, so j in 0..4)
        //     sc1 = get_scale(j*2), mn1 = get_min(j*2)
        //     sc2 = get_scale(j*2+1), mn2 = get_min(j*2+1)
        //     for l in 0..32:
        //       y[l+0]  = d * sc1 * (q[l] & 0xF) - dmin * mn1
        //       y[l+32] = d * sc2 * (q[l] >> 4)   - dmin * mn2
        //     q += 32, y += 64
        //
        // So it processes 64 values at a time using 32 bytes of qs.
        // Each byte holds two 4-bit values: low nibble and high nibble.
        let _ = (sc_val, mn_val); // will use below
    }

    // Re-implement following llama.cpp's actual loop structure
    let mut qi = 0; // index into qs
    let mut yi = 0; // index into output

    for j in 0..4 {
        let d_sc1 = d * sc[j * 2] as f32;
        let d_mn1 = dmin * mn[j * 2] as f32;
        let d_sc2 = d * sc[j * 2 + 1] as f32;
        let d_mn2 = dmin * mn[j * 2 + 1] as f32;

        for l in 0..32 {
            out[yi + l] = d_sc1 * (qs[qi + l] & 0xF) as f32 - d_mn1;
            out[yi + l + 32] = d_sc2 * (qs[qi + l] >> 4) as f32 - d_mn2;
        }
        qi += 32;
        yi += 64;
    }

    out
}

/// Dequantize a row of Q4_K_M blocks. `src` is raw bytes, `dst` is f32 output.
pub fn dequantize_q4_k_m_row(src: &[u8], dst: &mut [f32]) {
    let block_size = size_of::<BlockQ4KM>();
    let n_blocks = src.len() / block_size;
    debug_assert_eq!(src.len() % block_size, 0);
    debug_assert_eq!(dst.len(), n_blocks * 256);

    for i in 0..n_blocks {
        let block_bytes = &src[i * block_size..(i + 1) * block_size];
        let block = unsafe { &*(block_bytes.as_ptr() as *const BlockQ4KM) };
        let values = dequantize_q4_k_m_block(block);
        dst[i * 256..(i + 1) * 256].copy_from_slice(&values);
    }
}

/// Dot product of a Q4_K_M block with an f32 vector of length 256. Scalar version.
///
/// Ported from llama.cpp's ggml_vec_dot_q4_K_q8_K.
pub fn vec_dot_q4_k_m_f32_scalar(block: &BlockQ4KM, y: &[f32]) -> f32 {
    debug_assert_eq!(y.len(), 256);

    let d = f16_to_f32(block.d);
    let dmin = f16_to_f32(block.dmin);
    let (sc, mn) = decode_q4km_scales(&block.scales);
    let qs = &block.qs;

    let mut sumf = 0.0f32;
    let mut qi = 0usize;
    let mut yi = 0usize;

    for j in 0..4 {
        let sc1 = sc[j * 2] as f32;
        let mn1 = mn[j * 2] as f32;
        let sc2 = sc[j * 2 + 1] as f32;
        let mn2 = mn[j * 2 + 1] as f32;

        let mut sum1 = 0.0f32;
        let mut sum2 = 0.0f32;
        let mut sum_mn1 = 0.0f32;
        let mut sum_mn2 = 0.0f32;

        for l in 0..32 {
            sum1 += (qs[qi + l] & 0xF) as f32 * y[yi + l];
            sum2 += (qs[qi + l] >> 4) as f32 * y[yi + l + 32];
            sum_mn1 += y[yi + l];
            sum_mn2 += y[yi + l + 32];
        }

        sumf += d * (sc1 * sum1 + sc2 * sum2) - dmin * (mn1 * sum_mn1 + mn2 * sum_mn2);
        qi += 32;
        yi += 64;
    }

    sumf
}

// ── Q6_K dequantization ────────────────────────────────────────────────────

/// Dequantize a single Q6_K block to 256 f32 values.
///
/// Ported from llama.cpp's `dequantize_row_q6_K`. The 256 values are processed
/// in two passes of 128 values each. Within each pass, 32 iterations produce
/// 4 values each by reassembling 6-bit quants from ql (low 4 bits) and qh (high 2 bits).
pub fn dequantize_q6_k_block(block: &BlockQ6K) -> [f32; 256] {
    let d = f16_to_f32(block.d);
    let ql = &block.ql;
    let qh = &block.qh;
    let sc = &block.scales;

    let mut out = [0.0f32; 256];
    let mut ql_off = 0usize;
    let mut qh_off = 0usize;
    let mut sc_off = 0usize;
    let mut y_off = 0usize;

    // Two passes of 128 values (n = 0 and n = 128)
    for _n in 0..2 {
        for l in 0..32 {
            let is = l / 16;
            let q1 = ((ql[ql_off + l] & 0xF) | ((qh[qh_off + l] & 3) << 4)) as i8 - 32;
            let q2 = ((ql[ql_off + l + 32] & 0xF) | (((qh[qh_off + l] >> 2) & 3) << 4)) as i8 - 32;
            let q3 = ((ql[ql_off + l] >> 4) | (((qh[qh_off + l] >> 4) & 3) << 4)) as i8 - 32;
            let q4 = ((ql[ql_off + l + 32] >> 4) | (((qh[qh_off + l] >> 6) & 3) << 4)) as i8 - 32;
            out[y_off + l] = d * sc[sc_off + is] as f32 * q1 as f32;
            out[y_off + l + 32] = d * sc[sc_off + is + 2] as f32 * q2 as f32;
            out[y_off + l + 64] = d * sc[sc_off + is + 4] as f32 * q3 as f32;
            out[y_off + l + 96] = d * sc[sc_off + is + 6] as f32 * q4 as f32;
        }
        y_off += 128;
        ql_off += 64;
        qh_off += 32;
        sc_off += 8;
    }

    out
}

/// Dequantize a row of Q6_K blocks. `src` is raw bytes, `dst` is f32 output.
pub fn dequantize_q6_k_row(src: &[u8], dst: &mut [f32]) {
    let block_size = size_of::<BlockQ6K>();
    let n_blocks = src.len() / block_size;
    debug_assert_eq!(src.len() % block_size, 0);
    debug_assert_eq!(dst.len(), n_blocks * 256);

    for i in 0..n_blocks {
        let block_bytes = &src[i * block_size..(i + 1) * block_size];
        // SAFETY: BlockQ6K is repr(C, packed) and we've verified the slice length
        let block = unsafe { &*(block_bytes.as_ptr() as *const BlockQ6K) };
        let values = dequantize_q6_k_block(block);
        dst[i * 256..(i + 1) * 256].copy_from_slice(&values);
    }
}

/// Dot product of a Q6_K block with an f32 vector of length 256. Scalar version.
pub fn vec_dot_q6_k_f32_scalar(block: &BlockQ6K, y: &[f32]) -> f32 {
    debug_assert_eq!(y.len(), 256);
    let d = f16_to_f32(block.d);
    let ql = &block.ql;
    let qh = &block.qh;
    let sc = &block.scales;

    let mut sumf = 0.0f32;
    let mut ql_off = 0usize;
    let mut qh_off = 0usize;
    let mut sc_off = 0usize;
    let mut y_off = 0usize;

    for _n in 0..2 {
        for l in 0..32 {
            let is = l / 16;
            let q1 = ((ql[ql_off + l] & 0xF) | ((qh[qh_off + l] & 3) << 4)) as i8 - 32;
            let q2 = ((ql[ql_off + l + 32] & 0xF) | (((qh[qh_off + l] >> 2) & 3) << 4)) as i8 - 32;
            let q3 = ((ql[ql_off + l] >> 4) | (((qh[qh_off + l] >> 4) & 3) << 4)) as i8 - 32;
            let q4 = ((ql[ql_off + l + 32] >> 4) | (((qh[qh_off + l] >> 6) & 3) << 4)) as i8 - 32;
            sumf += sc[sc_off + is] as f32 * q1 as f32 * y[y_off + l];
            sumf += sc[sc_off + is + 2] as f32 * q2 as f32 * y[y_off + l + 32];
            sumf += sc[sc_off + is + 4] as f32 * q3 as f32 * y[y_off + l + 64];
            sumf += sc[sc_off + is + 6] as f32 * q4 as f32 * y[y_off + l + 96];
        }
        y_off += 128;
        ql_off += 64;
        qh_off += 32;
        sc_off += 8;
    }

    sumf * d
}

/// Dot product of a Q6_K block with an f32 vector. Dispatches to best available impl.
pub fn vec_dot_q6_k_f32(block: &BlockQ6K, y: &[f32]) -> f32 {
    crate::backend::simd::vec_dot_q6_k_f32(block, y)
}

// ── Q5_K dequantization ────────────────────────────────────────────────────

/// Dequantize a single Q5_K block to 256 f32 values.
///
/// Ported from llama.cpp's `dequantize_row_q5_K`. Q5_K shares Q4_K's 6-bit
/// scale/min packing (`decode_q4km_scales`); the extra `qh` plane supplies the
/// 5th bit of each quant. The 256 values are produced in 4 iterations of 64:
/// each iteration decodes two sub-blocks (low nibbles then high nibbles of the
/// same 32 `qs` bytes) and folds in `qh` via the `u1`/`u2` bit selectors, which
/// start at bit 0/1 and shift left by 2 each iteration so all 8 `qh` bits are
/// consumed across the 4×2 halves.
pub fn dequantize_q5_k_block(block: &BlockQ5K) -> [f32; 256] {
    let d = f16_to_f32(block.d);
    let dmin = f16_to_f32(block.dmin);
    let (sc, mn) = decode_q4km_scales(&block.scales);
    let ql = &block.qs; // low 4 bits
    let qh = &block.qh; // high (5th) bit

    let mut out = [0.0f32; 256];
    let mut qi = 0usize; // index into ql (qs), advances by 32 each iteration
    let mut yi = 0usize; // output index
    let mut u1: u8 = 1;
    let mut u2: u8 = 2;

    for j in 0..4 {
        let d1 = d * sc[j * 2] as f32;
        let m1 = dmin * mn[j * 2] as f32;
        let d2 = d * sc[j * 2 + 1] as f32;
        let m2 = dmin * mn[j * 2 + 1] as f32;

        for l in 0..32 {
            let hi = if qh[l] & u1 != 0 { 16.0 } else { 0.0 };
            out[yi + l] = d1 * ((ql[qi + l] & 0xF) as f32 + hi) - m1;
        }
        for l in 0..32 {
            let hi = if qh[l] & u2 != 0 { 16.0 } else { 0.0 };
            out[yi + l + 32] = d2 * ((ql[qi + l] >> 4) as f32 + hi) - m2;
        }
        qi += 32;
        yi += 64;
        u1 <<= 2;
        u2 <<= 2;
    }

    out
}

/// Dequantize a row of Q5_K blocks. `src` is raw bytes, `dst` is f32 output.
pub fn dequantize_q5_k_row(src: &[u8], dst: &mut [f32]) {
    let block_size = size_of::<BlockQ5K>();
    let n_blocks = src.len() / block_size;
    debug_assert_eq!(src.len() % block_size, 0);
    debug_assert_eq!(dst.len(), n_blocks * 256);

    for i in 0..n_blocks {
        let block_bytes = &src[i * block_size..(i + 1) * block_size];
        // SAFETY: BlockQ5K is repr(C, packed) and we've verified the slice length
        let block = unsafe { &*(block_bytes.as_ptr() as *const BlockQ5K) };
        let values = dequantize_q5_k_block(block);
        dst[i * 256..(i + 1) * 256].copy_from_slice(&values);
    }
}

/// Dot product of a Q5_K block with an f32 vector of length 256. Scalar version.
///
/// Same accumulation structure as `vec_dot_q4_k_m_f32_scalar`, extended with
/// the `qh` 5th-bit plane. Mathematically equal to `dot(dequant(block), y)`.
pub fn vec_dot_q5_k_f32_scalar(block: &BlockQ5K, y: &[f32]) -> f32 {
    debug_assert_eq!(y.len(), 256);

    let d = f16_to_f32(block.d);
    let dmin = f16_to_f32(block.dmin);
    let (sc, mn) = decode_q4km_scales(&block.scales);
    let ql = &block.qs;
    let qh = &block.qh;

    let mut sumf = 0.0f32;
    let mut qi = 0usize;
    let mut yi = 0usize;
    let mut u1: u8 = 1;
    let mut u2: u8 = 2;

    for j in 0..4 {
        let sc1 = sc[j * 2] as f32;
        let mn1 = mn[j * 2] as f32;
        let sc2 = sc[j * 2 + 1] as f32;
        let mn2 = mn[j * 2 + 1] as f32;

        let mut sum1 = 0.0f32;
        let mut sum2 = 0.0f32;
        let mut sum_mn1 = 0.0f32;
        let mut sum_mn2 = 0.0f32;

        for l in 0..32 {
            let hi1 = if qh[l] & u1 != 0 { 16.0 } else { 0.0 };
            let hi2 = if qh[l] & u2 != 0 { 16.0 } else { 0.0 };
            let q1 = (ql[qi + l] & 0xF) as f32 + hi1;
            let q2 = (ql[qi + l] >> 4) as f32 + hi2;
            sum1 += q1 * y[yi + l];
            sum2 += q2 * y[yi + l + 32];
            sum_mn1 += y[yi + l];
            sum_mn2 += y[yi + l + 32];
        }

        sumf += d * (sc1 * sum1 + sc2 * sum2) - dmin * (mn1 * sum_mn1 + mn2 * sum_mn2);
        qi += 32;
        yi += 64;
        u1 <<= 2;
        u2 <<= 2;
    }

    sumf
}

/// Dot product of a Q5_K block with an f32 vector. Dispatches to best available impl.
pub fn vec_dot_q5_k_f32(block: &BlockQ5K, y: &[f32]) -> f32 {
    crate::backend::simd::vec_dot_q5_k_f32(block, y)
}

// ── Dispatch functions ──────────────────────────────────────────────────────

/// Dot product of a Q4_0 block with an f32 vector. Dispatches to best available impl.
pub fn vec_dot_q4_0_f32(block: &BlockQ4_0, y: &[f32]) -> f32 {
    crate::backend::simd::vec_dot_q4_0_f32(block, y)
}

/// Dot product of a Q8_0 block with an f32 vector. Dispatches to best available impl.
pub fn vec_dot_q8_0_f32(block: &BlockQ8_0, y: &[f32]) -> f32 {
    crate::backend::simd::vec_dot_q8_0_f32(block, y)
}

/// Dot product of a Q4_K_M block with an f32 vector. Dispatches to best available impl.
pub fn vec_dot_q4_k_m_f32(block: &BlockQ4KM, y: &[f32]) -> f32 {
    crate::backend::simd::vec_dot_q4_k_m_f32(block, y)
}

// ── Tests ───────────────────────────────────────────────────────────────────

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

    /// `f32_to_f16` must round exactly like `half::f16::from_f32`.
    ///
    /// The full 2^32 input space was swept once against the `half` crate and
    /// came back with zero mismatches; that run takes far too long for CI, so
    /// what stays here is the structure that sweep was checking. Every f16 is
    /// covered by round-tripping all 65536 patterns, and the classes a
    /// hand-rolled rounder actually gets wrong are enumerated explicitly:
    /// tie-to-even at the rounding boundary, the carry that pushes a rounded
    /// significand up into the next exponent (and out to infinity), the
    /// normal-to-subnormal seam, underflow, and NaN. NaN is the one class the
    /// round-trip sweep cannot reach on its own: widening a half NaN and
    /// narrowing it back is a fixed point under any quieting rule, so the
    /// explicit f32 patterns below are what pin it.
    #[test]
    fn f32_to_f16_matches_half_crate() {
        // Every representable half, widened and narrowed back. NaN is compared
        // by bits rather than skipped: `exp == 0xFF && mant != 0` is the one arm
        // where narrowing is not a shift-and-round, so skipping it left the only
        // interesting case untested.
        for bits in 0..=u16::MAX {
            let v = f16_to_f32(bits);
            assert_eq!(
                f32_to_f16(v),
                f16::from_f32(v).to_bits(),
                "round-trip of half {bits:#06x} ({v})"
            );
        }

        // Both NaN kinds narrowed directly from f32, including a payload that
        // is entirely in the low bits the half cannot keep: `half` maps that to
        // a quiet NaN rather than to infinity, and so must this.
        for &nan_bits in &[
            0x7FC0_0000u32, // qNaN
            0xFFC0_0000,    // qNaN, negative
            0x7F80_0001,    // sNaN, minimal payload
            0x7F80_1000,    // sNaN, payload only in bits the half drops
            0xFF80_0001,    // sNaN, negative
        ] {
            let v = f32::from_bits(nan_bits);
            assert_eq!(
                f32_to_f16(v),
                f16::from_f32(v).to_bits(),
                "narrowing f32 NaN {nan_bits:#010x}"
            );
        }

        let cases: &[f32] = &[
            0.0,
            -0.0,
            1.0,
            -1.0,
            f32::INFINITY,
            f32::NEG_INFINITY,
            65504.0, // largest finite half
            65519.0, // rounds up to 65504
            65520.0, // first value that rounds to infinity
            -65520.0,
            1.0e30,         // plain overflow
            6.103_516e-5,   // smallest normal half
            6.097_555e-5,   // just below it, becomes subnormal
            5.960_464_5e-8, // smallest subnormal half
            2.980_232_2e-8, // exactly half of it, ties to even -> zero
            2.980_232_5e-8, // just above the tie, rounds up
            1.0e-10,        // underflows to zero
            -1.0e-10,
            1.000_976_6, // exact at half precision
            1.000_488_3, // exact tie between two halves, ties to even
            1.001_464_8, // tie the other way
            2048.0,      // integers around the half integer limit
            2049.0,
            4098.0,
        ];
        for &v in cases {
            assert_eq!(
                f32_to_f16(v),
                f16::from_f32(v).to_bits(),
                "boundary case {v:e}"
            );
        }

        // Deterministic sweep across the whole exponent range, including the
        // subnormal and overflow shoulders.
        let mut st = 0x1234_5678u32;
        for _ in 0..200_000 {
            st = st.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
            let v = f32::from_bits(st);
            assert_eq!(
                f32_to_f16(v),
                f16::from_f32(v).to_bits(),
                "random f32 {:#010x} ({v:e})",
                st
            );
        }
    }

    /// Same contract as the f16 twin below, over bf16's 65536 patterns.
    ///
    /// Cheap enough to be exhaustive, and worth being exhaustive about: the
    /// interesting inputs are exactly the ones a random sweep would miss, since
    /// the sNaN quieting is the only case where `bf16_to_f32` is not a plain
    /// shift and there are 127 such patterns per sign out of 65536.
    #[test]
    fn bf16_widen_matches_half_crate_exhaustively() {
        for bits in 0..=u16::MAX {
            // `half`, deliberately, and for the reason spelled out below.
            let want = half::bf16::from_bits(bits).to_f32();
            let got = bf16_to_f32(bits);
            if want.is_nan() {
                assert!(got.is_nan(), "bits {bits:#06x}: want NaN, got {got}");
                assert_eq!(
                    got.to_bits(),
                    want.to_bits(),
                    "bits {bits:#06x}: NaN payload/sign differs"
                );
            } else {
                assert_eq!(
                    got.to_bits(),
                    want.to_bits(),
                    "bits {bits:#06x}: want {want}, got {got}"
                );
            }
        }
    }

    /// `f16_to_f32` is open-coded, so it must equal `half`'s conversion for
    /// every one of the 65536 half patterns, exactly.
    ///
    /// Exhaustive rather than sampled because the interesting inputs are the
    /// rare ones: the subnormal renormalization is a different arm of the match
    /// from normals, and random draws would essentially never hit it. This also
    /// covers the two boundaries the arms meet at (`exp == 0` and `exp == 0x1F`)
    /// without having to enumerate them by hand.
    ///
    /// The kernel tests cannot stand in for this: they compare the f16 kernels
    /// against a scalar dot over values widened by *this same function*, so a
    /// wrong conversion cancels out on both sides and they still pass.
    #[test]
    fn f16_widen_matches_half_crate_exhaustively() {
        for bits in 0..=u16::MAX {
            // The `half` crate, deliberately: an open-coded conversion has to be
            // checked against an *independent* implementation. A mechanical pass
            // rewrote this reference to our own function once, which made the
            // test compare `f16_to_f32` to itself and pass unconditionally.
            let want = half::f16::from_bits(bits).to_f32();
            let got = f16_to_f32(bits);
            if want.is_nan() {
                assert!(got.is_nan(), "bits {bits:#06x}: want NaN, got {got}");
                // The payload rides along, so a NaN must not silently become a
                // different NaN (or an infinity with the sign flipped).
                assert_eq!(
                    got.to_bits() & 0x807F_FFFF,
                    want.to_bits() & 0x807F_FFFF,
                    "bits {bits:#06x}: NaN payload/sign differs"
                );
            } else {
                assert_eq!(
                    got.to_bits(),
                    want.to_bits(),
                    "bits {bits:#06x}: want {want}, got {got}"
                );
            }
        }
    }

    /// Build a Q8_0 block from known values for testing.
    fn make_q8_0_block(scale: f32, quants: [i8; 32]) -> BlockQ8_0 {
        BlockQ8_0 {
            delta: f16::from_f32(scale).to_bits(),
            quants,
        }
    }

    #[test]
    fn test_dequantize_q4_1_matches_ggml_formula() {
        // Reference is llama.cpp's dequantize_row_q4_1:
        //   x0 = qs[j] & 0x0F;  y[j]      = x0*d + m
        //   x1 = qs[j] >>   4;  y[j+qk/2] = x1*d + m
        // Note there is no `- 8`: the minimum replaces the recentering.
        let mut qs = [0u8; 16];
        for (i, qsi) in qs.iter_mut().enumerate() {
            *qsi = (i as u8) | (((15 - i) as u8) << 4);
        }
        let d = 0.25f32;
        let m = -1.5f32;
        let block = BlockQ4_1 {
            d: f16::from_f32(d).to_bits(),
            m: f16::from_f32(m).to_bits(),
            qs,
        };
        let out = dequantize_q4_1_block(&block);
        let d = f16::from_f32(d).to_f32();
        let m = f16::from_f32(m).to_f32();
        for i in 0..16 {
            let want_lo = i as f32 * d + m;
            let want_hi = (15 - i) as f32 * d + m;
            assert!(
                (out[i] - want_lo).abs() < 1e-5,
                "lo[{i}]: got {} want {want_lo}",
                out[i]
            );
            assert!(
                (out[i + 16] - want_hi).abs() < 1e-5,
                "hi[{i}]: got {} want {want_hi}",
                out[i + 16]
            );
        }
    }

    /// The minimum is a per-block constant, so `vec_dot` factors it out as
    /// `m * sum(y)` instead of adding it per element. That is an algebraic
    /// rearrangement, not the same operation — check it against the literal
    /// dequantize-then-dot.
    #[test]
    fn test_vec_dot_q4_1_matches_dequantize() {
        let mut st = 0x9e37_79b9u64;
        let mut lcg = || {
            st = st.wrapping_mul(6364136223846793005).wrapping_add(1);
            ((st >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0
        };
        for trial in 0..8 {
            let mut qs = [0u8; 16];
            for qsi in qs.iter_mut() {
                *qsi = ((lcg() + 1.0) * 127.0) as u8;
            }
            let block = BlockQ4_1 {
                d: f16::from_f32(0.1 + trial as f32 * 0.05).to_bits(),
                m: f16::from_f32(lcg()).to_bits(),
                qs,
            };
            let y: Vec<f32> = (0..32).map(|_| lcg()).collect();

            let want: f32 = dequantize_q4_1_block(&block)
                .iter()
                .zip(&y)
                .map(|(a, b)| a * b)
                .sum();
            let got = vec_dot_q4_1_f32(&block, &y);
            assert!(
                (got - want).abs() <= 1e-4 * (1.0 + want.abs()),
                "trial {trial}: got {got} want {want}"
            );
        }
    }

    /// A zero `d` with a non-zero `m` is a legal block (a constant row); the
    /// minimum must survive rather than being multiplied away.
    #[test]
    fn test_dequantize_q4_1_zero_scale_keeps_min() {
        let block = BlockQ4_1 {
            d: f16::from_f32(0.0).to_bits(),
            m: f16::from_f32(2.5).to_bits(),
            qs: [0xAB; 16],
        };
        let out = dequantize_q4_1_block(&block);
        assert!(out.iter().all(|v| (v - 2.5).abs() < 1e-5), "{out:?}");
    }

    #[test]
    fn test_dequantize_q4_1_row_matches_block() {
        let mut bytes = Vec::new();
        for b in 0..3u16 {
            bytes.extend_from_slice(&f16::from_f32(0.5).to_bits().to_le_bytes());
            bytes.extend_from_slice(&f16::from_f32(-0.25).to_bits().to_le_bytes());
            bytes.extend_from_slice(&[b as u8 | 0x30; 16]);
        }
        let mut dst = vec![0.0f32; 96];
        dequantize_q4_1_row(&bytes, &mut dst);
        for b in 0..3usize {
            let block = BlockQ4_1 {
                d: f16::from_f32(0.5).to_bits(),
                m: f16::from_f32(-0.25).to_bits(),
                qs: [b as u8 | 0x30; 16],
            };
            let want = dequantize_q4_1_block(&block);
            assert_eq!(&dst[b * 32..(b + 1) * 32], &want[..], "block {b}");
        }
    }

    #[test]
    fn test_dequantize_q4_0_simple() {
        // All nibbles = 8 → offset to 0
        let block = BlockQ4_0 {
            d: f16::from_f32(1.0).to_bits(),
            qs: [0x88; 16], // lo=8, hi=8 → both (8-8)*1.0 = 0.0
        };
        let out = dequantize_q4_0_block(&block);
        for (i, &v) in out.iter().enumerate() {
            assert!(v.abs() < 1e-3, "expected 0.0 at {i}, got {v}");
        }
    }

    #[test]
    fn test_dequantize_q4_0_varied() {
        // lo nibbles: 0..16, hi nibbles: all 15
        let mut qs = [0u8; 16];
        for (i, qsi) in qs.iter_mut().enumerate() {
            *qsi = (i as u8) | (15 << 4);
        }
        let block = BlockQ4_0 {
            d: f16::from_f32(0.5).to_bits(),
            qs,
        };
        let out = dequantize_q4_0_block(&block);

        // First 16: (i - 8) * 0.5
        for (i, &v) in out.iter().enumerate().take(16) {
            let expected = (i as f32 - 8.0) * 0.5;
            assert!(
                (v - expected).abs() < 1e-3,
                "lo[{i}]: got {v}, expected {expected}"
            );
        }
        // Last 16: (15 - 8) * 0.5 = 3.5
        for (i, &v) in out.iter().enumerate().skip(16) {
            assert!((v - 3.5).abs() < 1e-3, "hi[{i}]: got {v}, expected 3.5");
        }
    }

    #[test]
    fn test_vec_dot_q4_0_matches_dequantize() {
        let mut qs = [0u8; 16];
        for (i, qsi) in qs.iter_mut().enumerate() {
            *qsi = ((i % 13) as u8) | (((i % 7) as u8) << 4);
        }
        let block = BlockQ4_0 {
            d: f16::from_f32(0.3).to_bits(),
            qs,
        };
        let y: Vec<f32> = (0..32).map(|i| (i as f32 - 16.0) * 0.1).collect();

        let dequantized = dequantize_q4_0_block(&block);
        let expected: f32 = dequantized.iter().zip(y.iter()).map(|(a, b)| a * b).sum();
        let got = vec_dot_q4_0_f32(&block, &y);

        assert!(
            (got - expected).abs() < 1e-3,
            "vec_dot Q4_0 mismatch: got {got}, expected {expected}"
        );
    }

    #[test]
    fn test_dequantize_q8_0_simple() {
        let block = make_q8_0_block(0.5, {
            let mut q = [0i8; 32];
            for (i, qi) in q.iter_mut().enumerate() {
                *qi = i as i8;
            }
            q
        });
        let out = dequantize_q8_0_block(&block);
        for (i, &v) in out.iter().enumerate() {
            let expected = i as f32 * 0.5;
            assert!(
                (v - expected).abs() < 1e-3,
                "mismatch at {i}: got {v}, expected {expected}"
            );
        }
    }

    #[test]
    fn test_dequantize_q8_0_row() {
        // Two blocks
        let block1 = make_q8_0_block(1.0, {
            let mut q = [0i8; 32];
            for (i, qi) in q.iter_mut().enumerate() {
                *qi = (i as i8) - 16;
            }
            q
        });
        let block2 = make_q8_0_block(0.25, [1i8; 32]);

        let mut src = vec![0u8; 68];
        unsafe {
            std::ptr::copy_nonoverlapping(&block1 as *const _ as *const u8, src.as_mut_ptr(), 34);
            std::ptr::copy_nonoverlapping(
                &block2 as *const _ as *const u8,
                src.as_mut_ptr().add(34),
                34,
            );
        }

        let mut dst = vec![0.0f32; 64];
        dequantize_q8_0_row(&src, &mut dst);

        // Check block1 values
        for (i, &v) in dst.iter().enumerate().take(32) {
            let expected = (i as f32 - 16.0) * 1.0;
            assert!(
                (v - expected).abs() < 1e-3,
                "block1[{i}]: got {v}, expected {expected}"
            );
        }
        // Check block2 values
        for i in 0..32 {
            let expected = 1.0 * 0.25;
            assert!(
                (dst[32 + i] - expected).abs() < 1e-3,
                "block2[{i}]: got {}, expected {expected}",
                dst[32 + i]
            );
        }
    }

    /// `dequantize_q5_k_matrix` must equal a loop of `dequantize_q5_k_row`,
    /// bit for bit.
    ///
    /// This is the only thing standing between a Q5_K model and silently wrong
    /// prefill: the BLAS route dequantizes the whole weight matrix and SGEMMs,
    /// so an off-by-one in the row striding would corrupt every batched token
    /// while decode, which uses the per-row path, stayed correct and hid it.
    /// `m` is above the parallel threshold on purpose, since the matrix helper
    /// splits rows across workers and the row helper does not.
    #[test]
    fn test_dequantize_q5_k_matrix_matches_row() {
        let m = 128;
        let k = 512; // 2 superblocks per row
        let blocks_per_row = k / 256;
        let row_bytes = blocks_per_row * size_of::<BlockQ5K>();

        let mut st = 0x5EED_1234u32;
        let mut byte = || {
            st = st.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
            (st >> 24) as u8
        };
        let mut src = vec![0u8; m * row_bytes];
        src.iter_mut().for_each(|b| *b = byte());
        // Keep the scale fields finite. Random bits can land on inf/NaN, and
        // since NaN never compares equal to itself the `assert_eq!` below would
        // then fail on rows that actually agree: a flaky red, not a false pass.
        (0..m)
            .flat_map(|row| (0..blocks_per_row).map(move |b| (row, b)))
            .for_each(|(row, b)| {
                let off = row * row_bytes + b * size_of::<BlockQ5K>();
                let d = f16::from_f32(0.05 + (row % 7) as f32 * 0.01).to_bits();
                let dmin = f16::from_f32(0.02 + (b % 3) as f32 * 0.01).to_bits();
                src[off..off + 2].copy_from_slice(&d.to_le_bytes());
                src[off + 2..off + 4].copy_from_slice(&dmin.to_le_bytes());
            });

        let mut via_matrix = vec![0.0f32; m * k];
        dequantize_q5_k_matrix(&src, m, k, &mut via_matrix);

        let mut via_rows = vec![0.0f32; m * k];
        via_rows
            .chunks_mut(k)
            .zip(src.chunks(row_bytes))
            .for_each(|(dst, s)| dequantize_q5_k_row(s, dst));

        assert_eq!(
            via_matrix, via_rows,
            "matrix and row dequantization diverged"
        );
    }

    #[test]
    fn test_dequantize_q4_0_matrix_matches_row() {
        // Build `m` rows of Q4_0 blocks with distinct content, dequantize via
        // both the matrix helper and a loop of `dequantize_q4_0_row` calls, and
        // verify they produce byte-identical output.
        let m = 128; // above the parallelization threshold
        let k = 64; // 2 blocks per row
        let blocks_per_row = k / 32;
        let row_bytes = blocks_per_row * size_of::<BlockQ4_0>();

        let mut src = vec![0u8; m * row_bytes];
        for row in 0..m {
            for b in 0..blocks_per_row {
                let block = BlockQ4_0 {
                    d: f16::from_f32(0.1 + (row as f32) * 0.01).to_bits(),
                    qs: {
                        let mut qs = [0u8; 16];
                        for (i, q) in qs.iter_mut().enumerate() {
                            *q = ((row + b * 7 + i * 3) as u8).wrapping_mul(17);
                        }
                        qs
                    },
                };
                let offset = row * row_bytes + b * size_of::<BlockQ4_0>();
                unsafe {
                    std::ptr::copy_nonoverlapping(
                        &block as *const _ as *const u8,
                        src.as_mut_ptr().add(offset),
                        size_of::<BlockQ4_0>(),
                    );
                }
            }
        }

        let mut matrix_out = vec![0.0f32; m * k];
        dequantize_q4_0_matrix(&src, m, k, &mut matrix_out);

        let mut expected = vec![0.0f32; m * k];
        for row in 0..m {
            let src_row = &src[row * row_bytes..(row + 1) * row_bytes];
            let dst_row = &mut expected[row * k..(row + 1) * k];
            dequantize_q4_0_row(src_row, dst_row);
        }

        assert_eq!(matrix_out, expected);
    }

    #[test]
    fn test_dequantize_q8_0_matrix_matches_row() {
        let m = 96;
        let k = 96; // 3 blocks per row
        let blocks_per_row = k / 32;
        let row_bytes = blocks_per_row * size_of::<BlockQ8_0>();

        let mut src = vec![0u8; m * row_bytes];
        for row in 0..m {
            for b in 0..blocks_per_row {
                let block = make_q8_0_block(0.05 * (1 + row) as f32 + 0.001 * b as f32, {
                    let mut q = [0i8; 32];
                    for (i, slot) in q.iter_mut().enumerate() {
                        *slot = ((row + b + i) as i8).wrapping_mul(5).wrapping_sub(64);
                    }
                    q
                });
                let offset = row * row_bytes + b * size_of::<BlockQ8_0>();
                unsafe {
                    std::ptr::copy_nonoverlapping(
                        &block as *const _ as *const u8,
                        src.as_mut_ptr().add(offset),
                        size_of::<BlockQ8_0>(),
                    );
                }
            }
        }

        let mut matrix_out = vec![0.0f32; m * k];
        dequantize_q8_0_matrix(&src, m, k, &mut matrix_out);

        let mut expected = vec![0.0f32; m * k];
        for row in 0..m {
            let src_row = &src[row * row_bytes..(row + 1) * row_bytes];
            let dst_row = &mut expected[row * k..(row + 1) * k];
            dequantize_q8_0_row(src_row, dst_row);
        }

        assert_eq!(matrix_out, expected);
    }

    #[test]
    fn test_vec_dot_q8_0() {
        let block = make_q8_0_block(0.1, {
            let mut q = [0i8; 32];
            for (i, qi) in q.iter_mut().enumerate() {
                *qi = (i as i8) * 2 - 31;
            }
            q
        });
        let y: Vec<f32> = (0..32).map(|i| i as f32 * 0.5).collect();

        // Compute expected via dequantize
        let dequantized = dequantize_q8_0_block(&block);
        let expected: f32 = dequantized.iter().zip(y.iter()).map(|(a, b)| a * b).sum();
        let got = vec_dot_q8_0_f32(&block, &y);

        assert!(
            (got - expected).abs() < 1e-3,
            "vec_dot mismatch: got {got}, expected {expected}"
        );
    }

    #[test]
    fn test_dequantize_q4_k_m_basic() {
        // Create a Q4_K_M block with known values
        let mut block = BlockQ4KM {
            d: f16::from_f32(1.0).to_bits(),
            dmin: f16::from_f32(0.0).to_bits(), // zero min for simplicity
            scales: [0u8; 12],
            qs: [0u8; 128],
        };
        // Set all sub-block scales to 1 (6-bit value)
        for i in 0..4 {
            block.scales[i] = 1; // sc[i] = 1, bits 6-7 = 0
        }
        for i in 4..8 {
            block.scales[i] = 0; // mn[0..4] = 0
        }
        // sc[4..8] and mn[4..8] come from bytes 8-11
        for i in 8..12 {
            block.scales[i] = 0x01; // sc[j] low nibble = 1, mn[j] high nibble = 0
        }

        // Set qs: all nibbles = 3
        for b in block.qs.iter_mut() {
            *b = 0x33; // low nibble = 3, high nibble = 3
        }

        let out = dequantize_q4_k_m_block(&block);
        // With d=1.0, dmin=0.0, sc=1, all nibbles=3:
        // value = 1.0 * 1 * 3 - 0.0 = 3.0
        for (i, &v) in out.iter().enumerate() {
            assert!(
                (v - 3.0).abs() < 1e-3,
                "mismatch at {i}: got {v}, expected 3.0"
            );
        }
    }

    /// The K-quant *matrix* dequantizers must agree with the per-row ones they wrap.
    ///
    /// These feed `try_blas_prefill_gemm`, which is the path that actually ships on
    /// Apple Silicon (`--features blas`) — and the NEON kernels every other test in
    /// this PR covers are compiled *out* of that build. A row-stride mistake here
    /// (144 bytes per Q4_K superblock, 210 per Q6_K) would silently misalign every
    /// weight row and produce wrong logits with the whole suite green.
    #[test]
    fn kquant_matrix_dequant_matches_row_dequant() {
        let mut st = 0x2468_1357u64;
        let next = |st: &mut u64| {
            *st = st.wrapping_mul(6364136223846793005).wrapping_add(1);
            (*st >> 33) as u8
        };

        // m rows of k=512 (2 superblocks per row) — enough that a bad stride shows.
        let (m, k) = (3usize, 512usize);
        let nb = k / 256;

        // Q4_K
        let mut src = Vec::new();
        for _ in 0..m * nb {
            let blk = BlockQ4KM {
                d: half::f16::from_f32(0.03).to_bits(),
                dmin: half::f16::from_f32(0.01).to_bits(),
                scales: std::array::from_fn(|_| next(&mut st)),
                qs: std::array::from_fn(|_| next(&mut st)),
            };
            src.extend_from_slice(unsafe {
                std::slice::from_raw_parts((&raw const blk) as *const u8, size_of::<BlockQ4KM>())
            });
        }
        let mut got = vec![0.0f32; m * k];
        dequantize_q4_k_m_matrix(&src, m, k, &mut got);
        let row_bytes = nb * size_of::<BlockQ4KM>();
        for i in 0..m {
            let mut want = vec![0.0f32; k];
            dequantize_q4_k_m_row(&src[i * row_bytes..(i + 1) * row_bytes], &mut want);
            assert_eq!(&got[i * k..(i + 1) * k], &want[..], "Q4_K row {i} mismatch");
        }

        // Q6_K
        let mut src = Vec::new();
        for _ in 0..m * nb {
            let blk = BlockQ6K {
                ql: std::array::from_fn(|_| next(&mut st)),
                qh: std::array::from_fn(|_| next(&mut st)),
                scales: std::array::from_fn(|_| next(&mut st) as i8),
                d: half::f16::from_f32(0.02).to_bits(),
            };
            src.extend_from_slice(unsafe {
                std::slice::from_raw_parts((&raw const blk) as *const u8, size_of::<BlockQ6K>())
            });
        }
        let mut got = vec![0.0f32; m * k];
        dequantize_q6_k_matrix(&src, m, k, &mut got);
        let row_bytes = nb * size_of::<BlockQ6K>();
        for i in 0..m {
            let mut want = vec![0.0f32; k];
            dequantize_q6_k_row(&src[i * row_bytes..(i + 1) * row_bytes], &mut want);
            assert_eq!(&got[i * k..(i + 1) * k], &want[..], "Q6_K row {i} mismatch");
        }
    }

    #[test]
    fn test_vec_dot_q4km_matches_dequantize() {
        // Create a block with varied values
        let mut block = BlockQ4KM {
            d: f16::from_f32(0.5).to_bits(),
            dmin: f16::from_f32(0.1).to_bits(),
            scales: [0u8; 12],
            qs: [0u8; 128],
        };
        // Set scales: sc=2, mn=1 for first 4 sub-blocks
        for i in 0..4 {
            block.scales[i] = 2;
        }
        for i in 4..8 {
            block.scales[i] = 1;
        }
        for i in 8..12 {
            block.scales[i] = 0x21; // sc low=1, mn high=2 -> sc[j]=1|(bits<<4), mn[j]=(2)|(bits<<4)
        }

        // Varied quantized values
        for (i, b) in block.qs.iter_mut().enumerate() {
            *b = ((i % 7) as u8) | (((i % 11) as u8) << 4);
        }

        let y: Vec<f32> = (0..256).map(|i| (i as f32 - 128.0) * 0.01).collect();

        // Compute expected via dequantize + dot
        let dequantized = dequantize_q4_k_m_block(&block);
        let expected: f32 = dequantized.iter().zip(y.iter()).map(|(a, b)| a * b).sum();
        let got = vec_dot_q4_k_m_f32(&block, &y);

        assert!(
            (got - expected).abs() < 1e-2,
            "vec_dot mismatch: got {got}, expected {expected}"
        );
    }

    #[test]
    fn test_dequantize_q6_k_basic() {
        // Create a Q6_K block where all quants reassemble to 0 (offset 32 → value -32+32=0)
        // and scale d=1.0, sub-block scales=1
        let mut block = BlockQ6K {
            ql: [0u8; 128],
            qh: [0u8; 64],
            scales: [1i8; 16],
            d: f16::from_f32(1.0).to_bits(),
        };
        // Set ql and qh so that all 6-bit values = 32 (which becomes 32-32 = 0)
        // 32 in 6 bits = 0b100000 → low 4 bits = 0, high 2 bits = 0b10 = 2
        // ql stores pairs: ql[l] low nibble for q1, ql[l+32] low nibble for q2
        //                  ql[l] high nibble for q3, ql[l+32] high nibble for q4
        // qh[l] bits 0-1 for q1, bits 2-3 for q2, bits 4-5 for q3, bits 6-7 for q4
        // For value 32: low 4 = 0, high 2 = 2
        // So ql = 0x00 (both nibbles = 0), qh = 0b10_10_10_10 = 0xAA
        for b in block.ql.iter_mut() {
            *b = 0x00;
        }
        for b in block.qh.iter_mut() {
            *b = 0xAA; // bits: 10_10_10_10
        }

        let out = dequantize_q6_k_block(&block);
        for (i, &v) in out.iter().enumerate() {
            assert!(v.abs() < 1e-5, "expected ~0.0 at {i}, got {v}");
        }
    }

    #[test]
    fn test_vec_dot_q6_k_matches_dequantize() {
        // Build a Q6_K block with varied values
        let mut block = BlockQ6K {
            ql: [0u8; 128],
            qh: [0u8; 64],
            scales: [0i8; 16],
            d: f16::from_f32(0.5).to_bits(),
        };
        // Set sub-block scales to small values
        for (i, s) in block.scales.iter_mut().enumerate() {
            *s = (i as i8 % 5) + 1;
        }
        // Set varied ql values
        for (i, b) in block.ql.iter_mut().enumerate() {
            *b = ((i % 13) as u8) | (((i % 9) as u8) << 4);
        }
        // Set varied qh values
        for (i, b) in block.qh.iter_mut().enumerate() {
            *b = (i % 256) as u8;
        }

        let y: Vec<f32> = (0..256).map(|i| (i as f32 - 128.0) * 0.01).collect();

        let dequantized = dequantize_q6_k_block(&block);
        let expected: f32 = dequantized.iter().zip(y.iter()).map(|(a, b)| a * b).sum();
        let got = vec_dot_q6_k_f32(&block, &y);

        assert!(
            (got - expected).abs() < 1e-2,
            "vec_dot Q6_K mismatch: got {got}, expected {expected}"
        );
    }

    #[test]
    fn test_q5_k_block_is_176_bytes() {
        // Guards the repr(C, packed) field order/size against silent drift —
        // the row/vec_dot code reinterprets raw GGUF bytes as BlockQ5K.
        assert_eq!(size_of::<BlockQ5K>(), 176);
    }

    #[test]
    fn test_dequantize_q5_k_all_zero_quants() {
        // qs=0 (low nibble 0), qh=0 (5th bit 0) → every quant is 0, so each
        // output = d1*0 - m1 = -min. With dmin=0 the whole block is 0.0.
        let mut block = BlockQ5K {
            d: f16::from_f32(1.0).to_bits(),
            dmin: f16::from_f32(0.0).to_bits(),
            scales: [0u8; 12],
            qh: [0u8; 32],
            qs: [0u8; 128],
        };
        // sc[0..4]=1, mn[0..4]=0 (bytes 0-3 hold sc low6, bytes 4-7 hold mn low6)
        for s in block.scales.iter_mut().take(4) {
            *s = 1;
        }
        let out = dequantize_q5_k_block(&block);
        for (i, &v) in out.iter().enumerate() {
            assert!(v.abs() < 1e-5, "expected ~0.0 at {i}, got {v}");
        }
    }

    #[test]
    fn test_dequantize_q5_k_high_bit_extends_range() {
        // A single quant with low nibble = 0xF and its qh bit set must decode to
        // (15 + 16) = 31, i.e. the 5-bit max — proving the qh plane is applied.
        // Value 0 is the first low-nibble sub-block (selector u1 = bit 0 of qh[0]).
        let mut block = BlockQ5K {
            d: f16::from_f32(1.0).to_bits(),
            dmin: f16::from_f32(0.0).to_bits(),
            scales: [0u8; 12],
            qh: [0u8; 32],
            qs: [0u8; 128],
        };
        block.scales[0] = 1; // sc[0] = 1
        block.qs[0] = 0x0F; // value 0: low nibble = 15
        block.qh[0] = 0x01; // value 0: 5th bit set → +16
        let out = dequantize_q5_k_block(&block);
        assert!(
            (out[0] - 31.0).abs() < 1e-4,
            "expected 31.0 (15 + 16) at index 0, got {}",
            out[0]
        );
        // Without the high bit, value 1 (low nibble of qs[1]=0) stays 0.
        assert!(
            out[1].abs() < 1e-4,
            "expected 0.0 at index 1, got {}",
            out[1]
        );
    }

    #[test]
    fn test_vec_dot_q5_k_matches_dequantize() {
        // vec_dot must equal dot(dequant(block), y) for varied scales/quants/qh.
        let mut block = BlockQ5K {
            d: f16::from_f32(0.5).to_bits(),
            dmin: f16::from_f32(0.125).to_bits(),
            scales: [0u8; 12],
            qh: [0u8; 32],
            qs: [0u8; 128],
        };
        // Varied 6-bit scales/mins across the 12 packed bytes.
        for (i, s) in block.scales.iter_mut().enumerate() {
            *s = ((i * 7 + 3) % 64) as u8;
        }
        for (i, b) in block.qs.iter_mut().enumerate() {
            *b = ((i % 7) as u8) | (((i % 11) as u8) << 4);
        }
        for (i, b) in block.qh.iter_mut().enumerate() {
            *b = ((i * 37) % 256) as u8;
        }

        let y: Vec<f32> = (0..256).map(|i| (i as f32 - 128.0) * 0.01).collect();

        let dequantized = dequantize_q5_k_block(&block);
        let expected: f32 = dequantized.iter().zip(y.iter()).map(|(a, b)| a * b).sum();
        let got = vec_dot_q5_k_f32(&block, &y);

        assert!(
            (got - expected).abs() < 1e-2,
            "vec_dot Q5_K mismatch: got {got}, expected {expected}"
        );
    }

    #[test]
    fn test_dequantize_q5_k_row_multiple_blocks() {
        // Two blocks back-to-back must dequantize independently into 512 floats.
        let mut bytes = vec![0u8; 2 * size_of::<BlockQ5K>()];
        // Block 0 d=1.0 at offset 0..2; block 1 d=2.0 at offset 176..178.
        bytes[0..2].copy_from_slice(&f16::from_f32(1.0).to_bits().to_le_bytes());
        let b1 = size_of::<BlockQ5K>();
        bytes[b1..b1 + 2].copy_from_slice(&f16::from_f32(2.0).to_bits().to_le_bytes());
        // sc[0]=1 for both blocks (scales byte 0 is at offset 4 within each block).
        bytes[4] = 1;
        bytes[b1 + 4] = 1;
        // One nonzero quant in block 1 (value 0, low nibble 3, no high bit).
        bytes[b1 + 4 + 12 + 32] = 0x03; // qs starts after d,dmin,scales,qh
        let mut dst = vec![0.0f32; 512];
        dequantize_q5_k_row(&bytes, &mut dst);
        // Block 1's value 0 = d(2.0) * sc(1) * 3 = 6.0.
        assert!(
            (dst[256] - 6.0).abs() < 1e-3,
            "expected 6.0 at block-1 value 0, got {}",
            dst[256]
        );
    }

    #[test]
    fn test_decode_q4km_scales_roundtrip() {
        // Test that known scale values decode correctly
        let mut scales = [0u8; 12];
        // Set sc[0]=5, sc[1]=10, sc[2]=15, sc[3]=20 (6-bit, low bits in bytes 0-3)
        scales[0] = 5;
        scales[1] = 10;
        scales[2] = 15;
        scales[3] = 20;
        // Set mn[0]=1, mn[1]=2, mn[2]=3, mn[3]=4 (6-bit, low bits in bytes 4-7)
        scales[4] = 1;
        scales[5] = 2;
        scales[6] = 3;
        scales[7] = 4;
        // sc[4..8] and mn[4..8]: bytes 8-11, with high bits from bytes 0-3 bits 6-7
        // For simplicity set bytes 8-11 to 0 and don't use high bits
        scales[8] = 0;
        scales[9] = 0;
        scales[10] = 0;
        scales[11] = 0;

        let (sc, mn) = decode_q4km_scales(&scales);
        assert_eq!(sc[0], 5);
        assert_eq!(sc[1], 10);
        assert_eq!(sc[2], 15);
        assert_eq!(sc[3], 20);
        assert_eq!(mn[0], 1);
        assert_eq!(mn[1], 2);
        assert_eq!(mn[2], 3);
        assert_eq!(mn[3], 4);
    }
}