hanzo_ml/quantized/

mod.rs

use crate::{
    backend::BackendStorage, CpuStorage, DType, Device, Result, Shape, Storage, Tensor, D,
};
use iq_quants::*;
use k_quants::*;
use std::borrow::Cow;
use std::sync::Arc;

#[cfg(target_feature = "avx2")]
pub mod avx;
pub mod dsv4_qat;
mod dummy_cuda;
mod dummy_metal;
pub mod expert_stream;
pub mod ggml_file;
pub mod gguf_file;
pub mod imatrix_file;
mod iq_grids;
pub mod iq_quants;
pub mod k_quants;
#[cfg(feature = "metal")]
pub mod metal;
#[cfg(not(target_arch = "wasm32"))]
pub mod tokenizer;
#[cfg(not(feature = "metal"))]
mod metal {
    pub use super::dummy_metal::*;
}
#[cfg(feature = "cuda")]
pub mod cuda;
#[cfg(feature = "cuda")]
pub mod fast_mmq;
#[cfg(feature = "cuda")]
pub mod fast_mmvq;
#[cfg(not(feature = "cuda"))]
mod cuda {
    pub use super::dummy_cuda::*;
}

#[cfg(target_feature = "neon")]
pub mod neon;
#[cfg(target_feature = "simd128")]
pub mod simd128;
pub mod utils;
// Declarative GGUF quant-type formatter (Cut 2). Additive: defines `quant_format!` and a
// `#[cfg(test)]` equivalence proof; does not alter any existing type or wiring.
pub mod quant_format;
use half::{bf16, f16};

pub use k_quants::GgmlType;

// Borrows `data` (does not consume it) so the returned slice stays valid for the caller's lifetime.
// Taking `Cow` by value here was a use-after-free: the Cow dropped at return, dangling the slice for
// Cow::Owned inputs (segfault on large quantized tensors; mmap'd Cow::Borrowed happened to survive).
fn as_t_slice<T>(data: &[u8]) -> &[T] {
    let size = std::mem::size_of::<T>();
    assert_eq!(
        data.len() % size,
        0,
        "Data length must be a multiple of T's size"
    );
    let ptr = data.as_ptr();
    assert_eq!(
        (ptr as usize) % std::mem::align_of::<T>(),
        0,
        "Data pointer must be aligned to T's alignment"
    );
    unsafe { std::slice::from_raw_parts(ptr as *const T, data.len() / size) }
}

// The GPU-resident copies of this tensor's expert bank. A resident bank is a value DERIVED from
// these exact quantized bytes, so it is owned by the tensor that owns the bytes and dies with them.
// One slot per backend; each fills at most once (first routed token) and is then shared by every
// later call through the `Arc<QTensor>` all consumers already hold.
#[derive(Default)]
struct ResidentBanks {
    #[cfg(feature = "rocm")]
    rocm: std::sync::OnceLock<std::sync::Arc<crate::RocmStorage>>,
    #[cfg(feature = "vulkan")]
    vulkan: std::sync::OnceLock<std::sync::Arc<crate::VulkanStorage>>,
    #[cfg(feature = "vulkan")]
    vulkan_split: std::sync::OnceLock<std::sync::Arc<crate::vulkan::MoeBankSplit>>,
    #[cfg(feature = "wgpu")]
    wgpu: std::sync::OnceLock<std::sync::Arc<crate::WgpuStorage>>,
}

pub struct QTensor {
    storage: QStorage,
    shape: Shape,
    // Every field of ResidentBanks is gated on rocm/vulkan/wgpu, so with none of
    // them enabled it is an empty struct and this field is genuinely never read —
    // which is what clippy reports on a default-feature build. It IS read on any
    // accelerator build (see the cache_or_upload calls below), so the allowance is
    // scoped to exactly the configuration where the field is dead rather than
    // silencing the lint everywhere.
    #[cfg_attr(
        not(any(feature = "rocm", feature = "vulkan", feature = "wgpu")),
        allow(dead_code)
    )]
    banks: ResidentBanks,
}

impl Device {
    fn qzeros(&self, elem_count: usize, dtype: GgmlDType) -> Result<QStorage> {
        match self {
            Device::Cpu => {
                let storage = dtype.cpu_zeros(elem_count);
                Ok(QStorage::Cpu(storage))
            }
            Device::Metal(metal) => {
                let storage = metal::QMetalStorage::zeros(metal, elem_count, dtype)?;
                Ok(QStorage::Metal(storage))
            }
            Device::Cuda(cuda) => {
                let storage = cuda::QCudaStorage::zeros(cuda, elem_count, dtype)?;
                Ok(QStorage::Cuda(storage))
            }
            #[cfg(feature = "rocm")]
            Device::Rocm(d) => {
                // Mirror Vulkan: keep quantized blocks in a CPU box + the device; QMatMul
                // dequantizes them to a dense ROCm tensor on use.
                let storage = dtype.cpu_zeros(elem_count);
                Ok(QStorage::Rocm(storage, d.clone()))
            }
            #[cfg(feature = "vulkan")]
            Device::Vulkan(d) => {
                // Keep quantized blocks in a CPU box + the device (same as from_data); QMatMul's
                // VulkanQuant path uploads/dequantizes them to the GPU on use.
                let storage = dtype.cpu_zeros(elem_count);
                Ok(QStorage::Vulkan(storage, d.clone()))
            }
            #[cfg(feature = "wgpu")]
            Device::Wgpu(d) => {
                // Same as Vulkan: keep the quantized blocks in a CPU box + the device; QMatMul's
                // WgpuQuant path uploads them to the GPU on use.
                let storage = dtype.cpu_zeros(elem_count);
                Ok(QStorage::Wgpu(storage, d.clone()))
            }
        }
    }
}

pub enum QStorage {
    Cpu(Box<dyn QuantizedType>),
    Metal(metal::QMetalStorage),
    Cuda(cuda::QCudaStorage),
    // ROCm mirrors the Vulkan path: quantized blocks held in a CPU box + the device; QMatMul
    // dequantizes them to a dense f32 ROCm tensor on demand (rocBLAS matmul). A native HIP quant
    // matmul is the bandwidth-win follow-up.
    #[cfg(feature = "rocm")]
    Rocm(Box<dyn QuantizedType>, crate::RocmDevice),
    // Vulkan keeps the quantized blocks in a CPU box and dequantizes to an f32 Vulkan tensor on
    // demand (QMatMul forces dequantize for Vulkan). Lets GGUF-quantized models run on the GPU;
    // a direct quantized matmul (the Q8 kernel) is the bandwidth-win follow-up.
    #[cfg(feature = "vulkan")]
    Vulkan(Box<dyn QuantizedType>, crate::VulkanDevice),
    // wgpu mirror of the Vulkan path: quantized blocks held in a CPU box + the device. QMatMul's
    // WgpuQuant path reads the GGML bytes straight in the native quant matvec kernel (decode), or
    // dequantizes to an f32 wgpu tensor on demand.
    #[cfg(feature = "wgpu")]
    Wgpu(Box<dyn QuantizedType>, crate::WgpuDevice),
    // Disk-streaming MoE expert bank: the stacked [n_experts, n, k] quantized weight lives on NVMe,
    // not resident. Only `indexed_moe_forward` consumes it -- it fetches one expert's [n, k] slice
    // through the pin/LRU cache per token (host, CPU matmul). The low-memory mode for huge MoEs.
    Stream(Arc<expert_stream::ExpertStreamBank>),
}

impl QStorage {
    pub fn from_data(data: Cow<'_, [u8]>, device: &Device, dtype: GgmlDType) -> Result<Self> {
        match device {
            Device::Cpu => Ok(Self::Cpu(dtype.from_data(data))),
            Device::Metal(d) => match dtype {
                GgmlDType::F32 => metal::load_quantized(d, as_t_slice::<f32>(&data)),
                GgmlDType::F16 => metal::load_quantized(d, as_t_slice::<f16>(&data)),
                GgmlDType::Q4_0 => metal::load_quantized(d, as_t_slice::<BlockQ4_0>(&data)),
                GgmlDType::Q4_1 => metal::load_quantized(d, as_t_slice::<BlockQ4_1>(&data)),
                GgmlDType::Q5_0 => metal::load_quantized(d, as_t_slice::<BlockQ5_0>(&data)),
                GgmlDType::Q5_1 => metal::load_quantized(d, as_t_slice::<BlockQ5_1>(&data)),
                GgmlDType::Q8_0 => metal::load_quantized(d, as_t_slice::<BlockQ8_0>(&data)),
                GgmlDType::Q8_1 => metal::load_quantized(d, as_t_slice::<BlockQ8_1>(&data)),
                GgmlDType::Q2K => metal::load_quantized(d, as_t_slice::<BlockQ2K>(&data)),
                GgmlDType::Q3K => metal::load_quantized(d, as_t_slice::<BlockQ3K>(&data)),
                GgmlDType::Q4K => metal::load_quantized(d, as_t_slice::<BlockQ4K>(&data)),
                GgmlDType::Q5K => metal::load_quantized(d, as_t_slice::<BlockQ5K>(&data)),
                GgmlDType::Q6K => metal::load_quantized(d, as_t_slice::<BlockQ6K>(&data)),
                GgmlDType::Q8K => metal::load_quantized(d, as_t_slice::<BlockQ8K>(&data)),
                GgmlDType::IQ4_NL => metal::load_quantized(d, as_t_slice::<BlockIQ4nl>(&data)),
                GgmlDType::IQ4_XS => metal::load_quantized(d, as_t_slice::<BlockIQ4xs>(&data)),
                GgmlDType::MXFP4 => metal::load_quantized(d, as_t_slice::<BlockMXFP4>(&data)),
                GgmlDType::BF16 => metal::load_quantized(d, as_t_slice::<bf16>(&data)),
                GgmlDType::I32 => metal::load_quantized(d, as_t_slice::<i32>(&data)),
                // i-quant codebook family: the GGML blocks upload to the Metal buffer byte-for-byte
                // like any other quant, and the native MSL matvec/matmul/mul_mv_id kernels
                // (kernel_mul_mv_iq*_f32 etc.) read the block format directly -- no CPU dequant on the
                // hot path. QMetalStorage::dequantize (CPU codebook decode) backs ISQ / dequant-to-f32.
                // Exhaustive on purpose: a future GgmlDType fails-closed at compile rather than at load.
                GgmlDType::IQ2_XXS => metal::load_quantized(d, as_t_slice::<BlockIQ2xxs>(&data)),
                GgmlDType::IQ2_XS => metal::load_quantized(d, as_t_slice::<BlockIQ2xs>(&data)),
                GgmlDType::IQ2_S => metal::load_quantized(d, as_t_slice::<BlockIQ2s>(&data)),
                GgmlDType::IQ3_XXS => metal::load_quantized(d, as_t_slice::<BlockIQ3xxs>(&data)),
                GgmlDType::IQ3_S => metal::load_quantized(d, as_t_slice::<BlockIQ3s>(&data)),
                GgmlDType::IQ1_S => metal::load_quantized(d, as_t_slice::<BlockIQ1s>(&data)),
                GgmlDType::IQ1_M => metal::load_quantized(d, as_t_slice::<BlockIQ1m>(&data)),
                // Ternary / NVFP4 have no native Metal kernel yet (CPU-decode only).
                other => crate::bail!("{other:?} is not supported on the Metal backend"),
            },
            Device::Cuda(d) => match dtype {
                GgmlDType::F32 => cuda::load_quantized(d, as_t_slice::<f32>(&data)),
                GgmlDType::F16 => cuda::load_quantized(d, as_t_slice::<f16>(&data)),
                GgmlDType::Q4_0 => cuda::load_quantized(d, as_t_slice::<BlockQ4_0>(&data)),
                GgmlDType::Q4_1 => cuda::load_quantized(d, as_t_slice::<BlockQ4_1>(&data)),
                GgmlDType::Q5_0 => cuda::load_quantized(d, as_t_slice::<BlockQ5_0>(&data)),
                GgmlDType::Q5_1 => cuda::load_quantized(d, as_t_slice::<BlockQ5_1>(&data)),
                GgmlDType::Q8_0 => cuda::load_quantized(d, as_t_slice::<BlockQ8_0>(&data)),
                GgmlDType::Q8_1 => cuda::load_quantized(d, as_t_slice::<BlockQ8_1>(&data)),
                GgmlDType::Q2K => cuda::load_quantized(d, as_t_slice::<BlockQ2K>(&data)),
                GgmlDType::Q3K => cuda::load_quantized(d, as_t_slice::<BlockQ3K>(&data)),
                GgmlDType::Q4K => cuda::load_quantized(d, as_t_slice::<BlockQ4K>(&data)),
                GgmlDType::Q5K => cuda::load_quantized(d, as_t_slice::<BlockQ5K>(&data)),
                GgmlDType::Q6K => cuda::load_quantized(d, as_t_slice::<BlockQ6K>(&data)),
                GgmlDType::Q8K => cuda::load_quantized(d, as_t_slice::<BlockQ8K>(&data)),
                GgmlDType::IQ4_NL => cuda::load_quantized(d, as_t_slice::<BlockIQ4nl>(&data)),
                GgmlDType::IQ4_XS => cuda::load_quantized(d, as_t_slice::<BlockIQ4xs>(&data)),
                GgmlDType::MXFP4 => cuda::load_quantized(d, as_t_slice::<BlockMXFP4>(&data)),
                GgmlDType::BF16 => cuda::load_quantized(d, as_t_slice::<bf16>(&data)),
                GgmlDType::I32 => cuda::load_quantized(d, as_t_slice::<i32>(&data)),
                // IQ / ternary / 1-bit / NVFP4 codec types: no native on-GPU quant-matmul kernel, but the
                // GGML blocks upload to VRAM byte-for-byte like any other quant. QCudaStorage::fwd then
                // dequantizes them to f32 (CPU codebook decode + upload) for a dense matmul -- see
                // `has_native_q8_1_matmul`. So an i-quant GGUF that used to bail here now LOADS and DECODES
                // on CUDA, exact w.r.t. the CPU reference; a native mmvq kernel is the bandwidth follow-up.
                // Exhaustive on purpose: a future GgmlDType must be wired here (fail-closed at compile) rather
                // than silently bailing at load.
                GgmlDType::IQ2_XXS => cuda::load_quantized(d, as_t_slice::<BlockIQ2xxs>(&data)),
                GgmlDType::IQ2_XS => cuda::load_quantized(d, as_t_slice::<BlockIQ2xs>(&data)),
                GgmlDType::IQ2_S => cuda::load_quantized(d, as_t_slice::<BlockIQ2s>(&data)),
                GgmlDType::IQ3_XXS => cuda::load_quantized(d, as_t_slice::<BlockIQ3xxs>(&data)),
                GgmlDType::IQ3_S => cuda::load_quantized(d, as_t_slice::<BlockIQ3s>(&data)),
                GgmlDType::IQ1_S => cuda::load_quantized(d, as_t_slice::<BlockIQ1s>(&data)),
                GgmlDType::IQ1_M => cuda::load_quantized(d, as_t_slice::<BlockIQ1m>(&data)),
                GgmlDType::TQ1_0 => cuda::load_quantized(d, as_t_slice::<BlockTQ1_0>(&data)),
                GgmlDType::TQ2_0 => cuda::load_quantized(d, as_t_slice::<BlockTQ2_0>(&data)),
                GgmlDType::NVFP4 => cuda::load_quantized(d, as_t_slice::<BlockNVFP4>(&data)),
                GgmlDType::Q1_0 => cuda::load_quantized(d, as_t_slice::<BlockQ1_0>(&data)),
            },
            #[cfg(feature = "rocm")]
            Device::Rocm(d) => Ok(Self::Rocm(dtype.from_data(data), d.clone())),
            #[cfg(feature = "vulkan")]
            Device::Vulkan(d) => Ok(Self::Vulkan(dtype.from_data(data), d.clone())),
            #[cfg(feature = "wgpu")]
            Device::Wgpu(d) => Ok(Self::Wgpu(dtype.from_data(data), d.clone())),
        }
    }

    fn block_size(&self) -> usize {
        match self {
            QStorage::Cpu(storage) => storage.block_size(),
            QStorage::Metal(storage) => storage.dtype().block_size(),
            QStorage::Cuda(storage) => storage.dtype().block_size(),
            #[cfg(feature = "rocm")]
            QStorage::Rocm(storage, _) => storage.block_size(),
            #[cfg(feature = "vulkan")]
            QStorage::Vulkan(storage, _) => storage.block_size(),
            #[cfg(feature = "wgpu")]
            QStorage::Wgpu(storage, _) => storage.block_size(),
            QStorage::Stream(bank) => bank.dtype().block_size(),
        }
    }

    fn dtype(&self) -> GgmlDType {
        match self {
            QStorage::Cpu(storage) => storage.dtype(),
            QStorage::Metal(storage) => storage.dtype(),
            QStorage::Cuda(storage) => storage.dtype(),
            #[cfg(feature = "rocm")]
            QStorage::Rocm(storage, _) => storage.dtype(),
            #[cfg(feature = "vulkan")]
            QStorage::Vulkan(storage, _) => storage.dtype(),
            #[cfg(feature = "wgpu")]
            QStorage::Wgpu(storage, _) => storage.dtype(),
            QStorage::Stream(bank) => bank.dtype(),
        }
    }

    fn device(&self) -> Device {
        match self {
            QStorage::Cpu(_storage) => Device::Cpu,
            QStorage::Metal(storage) => Device::Metal(storage.device().clone()),
            QStorage::Cuda(storage) => Device::Cuda(storage.device().clone()),
            #[cfg(feature = "rocm")]
            QStorage::Rocm(_storage, device) => Device::Rocm(device.clone()),
            #[cfg(feature = "vulkan")]
            QStorage::Vulkan(_storage, device) => Device::Vulkan(device.clone()),
            #[cfg(feature = "wgpu")]
            QStorage::Wgpu(_storage, device) => Device::Wgpu(device.clone()),
            QStorage::Stream(_) => Device::Cpu,
        }
    }

