cera 0.3.1

use half::f16;

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

// ── 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 {
    pub d: u16, // f16 stored as raw bits
    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 {
    pub d: u16, // f16 stored as raw bits
    pub m: u16, // f16 stored as raw bits
    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 {
    pub delta: u16, // f16 stored as raw bits
    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 {
    pub d: u16,    // f16 stored as raw bits
    pub dmin: u16, // f16 stored as raw bits
    pub scales: [u8; 12],
    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 {
    pub ql: [u8; 128],
    pub qh: [u8; 64],
    pub scales: [i8; 16],
    pub d: u16, // f16 stored as raw bits
}

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 {
    pub d: u16,    // f16 stored as raw bits
    pub dmin: u16, // f16 stored as raw bits
    pub scales: [u8; 12],
    pub qh: [u8; 32],
    pub qs: [u8; 128],
}

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

// ── 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::from_bits(block.d).to_f32();
    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.
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);

    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 a Q4_0 matrix of shape `[m, k]` (row-major) to `out`.
///
/// `src` is the raw packed block bytes, `out` must have space for `m * k` f32s.
/// Rows are dequantized in parallel with rayon (rayon's split-on-demand handles
/// tiny inputs by running them on a single worker, so no manual cutoff needed).
pub fn dequantize_q4_0_matrix(src: &[u8], m: usize, k: usize, out: &mut [f32]) {
    debug_assert_eq!(
        k % 32,
        0,
        "dequantize_q4_0_matrix: k must be a multiple of 32"
    );
    let row_bytes = (k / 32) * size_of::<BlockQ4_0>();
    debug_assert_eq!(
        src.len(),
        m * row_bytes,
        "dequantize_q4_0_matrix: src length mismatch"
    );
    debug_assert_eq!(
        out.len(),
        m * k,
        "dequantize_q4_0_matrix: out length mismatch"
    );

    out.par_chunks_mut(k)
        .zip(src.par_chunks(row_bytes))
        .for_each(|(dst_row, src_row)| dequantize_q4_0_row(src_row, dst_row));
}

/// 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::from_bits(block.d).to_f32();
    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::from_bits(block.d).to_f32();
    let m = f16::from_bits(block.m).to_f32();
    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 a Q4_1 matrix of shape `[m, k]` (row-major) to `out`.
pub fn dequantize_q4_1_matrix(src: &[u8], m: usize, k: usize, out: &mut [f32]) {
    debug_assert_eq!(
        k % 32,
        0,
        "dequantize_q4_1_matrix: k must be a multiple of 32"
    );
    let row_bytes = (k / 32) * size_of::<BlockQ4_1>();
    debug_assert_eq!(
        src.len(),
        m * row_bytes,
        "dequantize_q4_1_matrix: src length mismatch"
    );
    debug_assert_eq!(
        out.len(),
        m * k,
        "dequantize_q4_1_matrix: out length mismatch"
    );

    out.par_chunks_mut(k)
        .zip(src.par_chunks(row_bytes))
        .for_each(|(dst_row, src_row)| dequantize_q4_1_row(src_row, dst_row));
}

/// Dot product of a Q4_1 block with an f32 vector of length 32.
///
/// Scalar only: Q4_1 is a legacy format with no SIMD or GPU kernels in this
/// tree, so this is the single implementation rather than a reference the
/// vectorized paths are checked against. Note that the `m` offset would not
/// carry over to the int8 kernels unchanged — the `dpbusd` sign trick the Q4_0
/// path relies on assumes a zero-centred quant, so a Q4_1 VNNI kernel would
/// need a separate correction term.
///
/// `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::from_bits(block.d).to_f32();
    let m = f16::from_bits(block.m).to_f32();
    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::from_bits(block.delta).to_f32();
    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.
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);

    for i in 0..n_blocks {
        let block_bytes = &src[i * block_size..(i + 1) * block_size];
        // SAFETY: BlockQ8_0 is repr(C, packed) and we've verified the slice length
        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 a Q8_0 matrix of shape `[m, k]` (row-major) to `out`.
///
/// `src` is the raw packed block bytes, `out` must have space for `m * k` f32s.
/// Rows are dequantized in parallel with rayon (rayon's split-on-demand handles
/// tiny inputs by running them on a single worker, so no manual cutoff needed).
pub fn dequantize_q8_0_matrix(src: &[u8], m: usize, k: usize, out: &mut [f32]) {
    debug_assert_eq!(
        k % 32,
        0,
        "dequantize_q8_0_matrix: k must be a multiple of 32"
    );
    let row_bytes = (k / 32) * size_of::<BlockQ8_0>();
    debug_assert_eq!(
        src.len(),
        m * row_bytes,
        "dequantize_q8_0_matrix: src length mismatch"
    );
    debug_assert_eq!(
        out.len(),
        m * k,
        "dequantize_q8_0_matrix: out length mismatch"
    );

    out.par_chunks_mut(k)
        .zip(src.par_chunks(row_bytes))
        .for_each(|(dst_row, src_row)| dequantize_q8_0_row(src_row, dst_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).
pub fn dequantize_q4_k_m_matrix(src: &[u8], m: usize, k: usize, out: &mut [f32]) {
    debug_assert_eq!(
        k % 256,
        0,
        "dequantize_q4_k_m_matrix: k must be a multiple of 256"
    );
    let row_bytes = (k / 256) * size_of::<BlockQ4KM>();
    debug_assert_eq!(
        src.len(),
        m * row_bytes,
        "dequantize_q4_k_m_matrix: src length mismatch"
    );
    debug_assert_eq!(
        out.len(),
        m * k,
        "dequantize_q4_k_m_matrix: out length mismatch"
    );

    out.par_chunks_mut(k)
        .zip(src.par_chunks(row_bytes))
        .for_each(|(dst_row, src_row)| dequantize_q4_k_m_row(src_row, dst_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).
pub fn dequantize_q6_k_matrix(src: &[u8], m: usize, k: usize, out: &mut [f32]) {
    debug_assert_eq!(
        k % 256,
        0,
        "dequantize_q6_k_matrix: k must be a multiple of 256"
    );
    let row_bytes = (k / 256) * size_of::<BlockQ6K>();
    debug_assert_eq!(
        src.len(),
        m * row_bytes,
        "dequantize_q6_k_matrix: src length mismatch"
    );
    debug_assert_eq!(
        out.len(),
        m * k,
        "dequantize_q6_k_matrix: out length mismatch"
    );

    out.par_chunks_mut(k)
        .zip(src.par_chunks(row_bytes))
        .for_each(|(dst_row, src_row)| dequantize_q6_k_row(src_row, dst_row));
}

/// 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::from_bits(block.delta).to_f32();
    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::from_bits(block.d).to_f32();
    let dmin = f16::from_bits(block.dmin).to_f32();
    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::from_bits(block.d).to_f32();
    let dmin = f16::from_bits(block.dmin).to_f32();
    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::from_bits(block.d).to_f32();
    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::from_bits(block.d).to_f32();
    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 {
    vec_dot_q6_k_f32_scalar(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::from_bits(block.d).to_f32();
    let dmin = f16::from_bits(block.dmin).to_f32();
    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::from_bits(block.d).to_f32();
    let dmin = f16::from_bits(block.dmin).to_f32();
    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 {
    vec_dot_q5_k_f32_scalar(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::*;

    /// 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]
            );
        }
    }

    #[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);
    }
}