hanzo-ml 0.11.71

//! Numerical validation of the native-GGML Vulkan quantized matvec kernels (Q4_0 / Q8_0 / Q4_K)
//! against the CPU reference, on the real Vulkan GPU.
//!
//! Each test:
//!   1. builds an f32 weight `W[nout, k]` and quantizes it to a GGML dtype on the CPU,
//!   2. dequantizes that QTensor back to f32 — the EXACT quantized weights the GPU kernel reads,
//!   3. uploads the raw GGML block bytes to the GPU (`upload_qweight`) and runs the matvec kernel,
//!   4. compares the GPU output to a host f64 matvec over the dequantized weights.
//!
//! The kernel reads the same quantized bytes as the reference dequant, so the only difference is
//! float accumulation order: agreement must be ~1e-3 relative. We also report the magnitude of the
//! quantization error itself (GPU vs the ORIGINAL f32 weight) for context.
//!
//! Skips cleanly (prints + returns) when no Vulkan GPU is present, so it is safe in CI without a
//! device. Requires the `vulkan` feature.
#![cfg(feature = "vulkan")]

use hanzo_ml::quantized::{GgmlDType, QMatMul, QStorage, QTensor};
use hanzo_ml::{Device, Module, Tensor};
use std::sync::Arc;

// Tests that force a kernel via a process-global env var (VK_DP4A_DECODE_OFF, VK_Q4K_{LEGACY,
// TILED}) must not run concurrently: cargo runs tests in parallel threads and the env set/remove races
// otherwise (flaky pass/fail at varying shapes). Each such test holds this lock for its whole body.
static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());

// Deterministic pseudo-random f32 in [-1, 1) from a counter (no rng dep; reproducible).
fn pseudo(i: usize) -> f32 {
    // splitmix64-ish; take 24 mantissa bits.
    let mut z = (i as u64).wrapping_add(0x9E37_79B9_7F4A_7C15);
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    z ^= z >> 31;
    ((z >> 40) as f32 / (1u32 << 24) as f32) * 2.0 - 1.0
}

struct ErrStats {
    max_abs: f32,
    max_rel: f32,
    rms: f32,
    quant_max_abs: f32, // GPU vs original f32 weight (the quantization error itself)
}

// Run one (dtype, nout, k) case end to end on the GPU and return error stats vs the CPU reference.
// True for GGML types with a CPU quantizer (`from_float`). The decode-only types (IQ4_NL / IQ4_XS /
// TQ2_0) have no `from_float`, so their test weights are synthesized as raw blocks instead.
fn quantizable(dtype: GgmlDType) -> bool {
    !matches!(
        dtype,
        GgmlDType::IQ4_NL
            | GgmlDType::IQ4_XS
            | GgmlDType::TQ2_0
            | GgmlDType::IQ2_XXS
            | GgmlDType::IQ2_S
            | GgmlDType::IQ3_XXS
            | GgmlDType::IQ3_S
            | GgmlDType::IQ1_S
            | GgmlDType::IQ1_M
            | GgmlDType::IQ2_XS
    )
}

// Deterministic pseudo-random byte (same splitmix source as `pseudo`).
fn pbyte(i: usize) -> u8 {
    (((pseudo(i) * 0.5 + 0.5) * 256.0) as i32).clamp(0, 255) as u8
}

// Synthesize valid raw GGML block bytes for a decode-only type (no `from_float` quantizer). Every
// nibble / 2-bit field is a valid index/value; the f16 scale magnitude is chosen so the dequantized
// weights stay O(0.5) -- the same range as the quantizable cases -- so the GPU-vs-CPU tolerance
// (max_abs < 1e-3) applies unchanged. Block layouts are bit-for-bit the CPU block structs.
fn synth_decode_only(dtype: GgmlDType, nout: usize, k: usize) -> Vec<u8> {
    use half::f16;
    let mut out: Vec<u8> = Vec::new();
    let mut c = 0usize;
    match dtype {
        // block_iq4_nl (18 B): f16 d, qs[16]. block of 32.
        GgmlDType::IQ4_NL => {
            for _ in 0..nout * (k / 32) {
                out.extend_from_slice(&f16::from_f32(pseudo(c) * 0.003).to_le_bytes());
                c += 1;
                for _ in 0..16 {
                    out.push(pbyte(c));
                    c += 1;
                }
            }
        }
        // block_iq4_xs (136 B): f16 d, u16 scales_h, scales_l[4], qs[128]. block of 256.
        GgmlDType::IQ4_XS => {
            for _ in 0..nout * (k / 256) {
                out.extend_from_slice(&f16::from_f32(pseudo(c) * 1e-4).to_le_bytes());
                c += 1;
                for _ in 0..6 {
                    out.push(pbyte(c)); // scales_h (u16) + scales_l[4]
                    c += 1;
                }
                for _ in 0..128 {
                    out.push(pbyte(c)); // qs[128]
                    c += 1;
                }
            }
        }
        // block_tq2_0 (66 B): qs[64] THEN f16 d. block of 256.
        GgmlDType::TQ2_0 => {
            for _ in 0..nout * (k / 256) {
                for _ in 0..64 {
                    out.push(pbyte(c)); // qs[64]
                    c += 1;
                }
                out.extend_from_slice(&f16::from_f32(pseudo(c) * 0.3).to_le_bytes());
                c += 1;
            }
        }
        // block_iq2_xxs (66 B): f16 d THEN qs[32] (u16 = 64 B). block of 256. Any byte pattern is a
        // valid codebook block (grid index 0..255, 7-bit sign group, 4-bit scale all in-range).
        GgmlDType::IQ2_XXS => {
            for _ in 0..nout * (k / 256) {
                out.extend_from_slice(&f16::from_f32(pseudo(c) * 0.05).to_le_bytes());
                c += 1;
                for _ in 0..64 {
                    out.push(pbyte(c)); // qs[32] u16
                    c += 1;
                }
            }
        }
        // block_iq2_xs (74 B): f16 d + qs[32] (u16 = 64 B) + scales[8]. block of 256. Any byte
        // pattern is valid (grid index 0..511 = qv&511, 7-bit sign = qv>>9, nibble scales in-range).
        GgmlDType::IQ2_XS => {
            for _ in 0..nout * (k / 256) {
                out.extend_from_slice(&f16::from_f32(pseudo(c) * 0.05).to_le_bytes());
                c += 1;
                for _ in 0..72 {
                    out.push(pbyte(c)); // qs[32] u16 (64 B) + scales[8]
                    c += 1;
                }
            }
        }
        // block_iq2s (82 B): f16 d THEN 80 bytes (qs/qh/scales). block of 256. Any byte
        // pattern is a valid codebook block.
        GgmlDType::IQ2_S => {
            for _ in 0..nout * (k / 256) {
                out.extend_from_slice(&f16::from_f32(pseudo(c) * 0.05).to_le_bytes());
                c += 1;
                for _ in 0..80 {
                    out.push(pbyte(c));
                    c += 1;
                }
            }
        }
        // block_iq3xxs (98 B): f16 d THEN 96 bytes (qs/qh/scales). block of 256. Any byte
        // pattern is a valid codebook block.
        GgmlDType::IQ3_XXS => {
            for _ in 0..nout * (k / 256) {
                out.extend_from_slice(&f16::from_f32(pseudo(c) * 0.02).to_le_bytes());
                c += 1;
                for _ in 0..96 {
                    out.push(pbyte(c));
                    c += 1;
                }
            }
        }
        // block_iq3s (110 B): f16 d THEN 108 bytes (qs/qh/scales). block of 256. Any byte
        // pattern is a valid codebook block.
        GgmlDType::IQ3_S => {
            for _ in 0..nout * (k / 256) {
                out.extend_from_slice(&f16::from_f32(pseudo(c) * 0.004).to_le_bytes());
                c += 1;
                for _ in 0..108 {
                    out.push(pbyte(c));
                    c += 1;
                }
            }
        }
        // block_iq1s (50 B): f16 d THEN 48 bytes (qs/qh/scales). block of 256. Any byte
        // pattern is a valid codebook block.
        GgmlDType::IQ1_S => {
            for _ in 0..nout * (k / 256) {
                out.extend_from_slice(&f16::from_f32(pseudo(c) * 0.05).to_le_bytes());
                c += 1;
                for _ in 0..48 {
                    out.push(pbyte(c));
                    c += 1;
                }
            }
        }
        // block_iq1m (56 B): qs[32] + qh[16] + scales[8]. NO leading d -- d is reconstructed from
        // the high nibbles of the 4 scale u16. Set scales so scale_u16 = 0x2C00 (f16 = 0.0625) and the
        // per-sub-block 3-bit scale fields are 0 (dl = d); qs/qh stay random (valid indices + signs).
        GgmlDType::IQ1_M => {
            for _ in 0..nout * (k / 256) {
                for _ in 0..48 {
                    out.push(pbyte(c)); // qs[32] + qh[16]
                    c += 1;
                }
                // scales[8] = sc0=0, sc1=0, sc2=0xC000, sc3=0x2000 -> scale_u16=0x2C00, dl fields 0.
                out.extend_from_slice(&[0u8, 0u8, 0u8, 0u8, 0u8, 0xC0u8, 0u8, 0x20u8]);
            }
        }
        _ => panic!("synth_decode_only: {dtype:?} is not a decode-only type"),
    }
    out
}