    fn size_in_bytes(&self) -> usize {
        match self {
            QStorage::Cpu(storage) => storage.storage_size_in_bytes(),
            QStorage::Metal(storage) => storage.storage_size_in_bytes(),
            QStorage::Cuda(storage) => storage.storage_size_in_bytes(),
            #[cfg(feature = "rocm")]
            QStorage::Rocm(storage, _) => storage.storage_size_in_bytes(),
            #[cfg(feature = "vulkan")]
            QStorage::Vulkan(storage, _) => storage.storage_size_in_bytes(),
            #[cfg(feature = "wgpu")]
            QStorage::Wgpu(storage, _) => storage.storage_size_in_bytes(),
            QStorage::Stream(bank) => bank.logical_bytes(),
        }
    }

    fn quantize(&mut self, src: &Storage) -> Result<()> {
        match (self, src) {
            (QStorage::Cpu(storage), Storage::Cpu(src)) => {
                storage.from_float(src.as_slice::<f32>()?);
            }
            (QStorage::Metal(storage), Storage::Metal(src)) => storage.quantize(src)?,
            (QStorage::Cuda(storage), Storage::Cuda(src)) => storage.quantize(src)?,
            _ => crate::bail!("Invalid quantize storage locations do not match"),
        }
        Ok(())
    }

    fn quantize_imatrix(
        &mut self,
        src: &Storage,
        imatrix_weights: &[f32],
        n_per_row: usize,
    ) -> Result<()> {
        match (self, src) {
            (QStorage::Cpu(storage), Storage::Cpu(src)) => {
                storage.from_float_imatrix(src.as_slice::<f32>()?, imatrix_weights, n_per_row);
            }
            (QStorage::Metal(storage), Storage::Metal(src)) => {
                storage.quantize_imatrix(src, imatrix_weights, n_per_row)?
            }
            (QStorage::Cuda(storage), Storage::Cuda(src)) => {
                storage.quantize_imatrix(src, imatrix_weights, n_per_row)?
            }
            _ => crate::bail!("Invalid quantize storage locations do not match"),
        }
        Ok(())
    }

    fn quantize_onto(&mut self, src: &Storage) -> Result<()> {
        match (self, src) {
            (QStorage::Cpu(storage), Storage::Cpu(src)) => {
                storage.from_float(src.as_slice::<f32>()?);
            }
            (QStorage::Metal(storage), Storage::Cpu(src)) => storage.quantize_onto(src)?,
            (QStorage::Cuda(storage), Storage::Cpu(src)) => storage.quantize_onto(src)?,
            _ => crate::bail!("Invalid quantize source storage locations: not on cpu"),
        }
        Ok(())
    }

    fn quantize_imatrix_onto(
        &mut self,
        src: &Storage,
        imatrix_weights: &[f32],
        n_per_row: usize,
    ) -> Result<()> {
        match (self, src) {
            (QStorage::Cpu(storage), Storage::Cpu(src)) => {
                storage.from_float_imatrix(src.as_slice::<f32>()?, imatrix_weights, n_per_row);
            }
            (QStorage::Metal(storage), Storage::Cpu(src)) => {
                storage.quantize_imatrix_onto(src, imatrix_weights, n_per_row)?
            }
            (QStorage::Cuda(storage), Storage::Cpu(src)) => {
                storage.quantize_imatrix_onto(src, imatrix_weights, n_per_row)?
            }
            _ => crate::bail!("Invalid quantize storage locations do not match"),
        }
        Ok(())
    }

    fn dequantize(&self, elem_count: usize) -> Result<Storage> {
        match self {
            QStorage::Cpu(storage) => Ok(Storage::Cpu(storage.dequantize(elem_count)?)),
            QStorage::Metal(storage) => Ok(Storage::Metal(storage.dequantize(elem_count)?)),
            QStorage::Cuda(storage) => Ok(Storage::Cuda(storage.dequantize(elem_count)?)),
            #[cfg(feature = "rocm")]
            QStorage::Rocm(storage, device) => {
                // Dequantize on the CPU, then upload the dense f32 weights to the ROCm device.
                use crate::backend::BackendDevice;
                let cpu = storage.dequantize(elem_count)?;
                Ok(Storage::Rocm(device.storage_from_cpu_storage(&cpu)?))
            }
            #[cfg(feature = "vulkan")]
            QStorage::Vulkan(storage, device) => {
                // Dequantize on the CPU, then upload the f32 weights to the GPU.
                let cpu = storage.dequantize(elem_count)?;
                Ok(Storage::Vulkan(device.upload_f32(cpu.as_slice::<f32>()?)?))
            }
            #[cfg(feature = "wgpu")]
            QStorage::Wgpu(storage, device) => {
                // Dequantize on the CPU, then upload the f32 weights to the GPU.
                let cpu = storage.dequantize(elem_count)?;
                Ok(Storage::Wgpu(device.upload_f32(cpu.as_slice::<f32>()?)?))
            }
            QStorage::Stream(_) => {
                crate::bail!("streaming expert bank has no whole-tensor dequantize; consume it via indexed_moe_forward")
            }
        }
    }

    fn data(&self) -> Result<Cow<'_, [u8]>> {
        match self {
            QStorage::Cpu(storage) => {
                let data_ptr = storage.as_ptr();
                let size_in_bytes = storage.storage_size_in_bytes();
                let data = unsafe { std::slice::from_raw_parts(data_ptr, size_in_bytes) };
                Ok(Cow::from(data))
            }
            QStorage::Cuda(storage) => Ok(Cow::from(storage.data()?)),
            QStorage::Metal(storage) => Ok(Cow::from(storage.data()?)),
            #[cfg(feature = "rocm")]
            QStorage::Rocm(storage, _) => {
                let data_ptr = storage.as_ptr();
                let size_in_bytes = storage.storage_size_in_bytes();
                let data = unsafe { std::slice::from_raw_parts(data_ptr, size_in_bytes) };
                Ok(Cow::from(data))
            }
            #[cfg(feature = "vulkan")]
            QStorage::Vulkan(storage, _) => {
                let data_ptr = storage.as_ptr();
                let size_in_bytes = storage.storage_size_in_bytes();
                let data = unsafe { std::slice::from_raw_parts(data_ptr, size_in_bytes) };
                Ok(Cow::from(data))
            }
            #[cfg(feature = "wgpu")]
            QStorage::Wgpu(storage, _) => {
                let data_ptr = storage.as_ptr();
                let size_in_bytes = storage.storage_size_in_bytes();
                let data = unsafe { std::slice::from_raw_parts(data_ptr, size_in_bytes) };
                Ok(Cow::from(data))
            }
            QStorage::Stream(_) => {
                crate::bail!(
                    "streaming expert bank is not resident; consume it via indexed_moe_forward"
                )
            }
        }
    }

    pub fn device_ptr(&self) -> Result<*const u8> {
        match self {
            QStorage::Cuda(storage) => storage.device_ptr(),
            #[cfg(feature = "rocm")]
            QStorage::Rocm(..) => crate::bail!("not implemented"),
            #[cfg(feature = "vulkan")]
            QStorage::Vulkan(..) => crate::bail!("not implemented"),
            #[cfg(feature = "wgpu")]
            QStorage::Wgpu(..) => crate::bail!("not implemented"),
            QStorage::Metal(_) | QStorage::Cpu(_) | QStorage::Stream(_) => {
                crate::bail!("not implemented");
            }
        }
    }

    #[cfg(feature = "cuda")]
    pub fn device_ptr_with_guard<'a>(
        &'a self,
        stream: &'a crate::cuda_backend::cudarc::driver::CudaStream,
    ) -> Result<(
        *const u8,
        crate::cuda_backend::cudarc::driver::SyncOnDrop<'a>,
    )> {
        match self {
            QStorage::Cuda(storage) => storage.device_ptr_with_guard(stream),
            QStorage::Metal(_) | QStorage::Cpu(_) | QStorage::Stream(_) => {
                crate::bail!("not implemented");
            }
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum GgmlDType {
    F32,
    F16,
    BF16,
    /// Raw int32 (GGML type 26) — integer side tables like DeepSeek-V4's
    /// `ffn_gate_tid2eid` hash-routing table. Dequantizes by casting to f32.
    I32,
    Q4_0,
    Q4_1,
    Q5_0,
    Q5_1,
    Q8_0,
    Q8_1,
    Q2K,
    Q3K,
    Q4K,
    Q5K,
    Q6K,
    Q8K,
    #[allow(non_camel_case_types)]
    IQ4_NL,
    #[allow(non_camel_case_types)]
    IQ4_XS,
    // 4-bit microscaling float (MXFP4), ggml type 39: 32 elems / block, E8M0 scale + 16 nibble-pairs.
    MXFP4,
    // IQ / ternary / 1-bit / NVFP4 codec types (decode-only; dequant->matmul on every backend).
    #[allow(non_camel_case_types)]
    IQ2_XXS,
    #[allow(non_camel_case_types)]
    IQ2_XS,
    #[allow(non_camel_case_types)]
    IQ3_XXS,
    #[allow(non_camel_case_types)]
    IQ1_S,
    #[allow(non_camel_case_types)]
    IQ3_S,
    #[allow(non_camel_case_types)]
    IQ2_S,
    #[allow(non_camel_case_types)]
    IQ1_M,
    TQ1_0,
    TQ2_0,
    NVFP4,
    Q1_0,
}

// --- Single-source-of-truth wiring (Cut 2) -----------------------------------------
//
// The five purely 1:1-per-block `GgmlDType` methods (`from_u32`, `to_u32`, `cpu_zeros`,
// `from_data`, `type_size`) are generated from the ONE `for_each_quant!` table in
// `quant_format.rs` instead of six hand-maintained match blocks. Each generator below is
// handed the whole `Variant => Block @ ggml_id` list and emits the block arms; the three
// non-block pseudo-types (F32/F16/BF16) that the table intentionally omits keep their
// bespoke arms inline. Behavior is byte-identical to the previous hand-written matches —
// the table's ids/blocks were verified equal to them. (`block_size` is left hand-written:
// the table carries no block-size data and its grouped form — F32=1, Q1_0/NVFP4 special,
// the big QK_K group — is already minimal.)
use crate::for_each_quant;

macro_rules! gen_from_u32 {
    ($($v:ident => $b:ident @ $id:literal),+ $(,)?) => {
        pub(crate) fn from_u32(u: u32) -> Result<Self> {
            let dtype = match u {
                0 => Self::F32,
                1 => Self::F16,
                30 => Self::BF16,
                26 => Self::I32,
                $( $id => Self::$v, )+
                _ => crate::bail!("unknown dtype for tensor {u}"),
            };
            Ok(dtype)
        }
    };
}

macro_rules! gen_to_u32 {
    ($($v:ident => $b:ident @ $id:literal),+ $(,)?) => {
        /// GGML type id for this dtype. Single source of truth (generated from the
        /// `for_each_quant!` table); cross-crate callers (e.g. hanzo-quant UQFF
        /// serialization) use this instead of hand-rolled id maps that drift.
        pub fn to_u32(self) -> u32 {
            match self {
                Self::F32 => 0,
                Self::F16 => 1,
                Self::BF16 => 30,
                Self::I32 => 26,
                $( Self::$v => $id, )+
            }
        }
    };
}

macro_rules! gen_cpu_zeros {
    ($($v:ident => $b:ident @ $id:literal),+ $(,)?) => {
        /// The block dtype
        pub fn cpu_zeros(&self, elem_count: usize) -> Box<dyn QuantizedType> {
            match self {
                Self::F32 => Box::new(vec![f32::zeros(); elem_count]),
                Self::F16 => Box::new(vec![f16::zeros(); elem_count]),
                Self::BF16 => Box::new(vec![bf16::zeros(); elem_count]),
                Self::I32 => Box::new(vec![0i32; elem_count]),
                $( Self::$v => Box::new(vec![<$b>::zeros(); elem_count / <$b>::BLCK_SIZE]), )+
            }
        }
    };
}

macro_rules! gen_from_data {
    ($($v:ident => $b:ident @ $id:literal),+ $(,)?) => {
        pub fn from_data(&self, data: Cow<'_, [u8]>) -> Box<dyn QuantizedType> {
            match self {
                Self::F32 => Box::new(as_t_slice::<f32>(&data).to_vec()),
                Self::F16 => Box::new(as_t_slice::<f16>(&data).to_vec()),
                Self::BF16 => Box::new(as_t_slice::<bf16>(&data).to_vec()),
                Self::I32 => Box::new(as_t_slice::<i32>(&data).to_vec()),
                $( Self::$v => Box::new(as_t_slice::<$b>(&data).to_vec()), )+
            }
        }
    };
}

macro_rules! gen_type_size {
    ($($v:ident => $b:ident @ $id:literal),+ $(,)?) => {
        /// The type size for blocks in bytes.
        pub fn type_size(&self) -> usize {
            use k_quants::*;
            match self {
                Self::F32 => 4,
                Self::F16 | Self::BF16 => 2,
                Self::I32 => 4,
                $( Self::$v => std::mem::size_of::<$b>(), )+
            }
        }
    };
}

macro_rules! gen_type_align {
    ($($v:ident => $b:ident @ $id:literal),+ $(,)?) => {
        /// Alignment (bytes) required to reinterpret raw GGUF bytes as this dtype's block type.
        /// The mmap loader uses it to decide whether a mapped region can be referenced no-copy
        /// (`from_mmap`) or must fall back to an owned, naturally-aligned copy. Single source of
        /// truth (generated from the `for_each_quant!` table), so it can't drift from `type_size`.
        pub fn type_align(&self) -> usize {
            use k_quants::*;
            match self {
                Self::F32 => std::mem::align_of::<f32>(),
                Self::F16 | Self::BF16 => std::mem::align_of::<f16>(),
                Self::I32 => std::mem::align_of::<i32>(),
                $( Self::$v => std::mem::align_of::<$b>(), )+
            }
        }
    };
}

macro_rules! gen_from_mmap {
    ($($v:ident => $b:ident @ $id:literal),+ $(,)?) => {
        /// No-copy CPU constructor: wrap a tensor's blocks *in place* inside the mmap'd GGUF region
        /// (`QMmap`) instead of copying them into an owned `Vec` (the `to_vec` in `from_data`). The
        /// returned store references the mapped pages directly, so the weight bytes stay on disk and
        /// the OS pages them in/out under memory pressure. Mirrors `from_data` arm-for-arm so the
        /// dtype -> block-type mapping can never disagree between the resident and mmap paths.
        #[allow(clippy::wrong_self_convention)] // dispatches on the block-type value, mirroring from_data
        pub(crate) fn from_mmap(
            &self,
            mmap: Arc<memmap2::Mmap>,
            offset: usize,
            n_blocks: usize,
        ) -> Box<dyn QuantizedType> {
            match self {
                Self::F32 => Box::new(QMmap::<f32>::new(mmap, offset, n_blocks)),
                Self::F16 => Box::new(QMmap::<f16>::new(mmap, offset, n_blocks)),
                Self::BF16 => Box::new(QMmap::<bf16>::new(mmap, offset, n_blocks)),
                Self::I32 => Box::new(QMmap::<i32>::new(mmap, offset, n_blocks)),
                $( Self::$v => Box::new(QMmap::<$b>::new(mmap, offset, n_blocks)), )+
            }
        }
    };
}

impl GgmlDType {
    for_each_quant!(gen_from_u32);
    for_each_quant!(gen_to_u32);
    for_each_quant!(gen_cpu_zeros);
    for_each_quant!(gen_from_data);
    for_each_quant!(gen_from_mmap);
    for_each_quant!(gen_type_size);
    for_each_quant!(gen_type_align);

    /// The block size, i.e. the number of elements stored in each block.
    pub fn block_size(&self) -> usize {
        match self {
            Self::F32 => 1,
            Self::F16 | Self::BF16 => 1,
            Self::I32 => 1,
            Self::Q4_0 => k_quants::QK4_0,
            Self::Q4_1 => k_quants::QK4_1,
            Self::Q5_0 => k_quants::QK5_0,
            Self::Q5_1 => k_quants::QK5_1,
            Self::Q8_0 => k_quants::QK8_0,
            Self::Q8_1 => k_quants::QK8_1,
            Self::IQ4_NL => k_quants::QK4_NL,
            Self::MXFP4 => k_quants::QK_MXFP4,
            Self::Q1_0 => iq_quants::QK1_0,
            Self::NVFP4 => iq_quants::QK_NVFP4,
            Self::Q2K
            | Self::Q3K
            | Self::Q4K
            | Self::Q5K
            | Self::Q6K
            | Self::Q8K
            | Self::IQ4_XS
            | Self::IQ2_XXS
            | Self::IQ2_XS
            | Self::IQ3_XXS
            | Self::IQ1_S
            | Self::IQ3_S
            | Self::IQ2_S
            | Self::IQ1_M
            | Self::TQ1_0
            | Self::TQ2_0 => k_quants::QK_K,
        }
    }
}

// A version of GgmlType without `vec_dot` so that it can be dyn boxed.
pub trait QuantizedType: Send + Sync {
    fn dtype(&self) -> GgmlDType;
    fn matmul_t(&self, mkn: (usize, usize, usize), lhs: &[f32], dst: &mut [f32]) -> Result<()>;
    fn matmul_t_f16(&self, mkn: (usize, usize, usize), lhs: &[f16], dst: &mut [f16]) -> Result<()>;
    fn dequantize(&self, elem_count: usize) -> Result<CpuStorage>;
    fn storage_size_in_bytes(&self) -> usize;
    fn as_ptr(&self) -> *const u8;
    fn block_size(&self) -> usize;
    #[allow(clippy::wrong_self_convention)]
    fn from_float(&mut self, xs: &[f32]);
    #[allow(clippy::wrong_self_convention)]
    fn from_float_imatrix(&mut self, xs: &[f32], imatrix_weights: &[f32], n_per_row: usize);
    fn size(&self) -> usize;
}

impl<T: k_quants::GgmlType + Send + Sync> QuantizedType for Vec<T> {
    fn matmul_t(&self, mkn: (usize, usize, usize), lhs: &[f32], dst: &mut [f32]) -> Result<()> {
        k_quants::matmul(mkn, lhs, self.as_slice(), dst)
    }
    fn matmul_t_f16(&self, mkn: (usize, usize, usize), lhs: &[f16], dst: &mut [f16]) -> Result<()> {
        k_quants::matmul_f16(mkn, lhs, self.as_slice(), dst)
    }

    fn size(&self) -> usize {
        self.len() * core::mem::size_of::<T>()
    }

    fn from_float(&mut self, xs: &[f32]) {
        T::from_float(xs, self)
    }

    fn from_float_imatrix(&mut self, xs: &[f32], imatrix_weights: &[f32], n_per_row: usize) {
        T::from_float_imatrix(xs, self, imatrix_weights, n_per_row)
    }

    fn dtype(&self) -> GgmlDType {
        T::DTYPE
    }

    fn block_size(&self) -> usize {
        T::BLCK_SIZE
    }

    fn dequantize(&self, elem_count: usize) -> Result<CpuStorage> {
        let mut ys = vec![0.0f32; elem_count];
        T::to_float(self.as_slice(), &mut ys);
        Ok(CpuStorage::F32(ys))
    }

    fn storage_size_in_bytes(&self) -> usize {
        self.len() * std::mem::size_of::<T>()
    }

    fn as_ptr(&self) -> *const u8 {
        self.as_ptr() as *const u8
    }
}

/// A CPU quantized store whose blocks live in an mmap'd GGUF region rather than an owned `Vec`.
///
/// Holds an `Arc<memmap2::Mmap>` plus the byte offset and block count of one tensor inside it, and
/// hands the dequant/matmul path a `&[T]` that points *directly* into the mapped file -- so the
/// quantized weight bytes are never copied resident; the OS pages them in on access and reclaims
/// them under memory pressure (page cache = "RAM as a speed spectrum", antirez ds4_ssd). The `Arc`
/// keeps the mapping alive for as long as any tensor references it.
///
/// This is the no-copy twin of `QuantizedType for Vec<T>`: every method forwards the same
/// `k_quants` routine over `self.as_slice()` instead of over a `Vec`. It is read-only -- the two
/// `from_float*` (quantize) entry points are unreachable in the load path and panic if ever called
/// (you cannot quantize *into* a memory-mapped, read-only weight file).
pub struct QMmap<T> {
    mmap: Arc<memmap2::Mmap>,
    /// Byte offset of this tensor's first block within the mapping.
    offset: usize,
    /// Number of `T` blocks.
    n_blocks: usize,
    _t: std::marker::PhantomData<T>,
}

impl<T> QMmap<T> {
    fn new(mmap: Arc<memmap2::Mmap>, offset: usize, n_blocks: usize) -> Self {
        Self {
            mmap,
            offset,
            n_blocks,
            _t: std::marker::PhantomData,
        }
    }

    /// The tensor's blocks as a slice into the live mapping. No copy: the pointer is inside the
    /// mmap'd region, valid for as long as `self` (and thus the `Arc<Mmap>`) is alive. Alignment is
    /// guaranteed by the caller (`TensorInfo::read_mmap` falls back to an owned copy for the rare
    /// misaligned region); `as_t_slice`'s asserts are the in-place safety net.
    #[inline]
    fn as_slice(&self) -> &[T] {
        let len = self.n_blocks * std::mem::size_of::<T>();
        as_t_slice::<T>(&self.mmap[self.offset..self.offset + len])
    }
}

impl<T: k_quants::GgmlType + Send + Sync> QuantizedType for QMmap<T> {
    fn matmul_t(&self, mkn: (usize, usize, usize), lhs: &[f32], dst: &mut [f32]) -> Result<()> {
        k_quants::matmul(mkn, lhs, self.as_slice(), dst)
    }

    fn matmul_t_f16(&self, mkn: (usize, usize, usize), lhs: &[f16], dst: &mut [f16]) -> Result<()> {
        k_quants::matmul_f16(mkn, lhs, self.as_slice(), dst)
    }

    fn size(&self) -> usize {
        self.n_blocks * std::mem::size_of::<T>()
    }

    fn from_float(&mut self, _xs: &[f32]) {
        panic!("QMmap is read-only: cannot quantize into a memory-mapped weight region")
    }

    fn from_float_imatrix(&mut self, _xs: &[f32], _imatrix_weights: &[f32], _n_per_row: usize) {
        panic!("QMmap is read-only: cannot quantize into a memory-mapped weight region")
    }

    fn dtype(&self) -> GgmlDType {
        T::DTYPE
    }

    fn block_size(&self) -> usize {
        T::BLCK_SIZE
    }

    fn dequantize(&self, elem_count: usize) -> Result<CpuStorage> {
        let mut ys = vec![0.0f32; elem_count];
        T::to_float(self.as_slice(), &mut ys);
        Ok(CpuStorage::F32(ys))
    }

    fn storage_size_in_bytes(&self) -> usize {
        self.n_blocks * std::mem::size_of::<T>()
    }

    fn as_ptr(&self) -> *const u8 {
        self.as_slice().as_ptr() as *const u8
    }
}

impl std::fmt::Debug for QTensor {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "QTensor[{:?}; {:?}]", self.shape, self.dtype())
    }
}

fn check_shape(shape: &Shape, block_size: usize) -> Result<()> {
    let dims = shape.dims();
    if dims.is_empty() {
        crate::bail!("scalar tensor cannot be quantized {shape:?}")
    }
    if !dims[dims.len() - 1].is_multiple_of(block_size) {
        crate::bail!(
            "quantized tensor must have their last dim divisible by block size {shape:?} {}",
            block_size
        )
    }
    Ok(())
}

impl QTensor {
    // The ONE way to build a QTensor from its parts: every constructor funnels here so a tensor's
    // resident-bank slots start empty exactly once, alongside the bytes they are derived from.
    fn make(storage: QStorage, shape: Shape) -> Self {
        Self {
            storage,
            shape,
            banks: ResidentBanks::default(),
        }
    }

    pub fn new<S: Into<Shape>>(storage: QStorage, shape: S) -> Result<Self> {
        let shape = shape.into();
        check_shape(&shape, storage.block_size())?;
        Ok(Self::make(storage, shape))
    }

    pub fn quantize(src: &Tensor, dtype: GgmlDType) -> Result<Self> {
        let shape = src.shape();
        let block_size = dtype.block_size();
        check_shape(shape, block_size)?;
        let src = src.to_dtype(crate::DType::F32)?.flatten_all()?;
        let elem_count = shape.elem_count();
        if !elem_count.is_multiple_of(block_size) {
            crate::bail!(
                "tensor size ({shape:?}) is not divisible by block size {}",
                block_size
            )
        }
        let mut storage = src.device().qzeros(elem_count, dtype)?;
        storage.quantize(&src.storage())?;
        Ok(Self::make(storage, shape.clone()))
    }

    pub fn quantize_imatrix(
        src: &Tensor,
        imatrix_weights: &[f32],
        dtype: GgmlDType,
    ) -> Result<Self> {
        // (n_per_row/QK_K-1)*QK_K+(QK_K/32-1)*32+32=n_per_row
        // Size of imatrix == last dim of tensor
        let n_per_row = src.dim(D::Minus1)?;
        if imatrix_weights.len() != n_per_row {
            crate::bail!(
                "imatrix weights must have the same length {} as the last dim of src {}",
                imatrix_weights.len(),
                src.dim(D::Minus1)?
            );
        }

        let shape = src.shape();
        let block_size = dtype.block_size();
        check_shape(shape, block_size)?;
        let src = src.to_dtype(crate::DType::F32)?.flatten_all()?;
        let elem_count = shape.elem_count();
        if !elem_count.is_multiple_of(block_size) {
            crate::bail!(
                "tensor size ({shape:?}) is not divisible by block size {}",
                block_size
            );
        }
        let mut storage = src.device().qzeros(elem_count, dtype)?;
        storage.quantize_imatrix(&src.storage(), imatrix_weights, n_per_row)?;
        Ok(Self::make(storage, shape.clone()))
    }

    /// Quantize `src` (currently on the CPU) to a QTensor on `dev`
    pub fn quantize_imatrix_onto(
        src: &Tensor,
        imatrix_weights: &[f32],
        dtype: GgmlDType,
        dev: &Device,
    ) -> Result<Self> {
        if !src.device().is_cpu() {
            crate::bail!(
                "`quantize_onto` expects a `src` to be on the cpu, got {:?}.",
                src.device()
            )
        }
        // (n_per_row/QK_K-1)*QK_K+(QK_K/32-1)*32+32=n_per_row
        // Size of imatrix == last dim of tensor
        let n_per_row = src.dim(D::Minus1)?;
        if imatrix_weights.len() != n_per_row {
            crate::bail!(
                "imatrix weights must have the same length {} as the last dim of src {}",
                imatrix_weights.len(),
                src.dim(D::Minus1)?
            );
        }
        let shape = src.shape();
        let block_size = dtype.block_size();
        check_shape(shape, block_size)?;
        let src = src.to_dtype(crate::DType::F32)?.flatten_all()?;
        let elem_count = shape.elem_count();
        if !elem_count.is_multiple_of(block_size) {
            crate::bail!(
                "tensor size ({shape:?}) is not divisible by block size {}",
                block_size
            )
        }
        // storage is on the `dev`, src is on `cpu`
        let mut storage = dev.qzeros(elem_count, dtype)?;
        storage.quantize_imatrix_onto(&src.storage(), imatrix_weights, n_per_row)?;
        Ok(Self::make(storage, shape.clone()))
    }

    /// Quantize `src` (currently on the CPU) to a QTensor on `dev`
    pub fn quantize_onto(src: &Tensor, dtype: GgmlDType, dev: &Device) -> Result<Self> {
        if !src.device().is_cpu() {
            crate::bail!(
                "`quantize_onto` expects a `src` to be on the cpu, got {:?}.",
                src.device()
            )
        }
        let shape = src.shape();
        let block_size = dtype.block_size();
        check_shape(shape, block_size)?;
        let src = src.to_dtype(crate::DType::F32)?.flatten_all()?;
        let elem_count = shape.elem_count();
        if !elem_count.is_multiple_of(block_size) {
            crate::bail!(
                "tensor size ({shape:?}) is not divisible by block size {}",
                block_size
            )
        }
        // storage is on the `dev`, src is on `cpu`
        let mut storage = dev.qzeros(elem_count, dtype)?;
        storage.quantize_onto(&src.storage())?;
        Ok(Self::make(storage, shape.clone()))
    }

    pub fn dtype(&self) -> GgmlDType {
        self.storage.dtype()
    }

    pub fn device(&self) -> Device {
        self.storage.device()
    }

    pub fn rank(&self) -> usize {
        self.shape.rank()
    }

    pub fn shape(&self) -> &Shape {
        &self.shape
    }

    pub fn dequantize(&self, device: &Device) -> Result<Tensor> {
        let storage = self.storage.dequantize(self.shape.elem_count())?;
        let none = crate::op::BackpropOp::none();
        crate::tensor::from_storage(storage, self.shape.clone(), none, false).to_device(device)
    }

    pub fn dequantize_f16(&self, device: &Device) -> Result<Tensor> {
        // In the CUDA case, we have a specialized kernel as this can be useful for volta
        // architectures. https://github.com/hanzoai/ml/issues/2136
        match &self.storage {
            QStorage::Cuda(s) => {
                let s = s.dequantize_f16(self.shape.elem_count())?;
                let none = crate::op::BackpropOp::none();
                crate::tensor::from_storage(Storage::Cuda(s), self.shape.clone(), none, false)
                    .to_device(device)
            }
            _ => {
                let s = self.dequantize(device)?.to_dtype(crate::DType::F16)?;
                Ok(s)
            }
        }
    }

    pub fn storage_size_in_bytes(&self) -> usize {
        self.storage.size_in_bytes()
    }