// Raw GGML block bytes for `dtype`, the exact f32 weights they dequantize to (`w_deq`, the kernel's
// ground truth), and the pre-quantization f32 source (`w_orig`, == w_deq for synthesized blocks).
fn weight_bytes(
    dtype: GgmlDType,
    nout: usize,
    k: usize,
) -> hanzo_ml::Result<(Vec<u8>, Vec<f32>, Vec<f32>)> {
    let cpu = Device::Cpu;
    if quantizable(dtype) {
        let w_host: Vec<f32> = (0..nout * k).map(|i| pseudo(i) * 0.5).collect();
        let w_t = Tensor::from_vec(w_host.clone(), (nout, k), &cpu)?;
        let q = QTensor::quantize(&w_t, dtype)?;
        let w_deq: Vec<f32> = q.dequantize(&cpu)?.flatten_all()?.to_vec1::<f32>()?;
        Ok((q.data()?.into_owned(), w_deq, w_host))
    } else {
        let raw = synth_decode_only(dtype, nout, k);
        let qs = QStorage::from_data(std::borrow::Cow::Owned(raw.clone()), &cpu, dtype)?;
        let q = QTensor::new(qs, (nout, k))?;
        let w_deq: Vec<f32> = q.dequantize(&cpu)?.flatten_all()?.to_vec1::<f32>()?;
        Ok((raw, w_deq.clone(), w_deq))
    }
}

fn run_case(dev: &Device, dtype: GgmlDType, nout: usize, k: usize) -> hanzo_ml::Result<ErrStats> {
    let x_host: Vec<f32> = (0..k).map(|i| pseudo(i + 1_000_003)).collect();
    // Raw GGML bytes + the EXACT dequantized weights the kernel reads + the original f32 source.
    let (raw, w_deq, w_host) = weight_bytes(dtype, nout, k)?;

    // Upload bytes + run the GPU kernel. Q6_K / Q3_K repack to a padded u32 stride on upload; every
    // other native type byte- or word-addresses the raw GGML bytes directly.
    let vk = dev.as_vulkan_device()?;
    let wq = match dtype {
        GgmlDType::Q6K => vk.quantize_q6k(&raw, nout, k)?,
        GgmlDType::Q3K => vk.quantize_q3k(&raw, nout, k)?,
        GgmlDType::IQ2_XXS => vk.quantize_iq2xxs(&raw, nout, k)?,
        GgmlDType::IQ2_XS => vk.quantize_iq2xs(&raw, nout, k)?,
        GgmlDType::IQ1_M => vk.quantize_iq1m(&raw, nout, k)?,
        GgmlDType::IQ1_S => vk.quantize_iq1s(&raw, nout, k)?,
        GgmlDType::IQ3_S => vk.quantize_iq3s(&raw, nout, k)?,
        GgmlDType::IQ3_XXS => vk.quantize_iq3xxs(&raw, nout, k)?,
        GgmlDType::IQ2_S => vk.quantize_iq2s(&raw, nout, k)?,
        _ => vk.upload_qweight(&raw)?,
    };
    let y_gpu: Vec<f32> = match dtype {
        GgmlDType::Q4_0 => vk.matvec_q4_0(&wq, &x_host, nout, k)?,
        GgmlDType::Q8_0 => vk.matvec_q8_0(&wq, &x_host, nout, k)?,
        GgmlDType::Q4K => vk.matvec_q4k_scalar(&wq, &x_host, nout, k)?,
        GgmlDType::Q5K => vk.matvec_q5k(&wq, &x_host, nout, k)?,
        GgmlDType::Q6K => vk.matvec_q6k(&wq, &x_host, nout, k)?,
        GgmlDType::Q2K => vk.matvec_q2k(&wq, &x_host, nout, k)?,
        GgmlDType::Q3K => vk.matvec_q3k(&wq, &x_host, nout, k)?,
        GgmlDType::IQ4_XS => vk.matvec_iq4xs(&wq, &x_host, nout, k)?,
        GgmlDType::IQ4_NL => vk.matvec_iq4nl(&wq, &x_host, nout, k)?,
        GgmlDType::IQ2_XXS => vk.matvec_iq2xxs(&wq, &x_host, nout, k)?,
        GgmlDType::IQ2_XS => vk.matvec_iq2xs(&wq, &x_host, nout, k)?,
        GgmlDType::IQ1_M => vk.matvec_iq1m(&wq, &x_host, nout, k)?,
        GgmlDType::IQ1_S => vk.matvec_iq1s(&wq, &x_host, nout, k)?,
        GgmlDType::IQ3_S => vk.matvec_iq3s(&wq, &x_host, nout, k)?,
        GgmlDType::IQ3_XXS => vk.matvec_iq3xxs(&wq, &x_host, nout, k)?,
        GgmlDType::IQ2_S => vk.matvec_iq2s(&wq, &x_host, nout, k)?,
        GgmlDType::TQ2_0 => vk.matvec_tq2_0(&wq, &x_host, nout, k)?,
        _ => panic!("unsupported dtype in run_case: {dtype:?}"),
    };
    assert_eq!(y_gpu.len(), nout);

    // 4a. Reference: host f64 matvec over the DEQUANTIZED weights (what the kernel approximates).
    // 4b. Also matvec over the ORIGINAL weights to report the quantization error magnitude.
    let mut max_abs = 0f32;
    let mut sse = 0f64;
    let mut ref_sq = 0f64; // sum of ref^2, for an RMS-normalized relative error
    let mut quant_max_abs = 0f32;
    for n in 0..nout {
        let mut ref_deq = 0f64;
        let mut ref_orig = 0f64;
        for j in 0..k {
            ref_deq += w_deq[n * k + j] as f64 * x_host[j] as f64;
            ref_orig += w_host[n * k + j] as f64 * x_host[j] as f64;
        }
        let g = y_gpu[n] as f64;
        max_abs = max_abs.max((g - ref_deq).abs() as f32);
        sse += (g - ref_deq) * (g - ref_deq);
        ref_sq += ref_deq * ref_deq;
        quant_max_abs = quant_max_abs.max((g - ref_orig).abs() as f32);
    }
    // Relative error normalized by the RMS of the reference outputs (the robust metric: a per-
    // element relative error explodes whenever a single dot-product lands near zero, which is a
    // property of the data, not the kernel). Both that and max_abs must be tiny.
    let ref_rms = (ref_sq / nout as f64).sqrt();
    let err_rms = (sse / nout as f64).sqrt();
    let max_rel = (err_rms / ref_rms.max(1e-9)) as f32;
    Ok(ErrStats {
        max_abs,
        max_rel,
        rms: err_rms as f32,
        quant_max_abs,
    })
}