    pub fn data(&self) -> Result<Cow<'_, [u8]>> {
        self.storage.data()
    }

    // Upload this expert bank's GGML bytes to VRAM ONCE and keep it resident, keyed by the stable
    // (ptr,len) of the QTensor's CPU bytes. Re-uploading the multi-GB bank per token/layer would
    // dominate decode, so the cache makes MoE bandwidth-bound on the quant matvec, not the H2D copy.
    #[cfg(feature = "rocm")]
    fn rocm_moe_bank(&self, dev: &crate::RocmDevice) -> Result<std::sync::Arc<crate::RocmStorage>> {
        use crate::backend::BackendDevice;
        let bank = self.data()?;
        cache_or_upload(&self.banks.rocm, bank.as_ref(), |b| {
            dev.storage_from_slice(b)
        })
    }

    // Resident Vulkan MoE expert bank (twin of `rocm_moe_bank`). Q8_0 repacks to the 9-u32/block
    // layout and Q6_K to the padded 53-u32 super-block layout their MoE shaders read (mirrors the 2D
    // decode paths); Q4_0/Q4_K shaders byte-address the raw GGML bytes, so a plain upload suffices.
    #[cfg(feature = "vulkan")]
    fn vulkan_moe_bank(
        &self,
        dev: &crate::VulkanDevice,
        e_cnt: usize,
        n: usize,
        k: usize,
    ) -> Result<std::sync::Arc<crate::VulkanStorage>> {
        let bank = self.data()?;
        let dt = self.storage.dtype();
        cache_or_upload(&self.banks.vulkan, bank.as_ref(), |b| match dt {
            GgmlDType::Q8_0 => dev.quantize_q8_blocks(b, e_cnt * n, k),
            GgmlDType::Q6K => dev.quantize_q6k(b, e_cnt * n, k),
            _ => dev.upload_qweight(b),
        })
    }

    // Resident PLANAR expert bank for the DSL block-reduced MoE kernels: the packed GGML bank is
    // de-interleaved into per-field device arrays once (at first sight of its CPU bytes) and reused
    // every token. Only Q4_K/Q6_K -- the dtypes with a committed `moe_matvec_q*k_blk_*` .spv.
    #[cfg(feature = "vulkan")]
    fn vulkan_moe_bank_split(
        &self,
        dev: &crate::VulkanDevice,
        e_cnt: usize,
        n: usize,
        k: usize,
    ) -> Result<std::sync::Arc<crate::vulkan::MoeBankSplit>> {
        let bank = self.data()?;
        let dt = self.storage.dtype();
        cache_or_upload(&self.banks.vulkan_split, bank.as_ref(), |b| match dt {
            GgmlDType::Q4K => dev.quantize_q4k_split(b, e_cnt * n, k),
            GgmlDType::Q6K => dev.quantize_q6k_split(b, e_cnt * n, k),
            _ => crate::bail!("vulkan_moe_bank_split: unsupported dtype {dt:?}"),
        })
    }

    // Resident wgpu MoE expert bank (twin of `vulkan_moe_bank`). The wgpu MoE shaders byte-address
    // the raw GGML bytes for every native type, so one upload path covers Q4_0/Q8_0/Q4K.
    #[cfg(feature = "wgpu")]
    fn wgpu_moe_bank(&self, dev: &crate::WgpuDevice) -> Result<std::sync::Arc<crate::WgpuStorage>> {
        let bank = self.data()?;
        cache_or_upload(&self.banks.wgpu, bank.as_ref(), |b| dev.upload_qweight(b))
    }

    pub fn indexed_moe_forward(&self, x: &Tensor, ids: &Tensor) -> Result<Tensor> {
        // The fused CUDA path reads ids as a flat row-major [batch*topk] buffer (rank-agnostic since
        // 0.11.17); force dense strides so a non-contiguous router output can't misindex the kernel.
        let ids = &ids.contiguous()?;
        match &self.storage {
            // Only dtypes with a fused CUDA indexed-MoE kernel take the fast path; others (e.g. MXFP4,
            // i-quant/ternary) fall through to the generic per-expert path below, which dequantizes via
            // QMatMul. The supported set is derived from the ONE kernel-name table (cuda.rs), so this gate
            // can't drift from the kernels that actually exist -- which is exactly what had stranded
            // Q4_0/Q4_1/Q5_0/Q5_1 on the CPU even though their fused kernels are now compiled.
            QStorage::Cuda(s) if cuda::QCudaStorage::supports_indexed_moe(s.dtype()) => {
                // The fused q8_1 MoE kernel reads the activation as f32 (as_cuda_slice::<f32>), so a
                // BF16/F16 compute dtype must ride f32 into the kernel and restore on the way out --
                // the same reconciliation the i-quant and Vulkan sibling branches already perform. A
                // f32 model no-ops both casts.
                let out_dtype = x.dtype();
                let x = x.to_dtype(crate::DType::F32)?.contiguous()?;
                // Bind the storage guards to named locals (declared after `x`, dropped before it)
                // so the borrow can't outlive the reconciled activation -- mirrors the Metal branch.
                let (x_guard, x_l) = x.storage_and_layout();
                let (ids_guard, ids_l) = ids.storage_and_layout();
                match (&*x_guard, &*ids_guard) {
                    (Storage::Cuda(x_storage), Storage::Cuda(ids_storage)) => {
                        let (storage, out_shape) = s.indexed_moe_forward(
                            self.shape(),
                            x_storage,
                            x_l,
                            ids_storage,
                            ids_l,
                        )?;
                        crate::tensor::from_storage(
                            Storage::Cuda(storage),
                            out_shape,
                            crate::op::BackpropOp::none(),
                            false,
                        )
                        .to_dtype(out_dtype)
                    }
                    _ => {
                        panic!("Non-cuda indexed_moe_forward is not implemented!");
                    }
                }
            }
            // Native CUDA i-quant MoE: the i-quant codebook types have no Blue-C fused q8_1 MoE kernel,
            // but a native dp4a MoE-decode kernel (moe_qmatvec_dp4a_<iq*>). The [E,n,k] bank stays
            // RESIDENT in VRAM and the router gather runs on-device -- the dp4a twin of the ROCm path.
            // This intercepts i-quant MoE BEFORE the generic fallback below, which would DtoH the whole
            // expert bank to host (self.data()) and re-upload every selected expert PER TOKEN.
            #[cfg(feature = "cuda")]
            QStorage::Cuda(s) if cuda::QCudaStorage::supports_iquant_moe(s.dtype()) => {
                let out_dtype = x.dtype();
                let (_e_cnt, n, k) = self.shape().dims3()?;
                let (t, topk) = ids.dims2()?;
                let nrows = t * topk;

                // PREFILL (t>1): expert-grouped int8-WMMA MMQ (qmmq) -- stage each expert's weight ONCE
                // and amortize it over all its routed tokens via the tensor cores (llama mul_mat_id),
                // instead of the per-slot dp4a re-streaming the weight per token. Uses the RAW [t,in1,k]
                // input (indexed_moe_grouped broadcasts/gathers internally). IQ1_M (no MMQ kernel) returns
                // None -> the per-slot dp4a below.
                if t > 1 {
                    let x_f32 = x.to_dtype(crate::DType::F32)?.contiguous()?;
                    let ids_u32 = ids.to_dtype(crate::DType::U32)?.contiguous()?;
                    let (xs, _) = x_f32.storage_and_layout();
                    let xc = match &*xs {
                        Storage::Cuda(c) => c,
                        _ => crate::bail!("cuda i-quant MoE: x not on cuda after contiguous()"),
                    };
                    let (ids_s, _) = ids_u32.storage_and_layout();
                    let idc = match &*ids_s {
                        Storage::Cuda(c) => c,
                        _ => crate::bail!("cuda i-quant MoE: ids not on cuda"),
                    };
                    if let Some((st, sh)) = s.moe_iquant_qmmq(
                        self.shape(),
                        xc.as_cuda_slice::<f32>()?,
                        x.shape(),
                        &idc.as_cuda_slice::<u32>()?.slice(0..),
                        ids.shape(),
                    )? {
                        return crate::tensor::from_storage(
                            Storage::Cuda(st),
                            sh,
                            crate::op::BackpropOp::none(),
                            false,
                        )
                        .to_dtype(out_dtype);
                    }
                }

                // DECODE (t==1) or qmmq-unsupported (IQ1_M): per-slot dp4a. Broadcast the shared gate/up
                // input across topk to a per-slot [nrows,k] activation. f32-native -> the matvec returns
                // F32 and the to_dtype below is a no-op for f32 models.
                let sdim = x.dim(1)?; // 1 (gate/up: shared input) or topk (down: per-slot)
                let x_exp = if sdim == topk {
                    x.clone()
                } else {
                    x.broadcast_as((t, topk, k))?
                };
                let x_flat = x_exp
                    .reshape((nrows, k))?
                    .to_dtype(crate::DType::F32)?
                    .contiguous()?;
                let ids_flat = ids
                    .reshape((nrows,))?
                    .to_dtype(crate::DType::U32)?
                    .contiguous()?;
                let (xstore, _) = x_flat.storage_and_layout();
                let xc = match &*xstore {
                    Storage::Cuda(c) => c,
                    _ => crate::bail!("cuda i-quant MoE: x not on cuda after contiguous()"),
                };
                let (idstore, _) = ids_flat.storage_and_layout();
                let idc = match &*idstore {
                    Storage::Cuda(c) => c,
                    _ => crate::bail!("cuda i-quant MoE: ids not on cuda"),
                };
                let y = s.moe_iquant_dp4a(
                    &xc.as_cuda_slice::<f32>()?.slice(0..),
                    &idc.as_cuda_slice::<u32>()?.slice(0..),
                    nrows,
                    n,
                    k,
                )?;
                let out = crate::tensor::from_storage(
                    Storage::Cuda(y),
                    (nrows, n),
                    crate::op::BackpropOp::none(),
                    false,
                );
                out.reshape((t, topk, n))?.to_dtype(out_dtype)
            }
            // Native Vulkan MoE: one fused grouped quant matvec dispatch reads the per-expert slice
            // out of the GGML weight bank [E, n, k] resident in VRAM and gathers by the router ids --
            // the whole expert compute runs on the GPU (no CPU expert loop; the CPU fallback below
            // would also hit the unimplemented Vulkan index_add). Supported for Q4_0/Q8_0/Q4_K; other
            // quant dtypes fall through to the (CPU-bound) generic path.
            #[cfg(feature = "vulkan")]
            QStorage::Vulkan(_, vk_dev) if vk_moe_kernel(self.storage.dtype()).is_some() => {
                let out_dtype = x.dtype();
                let (e_cnt, n, k) = self.shape().dims3()?;
                let (t, topk) = ids.dims2()?;
                let s = x.dim(1)?; // 1 (gate/up: shared input) or topk (down: per-slot)
                let x_exp = if s == topk {
                    x.clone()
                } else {
                    x.broadcast_as((t, topk, k))?
                };
                // [S, k] contiguous f32 on the Vulkan device; S = t*topk routed slots.
                let nrows = t * topk;
                let x_flat = x_exp
                    .reshape((nrows, k))?
                    .to_dtype(crate::DType::F32)?
                    .contiguous()?;
                // Keep routing ids on the GPU. They came from `moe_route`'s on-device sort, so reading
                // them to host (`to_vec1`) and re-uploading forced a GPU->CPU->GPU sync on EVERY expert
                // call -- 3x/layer (gate, up, down), ~144 fence stalls/token that serialized the whole
                // forward (the GPU sat idle-waiting). Bind the resident U32 buffer directly; the router
                // selects topk over exactly `e_cnt` logits, so ids are in-range by construction (the old
                // host-side OOB scan was the round-trip's only excuse).
                let dt = self.storage.dtype();
                let ids_u32 = ids
                    .reshape((nrows,))?
                    .to_dtype(crate::DType::U32)?
                    .contiguous()?;
                let y = {
                    let (store, _) = x_flat.storage_and_layout();
                    let xv = match &*store {
                        Storage::Vulkan(v) => v,
                        _ => crate::bail!("vulkan MoE: x not on vulkan after contiguous()"),
                    };
                    let (ids_store, _) = ids_u32.storage_and_layout();
                    let ids_v = match &*ids_store {
                        Storage::Vulkan(v) => v,
                        _ => crate::bail!("vulkan MoE: ids not on vulkan after contiguous()"),
                    };
                    // PREFILL (t > 1) routes to the expert-grouped MMQ on the matrix cores: each
                    // expert's weight is streamed ONCE and amortised over all of its routed tokens,
                    // where the per-slot matvec re-streams it per token and leaves the matrix cores
                    // idle. Twin of the CUDA/ROCm/Metal `t > 1` arms above. DECODE (t == 1) stays on
                    // the matvec -- one slot per expert means there is no reuse to win, and that path
                    // already runs at the memory wall. Q4_K only: the Q6_K `_dn` half of a Q4_K_M MoE
                    // has no MMQ kernel yet and keeps the matvec, which is correct at any token count.
                    // VK_MOE_PREFILL_GEMM_OFF forces the matvec for the A/B.
                    if t > 1
                        && dt == GgmlDType::Q4K
                        && vk_dev.has_int_dot8()
                        && std::env::var_os("VK_MOE_PREFILL_GEMM_OFF").is_none()
                    {
                        let bank = self.vulkan_moe_bank_split(vk_dev, e_cnt, n, k)?;
                        // Gate and up share this routed activation; quantizing here keeps the q8
                        // conversion identical to the decode matvec's, so the two paths differ only
                        // in how the weight is streamed.
                        let (xq, xsq, xsum) = vk_dev.quantize_act_q8(xv, nrows, k)?;
                        // cap = t: a token's top-k experts are distinct, so an expert claims at most
                        // one of that token's slots.
                        vk_dev.mmq_q4k_id_gpu(
                            bank.as_ref(),
                            &xq,
                            &xsq,
                            &xsum,
                            ids_v,
                            nrows,
                            e_cnt,
                            t,
                            n,
                            k,
                        )?
                    } else {
                        // Prefer the dp4a (int8 OpSDot) block kernel when the device supports integer
                        // dot-product and a specialized .spv exists (~1.4-1.6x the f32-decode block kernel,
                        // within activation-quant tolerance). Else the f32 DSL block kernel (planar bank,
                        // one workgroup/output, ~2-3x packed). Else the packed `vk_moe_kernel`. Same split
                        // bank feeds both block paths; the dp4a path q8-quantizes the activation once.
                        match vk_moe_blk_dp4a_kernel(dt, n, k).filter(|_| vk_dev.has_int_dot8()) {
                            Some((blk, with_xsum)) => {
                                let bank = self.vulkan_moe_bank_split(vk_dev, e_cnt, n, k)?;
                                vk_dev.moe_matvec_blk_dp4a_gpu(
                                    blk,
                                    with_xsum,
                                    bank.as_ref(),
                                    xv,
                                    ids_v,
                                    nrows,
                                    n,
                                    k,
                                )?
                            }
                            None => match vk_moe_blk_kernel(dt, n, k) {
                                Some(blk) => {
                                    let bank = self.vulkan_moe_bank_split(vk_dev, e_cnt, n, k)?;
                                    vk_dev.moe_matvec_blk_gpu(
                                        blk,
                                        bank.as_ref(),
                                        xv,
                                        ids_v,
                                        nrows,
                                        n,
                                        k,
                                    )?
                                }
                                None => {
                                    // Guarded by `vk_moe_kernel(..).is_some()`, so the packed kernel is present.
                                    let kernel = vk_moe_kernel(dt).unwrap();
                                    let wbank = self.vulkan_moe_bank(vk_dev, e_cnt, n, k)?;
                                    vk_dev.moe_matvec_gpu(
                                        kernel,
                                        wbank.as_ref(),
                                        xv,
                                        ids_v,
                                        nrows,
                                        n,
                                        k,
                                    )?
                                }
                            },
                        }
                    }
                };
                let out = crate::tensor::from_storage(
                    Storage::Vulkan(y),
                    (nrows, n),
                    crate::op::BackpropOp::none(),
                    false,
                );
                out.reshape((t, topk, n))?.to_dtype(out_dtype)
            }
            // Native wgpu MoE: mirror of the Vulkan fused grouped quant matvec dispatch. The GGML
            // weight bank [E, n, k] is uploaded once and the router gather + per-expert GEMM run in
            // one WGSL dispatch. Supported for Q4_0/Q8_0/Q4_K.
            #[cfg(feature = "wgpu")]
            QStorage::Wgpu(_, wgpu_dev) if wgpu_moe_kernel(self.storage.dtype()).is_some() => {
                let out_dtype = x.dtype();
                let (e_cnt, n, k) = self.shape().dims3()?;
                let (t, topk) = ids.dims2()?;
                let s = x.dim(1)?; // 1 (gate/up: shared input) or topk (down: per-slot)
                let x_exp = if s == topk {
                    x.clone()
                } else {
                    x.broadcast_as((t, topk, k))?
                };
                let nrows = t * topk;
                let x_flat = x_exp
                    .reshape((nrows, k))?
                    .to_dtype(crate::DType::F32)?
                    .contiguous()?;
                let ids_vec = ids
                    .reshape((nrows,))?
                    .to_dtype(crate::DType::U32)?
                    .to_vec1::<u32>()?;
                if let Some(&bad) = ids_vec.iter().find(|&&e| e as usize >= e_cnt) {
                    crate::bail!("indexed_moe_forward: expert id {bad} >= num_experts {e_cnt}");
                }
                // Guarded by `wgpu_moe_kernel(..).is_some()` above, so the kernel is always present.
                let kernel = wgpu_moe_kernel(self.storage.dtype()).unwrap();
                // Resident bank: uploaded once, reused every token (see `wgpu_moe_bank`).
                let wbank = self.wgpu_moe_bank(wgpu_dev)?;
                let ids_buf = wgpu_dev.upload_ids(&ids_vec)?;
                let y = {
                    let (store, _) = x_flat.storage_and_layout();
                    let xv = match &*store {
                        Storage::Wgpu(v) => v,
                        _ => crate::bail!("wgpu MoE: x not on wgpu after contiguous()"),
                    };
                    wgpu_dev.moe_matvec_gpu(kernel, wbank.as_ref(), xv, &ids_buf, nrows, n, k)?
                };
                let out = crate::tensor::from_storage(
                    Storage::Wgpu(y),
                    (nrows, n),
                    crate::op::BackpropOp::none(),
                    false,
                );
                out.reshape((t, topk, n))?.to_dtype(out_dtype)
            }
            // Native ROCm MoE: the GGML expert bank [E,n,k] is uploaded once and each routed slot
            // is dispatched through the SAME unified qmatvec_core<WTYPE> as ordinary decode (no
            // MoE-per-quant kernel; works for every wired quant). Avoids ROCm's missing index_add:
            // each routed slot writes exactly one output row, placed directly by slot index.
            #[cfg(feature = "rocm")]
            QStorage::Rocm(_, rocm_dev)
                if crate::RocmQuantType::from_ggml(self.storage.dtype()).is_some() =>
            {
                let qt = crate::RocmQuantType::from_ggml(self.storage.dtype()).unwrap();
                let out_dtype = x.dtype();
                // e_cnt is not read: ids stay on-device and the router guarantees the bound, so the
                // old host bounds check (and its DtoH `to_vec1`) is gone.
                let (_e_cnt, n, k) = self.shape().dims3()?;
                let (t, topk) = ids.dims2()?;
                let s = x.dim(1)?; // 1 (gate/up: shared input) or topk (down: per-slot)
                let x_exp = if s == topk {
                    x.clone()
                } else {
                    x.broadcast_as((t, topk, k))?
                };
                let nrows = t * topk;
                // PREFILL (t>1) routes to the fused expert-grouped WMMA GEMM (f16 activations); DECODE
                // (t==1) keeps the model's native bf16/f16 on the capture-clean matvec. See the twin in
                // QStorage::indexed_moe_forward.
                // Decode-only types (no qmmq kernel) ride the per-slot matvec core for prefill too
                // (correct at any token count). ONE predicate gates every prefill site.
                let use_qmmq = t > 1 && qt.qmmq_capable();
                let x_flat = match x_exp.dtype() {
                    // qmmq quantizes f16/f32 activations natively, so keep the model's dtype and skip
                    // the f32->f16 cast (a 16.7M-elem read+write per gate/up). Other dtypes (bf16 with
                    // a symmetric expert type) still cast to f16.
                    DType::F16 | DType::F32 if use_qmmq => {
                        x_exp.reshape((nrows, k))?.contiguous()?
                    }
                    _ if use_qmmq => x_exp
                        .reshape((nrows, k))?
                        .to_dtype(DType::F16)?
                        .contiguous()?,
                    DType::BF16 | DType::F16 => x_exp.reshape((nrows, k))?.contiguous()?,
                    // DECODE f32-native: dp4a experts quantize q8_1 from f32 and store f32, so an F32
                    // routed activation stays F32 (the matvec returns F32 -> the .to_dtype(out_dtype)
                    // below is a no-op), removing the cast pair that wrapped each gate/up/down matvec.
                    DType::F32 if qt.dp4a_active() => x_exp.reshape((nrows, k))?.contiguous()?,
                    _ => x_exp
                        .reshape((nrows, k))?
                        .to_dtype(DType::F16)?
                        .contiguous()?,
                };
                let wbank = self.rocm_moe_bank(rocm_dev)?;
                // Keep router ids ON the GPU for EVERY wired quant type and run ONE batched launch
                // (experts on grid.y, ids read on-device). No `to_vec1` DtoH sync -- that host round-
                // trip (3 per layer x 48 layers per token) was both the dominant WSL decode stall AND
                // what made HIP stream capture illegal (hipErrorStreamCaptureImplicit -> the graph-path
                // SIGSEGV). The router emits a top-k over the e_cnt expert logits, so 0 <= id < e_cnt
                // by construction; the prior host bounds check is dropped to stay capture-clean.
                let ids_u32 = ids
                    .reshape((nrows,))?
                    .to_dtype(crate::DType::U32)?
                    .contiguous()?;
                let (store, _) = x_flat.storage_and_layout();
                let xr = match &*store {
                    crate::Storage::Rocm(r) => r,
                    _ => crate::bail!("rocm MoE: x not on rocm after contiguous()"),
                };
                let (idstore, _) = ids_u32.storage_and_layout();
                let idr = match &*idstore {
                    crate::Storage::Rocm(r) => r,
                    _ => crate::bail!("rocm MoE: ids not on rocm"),
                };
                let y = if use_qmmq {
                    rocm_dev.moe_qmmq_quant(qt, wbank.as_ref(), xr, idr, nrows, n, k)?
                } else {
                    rocm_dev.moe_matvec_quant(qt, wbank.as_ref(), xr, idr, nrows, n, k)?
                };
                let out = crate::tensor::from_storage(
                    crate::Storage::Rocm(y),
                    (nrows, n),
                    crate::op::BackpropOp::none(),
                    false,
                );
                out.reshape((t, topk, n))?.to_dtype(out_dtype)
            }
            // Native Metal MoE: `QMetalStorage::indexed_moe_forward` runs the whole expert compute in
            // ONE fused dispatch straight out of the resident quantized bank -- decode (t == 1) via
            // `mul_mv_id` (per-slot matvec), prefill (t > 1) via `mul_mm_id` (expert-grouped GEMM, each
            // expert's weight read once and amortized over its tokens). Expert id read per row
            // on-device, no per-expert host loop, no `ids` DtoH sync. Same low buffer churn on both
            // paths (x_f32 + ids_u32 + dst), unlike the generic per-expert fallback whose per-expert
            // `from_data`+forward churned the Metal pool. Guarded to Metal-resident x/ids.
            #[cfg(feature = "metal")]
            QStorage::Metal(s)
                if matches!(&*x.storage(), Storage::Metal(_))
                    && matches!(&*ids.storage(), Storage::Metal(_)) =>
            {
                let out_dtype = x.dtype();
                let x = x.contiguous()?;
                let (xs_guard, x_l) = x.storage_and_layout();
                let (ids_guard, ids_l) = ids.storage_and_layout();
                let (Storage::Metal(x_storage), Storage::Metal(ids_storage)) =
                    (&*xs_guard, &*ids_guard)
                else {
                    unreachable!("metal MoE arm is guarded on Metal x/ids storage");
                };
                let (storage, out_shape) =
                    s.indexed_moe_forward(self.shape(), x_storage, x_l, ids_storage, ids_l)?;
                let out = crate::tensor::from_storage(
                    Storage::Metal(storage),
                    out_shape,
                    crate::op::BackpropOp::none(),
                    false,
                );
                out.to_dtype(out_dtype)
            }
            // Disk-streaming bank: same per-expert quantized matmul as the resident fallback below,
            // but each selected expert's [n, k] slice is fetched through the pin/LRU cache (disk)
            // instead of sliced out of a resident blob. Bit-identical -- only the fetch path differs.
            QStorage::Stream(bank) => {
                let (e_cnt, n, k) = self.shape().dims3()?;
                let dtype = bank.dtype();
                moe_grouped_per_expert(x, ids, n, k, |eid, device| {
                    if eid as usize >= e_cnt {
                        crate::bail!("indexed_moe_forward: expert id {eid} >= num_experts {e_cnt}");
                    }
                    let bytes = bank.fetch(eid)?;
                    QStorage::from_data(std::borrow::Cow::Borrowed(&bytes), device, dtype)
                })
            }
            _ => {
                // CPU / non-CUDA fallback: per-expert quantized matmul. The packed expert bank
                // [E, n, k] is sliced into equal, contiguous per-expert quantized blocks; for
                // each expert that is actually selected we run hanzo-ml's native quantized matmul
                // on just the tokens routed to it. Nothing is dequantized, so quantized MoE runs
                // on any backend (CPU, Metal, ...) at a cost proportional to the active experts.
                let (e_cnt, n, k) = self.shape().dims3()?;
                let dtype = self.storage.dtype();
                let all_bytes = self.data()?;
                let expert_bytes = all_bytes.len() / e_cnt;
                moe_grouped_per_expert(x, ids, n, k, |eid, device| {
                    let off = eid as usize * expert_bytes;
                    QStorage::from_data(
                        std::borrow::Cow::Borrowed(&all_bytes[off..off + expert_bytes]),
                        device,
                        dtype,
                    )
                })
            }
        }
    }

    pub fn device_ptr(&self) -> Result<*const u8> {
        match &self.storage {
            QStorage::Cuda(storage) => storage.device_ptr(),
            #[cfg(feature = "rocm")]
            QStorage::Rocm(..) => crate::bail!("not implemented"),
            #[cfg(feature = "vulkan")]
            QStorage::Vulkan(..) => crate::bail!("not implemented"),
            #[cfg(feature = "wgpu")]
            QStorage::Wgpu(..) => crate::bail!("not implemented"),
            QStorage::Metal(_) | QStorage::Cpu(_) | QStorage::Stream(_) => {
                crate::bail!("not implemented");
            }
        }
    }

    #[cfg(feature = "cuda")]
    pub fn device_ptr_with_guard<'a>(
        &'a self,
        stream: &'a crate::cuda_backend::cudarc::driver::CudaStream,
    ) -> Result<(
        *const u8,
        crate::cuda_backend::cudarc::driver::SyncOnDrop<'a>,
    )> {
        self.storage.device_ptr_with_guard(stream)
    }
}

#[derive(Clone, Debug)]
pub enum QMatMul {
    QTensor(std::sync::Arc<QTensor>),
    Tensor(Tensor),
    TensorF16(Tensor),
    // Native Vulkan quantized weight: the GGML quantized blocks live in VRAM (Q4_0/Q4_K ~0.5 B/elem,
    // Q8_0 ~1.06 B/elem). Decode (1 row) runs the matching on-GPU quant matvec kernel directly out of
    // the block format (no CPU dequant, no re-pack) -- the bandwidth lever for memory-bound decode.
    // Prefill (>1 row) dequantizes the original `qtensor` to a temporary f32 weight. `dtype` selects
    // the kernel; `n`/`k` are the weight dims.
    #[cfg(feature = "vulkan")]
    VulkanQuant {
        qtensor: std::sync::Arc<QTensor>,
        wq: std::sync::Arc<crate::VulkanStorage>,
        dtype: GgmlDType,
        n: usize,
        k: usize,
    },
    // wgpu mirror of VulkanQuant: GGML quantized blocks live in VRAM and decode (1 row) runs the
    // matching native-GGML quant matvec WGSL kernel straight out of the block format. Prefill (>1
    // row) dequantizes to a temporary f32 weight. `dtype` selects the kernel; `n`/`k` are the dims.
    #[cfg(feature = "wgpu")]
    WgpuQuant {
        qtensor: std::sync::Arc<QTensor>,
        wq: std::sync::Arc<crate::WgpuStorage>,
        dtype: GgmlDType,
        n: usize,
        k: usize,
    },
    // Native ROCm quantized weight: the GGML blocks live in VRAM. Decode (1 row) runs the ONE
    // unified on-GPU quant matvec (qmatvec_core<WTYPE>) straight out of the block format; prefill
    // (>1 row) runs the ONE unified int8 WMMA GEMM (qmmq_core<WTYPE>) -- both for the full wired
    // spread (Q8_0/Q4_0/Q4_K/Q6_K/IQ4_XS/TQ2_0). Unwired types dequantize to a temporary f16
    // weight (RDNA matrix-core matmul). `n`/`k` are the weight dims.
    #[cfg(feature = "rocm")]
    RocmQuant {
        qtensor: std::sync::Arc<QTensor>,
        wq: std::sync::Arc<crate::RocmStorage>,
        dtype: GgmlDType,
        n: usize,
        k: usize,
    },
}

// Upload `bank` through `upload` into this tensor's own resident `slot` on first call, then hand back
// that copy on every later call. The whole point: the H2D copy happens once, not per token.
//
// The slot lives in the owning QTensor, so a resident bank is reachable only from the bytes it was
// built from and is freed with them. A global map keyed by the CPU buffer's (ptr,len) cannot express
// that: the key stays live after the tensor drops, and the allocator reuses the address for the next
// same-size bank, so a later tensor reads an earlier one's weights. That aliased silently -- Q4_0 and
// Q4_K are both 0.5625 bytes/weight, so their banks are byte-identical in LENGTH for a given shape,
// and dtype was not part of the key.
#[cfg(any(feature = "rocm", feature = "vulkan", feature = "wgpu"))]
fn cache_or_upload<S>(
    slot: &std::sync::OnceLock<std::sync::Arc<S>>,
    bank: &[u8],
    upload: impl FnOnce(&[u8]) -> Result<S>,
) -> Result<std::sync::Arc<S>> {
    if let Some(w) = slot.get() {
        return Ok(w.clone());
    }
    let w = std::sync::Arc::new(upload(bank)?);
    // A racing caller may have filled the slot meanwhile; whoever lands first wins and both return
    // the same resident bank (the loser's upload drops).
    Ok(slot.get_or_init(|| w).clone())
}

// GGML dtype -> native fused grouped quant-matvec MoE kernel name on the Vulkan backend. ONE source
// of truth: the construction gate keeps the [E,n,k] bank quantized iff this returns Some, the
// dispatch guard fires on Some, and the branch selects the returned kernel. `None` dtypes fall
// through to dequantize. Q4_0/Q8_0/Q4K read the raw (or Q8_0-repacked) GGML bytes; Q6_K reads the
// padded 53-u32 super-block layout from `quantize_q6k`.
#[cfg(feature = "vulkan")]
fn vk_moe_kernel(dt: GgmlDType) -> Option<&'static str> {
    match dt {
        GgmlDType::Q4_0 => Some("moe_matvec_q4_0"),
        GgmlDType::Q8_0 => Some("moe_matvec_q8_0"),
        GgmlDType::Q4K => Some("moe_matvec_q4k"),
        GgmlDType::Q6K => Some("moe_matvec_q6k"),
        _ => None,
    }
}

// (dtype, n, k) -> committed DSL block-reduced MoE .spv. `n,k` are baked into each artifact at dump
// time (comptime -> full unroll + magic-multiply divide), so selection is per shape: gate/up (k=2048,
// nt=64) and down (k=768, nt=32) for Qwen3-30B-A3B. Shapes without a committed .spv return None and
// take the packed `vk_moe_kernel` path -- add a shape by dumping it in `matvec-check dump`. The split
// bank (`vulkan_moe_bank_split`) is uploaded iff this returns Some, keeping the two paths orthogonal.
#[cfg(feature = "vulkan")]
fn vk_moe_blk_kernel(dt: GgmlDType, n: usize, k: usize) -> Option<&'static str> {
    // A/B escape hatch: VK_MOE_PACKED forces the packed naive path so the block kernels' end-to-end
    // contribution is measurable under identical conditions (no cheating -- show the delta).
    if std::env::var_os("VK_MOE_PACKED").is_some() {
        return None;
    }
    match (dt, n, k) {
        (GgmlDType::Q4K, 768, 2048) => Some("moe_matvec_q4k_blk_gu"),
        (GgmlDType::Q4K, 2048, 768) => Some("moe_matvec_q4k_blk_dn"),
        (GgmlDType::Q6K, 2048, 768) => Some("moe_matvec_q6k_blk_dn"),
        _ => None,
    }
}

// dp4a (int8 OpSDot) block kernels for the shapes with a committed .spv. Selected before the f32 block
// kernel when the device advertises integer dot-product (see has_int_dot8). VK_MOE_DP4A_OFF forces the
// f32 path for an A/B. The bool is the kernel's activation-binding contract: whether it binds the
// per-32 q8 sums (`xsum`; Q4_K folds dmin against them) or derives its own half-block sums in-register
// (Q6_K's −32 fold needs per-16 sums, which per-32 xsum cannot express).
// Gated exactly like its sibling vk_moe_blk_kernel above. Both call sites are
// inside `#[cfg(feature = "vulkan")]` blocks, so without the feature this compiled
// with no callers and clippy correctly called it dead.
#[cfg(feature = "vulkan")]
fn vk_moe_blk_dp4a_kernel(dt: GgmlDType, n: usize, k: usize) -> Option<(&'static str, bool)> {
    if std::env::var_os("VK_MOE_PACKED").is_some() || std::env::var_os("VK_MOE_DP4A_OFF").is_some()
    {
        return None;
    }
    match (dt, n, k) {
        (GgmlDType::Q4K, 768, 2048) => Some(("moe_matvec_q4k_dp4a_blk_gu", true)),
        (GgmlDType::Q4K, 2048, 768) => Some(("moe_matvec_q4k_dp4a_blk_dn", true)),
        (GgmlDType::Q6K, 2048, 768) => Some(("moe_matvec_q6k_dp4a_blk_dn", false)),
        _ => None,
    }
}

// wgpu twin of `vk_moe_kernel`. The WGSL MoE shaders cover Q4_0/Q8_0/Q4K; Q6_K is Vulkan-only so far.
#[cfg(feature = "wgpu")]
fn wgpu_moe_kernel(dt: GgmlDType) -> Option<&'static str> {
    match dt {
        GgmlDType::Q4_0 => Some("moe_matvec_q4_0"),
        GgmlDType::Q8_0 => Some("moe_matvec_q8_0"),
        GgmlDType::Q4K => Some("moe_matvec_q4k"),
        _ => None,
    }
}

/// Per-expert quantized MoE matmul, shared by the resident and disk-streaming paths. Groups the
/// routed slots by expert id, and for each active expert builds a `QTensor` from `make_storage`
/// (resident slice or streamed slab), runs the native quantized matmul on just that expert's tokens,
/// and scatters the result back. The ONLY difference between resident and streaming is the closure.
fn moe_grouped_per_expert(
    x: &Tensor,
    ids: &Tensor,
    n: usize,
    k: usize,
    mut make_storage: impl FnMut(u32, &Device) -> Result<QStorage>,
) -> Result<Tensor> {
    use crate::Module; // brings QMatMul::forward into scope
    use std::collections::HashMap;
    use std::sync::Arc;
    let device = x.device();
    let out_dtype = x.dtype();
    let (t, topk) = ids.dims2()?;
    let s = x.dim(1)?; // 1 (gate/up: shared input) or topk (down: per-slot)
    let x_exp = if s == topk {
        x.clone()
    } else {
        x.broadcast_as((t, topk, k))?
    };
    let x_flat = x_exp
        .reshape((t * topk, k))?
        .to_dtype(DType::F32)?
        .contiguous()?;
    let ids_flat = ids.reshape((t * topk,))?.to_dtype(DType::U32)?;
    let ids_vec = ids_flat.to_vec1::<u32>()?;
    let mut groups: HashMap<u32, Vec<u32>> = HashMap::new();
    for (slot, eid) in ids_vec.iter().enumerate() {
        groups.entry(*eid).or_default().push(slot as u32);
    }
    let mut out_flat = Tensor::zeros((t * topk, n), DType::F32, device)?;
    for (eid, slots) in groups.into_iter() {
        let qs = make_storage(eid, device)?;
        let shape: crate::Shape = (n, k).into();
        let w_e = QTensor::make(qs, shape);
        let qm = QMatMul::from_arc(Arc::new(w_e))?;
        let m = slots.len();
        let idx = Tensor::from_vec(slots, (m,), device)?;
        let x_e = x_flat.index_select(&idx, 0)?; // [m, k]
        let y_e = qm.forward(&x_e)?.to_dtype(DType::F32)?; // [m, n]
        out_flat = out_flat.index_add(&idx, &y_e, 0)?;
    }
    out_flat.reshape((t, topk, n))?.to_dtype(out_dtype)
}

/// FUSED MoE expert-combine: `out[i,j] = sum_e scores[i,e] * ys[i,e,j]`, reducing the per-expert
/// outputs `ys` [t, topk, n] by the router weights `scores` [t, topk] into [t, n]. On ROCm this is
/// ONE fused kernel (`RocmDevice::moe_combine`) instead of `ys.broadcast_mul(scores).sum(Minus2)`,
/// which cast ys -> f32, wrote a [t,topk,n] f32 product temp, and ran an 8-wide strided reduce.
/// Other backends keep the generic broadcast-mul + sum.
#[cfg_attr(not(feature = "rocm"), allow(unused_variables))]
pub fn moe_combine(ys: &Tensor, scores: &Tensor) -> Result<Tensor> {
    let (t, topk, n) = ys.dims3()?;
    #[cfg(feature = "rocm")]
    if let Device::Rocm(dev) = ys.device() {
        let ys_c = ys.contiguous()?;
        let scores_c = scores.to_dtype(DType::F32)?.contiguous()?;
        let (ys_store, _) = ys_c.storage_and_layout();
        let yr = match &*ys_store {
            Storage::Rocm(r) => r,
            _ => crate::bail!("moe_combine: ys not on rocm after contiguous()"),
        };
        let (sc_store, _) = scores_c.storage_and_layout();
        let sr = match &*sc_store {
            Storage::Rocm(r) => r,
            _ => crate::bail!("moe_combine: scores not on rocm after contiguous()"),
        };
        let out = dev.moe_combine(yr, sr, t, topk, n)?;
        return Ok(crate::tensor::from_storage(
            Storage::Rocm(out),
            (t, n),
            crate::op::BackpropOp::none(),
            false,
        ));
    }
    // scores come from moe_route in f32; ys carries the model compute dtype, which on a backend
    // without a native moe_combine kernel (Metal) is bf16/f16 -> a raw broadcast_mul would hit a
    // dtype mismatch. Accumulate in f32 (as the rocm/cuda kernels do), then restore ys's dtype.
    // No-op for the f32 path, so CPU/CUDA stay byte-identical.
    let out_dtype = ys.dtype();
    ys.to_dtype(DType::F32)?
        .broadcast_mul(&scores.to_dtype(DType::F32)?.unsqueeze(D::Minus1)?)?
        .sum(D::Minus2)?
        .to_dtype(out_dtype)
}