// Open the first Vulkan GPU, or None (skip) if there is no device.
fn gpu() -> Option<Device> {
    match Device::new_vulkan(0) {
        Ok(d) => Some(d),
        Err(e) => {
            eprintln!("[vulkan_quant_tests] no Vulkan GPU ({e}); skipping");
            None
        }
    }
}

// Representative Qwen3 hidden sizes. k must be a multiple of the block size; these all are.
const SHAPES: &[(usize, usize)] = &[
    (2048, 2048), // Qwen3 hidden
    (4096, 2048), // up-proj-ish (nout != k)
    (2048, 4096), // down-proj-ish
    (512, 256),   // small, exercises k == one Q4_K super-block
];

#[test]
fn vulkan_matvec_q4_0_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::Q4_0, nout, k)?;
        println!(
            "Q4_0  nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}  (quant err vs f32: {:.3e})",
            s.max_abs, s.max_rel, s.rms, s.quant_max_abs
        );
        // GPU vs CPU-dequant reference: same bytes, only fp-order differences.
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "Q4_0 GPU/CPU mismatch too large: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_q8_0_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::Q8_0, nout, k)?;
        println!(
            "Q8_0  nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}  (quant err vs f32: {:.3e})",
            s.max_abs, s.max_rel, s.rms, s.quant_max_abs
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "Q8_0 GPU/CPU mismatch too large: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_q4k_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::Q4K, nout, k)?;
        println!(
            "Q4_K  nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}  (quant err vs f32: {:.3e})",
            s.max_abs, s.max_rel, s.rms, s.quant_max_abs
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "Q4_K GPU/CPU mismatch too large: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_q5k_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::Q5K, nout, k)?;
        println!(
            "Q5_K  nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}  (quant err vs f32: {:.3e})",
            s.max_abs, s.max_rel, s.rms, s.quant_max_abs
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "Q5_K GPU/CPU mismatch too large: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_q6k_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::Q6K, nout, k)?;
        println!(
            "Q6_K  nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}  (quant err vs f32: {:.3e})",
            s.max_abs, s.max_rel, s.rms, s.quant_max_abs
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "Q6_K GPU/CPU mismatch too large: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_q2k_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::Q2K, nout, k)?;
        println!(
            "Q2_K  nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}  (quant err vs f32: {:.3e})",
            s.max_abs, s.max_rel, s.rms, s.quant_max_abs
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "Q2_K GPU/CPU mismatch too large: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_q3k_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::Q3K, nout, k)?;
        println!(
            "Q3_K  nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}  (quant err vs f32: {:.3e})",
            s.max_abs, s.max_rel, s.rms, s.quant_max_abs
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "Q3_K GPU/CPU mismatch too large: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_iq4xs_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::IQ4_XS, nout, k)?;
        println!(
            "IQ4_XS nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}  (quant err vs f32: {:.3e})",
            s.max_abs, s.max_rel, s.rms, s.quant_max_abs
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "IQ4_XS GPU/CPU mismatch too large: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_iq2xxs_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::IQ2_XXS, nout, k)?;
        println!(
            "IQ2_XXS nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}  (quant err vs f32: {:.3e})",
            s.max_abs, s.max_rel, s.rms, s.quant_max_abs
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "IQ2_XXS GPU/CPU mismatch too large: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_iq2xs_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::IQ2_XS, nout, k)?;
        println!(
            "IQ2_XS  nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}  (quant err vs f32: {:.3e})",
            s.max_abs, s.max_rel, s.rms, s.quant_max_abs
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "IQ2_XS GPU/CPU mismatch too large: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_iq4nl_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::IQ4_NL, nout, k)?;
        println!(
            "IQ4_NL nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}  (quant err vs f32: {:.3e})",
            s.max_abs, s.max_rel, s.rms, s.quant_max_abs
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "IQ4_NL GPU/CPU mismatch too large: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_tq2_0_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::TQ2_0, nout, k)?;
        println!(
            "TQ2_0  nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}  (quant err vs f32: {:.3e})",
            s.max_abs, s.max_rel, s.rms, s.quant_max_abs
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "TQ2_0 GPU/CPU mismatch too large: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_iq2s_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::IQ2_S, nout, k)?;
        println!(
            "IQ2_S nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}",
            s.max_abs, s.max_rel, s.rms
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "IQ2_S GPU/CPU mismatch: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_iq3xxs_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::IQ3_XXS, nout, k)?;
        println!(
            "IQ3_XXS nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}",
            s.max_abs, s.max_rel, s.rms
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "IQ3_XXS GPU/CPU mismatch: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_iq3s_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::IQ3_S, nout, k)?;
        println!(
            "IQ3_S nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}",
            s.max_abs, s.max_rel, s.rms
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "IQ3_S GPU/CPU mismatch: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_iq1s_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::IQ1_S, nout, k)?;
        println!(
            "IQ1_S nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}",
            s.max_abs, s.max_rel, s.rms
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "IQ1_S GPU/CPU mismatch: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}

#[test]
fn vulkan_matvec_iq1m_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    for &(nout, k) in SHAPES {
        let s = run_case(&dev, GgmlDType::IQ1_M, nout, k)?;
        println!(
            "IQ1_M nout={nout:5} k={k:5}  max_abs={:.3e} max_rel={:.3e} rms={:.3e}",
            s.max_abs, s.max_rel, s.rms
        );
        assert!(
            s.max_rel < 1e-3 && s.max_abs < 1e-3,
            "IQ1_M GPU/CPU mismatch: max_abs={} max_rel={}",
            s.max_abs,
            s.max_rel
        );
    }
    Ok(())
}
// ---------------------------------------------------------------------------------------------
// End-to-end// ---------------------------------------------------------------------------------------------
// End-to-end// ---------------------------------------------------------------------------------------------
// End-to-end// ---------------------------------------------------------------------------------------------
// End-to-end// ---------------------------------------------------------------------------------------------
// End-to-end// ---------------------------------------------------------------------------------------------
// End-to-end: the actual model path. Build a quantized QTensor ON the Vulkan device (exactly like
// the GGUF loader: QStorage::from_data on Device::Vulkan -> QTensor::new), wrap it in a QMatMul
// (which now routes Q4_0/Q8_0/Q4_K to the native VulkanQuant kernels), and run QMatMul::forward on
// a single-row Vulkan activation -- the decode hot path. Compare to the CPU QMatMul forward over
// the same quantized weights. This validates the wiring, not just the raw kernel.
fn end_to_end_case(dev: &Device, dtype: GgmlDType, nout: usize, k: usize) -> hanzo_ml::Result<f32> {
    let x_host: Vec<f32> = (0..k).map(|i| pseudo(i + 7)).collect();

    // Raw GGML bytes + the exact dequantized weights. Quantizable types quantize on CPU; decode-only
    // types (IQ4_NL/IQ4_XS/TQ2_0) synthesize raw blocks. The dequantized weights are the f64 ground
    // truth the GPU kernel targets (dequantize-then-dot). NOTE we deliberately do NOT use the CPU
    // `QMatMul::forward`: for k-quants it runs `vec_dot`, which quantizes the ACTIVATION to int8 per
    // block first -- a lossier algorithm than the GPU's f32-activation dequant matvec (differs ~1e-1).
    let (bytes, w_deq, _) = weight_bytes(dtype, nout, k)?;
    let mut y_ref = vec![0f64; nout];
    for n in 0..nout {
        let mut acc = 0f64;
        for j in 0..k {
            acc += w_deq[n * k + j] as f64 * x_host[j] as f64;
        }
        y_ref[n] = acc;
    }

    // Vulkan quantized weight built the loader way, then QMatMul (native VulkanQuant path).
    let qs_vk = QStorage::from_data(std::borrow::Cow::Owned(bytes), dev, dtype)?;
    let q_vk = QTensor::new(qs_vk, (nout, k))?;
    let qm_vk = QMatMul::from_qtensor(q_vk)?;
    // The native decode kernel triggers on a single-row input on the Vulkan device.
    let x_vk = Tensor::from_vec(x_host, (1, k), dev)?;
    let y_vk = qm_vk.forward(&x_vk)?.flatten_all()?.to_vec1::<f32>()?;

    assert_eq!(y_vk.len(), nout);
    let mut max_abs = 0f32;
    for n in 0..nout {
        max_abs = max_abs.max((y_vk[n] as f64 - y_ref[n]).abs() as f32);
    }
    Ok(max_abs)
}

#[test]
fn vulkan_qmatmul_forward_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    let _env_guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
    // Q4_K decode defaults to int8 dp4a (q8_1 activation quant, ~0.4%); force the scalar path so this
    // exact-vs-CPU check validates the decode math. The dp4a q8_1 path is gated separately by
    // vulkan_q4k_dp4a_decode_matches_scalar.
    std::env::set_var("VK_DP4A_DECODE_OFF", "1");
    // (nout, k); k divisible by 256 so every dtype is exercised on the same shapes.
    for &(nout, k) in &[(2048usize, 2048usize), (4096, 2048), (512, 256)] {
        for dt in [
            GgmlDType::Q4_0,
            GgmlDType::Q8_0,
            GgmlDType::Q4K,
            GgmlDType::Q5K,
            GgmlDType::Q6K,
            GgmlDType::Q2K,
            GgmlDType::Q3K,
            GgmlDType::IQ4_XS,
            GgmlDType::IQ4_NL,
            GgmlDType::TQ2_0,
        ] {
            let max_abs = end_to_end_case(&dev, dt, nout, k)?;
            println!("QMatMul::forward {dt:?}  nout={nout:5} k={k:5}  GPU-vs-(dequant ref) max_abs={max_abs:.3e}");
            // GPU dequant-matvec vs the f64 dequant reference over identical quantized weights:
            // pure fp32 accumulation noise.
            assert!(
                max_abs < 1e-3,
                "QMatMul::forward {dt:?} GPU/ref mismatch too large: {max_abs}"
            );
        }
    }
    std::env::remove_var("VK_DP4A_DECODE_OFF");
    Ok(())
}

// ---------------------------------------------------------------------------------------------
// End-to-end PREFILL (M>1): the prompt-processing GEMM path. Build a quantized QTensor on the Vulkan
// device, wrap it in a QMatMul, and run forward on an [M, k] activation -- which now dispatches the
// native quantized GEMM (`mul_mat_q*`: weights stay quantized in VRAM, each weight block decoded ONCE
// per output column and reused across all M rows) instead of dequantizing the whole weight to f32.
// Compare every output row to a host f64 matmul over the dequantized weights. M straddles the host's
// MAX_M=8 row-tiling (partial tile / exact tile / multi-tile + remainder) so the tiling is exercised.
fn prefill_case(
    dev: &Device,
    dtype: GgmlDType,
    m: usize,
    nout: usize,
    k: usize,
) -> hanzo_ml::Result<f32> {
    let cpu = Device::Cpu;
    let w_host: Vec<f32> = (0..nout * k).map(|i| pseudo(i) * 0.5).collect();
    let x_host: Vec<f32> = (0..m * k).map(|i| pseudo(i + 7)).collect();

    // CPU quantized weight; reuse it for both the GPU upload bytes and the dequant reference.
    let w_t = Tensor::from_vec(w_host, (nout, k), &cpu)?;
    let q_cpu = Arc::new(QTensor::quantize(&w_t, dtype)?);
    let bytes = q_cpu.data()?.into_owned();
    let w_deq: Vec<f32> = q_cpu.dequantize(&cpu)?.flatten_all()?.to_vec1::<f32>()?;

    // Reference: y[mi, n] = sum_k W_deq[n, k] * x[mi, k] in f64 (dequantize-then-dot ground truth).
    let mut y_ref = vec![0f64; m * nout];
    for mi in 0..m {
        for n in 0..nout {
            let mut acc = 0f64;
            for j in 0..k {
                acc += w_deq[n * k + j] as f64 * x_host[mi * k + j] as f64;
            }
            y_ref[mi * nout + n] = acc;
        }
    }

    // Vulkan quantized weight built the loader way, then QMatMul (native VulkanQuant prefill GEMM).
    let qs_vk = QStorage::from_data(std::borrow::Cow::Owned(bytes), dev, dtype)?;
    let q_vk = QTensor::new(qs_vk, (nout, k))?;
    let qm_vk = QMatMul::from_qtensor(q_vk)?;
    // A multi-row [M, k] input on the Vulkan device triggers the native prefill GEMM.
    let x_vk = Tensor::from_vec(x_host, (m, k), dev)?;
    let y_vk = qm_vk.forward(&x_vk)?.flatten_all()?.to_vec1::<f32>()?;

    assert_eq!(y_vk.len(), m * nout);
    let mut max_abs = 0f32;
    for i in 0..m * nout {
        max_abs = max_abs.max((y_vk[i] as f64 - y_ref[i]).abs() as f32);
    }
    Ok(max_abs)
}

#[test]
fn vulkan_qmatmul_prefill_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    let _env_guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
    // Q4_K prefill defaults to int8 dp4a-2D (q8_1 activation quant, ~0.4%); force the exact column path
    // so this exact-vs-CPU check validates the prefill math. The dp4a-2D path is gated by
    // vulkan_q4k_dp4a2d_matches_default. (Q5_K/Q6_K still exercise the GEMM/dequant-fallback split.)
    std::env::set_var("VK_Q4K_LEGACY", "1");
    // M straddles the host MAX_M=8 tiling AND the per-dtype GEMM gate (Q5_K/Q6_K: 64, else 128):
    // 5/8/17 hit the native GEMM for every dtype; 100 hits the GEMM for Q4_0/Q8_0/Q4_K but the dequant
    // fallback for Q5_K/Q6_K (exercising the split); 200 (> both gates) is the dequant path for all.
    // BOTH prefill paths are thus validated against the same f64 reference.
    // (nout, k) with k divisible by 256 so all five dtypes share the same shapes.
    for &m in &[5usize, 8, 17, 100, 200] {
        for &(nout, k) in &[(2048usize, 2048usize), (4096, 2048), (512, 256)] {
            for dt in [
                GgmlDType::Q4_0,
                GgmlDType::Q8_0,
                GgmlDType::Q4K,
                GgmlDType::Q5K,
                GgmlDType::Q6K,
            ] {
                let max_abs = prefill_case(&dev, dt, m, nout, k)?;
                println!(
                    "QMatMul::prefill {dt:?}  M={m:3} nout={nout:5} k={k:5}  GPU-vs-(dequant ref) max_abs={max_abs:.3e}"
                );
                // The GEMM decodes the SAME quantized bytes as the reference dequant and accumulates
                // per row exactly like the decode matvec; only fp32-vs-f64 accumulation order differs.
                assert!(
                    max_abs < 1e-3,
                    "QMatMul::prefill {dt:?} M={m} GPU/ref mismatch too large: {max_abs}"
                );
            }
        }
    }
    std::env::remove_var("VK_Q4K_LEGACY");
    Ok(())
}

// ---------------------------------------------------------------------------------------------
// MoE: the fused grouped quant matvec. Build a quantized expert bank [E, n, k] on the Vulkan device
// and call QTensor::indexed_moe_forward (the path Qwen3-MoE uses) -- which now runs one on-GPU
// dispatch that gathers each routed slot's expert and computes its matvec -- and compare to a host
// f64 reference over the dequantized per-expert weights.
fn moe_case(
    dev: &Device,
    dtype: GgmlDType,
    e_cnt: usize,
    n: usize,
    k: usize,
    t: usize,
    topk: usize,
) -> hanzo_ml::Result<f32> {
    let cpu = Device::Cpu;
    // Expert bank [E, n, k] and its dequantized f32 (the reference weights).
    let bank_host: Vec<f32> = (0..e_cnt * n * k).map(|i| pseudo(i) * 0.5).collect();
    let bank_t = Tensor::from_vec(bank_host, (e_cnt, n, k), &cpu)?;
    let q_bank = QTensor::quantize(&bank_t.reshape((e_cnt * n, k))?, dtype)?; // quantize per row
    let bank_deq: Vec<f32> = q_bank.dequantize(&cpu)?.flatten_all()?.to_vec1::<f32>()?; // [E*n*k]
    let bytes = q_bank.data()?.into_owned();

    // Activations [t, topk, k] (per-slot inputs) and router ids [t, topk].
    let x_host: Vec<f32> = (0..t * topk * k).map(|i| pseudo(i + 11) * 0.7).collect();
    // Deterministic expert assignment spread across all experts.
    let ids_host: Vec<u32> = (0..t * topk)
        .map(|i| ((i * 7 + 3) % e_cnt) as u32)
        .collect();

    // Build the Vulkan QTensor bank the loader way and run the fused MoE forward.
    let qs_vk = QStorage::from_data(std::borrow::Cow::Owned(bytes), dev, dtype)?;
    let q_vk = QTensor::new(qs_vk, (e_cnt, n, k))?;
    let x_vk = Tensor::from_vec(x_host.clone(), (t, topk, k), dev)?;
    let ids_vk = Tensor::from_vec(ids_host.clone(), (t, topk), dev)?;
    let y_vk = q_vk
        .indexed_moe_forward(&x_vk, &ids_vk)?
        .reshape((t * topk, n))?
        .to_vec2::<f32>()?;

    // Reference: y[slot, r] = sum_k W_deq[ids[slot], r, k] * x[slot, k].
    let mut max_abs = 0f32;
    for slot in 0..t * topk {
        let e = ids_host[slot] as usize;
        for r in 0..n {
            let wbase = (e * n + r) * k;
            let xbase = slot * k;
            let mut acc = 0f64;
            for j in 0..k {
                acc += bank_deq[wbase + j] as f64 * x_host[xbase + j] as f64;
            }
            max_abs = max_abs.max((y_vk[slot][r] as f64 - acc).abs() as f32);
        }
    }
    Ok(max_abs)
}

// ---------------------------------------------------------------------------------------------
// Flash-attention: fused QK^T -> softmax -> V (online softmax) vs an eager f64 reference. The GPU
// kernel never materializes the [Lq, Lk] score matrix; the reference computes it explicitly.
fn flash_case(
    dev: &Device,
    bh: usize,
    lq: usize,
    lk: usize,
    d: usize,
    causal: bool,
) -> hanzo_ml::Result<f32> {
    let scale = 1.0f32 / (d as f32).sqrt();
    let q: Vec<f32> = (0..bh * lq * d).map(|i| pseudo(i) * 0.5).collect();
    let k: Vec<f32> = (0..bh * lk * d).map(|i| pseudo(i + 5) * 0.5).collect();
    let v: Vec<f32> = (0..bh * lk * d).map(|i| pseudo(i + 9) * 0.5).collect();

    let vk = dev.as_vulkan_device()?;
    let out = vk.flash_attn(&q, &k, &v, bh, lq, lk, d, scale, causal)?;
    assert_eq!(out.len(), bh * lq * d);

    // Eager f64 reference.
    let mut max_abs = 0f32;
    for b in 0..bh {
        for qi in 0..lq {
            let last = if causal { qi + (lk - lq) + 1 } else { lk };
            let last = last.min(lk);
            // scores
            let mut sc = vec![0f64; last];
            let mut mx = f64::NEG_INFINITY;
            for (j, scj) in sc.iter_mut().enumerate() {
                let mut s = 0f64;
                for t in 0..d {
                    s += q[(b * lq + qi) * d + t] as f64 * k[(b * lk + j) * d + t] as f64;
                }
                *scj = s * scale as f64;
                mx = mx.max(*scj);
            }
            let mut denom = 0f64;
            for scj in sc.iter_mut() {
                *scj = (*scj - mx).exp();
                denom += *scj;
            }
            for t in 0..d {
                let mut acc = 0f64;
                for (j, &p) in sc.iter().enumerate() {
                    acc += p * v[(b * lk + j) * d + t] as f64;
                }
                acc /= denom;
                let g = out[(b * lq + qi) * d + t] as f64;
                max_abs = max_abs.max((g - acc).abs() as f32);
            }
        }
    }
    Ok(max_abs)
}

#[test]
fn vulkan_flash_attn_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    // (bh, lq, lk, d, causal). d = 128 is Qwen3's head dim.
    let cases = [
        (8usize, 16usize, 16usize, 128usize, false), // prefill, non-causal
        (8, 16, 16, 128, true),                      // prefill, causal
        (8, 1, 64, 128, true),                       // decode: single query over a 64-key cache
        (4, 7, 13, 64, false),                       // ragged Lq != Lk, smaller d
    ];
    for &(bh, lq, lk, d, causal) in &cases {
        let max_abs = flash_case(&dev, bh, lq, lk, d, causal)?;
        println!(
            "FlashAttn bh={bh} lq={lq:3} lk={lk:3} d={d:3} causal={causal}  GPU-vs-(eager ref) max_abs={max_abs:.3e}"
        );
        assert!(
            max_abs < 1e-4,
            "FlashAttn GPU/ref mismatch too large: {max_abs}"
        );
    }
    Ok(())
}

#[test]
fn vulkan_moe_forward_matches_cpu() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else { return Ok(()) };
    // Qwen3-MoE-ish: many experts, small expert intermediate n, hidden k, a few tokens, top-k.
    // (e_cnt, n, k, t, topk). k divisible by 256 so all three dtypes work on the same shapes.
    let cases = [
        (8usize, 256usize, 256usize, 2usize, 2usize),
        (16, 512, 256, 3, 4),
        (4, 768, 512, 1, 2), // decode-like: single token, top-2
    ];
    for &(e_cnt, n, k, t, topk) in &cases {
        for dt in [GgmlDType::Q4_0, GgmlDType::Q8_0, GgmlDType::Q4K] {
            let max_abs = moe_case(&dev, dt, e_cnt, n, k, t, topk)?;
            println!(
                "MoE {dt:?}  E={e_cnt:3} n={n:4} k={k:4} t={t} topk={topk}  GPU-vs-(dequant ref) max_abs={max_abs:.3e}"
            );
            assert!(
                max_abs < 1e-3,
                "MoE {dt:?} GPU/ref mismatch too large: {max_abs}"
            );
        }
    }
    Ok(())
}

// L1 2D-tiled prefill GEMM (mul_mm_q4k_tiled, VK_Q4K_TILED2D): 64x64 output tile, both operands
// staged in LDS per K-step. Decode per element is identical to mul_mat_q4k but accumulation is tiled
// (partial sums), so it matches the default within f32-reorder tolerance, not bit-exact. A decode/index
// bug is a systematic >>1e-3 divergence; f32 reassociation is <~1e-4.
#[test]
fn vulkan_q4k_tiled2d_matches_default() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else {
        return Ok(());
    };
    let _env_guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
    let vk = dev.as_vulkan_device()?;
    for &(nout, k) in &[(512usize, 256usize), (2048, 2048), (256, 4096), (320, 1024)] {
        for &m in &[1usize, 7, 64, 65, 512] {
            let (raw, _, _) = weight_bytes(GgmlDType::Q4K, nout, k)?;
            let wq = vk.upload_qweight(&raw)?;
            let x: Vec<f32> = (0..m * k).map(|i| pseudo(i + 11)).collect();
            std::env::set_var("VK_Q4K_LEGACY", "1");
            let y_ref = vk.matmul_q4k(&wq, &x, m, nout, k)?;
            std::env::remove_var("VK_Q4K_LEGACY");
            std::env::set_var("VK_Q4K_TILED2D", "1");
            let y_2d = vk.matmul_q4k(&wq, &x, m, nout, k)?;
            std::env::remove_var("VK_Q4K_TILED2D");
            assert_eq!(y_2d.len(), m * nout);
            let mut max_abs = 0f32;
            let mut max_ref = 0f32;
            for (a, b) in y_ref.iter().zip(y_2d.iter()) {
                max_abs = max_abs.max((a - b).abs());
                max_ref = max_ref.max(a.abs());
            }
            let rel = if max_ref > 0.0 {
                max_abs / max_ref
            } else {
                max_abs
            };
            assert!(
                rel < 2e-3,
                "2d-tiled != default (m={m}, nout={nout}, k={k}): max_abs={max_abs}, rel={rel}"
            );
        }
    }
    Ok(())
}

// L1 perf: default mul_mat_q4k vs 2D-tiled mul_mm_q4k_tiled at a realistic prefill shape (run --ignored).
#[test]
#[ignore]
fn vulkan_q4k_tiled2d_bench() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else {
        return Ok(());
    };
    let vk = dev.as_vulkan_device()?;
    let (m, nout, k) = (512usize, 4096usize, 4096usize);
    let (raw, _, _) = weight_bytes(GgmlDType::Q4K, nout, k)?;
    let wq = vk.upload_qweight(&raw)?;
    let x: Vec<f32> = (0..m * k).map(|i| pseudo(i + 7)).collect();
    let iters = 20;
    let bench = |var: Option<&str>| -> hanzo_ml::Result<f64> {
        std::env::remove_var("VK_Q4K_TILED2D");
        if let Some(v) = var {
            std::env::set_var(v, "1");
        }
        let _ = vk.matmul_q4k(&wq, &x, m, nout, k)?;
        let t = std::time::Instant::now();
        for _ in 0..iters {
            let _ = vk.matmul_q4k(&wq, &x, m, nout, k)?;
        }
        if let Some(v) = var {
            std::env::remove_var(v);
        }
        Ok(t.elapsed().as_secs_f64() * 1e3 / iters as f64)
    };
    let def_ms = bench(Some("VK_Q4K_LEGACY"))?;
    let tiled2d_ms = bench(Some("VK_Q4K_TILED2D"))?;
    eprintln!(
        "[L1 2D bench] m={m} nout={nout} k={k}: default={def_ms:.3} ms  tiled2d={tiled2d_ms:.3} ms  speedup={:.2}x",
        def_ms / tiled2d_ms
    );
    Ok(())
}

// L1 dp4a 2D-tiled GEMM (mul_mm_q4k_tiled_dp4a, VK_Q4K_DP4A): the 2D tile with an int8 dp4a inner
// loop -- activations are quantized to q8_1 (int8 + per-32-block scale/sum) so the dot is integer. That
// q8_1 quant adds ~0.5-1% error on top of the f32 tiling reorder, so the gate vs the exact f32 default
// is a ~1.5% relative bound (a decode/index/pack bug is a systematic >>1.5% garble).
#[test]
fn vulkan_q4k_dp4a2d_matches_default() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else {
        return Ok(());
    };
    let _env_guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
    let vk = dev.as_vulkan_device()?;
    for &(nout, k) in &[(512usize, 256usize), (2048, 2048), (256, 4096), (320, 1024)] {
        for &m in &[1usize, 7, 64, 65, 512] {
            let (raw, _, _) = weight_bytes(GgmlDType::Q4K, nout, k)?;
            let wq = vk.upload_qweight(&raw)?;
            let x: Vec<f32> = (0..m * k).map(|i| pseudo(i + 13)).collect();
            std::env::set_var("VK_Q4K_LEGACY", "1");
            let y_ref = vk.matmul_q4k(&wq, &x, m, nout, k)?;
            std::env::remove_var("VK_Q4K_LEGACY");
            std::env::set_var("VK_Q4K_DP4A", "1");
            let y_dp = vk.matmul_q4k(&wq, &x, m, nout, k)?;
            std::env::remove_var("VK_Q4K_DP4A");
            assert_eq!(y_dp.len(), m * nout);
            let mut max_abs = 0f32;
            let mut max_ref = 0f32;
            for (a, b) in y_ref.iter().zip(y_dp.iter()) {
                max_abs = max_abs.max((a - b).abs());
                max_ref = max_ref.max(a.abs());
            }
            let rel = if max_ref > 0.0 {
                max_abs / max_ref
            } else {
                max_abs
            };
            assert!(
                rel < 1.5e-2,
                "dp4a-2d != default (m={m}, nout={nout}, k={k}): max_abs={max_abs}, max_ref={max_ref}, rel={rel}"
            );
        }
    }
    Ok(())
}

// L1 perf: default vs 2D-f32 vs 2D-dp4a at a realistic prefill shape (run --ignored).
#[test]
#[ignore]
fn vulkan_q4k_dp4a2d_bench() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else {
        return Ok(());
    };
    let vk = dev.as_vulkan_device()?;
    let (m, nout, k) = (512usize, 4096usize, 4096usize);
    let (raw, _, _) = weight_bytes(GgmlDType::Q4K, nout, k)?;
    let wq = vk.upload_qweight(&raw)?;
    let x: Vec<f32> = (0..m * k).map(|i| pseudo(i + 7)).collect();
    let iters = 20;
    let bench = |var: Option<&str>| -> hanzo_ml::Result<f64> {
        for v in ["VK_Q4K_TILED2D", "VK_Q4K_DP4A"] {
            std::env::remove_var(v);
        }
        if let Some(v) = var {
            std::env::set_var(v, "1");
        }
        let _ = vk.matmul_q4k(&wq, &x, m, nout, k)?;
        let t = std::time::Instant::now();
        for _ in 0..iters {
            let _ = vk.matmul_q4k(&wq, &x, m, nout, k)?;
        }
        if let Some(v) = var {
            std::env::remove_var(v);
        }
        Ok(t.elapsed().as_secs_f64() * 1e3 / iters as f64)
    };
    let def_ms = bench(Some("VK_Q4K_LEGACY"))?;
    let f32_ms = bench(Some("VK_Q4K_TILED2D"))?;
    let dp4a_ms = bench(Some("VK_Q4K_DP4A"))?;
    eprintln!(
        "[L1 dp4a bench] m={m} nout={nout} k={k}: default={def_ms:.3} 2d_f32={f32_ms:.3} 2d_dp4a={dp4a_ms:.3} ms | dp4a vs default={:.2}x vs 2d_f32={:.2}x",
        def_ms / dp4a_ms,
        f32_ms / dp4a_ms
    );
    Ok(())
}

// L4 coopmat Q4_K prefill GEMM (mul_mm_q4k_coopmat, VK_Q4K_COOPMAT): tensor-core path (decode
// weight to f16 LDS tiles + coopMatMulAdd). f16 weight + f16 activation rounding (f32 accumulate), so
// the gate vs the exact legacy column kernel is ~1% relative (a decode/transpose/tile bug is >>1%).
// Skips if the device lacks coopmat (the env path is a no-op fallthrough -> equals the default).
#[test]
fn vulkan_q4k_coopmat_matches_default() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else {
        return Ok(());
    };
    let _env_guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
    let vk = dev.as_vulkan_device()?;
    if vk.coopmat_info().is_none() {
        eprintln!("[vulkan_quant_tests] no coopmat; skipping coopmat gate");
        return Ok(());
    }
    for &(nout, k) in &[(512usize, 256usize), (2048, 2048), (256, 4096), (320, 1024)] {
        for &m in &[2usize, 16, 17, 64, 512] {
            let (raw, _, _) = weight_bytes(GgmlDType::Q4K, nout, k)?;
            let wq = vk.upload_qweight(&raw)?;
            let x: Vec<f32> = (0..m * k).map(|i| pseudo(i + 17)).collect();
            std::env::set_var("VK_Q4K_LEGACY", "1");
            let y_ref = vk.matmul_q4k(&wq, &x, m, nout, k)?;
            std::env::remove_var("VK_Q4K_LEGACY");
            std::env::set_var("VK_Q4K_COOPMAT", "1");
            let y_cm = vk.matmul_q4k(&wq, &x, m, nout, k)?;
            std::env::remove_var("VK_Q4K_COOPMAT");
            let mut max_abs = 0f32;
            let mut max_ref = 0f32;
            for (a, b) in y_ref.iter().zip(y_cm.iter()) {
                max_abs = max_abs.max((a - b).abs());
                max_ref = max_ref.max(a.abs());
            }
            let rel = if max_ref > 0.0 {
                max_abs / max_ref
            } else {
                max_abs
            };
            assert!(
                rel < 1e-2,
                "coopmat != legacy (m={m}, nout={nout}, k={k}): max_abs={max_abs}, rel={rel}"
            );
        }
    }
    Ok(())
}

// L4 perf: default (dp4a) vs coopmat at a realistic prefill shape (run --ignored).
#[test]
#[ignore]
fn vulkan_q4k_coopmat_bench() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else {
        return Ok(());
    };
    let vk = dev.as_vulkan_device()?;
    if vk.coopmat_info().is_none() {
        eprintln!("[vulkan_quant_tests] no coopmat; skipping");
        return Ok(());
    }
    let (m, nout, k) = (512usize, 4096usize, 4096usize);
    let (raw, _, _) = weight_bytes(GgmlDType::Q4K, nout, k)?;
    let wq = vk.upload_qweight(&raw)?;
    let x: Vec<f32> = (0..m * k).map(|i| pseudo(i + 7)).collect();
    let iters = 20;
    let bench = |var: Option<&str>| -> hanzo_ml::Result<f64> {
        for v in ["VK_Q4K_COOPMAT", "VK_Q4K_LEGACY"] {
            std::env::remove_var(v);
        }
        if let Some(v) = var {
            std::env::set_var(v, "1");
        }
        let _ = vk.matmul_q4k(&wq, &x, m, nout, k)?;
        let t = std::time::Instant::now();
        for _ in 0..iters {
            let _ = vk.matmul_q4k(&wq, &x, m, nout, k)?;
        }
        if let Some(v) = var {
            std::env::remove_var(v);
        }
        Ok(t.elapsed().as_secs_f64() * 1e3 / iters as f64)
    };
    let dp4a_ms = bench(None)?; // default == dp4a on this device
    let legacy_ms = bench(Some("VK_Q4K_LEGACY"))?;
    let cm_ms = bench(Some("VK_Q4K_COOPMAT"))?;
    eprintln!(
        "[L4 coopmat bench] m={m} nout={nout} k={k}: legacy={legacy_ms:.3} dp4a={dp4a_ms:.3} coopmat={cm_ms:.3} ms | coopmat vs dp4a={:.2}x vs legacy={:.2}x",
        dp4a_ms / cm_ms,
        legacy_ms / cm_ms
    );
    Ok(())
}

// Kernel-isolated perf: the TRUE Q4_K prefill kernel time (GPU-resident operands, no per-iter host
// transfer) for legacy / f32-2D / dp4a / coopmat, plus the effective GFLOP/s. Run with --ignored.
#[test]
#[ignore]
fn vulkan_q4k_kernel_bench() -> hanzo_ml::Result<()> {
    let Some(dev) = gpu() else {
        return Ok(());
    };
    let vk = dev.as_vulkan_device()?;
    let (m, nout, k) = (512usize, 4096usize, 4096usize);
    let (raw, _, _) = weight_bytes(GgmlDType::Q4K, nout, k)?;
    let wq = vk.upload_qweight(&raw)?;
    let x: Vec<f32> = (0..m * k).map(|i| pseudo(i + 7)).collect();
    let bench = |var: Option<&str>| -> hanzo_ml::Result<f64> {
        for v in [
            "VK_Q4K_LEGACY",
            "VK_Q4K_TILED2D",
            "VK_Q4K_DP4A",
            "VK_Q4K_COOPMAT",
        ] {
            std::env::remove_var(v);
        }
        if let Some(v) = var {
            std::env::set_var(v, "1");
        }
        let r = vk.bench_matmul_q4k(&wq, &x, m, nout, k, 50);
        if let Some(v) = var {
            std::env::remove_var(v);
        }
        r
    };
    let leg = bench(Some("VK_Q4K_LEGACY"))?;
    let f32t = bench(Some("VK_Q4K_TILED2D"))?;
    let dp = bench(None)?; // default == dp4a on int_dot8 devices
    let cm = if vk.coopmat_info().is_some() {
        bench(Some("VK_Q4K_COOPMAT"))?
    } else {
        0.0
    };
    let gflop = 2.0 * m as f64 * nout as f64 * k as f64 / 1e9;
    eprintln!(
        "[KERNEL bench] m={m} nout={nout} k={k} ({gflop:.1} GFLOP): legacy={leg:.3} f32_2d={f32t:.3} dp4a={dp:.3} coopmat={cm:.3} ms | dp4a={:.0} GFLOP/s, {:.1}x vs legacy",
        gflop / (dp / 1e3),
        leg / dp
    );
    Ok(())
}

// Fused rope_norm (rms_norm + NeoX rope in ONE dispatch) must match the unfused two-op chain
// bit-exactly (same f32 ops, same order). First rung of the Vulkan decode op-fusion campaign:
// it removes the inter-op global barrier that serializes decode (see LLM.md / paper 06).
#[test]
fn vulkan_rope_norm_matches_unfused() {
    let Some(dev) = gpu() else {
        return;
    };
    let vk = dev.as_vulkan_device().unwrap();
    // (b, h, t, d): decode (t=1) + short prefill; q/k head counts; head dims 64 and 128.
    for &(b, h, t, d) in &[
        (1usize, 4usize, 1usize, 128usize),
        (1, 8, 5, 64),
        (2, 4, 3, 128),
    ] {
        let hd = d / 2;
        let n = b * h * t * d;
        let x: Vec<f32> = (0..n).map(|i| pseudo(i + 7)).collect();
        let weight: Vec<f32> = (0..d).map(|i| 0.5 + pseudo(i + 99).abs()).collect();
        let cs: Vec<f32> = (0..t * hd).map(|i| pseudo(i + 3).cos()).collect();
        let sn: Vec<f32> = (0..t * hd).map(|i| pseudo(i + 3).sin()).collect();
        let eps = 1e-6f32;
        // CPU reference: rms_norm (x / sqrt(mean(x^2)+eps) * weight) then NeoX rope.
        let mut refv = vec![0f32; n];
        for row in 0..b * h * t {
            let base = row * d;
            let mut ss = 0f32;
            for i in 0..d {
                let v = x[base + i];
                ss += v * v;
            }
            let denom = (ss / d as f32 + eps).sqrt();
            let i_t = row % t;
            for i_d in 0..hd {
                let (i1, i2) = (base + i_d, base + i_d + hd);
                let x1 = x[i1] / denom * weight[i_d];
                let x2 = x[i2] / denom * weight[i_d + hd];
                let (c, s) = (cs[i_t * hd + i_d], sn[i_t * hd + i_d]);
                refv[i1] = x1 * c - x2 * s;
                refv[i2] = x1 * s + x2 * c;
            }
        }
        let got = vk
            .rope_norm_f32(&x, &weight, &cs, &sn, b, h, t, d, eps)
            .unwrap();
        assert_eq!(got.len(), n);
        let mut maxabs = 0f32;
        for i in 0..n {
            maxabs = maxabs.max((got[i] - refv[i]).abs());
        }
        eprintln!("rope_norm b{b} h{h} t{t} d{d}: maxabs={maxabs:.3e}");
        assert!(
            maxabs < 1e-4,
            "rope_norm (b{b} h{h} t{t} d{d}) mismatch maxabs={maxabs}"
        );
    }
}

// Fused add_rmsnorm (residual-add + RMSNorm -> (s, y) in one dispatch) must match the unfused
// add-then-rms_norm chain bit-exactly. Rung 2 of the Vulkan decode op-fusion campaign.
#[test]
fn vulkan_add_rmsnorm_matches_unfused() {
    let Some(dev) = gpu() else {
        return;
    };
    let vk = dev.as_vulkan_device().unwrap();
    for &(nrows, m) in &[(1usize, 2048usize), (5, 64), (3, 4096)] {
        let n = nrows * m;
        let x: Vec<f32> = (0..n).map(|i| pseudo(i + 11)).collect();
        let res: Vec<f32> = (0..n).map(|i| pseudo(i + 222)).collect();
        let alpha: Vec<f32> = (0..m).map(|i| 0.5 + pseudo(i + 9).abs()).collect();
        let eps = 1e-5f32;
        // CPU reference: s = x + res ; y = s / sqrt(mean(s^2)+eps) * alpha.
        let mut s_ref = vec![0f32; n];
        let mut y_ref = vec![0f32; n];
        for row in 0..nrows {
            let base = row * m;
            let mut ss = 0f32;
            for i in 0..m {
                let v = x[base + i] + res[base + i];
                s_ref[base + i] = v;
                ss += v * v;
            }
            let denom = (ss / m as f32 + eps).sqrt();
            for i in 0..m {
                y_ref[base + i] = s_ref[base + i] / denom * alpha[i];
            }
        }
        let (s_gpu, y_gpu) = vk.add_rmsnorm_f32(&x, &res, &alpha, nrows, m, eps).unwrap();
        let (mut ms, mut my) = (0f32, 0f32);
        for i in 0..n {
            ms = ms.max((s_gpu[i] - s_ref[i]).abs());
            my = my.max((y_gpu[i] - y_ref[i]).abs());
        }
        eprintln!("add_rmsnorm r{nrows} m{m}: s_maxabs={ms:.3e} y_maxabs={my:.3e}");
        assert!(
            ms < 1e-5 && my < 1e-4,
            "add_rmsnorm r{nrows} m{m} mismatch s={ms} y={my}"
        );
    }
}

// The dp4a column decode matvec (matvec_q4k_dp4a, the shipped optimum) must match the SCALAR matvec
// within the q8_1 activation-quant tolerance. matvec_q4k_scalar forces the non-dp4a path (the exact
// reference) regardless of the dp4a default. Validates the Vulkan decode lever (1.8x).
#[test]
fn vulkan_q4k_dp4a_decode_matches_scalar() {
    let Some(dev) = gpu() else {
        return;
    };
    let vk = dev.as_vulkan_device().unwrap();
    for &(nout, k) in SHAPES {
        let x: Vec<f32> = (0..k).map(|i| pseudo(i + 5)).collect();
        let (raw, _w_deq, _w_host) = weight_bytes(GgmlDType::Q4K, nout, k).unwrap();
        let wq = vk.upload_qweight(&raw).unwrap();
        let scalar = vk.matvec_q4k_scalar(&wq, &x, nout, k).unwrap();
        let column = vk.matvec_q4k_dp4a(&wq, &x, nout, k).unwrap();
        for (name, got) in [("column", &column)] {
            let (mut sse, mut refsq) = (0f64, 0f64);
            for i in 0..nout {
                let d = (got[i] - scalar[i]) as f64;
                sse += d * d;
                refsq += (scalar[i] as f64) * (scalar[i] as f64);
            }
            let rel = (sse / refsq.max(1e-9)).sqrt();
            eprintln!("dp4a-{name} vs scalar nout={nout} k={k}: rel={rel:.3e}");
            assert!(
                rel < 2e-2,
                "dp4a {name} rel err {rel} too large (nout={nout} k={k})"
            );
        }
    }
}