/// Fused MoE router. Reduces the F32 router logits [ntok, n_experts] to the topk selected expert
/// ids [ntok, topk] (descending logit) and their softmax weights [ntok, topk]; `norm` renormalizes
/// the topk weights to sum 1 (norm_topk_prob). On ROCm this is ONE `moe_route` kernel replacing the
/// softmax->sort->narrow->sum->div chain (~6 launches/layer); elsewhere it is that chain via ml ops.
#[cfg_attr(not(feature = "rocm"), allow(unused_variables))]
pub fn moe_route(logits: &Tensor, topk: usize, norm: bool) -> Result<(Tensor, Tensor)> {
    let (ntok, n_experts) = logits.dims2()?;
    #[cfg(feature = "rocm")]
    if let Device::Rocm(dev) = logits.device() {
        let logits_c = logits.to_dtype(DType::F32)?.contiguous()?;
        let (lg_store, _) = logits_c.storage_and_layout();
        let lr = match &*lg_store {
            Storage::Rocm(r) => r,
            _ => crate::bail!("moe_route: logits not on rocm after contiguous()"),
        };
        let (ids, w) = dev.moe_route(lr, ntok, n_experts, topk, norm)?;
        let ids_t = crate::tensor::from_storage(
            Storage::Rocm(ids),
            (ntok, topk),
            crate::op::BackpropOp::none(),
            false,
        );
        let w_t = crate::tensor::from_storage(
            Storage::Rocm(w),
            (ntok, topk),
            crate::op::BackpropOp::none(),
            false,
        );
        return Ok((ids_t, w_t));
    }
    #[cfg(feature = "cuda")]
    if let Device::Cuda(cdev) = logits.device() {
        if n_experts <= 256 && topk <= 32 {
            let logits_c = logits.to_dtype(DType::F32)?.contiguous()?;
            let (lg_store, _) = logits_c.storage_and_layout();
            let lr = match &*lg_store {
                Storage::Cuda(c) => c,
                _ => crate::bail!("moe_route: logits not on cuda after contiguous()"),
            };
            let lview = lr.as_cuda_slice::<f32>()?.slice(0..);
            let (ids, w) = cuda::moe_route(&lview, ntok, n_experts, topk, norm, cdev)?;
            let ids_t = crate::tensor::from_storage(
                Storage::Cuda(ids),
                (ntok, topk),
                crate::op::BackpropOp::none(),
                false,
            );
            let w_t = crate::tensor::from_storage(
                Storage::Cuda(w),
                (ntok, topk),
                crate::op::BackpropOp::none(),
                false,
            );
            return Ok((ids_t, w_t));
        }
    }
    // Fused Vulkan router (one workgroup/token, softmax + top-k in shared mem) for the committed .spv
    // shape (E=128, top-8) with norm_topk_prob; replaces the generic softmax+sort+gather op-chain and
    // its per-op layout copies. Other shapes / norm=false fall through to the generic path below.
    #[cfg(feature = "vulkan")]
    if let Device::Vulkan(vdev) = logits.device() {
        if norm && n_experts == 128 && topk == 8 {
            let logits_c = logits.to_dtype(DType::F32)?.contiguous()?;
            let (lg_store, _) = logits_c.storage_and_layout();
            let lv = match &*lg_store {
                Storage::Vulkan(v) => v,
                _ => crate::bail!("moe_route: logits not on vulkan after contiguous()"),
            };
            let (ids, w) = vdev.moe_route_vk(lv, ntok, n_experts, topk)?;
            let ids_t = crate::tensor::from_storage(
                Storage::Vulkan(ids),
                (ntok, topk),
                crate::op::BackpropOp::none(),
                false,
            );
            let w_t = crate::tensor::from_storage(
                Storage::Vulkan(w),
                (ntok, topk),
                crate::op::BackpropOp::none(),
                false,
            );
            return Ok((ids_t, w_t));
        }
    }
    let lf = logits.to_dtype(DType::F32)?;
    let mx = lf.max_keepdim(D::Minus1)?;
    let e = lf.broadcast_sub(&mx)?.exp()?;
    let z = e.sum_keepdim(D::Minus1)?;
    let p = e.broadcast_div(&z)?;
    let (sv, si) = p.sort_last_dim(false)?;
    let ids = si.narrow(D::Minus1, 0, topk)?.contiguous()?;
    let mut w = sv.narrow(D::Minus1, 0, topk)?.contiguous()?;
    if norm {
        w = w.broadcast_div(&w.sum_keepdim(D::Minus1)?)?;
    }
    Ok((ids, w))
}

/// Fused MoE gate+up projections. Both expert banks consume the SAME routed token `x` [t,1,k], so the
/// shared input is broadcast + quantized ONCE and matvec'd against both banks (vs once per bank,
/// re-materializing + re-quantizing the identical activation). Returns the raw (gate_out, up_out)
/// [t,topk,n]; the caller applies silu(gate)*up. Each output is bit-identical to the unfused
/// `indexed_moe_forward` because the q8_1 activation is deterministic in `x`. ROCm decode/matvec path
/// only (the prefill qmmq path keeps its own per-bank quantize); every other case runs the two
/// unfused forwards.
pub fn moe_gate_up(
    x: &Tensor,
    ids: &Tensor,
    gate: &QMatMul,
    up: &QMatMul,
) -> Result<(Tensor, Tensor)> {
    #[cfg(feature = "rocm")]
    {
        if let (QMatMul::QTensor(gq), QMatMul::QTensor(uq)) = (gate, up) {
            if let (QStorage::Rocm(_, dev), QStorage::Rocm(..)) = (&gq.storage, &uq.storage) {
                let dt = gq.storage.dtype();
                if dt == uq.storage.dtype() {
                    if let Some(qt) = crate::RocmQuantType::from_ggml(dt) {
                        let (_e, n, k) = gq.shape().dims3()?;
                        let (t, topk) = ids.dims2()?;
                        let use_qmmq = t > 1 && qt.qmmq_capable();
                        if x.dim(1)? == 1 && !use_qmmq {
                            let nrows = t * topk;
                            let x_exp = x.broadcast_as((t, topk, k))?;
                            let x_flat = match x_exp.dtype() {
                                DType::BF16 | DType::F16 => {
                                    x_exp.reshape((nrows, k))?.contiguous()?
                                }
                                DType::F32 if qt.dp4a_active() => {
                                    x_exp.reshape((nrows, k))?.contiguous()?
                                }
                                _ => x_exp
                                    .reshape((nrows, k))?
                                    .to_dtype(DType::F16)?
                                    .contiguous()?,
                            };
                            let out_dtype = x.dtype();
                            let ids_u32 =
                                ids.reshape((nrows,))?.to_dtype(DType::U32)?.contiguous()?;
                            let gwb = gq.rocm_moe_bank(dev)?;
                            let uwb = uq.rocm_moe_bank(dev)?;
                            let (xstore, _) = x_flat.storage_and_layout();
                            let xr = match &*xstore {
                                Storage::Rocm(r) => r,
                                _ => crate::bail!("moe_gate_up: x not on rocm after contiguous()"),
                            };
                            let (idstore, _) = ids_u32.storage_and_layout();
                            let idr = match &*idstore {
                                Storage::Rocm(r) => r,
                                _ => crate::bail!("moe_gate_up: ids not on rocm"),
                            };
                            let (gy, uy) = dev.moe_matvec_pair(
                                qt,
                                gwb.as_ref(),
                                uwb.as_ref(),
                                xr,
                                idr,
                                nrows,
                                n,
                                k,
                            )?;
                            let g = crate::tensor::from_storage(
                                Storage::Rocm(gy),
                                (nrows, n),
                                crate::op::BackpropOp::none(),
                                false,
                            )
                            .reshape((t, topk, n))?
                            .to_dtype(out_dtype)?;
                            let u = crate::tensor::from_storage(
                                Storage::Rocm(uy),
                                (nrows, n),
                                crate::op::BackpropOp::none(),
                                false,
                            )
                            .reshape((t, topk, n))?
                            .to_dtype(out_dtype)?;
                            return Ok((g, u));
                        }
                    }
                }
            }
        }
    }
    // Vulkan dp4a twin of the ROCm block above: gate and up contract the SAME routed token, so the
    // q8 activation quantize is hoisted out and both matvecs dispatch against one copy. Falling
    // through to two `indexed_moe_forward` calls re-derives it per matvec -- byte-identical work,
    // twice. Only the quantize is shared; the dispatch is `moe_matvec_blk_dp4a_pre_gpu` either way.
    // A/B escape hatch, mirroring VK_MOE_PACKED / VK_MOE_DP4A_OFF: VK_MOE_GU_FUSE_OFF falls through to
    // the two-call path so the shared quantize can be measured against re-quantizing per matvec.
    #[cfg(feature = "vulkan")]
    if std::env::var_os("VK_MOE_GU_FUSE_OFF").is_none() {
        if let (QMatMul::QTensor(gq), QMatMul::QTensor(uq)) = (gate, up) {
            if let (QStorage::Vulkan(_, dev), QStorage::Vulkan(..)) = (&gq.storage, &uq.storage) {
                let dt = gq.storage.dtype();
                if dt == uq.storage.dtype() && dev.has_int_dot8() {
                    let (e_cnt, n, k) = gq.shape().dims3()?;
                    // gate/up share one input row per token; the per-slot (down) shape is not this path.
                    if uq.shape().dims3()? == (e_cnt, n, k) && x.dim(1)? == 1 {
                        if let Some((blk, with_xsum)) = vk_moe_blk_dp4a_kernel(dt, n, k) {
                            let (t, topk) = ids.dims2()?;
                            let nrows = t * topk;
                            let x_flat = x
                                .broadcast_as((t, topk, k))?
                                .reshape((nrows, k))?
                                .to_dtype(DType::F32)?
                                .contiguous()?;
                            let ids_u32 =
                                ids.reshape((nrows,))?.to_dtype(DType::U32)?.contiguous()?;
                            let (xstore, _) = x_flat.storage_and_layout();
                            let xv = match &*xstore {
                                Storage::Vulkan(v) => v,
                                _ => {
                                    crate::bail!("moe_gate_up: x not on vulkan after contiguous()")
                                }
                            };
                            let (idstore, _) = ids_u32.storage_and_layout();
                            let idv = match &*idstore {
                                Storage::Vulkan(v) => v,
                                _ => crate::bail!("moe_gate_up: ids not on vulkan"),
                            };
                            // The one quantize both projections read.
                            let (xq, xs, xsum) = dev.quantize_act_q8(xv, nrows, k)?;
                            let gbank = gq.vulkan_moe_bank_split(dev, e_cnt, n, k)?;
                            let ubank = uq.vulkan_moe_bank_split(dev, e_cnt, n, k)?;
                            let out_dtype = x.dtype();
                            // PREFILL (t > 1): expert-grouped MMQ on the matrix cores, the same
                            // t > 1 split `indexed_moe_forward` makes. Gate and up share the routing,
                            // so they also share ONE grouping pass. DECODE (t == 1) keeps the matvec.
                            let (gy, uy) = if t > 1
                                && dt == GgmlDType::Q4K
                                && std::env::var_os("VK_MOE_PREFILL_GEMM_OFF").is_none()
                            {
                                let (counts, rows) =
                                    dev.moe_expert_rows_vk(idv, nrows, e_cnt, t)?;
                                let gy = dev.mmq_q4k_id_pre_gpu(
                                    gbank.as_ref(),
                                    &xq,
                                    &xs,
                                    &xsum,
                                    &rows,
                                    &counts,
                                    nrows,
                                    e_cnt,
                                    t,
                                    n,
                                    k,
                                )?;
                                let uy = dev.mmq_q4k_id_pre_gpu(
                                    ubank.as_ref(),
                                    &xq,
                                    &xs,
                                    &xsum,
                                    &rows,
                                    &counts,
                                    nrows,
                                    e_cnt,
                                    t,
                                    n,
                                    k,
                                )?;
                                (gy, uy)
                            } else {
                                let gy = dev.moe_matvec_blk_dp4a_pre_gpu(
                                    blk,
                                    with_xsum,
                                    gbank.as_ref(),
                                    &xq,
                                    &xs,
                                    &xsum,
                                    idv,
                                    nrows,
                                    n,
                                )?;
                                let uy = dev.moe_matvec_blk_dp4a_pre_gpu(
                                    blk,
                                    with_xsum,
                                    ubank.as_ref(),
                                    &xq,
                                    &xs,
                                    &xsum,
                                    idv,
                                    nrows,
                                    n,
                                )?;
                                (gy, uy)
                            };
                            let shape = |o| -> Result<Tensor> {
                                crate::tensor::from_storage(
                                    Storage::Vulkan(o),
                                    (nrows, n),
                                    crate::op::BackpropOp::none(),
                                    false,
                                )
                                .reshape((t, topk, n))?
                                .to_dtype(out_dtype)
                            };
                            return Ok((shape(gy)?, shape(uy)?));
                        }
                    }
                }
            }
        }
    }
    Ok((
        gate.indexed_moe_forward(x, ids)?,
        up.indexed_moe_forward(x, ids)?,
    ))
}

impl QMatMul {
    pub fn from_arc(qtensor: std::sync::Arc<QTensor>) -> Result<Self> {
        // Native Vulkan quantized path: keep the GGML quantized blocks in VRAM and run the matching
        // on-GPU quant matvec for decode, instead of dequantizing the whole model to f32 (4x the
        // decode bandwidth). The kernel reads the GGML block format straight from the uploaded bytes
        // -- no CPU dequant, no re-pack -- so this is exact w.r.t. the CPU reference. Q4_0/Q4_K need
        // k a multiple of their block (32 / 256); Q8_0 needs a multiple of 32.
        #[cfg(feature = "vulkan")]
        {
            let dt = qtensor.dtype();
            let native_vk = matches!(
                dt,
                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
                    | GgmlDType::IQ2_XXS
                    | GgmlDType::IQ2_S
                    | GgmlDType::IQ3_XXS
                    | GgmlDType::IQ3_S
                    | GgmlDType::IQ1_S
                    | GgmlDType::IQ1_M
                    | GgmlDType::IQ2_XS
            );
            if native_vk {
                if let Device::Vulkan(d) = qtensor.device() {
                    if let Ok((n, k)) = qtensor.shape().dims2() {
                        let blk = dt.block_size();
                        if k % blk == 0 {
                            let bytes = qtensor.data()?;
                            // Q6_K (210 B) and Q3_K (110 B) blocks are not u32-aligned; their shaders
                            // read a padded u32 stride, so repack on upload. Q8_0 repacks to the
                            // 9-u32/block layout that BOTH its decode (mul_mat_vec_q8) and prefill GEMM
                            // (mul_mat_q8) read -- ONE layout, so decode + prefill + MoE all agree
                            // (mirrors how the MoE bank repacks Q8_0). Every other native type's
                            // shaders byte-address the raw GGML bytes directly.
                            let wq = match dt {
                                GgmlDType::Q6K => d.quantize_q6k(&bytes, n, k)?,
                                GgmlDType::Q3K => d.quantize_q3k(&bytes, n, k)?,
                                GgmlDType::Q8_0 => d.quantize_q8_blocks(&bytes, n, k)?,
                                GgmlDType::IQ2_XXS => d.quantize_iq2xxs(&bytes, n, k)?,
                                GgmlDType::IQ2_XS => d.quantize_iq2xs(&bytes, n, k)?,
                                GgmlDType::IQ1_M => d.quantize_iq1m(&bytes, n, k)?,
                                GgmlDType::IQ1_S => d.quantize_iq1s(&bytes, n, k)?,
                                GgmlDType::IQ3_S => d.quantize_iq3s(&bytes, n, k)?,
                                GgmlDType::IQ3_XXS => d.quantize_iq3xxs(&bytes, n, k)?,
                                GgmlDType::IQ2_S => d.quantize_iq2s(&bytes, n, k)?,
                                _ => d.upload_qweight(&bytes)?,
                            };
                            return Ok(Self::VulkanQuant {
                                qtensor,
                                wq: std::sync::Arc::new(wq),
                                dtype: dt,
                                n,
                                k,
                            });
                        }
                    }
                }
            }
        }
        // Native wgpu quantized path: same idea as Vulkan. Ships the Q4_0/Q8_0/Q4_K WGSL matvec
        // kernels; other dtypes fall through to the dequantize path below.
        #[cfg(feature = "wgpu")]
        {
            let dt = qtensor.dtype();
            let native_wgpu = matches!(dt, GgmlDType::Q4_0 | GgmlDType::Q8_0 | GgmlDType::Q4K);
            if native_wgpu {
                if let Device::Wgpu(d) = qtensor.device() {
                    if let Ok((n, k)) = qtensor.shape().dims2() {
                        let blk = dt.block_size();
                        if k % blk == 0 {
                            let bytes = qtensor.data()?;
                            let wq = d.upload_qweight(&bytes)?;
                            return Ok(Self::WgpuQuant {
                                qtensor,
                                wq: std::sync::Arc::new(wq),
                                dtype: dt,
                                n,
                                k,
                            });
                        }
                    }
                }
            }
        }
        // Native ROCm quantized path: keep the GGML blocks in VRAM and run the ONE unified on-GPU
        // quant decode core (qmatvec_core<WTYPE>; Q8_0+Q4_0 also have the int8 WMMA prefill gemm),
        // instead of dequantizing the whole model to dense f16 (2x+ the decode bandwidth). Reads the
        // block format straight from the uploaded bytes -- exact w.r.t. the CPU reference. The wired
        // set is exactly RocmQuantType::from_ggml; k must be a multiple of that type's block size.
        #[cfg(feature = "rocm")]
        {
            let dt = qtensor.dtype();
            // Decode AND prefill native iff the unified core has this type wired (one enum row in
            // RocmQuantType): decode rides qmatvec_core<WTYPE>, prefill (rows>1) rides the int8 WMMA
            // qmmq_core<WTYPE>. Both cover the SAME wired spread; adding a type is one enum row + the
            // in-kernel decode, no per-quant kernel. Unwired types dequantize-to-f16 in forward().
            if let Some(qt) = crate::RocmQuantType::from_ggml(dt) {
                if let Device::Rocm(d) = qtensor.device() {
                    if let Ok((n, k)) = qtensor.shape().dims2() {
                        let blk_ok = k % qt.block_elems() == 0;
                        if blk_ok {
                            use crate::backend::BackendDevice;
                            let bytes = qtensor.data()?;
                            let wq = d.storage_from_slice(bytes.as_ref())?;
                            return Ok(Self::RocmQuant {
                                qtensor,
                                wq: std::sync::Arc::new(wq),
                                dtype: dt,
                                n,
                                k,
                            });
                        }
                    }
                }
            }
        }
        // ROCm MoE bank: a 3D [E,n,k] expert bank of a wired quant type stays QUANTIZED (kept as
        // QTensor), so `indexed_moe_forward` runs each routed expert through the ONE unified
        // qmatvec_core (no per-expert kernel) instead of dequantizing the whole bank to dense f16
        // (which for a 30B-A3B model is many GB of resident f16 AND has no indexed_moe path). The 2D
        // RocmQuant decode/prefill path above already handles ordinary weights; this is the MoE case.
        #[cfg(feature = "rocm")]
        {
            if qtensor.device().is_rocm()
                && qtensor.shape().dims().len() == 3
                && crate::RocmQuantType::from_ggml(qtensor.dtype()).is_some()
            {
                return Ok(Self::QTensor(qtensor));
            }
        }
        // Vulkan/wgpu MoE bank: a 3D [E,n,k] expert bank whose dtype has a native fused grouped
        // quant-matvec kernel stays QUANTIZED (kept as QTensor), so `indexed_moe_forward` gathers
        // the routed experts on the GPU straight from the resident bank -- the twin of the ROCm
        // case above. Dequantizing here instead explodes the bank to dense f32 in VRAM (30B-A3B ->
        // 100+ GB) AND lands on the generic path, whose Vulkan/wgpu index_add is unwired.
        #[cfg(feature = "vulkan")]
        {
            if qtensor.device().is_vulkan()
                && qtensor.shape().dims().len() == 3
                && vk_moe_kernel(qtensor.dtype()).is_some()
            {
                return Ok(Self::QTensor(qtensor));
            }
        }
        #[cfg(feature = "wgpu")]
        {
            if qtensor.device().is_wgpu()
                && qtensor.shape().dims().len() == 3
                && wgpu_moe_kernel(qtensor.dtype()).is_some()
            {
                return Ok(Self::QTensor(qtensor));
            }
        }
        let dequantize = match qtensor.dtype() {
            GgmlDType::F32 | GgmlDType::F16 | GgmlDType::BF16 | GgmlDType::I32 => true,
            // The Vulkan/wgpu/ROCm backends have no generic native quantized matmul, so dequantize
            // to f32 here (once, at construction) and run the regular f32 GPU matmul.
            _ => {
                qtensor.device().is_vulkan()
                    || qtensor.device().is_wgpu()
                    || qtensor.device().is_rocm()
            }
        };
        let t = if dequantize {
            // ROCm: dequantize to f16 so the matmul hits RDNA3.5 matrix cores (WMMA). Dense f32
            // (sgemm) has no matrix-core path on RDNA and runs ~an order of magnitude slower, and
            // f16 also halves the resident weight memory.
            if qtensor.device().is_rocm() {
                Self::TensorF16(qtensor.dequantize_f16(&qtensor.device())?)
            } else {
                Self::Tensor(qtensor.dequantize(&qtensor.device())?)
            }
        } else {
            Self::QTensor(qtensor)
        };
        Ok(t)
    }

    pub fn from_qtensor(qtensor: QTensor) -> Result<Self> {
        Self::from_arc(std::sync::Arc::new(qtensor))
    }

    pub fn dequantize_f16(&self) -> Result<Tensor> {
        match self {
            Self::QTensor(t) => t.dequantize_f16(&t.device()),
            Self::Tensor(t) => t.to_dtype(DType::F16),
            Self::TensorF16(t) => Ok(t.clone()),
            #[cfg(feature = "rocm")]
            Self::RocmQuant { qtensor, .. } => qtensor.dequantize_f16(&qtensor.device()),
            #[cfg(feature = "vulkan")]
            Self::VulkanQuant { qtensor, .. } => qtensor.dequantize_f16(&qtensor.device()),
            #[cfg(feature = "wgpu")]
            Self::WgpuQuant { qtensor, .. } => qtensor.dequantize_f16(&qtensor.device()),
        }
    }

    pub fn forward_via_f16(&self, xs: &Tensor) -> Result<Tensor> {
        let w = self.dequantize_f16()?;
        let in_dtype = xs.dtype();
        let w = match *xs.dims() {
            [b1, b2, _, _] => w.broadcast_left((b1, b2))?.t()?,
            [bsize, _, _] => w.broadcast_left(bsize)?.t()?,
            _ => w.t()?,
        };
        xs.to_dtype(DType::F16)?.matmul(&w)?.to_dtype(in_dtype)
    }

    pub fn indexed_moe_forward(&self, x: &Tensor, ids: &Tensor) -> Result<Tensor> {
        match self {
            Self::QTensor(t) => t.indexed_moe_forward(x, ids),
            // Resident-bank MoE: `wq` already holds the [E,n,k] GGML blocks in VRAM (uploaded at
            // load), so route straight to the batched on-GPU quant matvec. Delegating to `qtensor`
            // (CPU-side) instead drops to the generic fallback that re-uploads every routed expert
            // every token -- the 20-50x decode cliff. `qtensor` is read only for its shape.
            #[cfg(feature = "rocm")]
            Self::RocmQuant {
                qtensor, wq, dtype, ..
            } if crate::RocmQuantType::from_ggml(*dtype).is_some() => {
                let qt = crate::RocmQuantType::from_ggml(*dtype).unwrap();
                let wbank = wq.as_ref();
                // e_cnt unused: ids stay on-device, router guarantees the bound (no host check).
                let (_e_cnt, n, k) = qtensor.shape().dims3()?;
                let (t, topk) = ids.dims2()?;
                let s = x.dim(1)?; // 1 (gate/up: shared input) or topk (down: per-slot)
                let x_exp = if s == topk {
                    x.clone()
                } else {
                    x.broadcast_as((t, topk, k))?
                };
                let nrows = t * topk;
                // PREFILL (t>1, never graph-captured) routes to the FUSED expert-grouped WMMA GEMM
                // (`moe_qmmq_quant`), which needs f16 activations; DECODE (t==1) stays on the
                // capture-clean dp4a/scalar matvec, which takes bf16/f16 natively.
                // Decode-only types (no qmmq kernel) ride the per-slot matvec core for prefill too
                // (correct at any token count). ONE predicate gates every prefill site.
                let use_qmmq = t > 1 && qt.qmmq_capable();
                let x_flat = match x_exp.dtype() {
                    // qmmq quantizes f16/f32 activations natively, so keep the model's dtype and skip
                    // the f32->f16 cast (a 16.7M-elem read+write per gate/up). Other dtypes (bf16 with
                    // a symmetric expert type) still cast to f16.
                    DType::F16 | DType::F32 if use_qmmq => {
                        x_exp.reshape((nrows, k))?.contiguous()?
                    }
                    _ if use_qmmq => x_exp
                        .reshape((nrows, k))?
                        .to_dtype(DType::F16)?
                        .contiguous()?,
                    DType::BF16 | DType::F16 => x_exp.reshape((nrows, k))?.contiguous()?,
                    // DECODE f32-native dp4a: keep F32 routed activation F32 end-to-end (matvec stores
                    // F32), eliding the cast pair around each expert matvec. See QStorage twin above.
                    DType::F32 if qt.dp4a_active() => x_exp.reshape((nrows, k))?.contiguous()?,
                    _ => x_exp
                        .reshape((nrows, k))?
                        .to_dtype(DType::F16)?
                        .contiguous()?,
                };
                let out_dtype = x.dtype();
                // Keep router ids ON the GPU for EVERY wired quant type: the batched kernels index
                // experts on-device, so there is no per-call `to_vec1` host round-trip. That DtoH
                // sync (3 per layer x 48 layers per token) was both the dominant decode stall on WSL
                // AND the HIP-graph capture breaker. Router top-k guarantees 0 <= id < e_cnt.
                let ids_u32 = ids
                    .reshape((nrows,))?
                    .to_dtype(crate::DType::U32)?
                    .contiguous()?;
                let (xstore, _) = x_flat.storage_and_layout();
                let xr = match &*xstore {
                    crate::Storage::Rocm(r) => r,
                    _ => crate::bail!("rocm MoE: x not on rocm after contiguous()"),
                };
                let (idstore, _) = ids_u32.storage_and_layout();
                let idr = match &*idstore {
                    crate::Storage::Rocm(r) => r,
                    _ => crate::bail!("rocm MoE: ids not on rocm"),
                };
                let y = if use_qmmq {
                    wbank
                        .device
                        .moe_qmmq_quant(qt, wbank, xr, idr, nrows, n, k)?
                } else {
                    wbank
                        .device
                        .moe_matvec_quant(qt, wbank, xr, idr, nrows, n, k)?
                };
                let out = crate::tensor::from_storage(
                    crate::Storage::Rocm(y),
                    (nrows, n),
                    crate::op::BackpropOp::none(),
                    false,
                );
                out.reshape((t, topk, n))?.to_dtype(out_dtype)
            }
            // Unwired ROCm quant dtypes (no on-GPU quant matvec): CPU per-expert fallback.
            #[cfg(feature = "rocm")]
            Self::RocmQuant { qtensor, .. } => qtensor.indexed_moe_forward(x, ids),
            #[cfg(feature = "vulkan")]
            Self::VulkanQuant { qtensor, .. } => qtensor.indexed_moe_forward(x, ids),
            #[cfg(feature = "wgpu")]
            Self::WgpuQuant { qtensor, .. } => qtensor.indexed_moe_forward(x, ids),
            _ => {
                panic!("Not implemented!")
            }
        }
    }
}

impl crate::CustomOp1 for QTensor {
    fn name(&self) -> &'static str {
        "qmatmul"
    }

    fn cpu_fwd(
        &self,
        storage: &crate::CpuStorage,
        layout: &crate::Layout,
    ) -> Result<(crate::CpuStorage, Shape)> {
        if !layout.is_contiguous() {
            crate::bail!("input tensor is not contiguous {layout:?}")
        }
        let src_shape = layout.shape();
        // self is transposed so n is first then k.
        let (n, k) = self.shape.dims2()?;
        if src_shape.rank() < 2 {
            crate::bail!("input tensor has only one dimension {layout:?}")
        }
        let mut dst_shape = src_shape.dims().to_vec();
        let last_k = dst_shape.pop().unwrap();
        if last_k != k {
            crate::bail!("input tensor {layout:?} incompatible with {:?}", self.shape)
        }
        dst_shape.push(n);
        let dst_shape = Shape::from(dst_shape);
        #[allow(clippy::infallible_destructuring_match)]
        let self_storage = match &self.storage {
            QStorage::Cpu(storage) => storage,
            #[cfg(feature = "rocm")]
            QStorage::Rocm(..) => crate::bail!("Invalid storage"),
            #[cfg(feature = "vulkan")]
            QStorage::Vulkan(..) => crate::bail!("Invalid storage"),
            #[cfg(feature = "wgpu")]
            QStorage::Wgpu(..) => crate::bail!("Invalid storage"),
            QStorage::Metal(_) | QStorage::Cuda(_) | QStorage::Stream(_) => {
                crate::bail!("Invalid storage")
            }
        };
        match storage.dtype() {
            DType::F32 => {
                let slice = storage.as_slice::<f32>()?;
                let slice =
                    &slice[layout.start_offset()..layout.start_offset() + src_shape.elem_count()];
                let mut dst_storage = vec![0f32; dst_shape.elem_count()];
                self_storage.matmul_t(
                    (dst_shape.elem_count() / n, k, n),
                    slice,
                    &mut dst_storage,
                )?;
                Ok((crate::CpuStorage::F32(dst_storage), dst_shape))
            }
            DType::F16 => {
                let slice = storage.as_slice::<f16>()?;
                let slice =
                    &slice[layout.start_offset()..layout.start_offset() + src_shape.elem_count()];
                let mut dst_storage = vec![f16::ZERO; dst_shape.elem_count()];
                self_storage.matmul_t_f16(
                    (dst_shape.elem_count() / n, k, n),
                    slice,
                    &mut dst_storage,
                )?;
                Ok((crate::CpuStorage::F16(dst_storage), dst_shape))
            }
            _ => crate::bail!("Expected f32/f16"),
        }
    }

    fn metal_fwd(
        &self,
        storage: &crate::MetalStorage,
        layout: &crate::Layout,
    ) -> Result<(crate::MetalStorage, Shape)> {
        let self_storage = match &self.storage {
            QStorage::Metal(metal) => metal,
            _ => unreachable!("Cannot call metal matmul on non metal QTensor"),
        };
        self_storage.fwd(&self.shape, storage, layout)
    }

    fn cuda_fwd(
        &self,
        storage: &crate::CudaStorage,
        layout: &crate::Layout,
    ) -> Result<(crate::CudaStorage, Shape)> {
        let self_storage = match &self.storage {
            QStorage::Cuda(cuda) => cuda,
            _ => unreachable!("Cannot call cuda matmul on non cuda QTensor"),
        };
        self_storage.fwd(&self.shape, storage, layout)
    }
}

/// Dense (non-quantized) matmul `xs @ w^T` for the `Tensor`/`TensorF16` `QMatMul` variants, where
/// the stored weight `w` is `[n, k]` and `xs` is `[.., k]`. On ROCm at decode (a single-row matvec)
/// this computes the result as `sum_k(xs[k] * w[n, k])` via pooled broadcast-mul + reduce instead of
/// rocBLAS `gemm_ex`. rocBLAS's GEMM dispatch records a vendor-specific PM4 indirect-buffer packet
/// that WSL's HSA thunk rejects on hipGraph replay (`VendorSpecificAqlToPm4` assert), so a captured
/// decode forward containing one (e.g. the MoE F32 router gate) corrupts/aborts on replay. The
/// reduce path uses only ops already exercised under capture (RMSNorm etc.), so it replays cleanly,
/// and at M=1 a GEMV-as-reduce is as cheap as the GEMM (it materializes only the `[n, k]` weight).
/// Prefill (rows > 1, never graph-captured) keeps the rocBLAS GEMM. Non-ROCm devices are unchanged.
fn dense_matmul(xs: &Tensor, w: &Tensor) -> Result<Tensor> {
    let k = *w.dims().last().unwrap();
    let rows = xs.elem_count() / k;
    if rows == 1 && xs.device().is_rocm() {
        let n = w.dim(0)?;
        #[cfg(feature = "rocm")]
        {
            // Dense decode GEMV: read the [n,k] weight ONCE (warp/row dot) instead of materializing
            // and re-reading the broadcast_mul product. The activation is matched to the weight dtype
            // (a [k] cast, negligible); the GEMV stays capture-clean (no rocBLAS).
            let d = match xs.device() {
                Device::Rocm(d) => d.clone(),
                _ => unreachable!(),
            };
            let xs1 = xs.reshape((k,))?.to_dtype(w.dtype())?.contiguous()?;
            let w = w.contiguous()?;
            let (wstore, _) = w.storage_and_layout();
            let wr = match &*wstore {
                crate::Storage::Rocm(r) => r,
                _ => crate::bail!("dense_matmul: weight not on rocm"),
            };
            let (xstore, _) = xs1.storage_and_layout();
            let xr = match &*xstore {
                crate::Storage::Rocm(r) => r,
                _ => crate::bail!("dense_matmul: x not on rocm"),
            };
            let y = d.dense_gemv(wr, xr, n, k)?;
            let mut dims = xs.dims().to_vec();
            *dims.last_mut().unwrap() = n;
            return crate::tensor::from_storage(
                crate::Storage::Rocm(y),
                dims,
                crate::op::BackpropOp::none(),
                false,
            )
            .to_dtype(xs.dtype());
        }
        #[cfg(not(feature = "rocm"))]
        {
            let out = xs.reshape((1, k))?.broadcast_mul(w)?.sum(D::Minus1)?;
            let mut dims = xs.dims().to_vec();
            *dims.last_mut().unwrap() = n;
            return out.reshape(dims);
        }
    }
    let w = match *xs.dims() {
        [b1, b2, _, _] => w.broadcast_left((b1, b2))?.t()?,
        [bsize, _, _] => w.broadcast_left(bsize)?.t()?,
        _ => w.t()?,
    };
    xs.matmul(&w)
}

// Prefill gate for the native Vulkan quantized GEMM. The `else` (dequant) path materializes a FULL f32
// copy of the weight (`QTensor::dequantize` -> `upload_f32` -> a fresh `amdgpu_bo_alloc`), ~100-235 MB
// per Linear weight. Under the deferred single command batch (`BATCH_CAP`) none of those f32 BOs free
// until the end-of-forward flush, so a dense prefill re-expands the whole model to f32 (~32 GB for an
// 8B) in fresh allocations; on the gfx1151 Strix Halo UMA the accumulated BOs exhaust VRAM+GTT and
// `GEM_CREATE` BLOCKS IN THE KERNEL waiting for a free that can only happen after the still-recording
// batch flushes -- a self-deadlock (engine hangs at the first long prefill). ROCm never hits it: its
// prefill rides the int8-WMMA `qmmq` GEMM, never dequantizes. So for every dtype that HAS a native
// quantized GEMM, ALWAYS use it (`usize::MAX`) -- liveness beats the at-large-M throughput edge the
// dequant path used to win (the GEMM re-reads the weight ceil(M/8) times so it is slower at big M, but
// it cannot deadlock; the future perf lever is a shared-memory-tiled int8 GEMM that reads the weight
// once, mirroring ROCm `qmmq`). Types WITHOUT a GEMM kernel return 0 so rows>1 dequantize (their only
// option; a handful of such weights allocate little and do not accumulate enough to deadlock). rows==1
// (decode) never reaches this gate -- it uses the native matvec straight out of the block format.
#[cfg(feature = "vulkan")]
fn vulkan_prefill_gemm_max_rows(dtype: GgmlDType) -> usize {
    match dtype {
        GgmlDType::Q4_0 | GgmlDType::Q8_0 | GgmlDType::Q4K | GgmlDType::Q5K | GgmlDType::Q6K => {
            usize::MAX
        }
        _ => 0,
    }
}

// The native quant matvec/matmul kernels read the activation from the storage buffer base; they take
// no per-tensor start_offset. A tensor narrowed out of a larger buffer (e.g. extract_logits' last-row
// slice: `x.narrow(1, seq_len-1, 1)`) is `is_contiguous()` yet carries a non-zero start_offset because
// its size-1 leading dims mask the stride mismatch -- so `contiguous()` returns it as a view and the
// kernel would read the WRONG rows (position 0 instead of the narrowed one). Materialize a genuine
// offset-0 buffer before the storage reaches the kernel.
#[cfg(feature = "vulkan")]
fn vulkan_act_offset0(xs: &Tensor) -> Result<Tensor> {
    let xs = xs.contiguous()?;
    if xs.layout().start_offset() == 0 {
        Ok(xs)
    } else {
        xs.force_contiguous()
    }
}

impl crate::Module for QMatMul {
    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        match self {
            #[cfg(feature = "rocm")]
            Self::RocmQuant {
                qtensor,
                wq,
                dtype,
                n,
                k,
            } => {
                // Device-residency guard: multi-token attention prefill can leave the activation
                // off-device (an upstream op leaked to host); recover instead of bailing so prefill
                // stays correct. The leak is the bug to fix for speed; this keeps us correct meanwhile.
                let xs_recovered = if xs.device().is_rocm() {
                    None
                } else {
                    Some(xs.to_device(&qtensor.device())?)
                };
                let xs = xs_recovered.as_ref().unwrap_or(xs);
                let rows: usize = xs.elem_count() / *k;
                // Table-driven unified decode: a type is decode-native iff the single
                // `qmatvec_core<WTYPE>` has a `decode_block` wired for it (RocmQuantType::from_ggml).
                // Q8_0/Q4_0/Q4_K/Q6_K/IQ4_XS/TQ2_0 today; adding a type is one enum row, no kernel.
                // The dp4a-vs-scalar routing for dp4a-capable types lives entirely inside `matvec_quant`
                // (dp4a_active) -- ONE path, one place; the type stays native either way.
                #[cfg(feature = "rocm")]
                let unified_qt = crate::RocmQuantType::from_ggml(*dtype);
                #[cfg(not(feature = "rocm"))]
                let unified_qt: Option<()> = None;
                // Native int8-WMMA prefill exists only for `qmmq_capable` types; the decode-only types
                // (Q2_K/Q3_K + every IQ*/TQ* codebook/fractional type) dequantize-to-f16 for rows>1 via
                // the `else` branch below -- correct, just not WMMA-accelerated. ONE predicate, read here
                // and at the two MoE `use_qmmq` sites.
                #[cfg(feature = "rocm")]
                let qmmq_ok = unified_qt.map(|qt| qt.qmmq_capable()).unwrap_or(false);
                #[cfg(not(feature = "rocm"))]
                let qmmq_ok = false;
                if rows == 1 && unified_qt.is_some() {
                    // Decode: weights stay quantized in VRAM; the ONE native on-GPU quant matvec core
                    // dequantizes per-block on-the-fly (no dense f16 copy). The matvec consumes
                    // bf16/f16 activations directly and returns the same dtype, so the model's working
                    // dtype (bf16) is kept end-to-end -- no bf16->f32->f16->bf16 cast detour. Only fall
                    // back to an f16 cast for exotic input dtypes. Every wired type (symmetric 8-bit
                    // through asymmetric super-block through sub-4-bit ternary) rides the same core.
                    // dp4a-capable types accept the F32 residual/norm stream DIRECTLY (q8_1 quantize
                    // from f32 + f32-store matvec), so an F32 activation stays F32 end-to-end with no
                    // f16 bounce -- this removes the cast_f32_f16-before / cast_f16_f32-after pair that
                    // wrapped every decode matvec. Non-dp4a (scalar) types keep the f16 cast.
                    #[cfg(feature = "rocm")]
                    let keep_f32 = unified_qt.map(|qt| qt.dp4a_active()).unwrap_or(false);
                    #[cfg(not(feature = "rocm"))]
                    let keep_f32 = false;
                    let xs = match xs.dtype() {
                        DType::BF16 | DType::F16 => xs.contiguous()?,
                        DType::F32 if keep_f32 => xs.contiguous()?,
                        _ => xs.to_dtype(DType::F16)?.contiguous()?,
                    };
                    let d = match xs.device() {
                        Device::Rocm(d) => d,
                        _ => crate::bail!("RocmQuant input not on rocm"),
                    };
                    let y = {
                        let (store, _) = xs.storage_and_layout();
                        let xr = match &*store {
                            crate::Storage::Rocm(r) => r,
                            _ => crate::bail!("RocmQuant expected rocm storage"),
                        };
                        #[cfg(feature = "rocm")]
                        {
                            d.matvec_quant(unified_qt.unwrap(), wq, xr, *n, *k)?
                        }
                        #[cfg(not(feature = "rocm"))]
                        {
                            crate::bail!("rocm feature disabled")
                        }
                    };
                    let mut dims = xs.dims().to_vec();
                    let last = dims.len() - 1;
                    dims[last] = *n;
                    Ok(crate::tensor::from_storage(
                        crate::Storage::Rocm(y),
                        dims,
                        crate::op::BackpropOp::none(),
                        false,
                    ))
                } else if let Some(qt) = unified_qt.filter(|_| qmmq_ok) {
                    // Prefill (rows>1): native int8 WMMA gemm through the ONE unified core
                    // (`qmmq_core<WTYPE>` in quant.hip). Weights stay quantized in VRAM (no resident
                    // dense f16, which would slow the memory-bound decode) and the MAC runs on the
                    // RDNA3 int8 matrix cores instead of rocBLAS. The SAME core covers the whole wired
                    // spread: Q8_0/Q4_0 (symmetric, proven), Q4_K (asymmetric -- min bias via the
                    // q8_1 block-sum), and the symmetric super-block / IQ / ternary types (Q6_K,
                    // IQ4_XS, TQ2_0). Selecting the type is one `RocmQuantType` row + the in-kernel
                    // decode; there is NO per-quant prefill kernel.
                    let xs = xs.to_dtype(DType::F16)?.contiguous()?;
                    let d = match xs.device() {
                        Device::Rocm(d) => d,
                        _ => crate::bail!("RocmQuant input not on rocm"),
                    };
                    let m = xs.elem_count() / *k;
                    let y = {
                        let (store, _) = xs.storage_and_layout();
                        let xr = match &*store {
                            crate::Storage::Rocm(r) => r,
                            _ => crate::bail!("RocmQuant expected rocm storage"),
                        };
                        #[cfg(feature = "rocm")]
                        {
                            d.qmmq_quant(qt, xr, wq, m, *n, *k)?
                        }
                        #[cfg(not(feature = "rocm"))]
                        {
                            let _ = qt;
                            crate::bail!("rocm feature disabled")
                        }
                    };
                    let mut dims = xs.dims().to_vec();
                    let last = dims.len() - 1;
                    dims[last] = *n;
                    Ok(crate::tensor::from_storage(
                        crate::Storage::Rocm(y),
                        dims,
                        crate::op::BackpropOp::none(),
                        false,
                    ))
                } else {
                    // Unwired-type prefill: dequantize to a temporary f16 weight (freed after; a
                    // persistent f16 copy would slow the memory-bound decode). Only reached for quants
                    // with no `qmmq_core<WTYPE>` wired (e.g. Q5_K/MXFP4).
                    let w = qtensor.dequantize_f16(&xs.device())?;
                    let w = match *xs.dims() {
                        [b1, b2, _, _] => w.broadcast_left((b1, b2))?.t()?,
                        [bsize, _, _] => w.broadcast_left(bsize)?.t()?,
                        _ => w.t()?,
                    };
                    xs.to_dtype(DType::F16)?.matmul(&w)
                }
            }
            #[cfg(feature = "vulkan")]
            Self::VulkanQuant {
                qtensor,
                wq,
                dtype,
                n,
                k,
            } => {
                // The activation can reach us on a different device than the weight -- e.g. a
                // CPU-resident prompt activation on the prefill path -- so place it on the weight's
                // Vulkan device rather than refusing. A no-op when it is already there, so the hot
                // decode path is unchanged; every sub-path below then sees a Vulkan input.
                let vdev = qtensor.device();
                let xs_on_vdev = if xs.device().same_device(&vdev) {
                    None
                } else {
                    Some(xs.to_device(&vdev)?)
                };
                let xs = xs_on_vdev.as_ref().unwrap_or(xs);
                let rows: usize = xs.elem_count() / *k;
                if rows == 1 {
                    // Decode: weights stay quantized in VRAM; the matching native-GGML quant matvec
                    // runs straight out of the block format (no dequant, no copy).
                    let xs = vulkan_act_offset0(xs)?;
                    let d = match xs.device() {
                        Device::Vulkan(d) => d,
                        _ => crate::bail!("VulkanQuant input not on vulkan"),
                    };
                    let y = {
                        let (store, _) = xs.storage_and_layout();
                        let xv = match &*store {
                            crate::Storage::Vulkan(v) => v,
                            _ => crate::bail!("VulkanQuant expected vulkan storage"),
                        };
                        match dtype {
                            GgmlDType::Q4_0 => d.matvec_q4_0_gpu(wq, xv, *n, *k)?,
                            // Q8_0 rides the 9-u32 repacked layout (mul_mat_vec_q8), the SAME blocks
                            // the prefill GEMM (mul_mat_q8) reads -- one layout for decode + prefill.
                            GgmlDType::Q8_0 => d.matvec_q8_gpu(wq, xv, *n, *k)?,
                            GgmlDType::Q4K => d.matvec_q4k_gpu(wq, xv, *n, *k)?,
                            GgmlDType::Q5K => d.matvec_q5k_gpu(wq, xv, *n, *k)?,
                            GgmlDType::Q6K => d.matvec_q6k_gpu(wq, xv, *n, *k)?,
                            GgmlDType::Q2K => d.matvec_q2k_gpu(wq, xv, *n, *k)?,
                            GgmlDType::Q3K => d.matvec_q3k_gpu(wq, xv, *n, *k)?,
                            GgmlDType::IQ4_XS => d.matvec_iq4xs_gpu(wq, xv, *n, *k)?,
                            GgmlDType::IQ4_NL => d.matvec_iq4nl_gpu(wq, xv, *n, *k)?,
                            GgmlDType::IQ2_XXS => d.matvec_iq2xxs_gpu(wq, xv, *n, *k)?,
                            GgmlDType::IQ2_XS => d.matvec_iq2xs_gpu(wq, xv, *n, *k)?,
                            GgmlDType::IQ1_M => d.matvec_iq1m_gpu(wq, xv, *n, *k)?,
                            GgmlDType::IQ1_S => d.matvec_iq1s_gpu(wq, xv, *n, *k)?,
                            GgmlDType::IQ3_S => d.matvec_iq3s_gpu(wq, xv, *n, *k)?,
                            GgmlDType::IQ3_XXS => d.matvec_iq3xxs_gpu(wq, xv, *n, *k)?,
                            GgmlDType::IQ2_S => d.matvec_iq2s_gpu(wq, xv, *n, *k)?,
                            GgmlDType::TQ2_0 => d.matvec_tq2_0_gpu(wq, xv, *n, *k)?,
                            other => crate::bail!("VulkanQuant: no native matvec for {other:?}"),
                        }
                    };
                    let mut dims = xs.dims().to_vec();
                    let last = dims.len() - 1;
                    dims[last] = *n;
                    Ok(crate::tensor::from_storage(
                        crate::Storage::Vulkan(y),
                        dims,
                        crate::op::BackpropOp::none(),
                        false,
                    ))
                } else if rows <= vulkan_prefill_gemm_max_rows(*dtype) {
                    // Prefill (small/moderate M): native quantized GEMM. Weights stay quantized in
                    // VRAM and each weight block is decoded ONCE per output column then reused across
                    // a tile of up to MATMUL_Q_MAX_M(=8) rows -- so the weight is re-read+re-decoded
                    // ceil(M/8) times, vs the dequant path's one-time f32 materialization. The GEMM
                    // therefore wins decisively while M is small (short / chunked prefill, batched
                    // decode) and would lose at large M, which the dtype-aware `rows` gate routes to
                    // the dequant path below. Same block layout + decode as the decode matvec above;
                    // one matmul_q*_gpu per native dtype. Leading batch dims flatten into M.
                    let m = rows;
                    let xs = vulkan_act_offset0(xs)?;
                    let d = match xs.device() {
                        Device::Vulkan(d) => d,
                        _ => crate::bail!("VulkanQuant input not on vulkan"),
                    };
                    let y = {
                        let (store, _) = xs.storage_and_layout();
                        let xv = match &*store {
                            crate::Storage::Vulkan(v) => v,
                            _ => crate::bail!("VulkanQuant expected vulkan storage"),
                        };
                        match dtype {
                            GgmlDType::Q4_0 => d.matmul_q4_0_gpu(wq, xv, m, *n, *k)?,
                            GgmlDType::Q8_0 => d.matmul_q8_gpu(wq, xv, m, *n, *k)?,
                            // Coopmat (tensor-core) Q4_K prefill GEMM. The mmq_q4k .spv bakes n/k, so
                            // only the committed shape routes here; the split bank is the 2D weight as
                            // a 1-expert MoE bank (cached), the activation q8_1-quantized once. Opt-in
                            // (VK_MMQ_Q4K) until measured against matmul_q4k_gpu with VK_PROFILE_GPU.
                            GgmlDType::Q4K
                                if *n == 2048
                                    && *k == 2048
                                    && std::env::var_os("VK_MMQ_Q4K").is_some() =>
                            {
                                let bank = qtensor.vulkan_moe_bank_split(d, 1, *n, *k)?;
                                let (xq, xsq, xsum) = d.quantize_act_q8(xv, m, *k)?;
                                d.mmq_q4k_gpu(&xq, &xsq, &xsum, bank.as_ref(), m, *n)?
                            }
                            GgmlDType::Q4K => d.matmul_q4k_gpu(wq, xv, m, *n, *k)?,
                            GgmlDType::Q5K => d.matmul_q5k_gpu(wq, xv, m, *n, *k)?,
                            GgmlDType::Q6K => d.matmul_q6k_gpu(wq, xv, m, *n, *k)?,
                            other => crate::bail!("VulkanQuant: no native matmul for {other:?}"),
                        }
                    };
                    let mut dims = xs.dims().to_vec();
                    let last = dims.len() - 1;
                    dims[last] = *n;
                    Ok(crate::tensor::from_storage(
                        crate::Storage::Vulkan(y),
                        dims,
                        crate::op::BackpropOp::none(),
                        false,
                    ))
                } else {
                    // Large dense prefill (M > the crossover): the column-per-invocation GEMM would
                    // re-read the weight ceil(M/8) times and lose to materializing the f32 weight once
                    // and running a dense matmul. Keep the dequant path here -- a strict non-regression
                    // until a shared-memory-tiled int8 GEMM (reads the weight once) removes the gate.
                    let w = qtensor.dequantize(&xs.device())?;
                    let w = match *xs.dims() {
                        [b1, b2, _, _] => w.broadcast_left((b1, b2))?.t()?,
                        [bsize, _, _] => w.broadcast_left(bsize)?.t()?,
                        _ => w.t()?,
                    };
                    xs.matmul(&w)
                }
            }
            #[cfg(feature = "wgpu")]
            Self::WgpuQuant {
                qtensor,
                wq,
                dtype,
                n,
                k,
            } => {
                // Same device coercion as the Vulkan arm: a CPU-resident prompt activation on the
                // prefill path is placed on the weight's Wgpu device rather than refused. A no-op
                // when it is already there, so the hot decode path is unchanged.
                let vdev = qtensor.device();
                let xs_on_vdev = if xs.device().same_device(&vdev) {
                    None
                } else {
                    Some(xs.to_device(&vdev)?)
                };
                let xs = xs_on_vdev.as_ref().unwrap_or(xs);
                let rows: usize = xs.elem_count() / *k;
                if rows == 1 {
                    // Decode: weights stay quantized in VRAM; the matching native-GGML quant matvec
                    // WGSL kernel runs straight out of the block format (no dequant, no copy).
                    let xs = xs.contiguous()?;
                    let d = match xs.device() {
                        Device::Wgpu(d) => d,
                        _ => crate::bail!("WgpuQuant input not on wgpu"),
                    };
                    let y = {
                        let (store, _) = xs.storage_and_layout();
                        let xv = match &*store {
                            crate::Storage::Wgpu(v) => v,
                            _ => crate::bail!("WgpuQuant expected wgpu storage"),
                        };
                        match dtype {
                            GgmlDType::Q4_0 => d.matvec_q4_0_gpu(wq, xv, *n, *k)?,
                            GgmlDType::Q8_0 => d.matvec_q8_0_gpu(wq, xv, *n, *k)?,
                            GgmlDType::Q4K => d.matvec_q4k_gpu(wq, xv, *n, *k)?,
                            other => crate::bail!("WgpuQuant: no native matvec for {other:?}"),
                        }
                    };
                    let mut dims = xs.dims().to_vec();
                    let last = dims.len() - 1;
                    dims[last] = *n;
                    Ok(crate::tensor::from_storage(
                        crate::Storage::Wgpu(y),
                        dims,
                        crate::op::BackpropOp::none(),
                        false,
                    ))
                } else {
                    // Prefill: dequantize to a temporary f32 weight (reuses the NT matmul path).
                    let w = qtensor.dequantize(&xs.device())?;
                    let w = match *xs.dims() {
                        [b1, b2, _, _] => w.broadcast_left((b1, b2))?.t()?,
                        [bsize, _, _] => w.broadcast_left(bsize)?.t()?,
                        _ => w.t()?,
                    };
                    xs.matmul(&w)
                }
            }
            Self::QTensor(t) => xs.apply_op1_no_bwd(t.as_ref()),
            Self::Tensor(w) => dense_matmul(xs, w),
            Self::TensorF16(w) => {
                let in_dtype = xs.dtype();
                dense_matmul(&xs.to_dtype(DType::F16)?, w)?.to_dtype(in_dtype)
            }
        }
    }
}