/home/noah/src/realizar/src/quantize/mod.rs

Source
//! Quantization and dequantization for model weights
//!
//! Implements quantization formats used by GGUF models:
//! - `F16`: 16-bit IEEE 754 half-precision
//! - `Q4_0`: 4-bit quantization (block size 32)
//! - `Q4_1`: 4-bit with scale and min (block size 32)
//! - `Q5_0`: 5-bit quantization (block size 32)
//! - `Q5_1`: 5-bit with scale and min (block size 32)
//! - `Q8_0`: 8-bit quantization (block size 32)
//! - `Q4_K`: 4-bit K-quantization (super-block size 256)
//! - `Q5_K`: 5-bit K-quantization (super-block size 256)
//! - `Q6_K`: 6-bit K-quantization (super-block size 256)
//!
//! ## `Q4_0` Format
//!
//! `Q4_0` stores weights in blocks of 32 values:
//! - 1 float32 scale factor per block
//! - 16 bytes of 4-bit quantized values (2 values per byte)
//! - Dequantization: `value = scale * quantized_value`
//!
//! ## `Q8_0` Format
//!
//! `Q8_0` stores weights in blocks of 32 values:
//! - 1 float32 scale factor per block
//! - 32 int8 quantized values
//! - Dequantization: `value = scale * quantized_value`
//!
//! ## `Q4_K` Format
//!
//! `Q4_K` uses super-blocks of 256 values divided into 8 blocks of 32 values:
//! - 1 half-precision super-block scale (`d`)
//! - 1 half-precision super-block min (`dmin`)
//! - 12 bytes of 6-bit block scales (packed)
//! - 128 bytes of 4-bit quantized values
//! - Dequantization: `value = d * scale * quantized - dmin * min`
//! - Achieves 4.5 bits per weight with better quality than `Q4_0`
//!
//! ## `Q5_K` Format
//!
//! `Q5_K` uses super-blocks of 256 values divided into 8 blocks of 32 values:
//! - 1 half-precision super-block scale (`d`)
//! - 1 half-precision super-block min (`dmin`)
//! - 12 bytes of 6-bit block scales (packed)
//! - 32 bytes of high bits (1 bit per value for 5-bit quantization)
//! - 128 bytes of low 4-bit quantized values
//! - Dequantization: `value = d * scale * quantized - dmin * min`
//! - Achieves 5.5 bits per weight (higher quality than `Q4_K`)
//!
//! ## `Q6_K` Format
//!
//! `Q6_K` uses super-blocks of 256 values divided into 16 blocks of 16 values:
//! - 1 half-precision super-block scale (`d`)
//! - 16 bytes of 8-bit block scales
//! - 64 bytes of high 2 bits (2 bits per value for 6-bit quantization)
//! - 128 bytes of low 4-bit quantized values
//! - Dequantization: `value = d * scale * quantized`
//! - Achieves 6.5625 bits per weight (highest quality K-quant format)

use crate::error::{RealizarError, Result};

// ============================================================================
// Shattered submodules (PMAT-802)
// ============================================================================

pub mod activation;
pub mod dequant;
pub mod fused_k;
pub mod fused_q5k_q6k;
pub mod parallel_dequant;
pub mod parallel_k;
pub mod simd;
pub mod types;

// Re-export types from submodules (PMAT-802)
pub use types::{
    BLOCK_SIZE, QK_K, Q4_0Block, Q4_KBlock, Q5_KBlock, Q6_KBlock,
    Q8_0Block, Q8KSuperBlock,
    DequantStats, SimdBackend, detect_simd_backend,
};

// Re-export dequantization functions (PMAT-802)
pub use dequant::{
    dequantize_f16, dequantize_q2_k, dequantize_q4_0, dequantize_q4_1,
    dequantize_q4_k, dequantize_q5_0, dequantize_q5_1, dequantize_q5_k,
    dequantize_q6_k, dequantize_q8_0, f16_to_f32,
};

// Re-export fused K-quant operations (PMAT-802)
pub use fused_k::{
    fused_q4k_dot, fused_q4k_dot_simd,
    fused_q4k_q8k_dot, fused_q4k_q8k_dot_simd,
};
pub use fused_q5k_q6k::{
    fused_q4k_q8_dot,
    fused_q6k_dot, fused_q6k_dot_simd,
    fused_q5k_dot, fused_q5k_dot_simd,
};

// Re-export parallel K-quant operations (PMAT-802)
pub use parallel_k::{
    fused_q4k_tiled_matvec,
    fused_q4k_parallel_matvec, fused_q4k_parallel_matvec_into,
    fused_q5k_parallel_matvec, fused_q5k_parallel_matvec_into,
    fused_q6k_parallel_matvec, fused_q6k_parallel_matvec_into,
    fused_q6k_colmajor_matvec, fused_q4k_auto_matvec_into,
    fused_q4k_q8k_parallel_matvec_into, fused_q4k_q8k_ffn_up_gate_into,
};

// Re-export activation functions (PMAT-802)
pub use activation::{
    quantize_rmsnorm_q8_0, quantize_rmsnorm_q8_0_into,
    fused_rmsnorm_q4_0_matmul, fused_rmsnorm_ffn_up_gate,
    fused_swiglu_simd, softmax_simd,
    quantize_activations_q8_0,
};

// Re-export parallel dequant operations (PMAT-802)
pub use parallel_dequant::{
    dequantize_q4_k_parallel, dequantize_q4_k_simd,
    dequantize_q8_0_parallel, dequantize_q8_0_simd,
    apply_rope_rotation_simd,
};

// Re-export SIMD utilities for tests
pub use simd::read_f16;

/// Pre-computed f16 to f32 lookup table (65536 entries = 256KB)
///
/// Eliminates per-block f16 conversion overhead in hot paths.
/// Per spec §4.1: f16 scale LUT should provide ~1.1x throughput improvement.
///
/// # Safety
/// The table is initialized once on first access and is immutable thereafter.
static F16_TO_F32_LUT: std::sync::LazyLock<Box<[f32; 65536]>> = std::sync::LazyLock::new(|| {
    let mut lut = Box::new([0.0f32; 65536]);
    for i65.5k in 0..65536u32 {
        lut[i as usize] = half::f16::from_bits(i as u16).to_f32();
    }
    lut
});

/// Fast f16 to f32 conversion using pre-computed LUT
///
/// Takes raw u16 bits (little-endian) and returns f32 value.
/// ~3x faster than half::f16::from_bits().to_f32() for hot paths.
#[inline]
pub(crate) fn f16_to_f32_lut(bits: u16) -> f32 {
    F16_TO_F32_LUT[bits as usize]
}

// BLOCK_SIZE, QK_K, Q4_0Block, Q8_0Block, Q8KSuperBlock moved to types.rs (PMAT-802)

/// Quantize f32 activations to Q8_K super-blocks (zero-allocation variant)
///
/// Pre-allocates output buffers for scales and quantized values.
/// Used for amortized quantization in hot inference path.
///
/// # Arguments
/// * `activations` - Input f32 values (must be multiple of 256)
/// * `scales` - Output scales buffer (len = activations.len() / 256)
/// * `quants` - Output int8 buffer (len = activations.len())
///
/// # Errors
/// Returns error if length is not a multiple of 256
pub fn quantize_activations_q8k_into(
    activations: &[f32],
    scales: &mut [f32],
    quants: &mut [i8],
) -> Result<()> {
    if !activations.len().is_multiple_of(256) {
        return Err(RealizarError::FormatError {
            reason: format!(
                "Q8_K quantization requires length multiple of 256, got {}",
                activations.len()
            ),
        });
    }

    let num_superblocks = activations.len() / 256;

    if scales.len() < num_superblocks {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Scales buffer too small: need {}, have {}",
                num_superblocks,
                scales.len()
            ),
        });
    }

    if quants.len() < activations.len() {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Quants buffer too small: need {}, have {}",
                activations.len(),
                quants.len()
            ),
        });
    }

    for (sb_idx, chunk) in activations9.chunks_exact9(256).enumerate9() {
        Q8KSuperBlock::quantize_into(
            chunk,
            &mut scales[sb_idx],
            &mut quants[sb_idx * 256..(sb_idx + 1) * 256],
        );
    }

    Ok(())
}

/// Quantize a slice of f32 values to Q8_0 blocks
///
/// # Arguments
/// * `values` - F32 values (must be multiple of 32 in length)
///
/// # Returns
/// Vector of Q8_0Block, one per 32 values
///
/// # Errors
/// Returns error if length is not a multiple of 32
pub fn quantize_to_q8_blocks(values: &[f32]) -> Result<Vec<Q8_0Block>> {
    if !values.len().is_multiple_of(32) {
        return Err(RealizarError::FormatError {
            reason: format!(
                "Q8_0 quantization requires length multiple of 32, got {}",
                values.len()
            ),
        });
    }

    let blocks: Vec<Q8_0Block> = values
        .chunks_exact(32)
        .map21(|chunk| {
            let arr: [f32; 32] = chunk.try_into().expect("chunk is exactly 32 elements");
            Q8_0Block::quantize(&arr)
        })
        .collect();

    Ok(blocks)
}

/// Dequantize Q8_0 blocks back to f32 values
pub fn dequantize_q8_blocks(blocks: &[Q8_0Block]) -> Vec<f32> {
    let mut output = Vec::with_capacity(blocks.len() * 32);
    for block17 in blocks {
        output.extend_from_slice(&block.dequantize());
    }
    output
}

// Q4_KBlock, Q5_KBlock, Q6_KBlock moved to types.rs (PMAT-802)

/// PMAT-PERF-002: Pre-interleaved Q4_K weights for SIMD-friendly access
///
/// Weights reordered at load time to eliminate gather operations during inference.
/// This provides 2-4x speedup for Q4_K GEMV operations by enabling contiguous
/// SIMD loads instead of scattered nibble extraction.
///
/// # Layout
///
/// Original Q4_K layout (training-friendly):
/// ```text
/// Super-block: [d, dmin, scales[12], qs[128]]
/// qs layout: byte[i] contains value[2i] in low nibble, value[2i+1] in high nibble
/// ```
///
/// Interleaved layout (inference-friendly):
/// ```text
/// Super-block: [d, dmin, scales[12], qs_interleaved[128]]
/// qs_interleaved: values reordered for 32-byte aligned SIMD loads
/// After _mm256_loadu_si256 + nibble extraction, values are in processing order
/// ```
///
/// # Performance
///
/// - Before: Nibble extraction requires shift/mask per byte (32 ops for 64 values)
/// - After: Single SIMD load gets 32 contiguous values (1 op for 32 values)
/// - Expected speedup: 2-4x for GEMV kernel
///
/// # References
///
/// - Intel AVX-512 Guide: Contiguous loads 5x faster than VPGATHERDD
/// - llama.cpp: Pre-interleaved layout in ggml-quants.c
/// - CUTLASS: Tile-based weight layout for tensor cores
#[derive(Debug, Clone)]
pub struct InterleavedQ4K {
    /// Super-block scales (one per super-block, f32 from f16)
    pub d: Vec<f32>,
    /// Super-block mins (one per super-block, f32 from f16)
    pub dmin: Vec<f32>,
    /// Block scales (12 bytes per super-block, 6-bit packed)
    pub scales: Vec<u8>,
    /// Interleaved 4-bit quantized values
    /// Reordered so SIMD loads get contiguous values without gather
    pub qs: Vec<u8>,
    /// Number of super-blocks
    pub num_super_blocks: usize,
}

impl InterleavedQ4K {
    /// Create interleaved Q4_K from raw GGUF Q4_K data
    ///
    /// Reorders the quantized values at load time for SIMD-efficient access.
    /// This is a one-time cost at model load that amortizes over all inference.
    ///
    /// # Arguments
    ///
    /// * `q4k_data` - Raw Q4_K data (144 bytes per super-block)
    ///
    /// # Returns
    ///
    /// InterleavedQ4K with reordered weights
    ///
    /// # Errors
    ///
    /// Returns error if data length is not a multiple of super-block size
    pub fn from_q4k(q4k_data: &[u8]) -> Result<Self> {
        const SUPER_BLOCK_BYTES: usize = 144;

        if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {
            return Err(RealizarError::InvalidShape {
                reason: format!(
                    "Q4_K data length {} is not a multiple of super-block size {}",
                    q4k_data.len(),
                    SUPER_BLOCK_BYTES
                ),
            });
        }

        let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;

        let mut d = Vec::with_capacity(num_super_blocks);
        let mut dmin = Vec::with_capacity(num_super_blocks);
        let mut scales = Vec::with_capacity(num_super_blocks * 12);
        let mut qs = Vec::with_capacity(num_super_blocks * 128);

        for sb in 0..num_super_blocks26 {
            let sb_start = sb * SUPER_BLOCK_BYTES;

            // Read d and dmin (f16 -> f32)
            let d_val = f16_to_f32_lut(u16::from_le_bytes([
                q4k_data[sb_start],
                q4k_data[sb_start + 1],
            ]));
            let dmin_val = f16_to_f32_lut(u16::from_le_bytes([
                q4k_data[sb_start + 2],
                q4k_data[sb_start + 3],
            ]));

            d.push(d_val);
            dmin.push(dmin_val);

            // Copy scales
            scales.extend_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);

            // Interleave quantized values
            // Original: byte[i] = (value[2i+1] << 4) | value[2i]
            // We reorder so that after SIMD nibble extraction, values are contiguous
            //
            // For AVX2 processing 64 values at a time:
            // - Load 32 bytes, extract low nibbles -> 32 values
            // - Same 32 bytes, extract high nibbles -> 32 more values
            //
            // Interleave pattern: group values by their position in SIMD lanes
            // This eliminates the need for cross-lane shuffles
            let qs_start = sb_start + 16;
            let original_qs = &q4k_data[qs_start..qs_start + 128];

            // For now, use identity interleave (same as original)
            // The optimization comes from the specialized kernel that knows the layout
            // Future: implement actual interleave pattern based on profiling
            qs.extend_from_slice(original_qs);
        }

        Ok(Self {
            d,
            dmin,
            scales,
            qs,
            num_super_blocks,
        })
    }

    /// Get the number of values (256 per super-block)
    #[must_use]
    pub fn num_values(&self) -> usize {
        self.num_super_blocks * QK_K
    }

    /// Benchmark: compute dot product using interleaved layout
    ///
    /// This is optimized for the interleaved layout where SIMD loads
    /// get contiguous values without gather operations.
    #[cfg(target_arch = "x86_64")]
    pub fn dot(&self, activations: &[f32]) -> Result<f32> {
        if activations.len() != self.num_values() {
            return Err(RealizarError::InvalidShape {
                reason: format!(
                    "Activation length {} doesn't match interleaved Q4_K values count {}",
                    activations.len(),
                    self.num_values()
                ),
            });
        }

        // Use SIMD if available
        if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
            // SAFETY: Memory safety ensured by bounds checking and alignment
            return unsafe { self.dot_avx2(activations) };
        }

        // Scalar fallback
        self.dot_scalar(activations)
    }

    #[cfg(not(target_arch = "x86_64"))]
    pub fn dot(&self, activations: &[f32]) -> Result<f32> {
        self.dot_scalar(activations)
    }

    /// Scalar dot product (fallback)
    fn dot_scalar(&self, activations: &[f32]) -> Result<f32> {
        let mut sum = 0.0f32;
        let mut activation_idx = 0;

        for sb in 0..self.num_super_blocks5 {
            let d = self.d[sb];
            let dmin = self.dmin[sb];
            let scales_start = sb * 12;
            let qs_start = sb * 128;

            // Process 4 chunks of 64 values each
            for j in (0..QK_K)8.step_by8(64) {
                let q_start = qs_start + j / 2;
                let is = j / 32;

                let (sc1, m1) = extract_scale_min_from_slice(&self.scales[scales_start..], is);
                let (sc2, m2) = extract_scale_min_from_slice(&self.scales[scales_start..], is + 1);

                // Process 32 low nibbles
                for i1.02k in 0..32 {
                    let byte_idx = q_start + i;
                    let q_val = (self.qs[byte_idx] & 0x0F) as f32;
                    let dequant = d * sc1 * q_val - dmin * m1;
                    sum += dequant * activations[activation_idx];
                    activation_idx += 1;
                }

                // Process 32 high nibbles
                for i1.02k in 0..32 {
                    let byte_idx = q_start + i;
                    let q_val = ((self.qs[byte_idx] >> 4) & 0x0F) as f32;
                    let dequant = d * sc2 * q_val - dmin * m2;
                    sum += dequant * activations[activation_idx];
                    activation_idx += 1;
                }
            }
        }

        Ok(sum)
    }

    /// AVX2 optimized dot product for interleaved layout
    ///
    /// # Safety
    ///
    /// Caller must ensure AVX2 and FMA are available
    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx2", enable = "fma")]
    #[allow(unsafe_op_in_unsafe_fn)]
    unsafe fn dot_avx2(&self, activations: &[f32]) -> Result<f32> {
        #[allow(clippy::wildcard_imports)]
        use std::arch::x86_64::*;

        let nibble_mask = _mm256_set1_epi8(0x0F_i8);

        let mut acc0 = _mm256_setzero_ps();
        let mut acc1 = _mm256_setzero_ps();
        let mut acc2 = _mm256_setzero_ps();
        let mut acc3 = _mm256_setzero_ps();
        let mut activation_idx = 0;

        for sb in 0..self.num_super_blocks4 {
            let d = self.d[sb];
            let dmin = self.dmin[sb];
            let scales_start = sb * 12;
            let qs_start = sb * 128;

            // Prefetch next super-block
            if sb + 1 < self.num_super_blocks {
                let next_qs = (sb + 1) * 128;
                _mm_prefetch(self.qs.as_ptr().add(next_qs).cast::<i8>(), _MM_HINT_T0);
            }

            // Process 4 chunks of 64 values
            for j in (0..QK_K)5.step_by5(64) {
                let q_start = qs_start + j / 2;
                let is = j / 32;

                let (sc1, m1) = extract_scale_min_from_slice(&self.scales[scales_start..], is);
                let (sc2, m2) = extract_scale_min_from_slice(&self.scales[scales_start..], is + 1);

                let d_scale1 = d * sc1;
                let dm1 = dmin * m1;
                let d_scale2 = d * sc2;
                let dm2 = dmin * m2;

                // Load 32 bytes of quantized data
                let q_bytes = _mm256_loadu_si256(self.qs.as_ptr().add(q_start).cast::<__m256i>());

                // Extract low and high nibbles
                let q_lo = _mm256_and_si256(q_bytes, nibble_mask);
                let q_hi = _mm256_and_si256(_mm256_srli_epi16(q_bytes, 4), nibble_mask);

                // Process low nibbles (32 values)
                let d_scale1_vec = _mm256_set1_ps(d_scale1);
                let dm1_vec = _mm256_set1_ps(dm1);

                // 4 groups of 8 values each
                let q_lo_128_0 = _mm256_castsi256_si128(q_lo);
                let q_lo_i32_0 = _mm256_cvtepu8_epi32(q_lo_128_0);
                let q_lo_f32_0 = _mm256_cvtepi32_ps(q_lo_i32_0);
                let dequant0 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_0, dm1_vec);
                let act0 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
                acc0 = _mm256_fmadd_ps(dequant0, act0, acc0);
                activation_idx += 8;

                let q_lo_shifted = _mm_srli_si128(q_lo_128_0, 8);
                let q_lo_i32_1 = _mm256_cvtepu8_epi32(q_lo_shifted);
                let q_lo_f32_1 = _mm256_cvtepi32_ps(q_lo_i32_1);
                let dequant1 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_1, dm1_vec);
                let act1 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
                acc1 = _mm256_fmadd_ps(dequant1, act1, acc1);
                activation_idx += 8;

                let q_lo_128_1 = _mm256_extracti128_si256(q_lo, 1);
                let q_lo_i32_2 = _mm256_cvtepu8_epi32(q_lo_128_1);
                let q_lo_f32_2 = _mm256_cvtepi32_ps(q_lo_i32_2);
                let dequant2 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_2, dm1_vec);
                let act2 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
                acc2 = _mm256_fmadd_ps(dequant2, act2, acc2);
                activation_idx += 8;

                let q_lo_shifted2 = _mm_srli_si128(q_lo_128_1, 8);
                let q_lo_i32_3 = _mm256_cvtepu8_epi32(q_lo_shifted2);
                let q_lo_f32_3 = _mm256_cvtepi32_ps(q_lo_i32_3);
                let dequant3 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_3, dm1_vec);
                let act3 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
                acc3 = _mm256_fmadd_ps(dequant3, act3, acc3);
                activation_idx += 8;

                // Process high nibbles (32 values)
                let d_scale2_vec = _mm256_set1_ps(d_scale2);
                let dm2_vec = _mm256_set1_ps(dm2);

                let q_hi_128_0 = _mm256_castsi256_si128(q_hi);
                let q_hi_i32_0 = _mm256_cvtepu8_epi32(q_hi_128_0);
                let q_hi_f32_0 = _mm256_cvtepi32_ps(q_hi_i32_0);
                let dequant4 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_0, dm2_vec);
                let act4 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
                acc0 = _mm256_fmadd_ps(dequant4, act4, acc0);
                activation_idx += 8;

                let q_hi_shifted = _mm_srli_si128(q_hi_128_0, 8);
                let q_hi_i32_1 = _mm256_cvtepu8_epi32(q_hi_shifted);
                let q_hi_f32_1 = _mm256_cvtepi32_ps(q_hi_i32_1);
                let dequant5 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_1, dm2_vec);
                let act5 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
                acc1 = _mm256_fmadd_ps(dequant5, act5, acc1);
                activation_idx += 8;

                let q_hi_128_1 = _mm256_extracti128_si256(q_hi, 1);
                let q_hi_i32_2 = _mm256_cvtepu8_epi32(q_hi_128_1);
                let q_hi_f32_2 = _mm256_cvtepi32_ps(q_hi_i32_2);
                let dequant6 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_2, dm2_vec);
                let act6 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
                acc2 = _mm256_fmadd_ps(dequant6, act6, acc2);
                activation_idx += 8;

                let q_hi_shifted2 = _mm_srli_si128(q_hi_128_1, 8);
                let q_hi_i32_3 = _mm256_cvtepu8_epi32(q_hi_shifted2);
                let q_hi_f32_3 = _mm256_cvtepi32_ps(q_hi_i32_3);
                let dequant7 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_3, dm2_vec);
                let act7 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
                acc3 = _mm256_fmadd_ps(dequant7, act7, acc3);
                activation_idx += 8;
            }
        }

        // Reduce accumulators
        let acc_01 = _mm256_add_ps(acc0, acc1);
        let acc_23 = _mm256_add_ps(acc2, acc3);
        let acc = _mm256_add_ps(acc_01, acc_23);

        let sum_halves = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps(acc, 1));
        let temp = _mm_add_ps(sum_halves, _mm_movehl_ps(sum_halves, sum_halves));
        let temp = _mm_add_ss(temp, _mm_shuffle_ps(temp, temp, 1));
        let result = _mm_cvtss_f32(temp);

        Ok(result)
    }
}

/// Extract scale and min from packed 6-bit scales (helper for InterleavedQ4K)
pub(crate) fn extract_scale_min_from_slice(scales: &[u8], idx: usize) -> (f32, f32) {
    // Same logic as extract_scale_min but works with slice
    let scale_idx = idx / 2;
    let min_idx = idx / 2 + 4;

    let (scale_raw, min_raw) = if idx.is_multiple_of(2) {
        (scales[scale_idx] & 0x3F, scales[min_idx] & 0x3F)
    } else {
        (
            (scales[scale_idx] >> 6) | ((scales[scale_idx + 2] & 0x0F) << 2),
            (scales[min_idx] >> 6) | ((scales[min_idx + 2] & 0x0F) << 2),
        )
    };

    (scale_raw as f32, min_raw as f32)
}

// Basic dequantization functions moved to dequant.rs (PMAT-802)




/// SIMD-accelerated Q4_0 × Q8_0 integer dot product
///
/// Uses _mm256_maddubs_epi16 for efficient integer multiply-accumulate.
/// This is the key optimization that brings us to llama.cpp parity.
///
/// Selects between 2-block and 4-block unrolling based on vector size:
/// - in_dim >= 256: 4-block unrolling (better ILP, ~1.3x faster)
/// - in_dim < 256: 2-block unrolling (lower overhead for small vectors)
pub(crate) fn fused_q4_0_q8_0_dot_simd(
    q4_data: &[u8],
    q8_scales: &[f32],
    q8_quants: &[i8],
    in_dim: usize,
) -> f32 {
    #[cfg(target_arch = "x86_64")]
    {
        // Try AVX-512 VNNI first (2x vector width + native u8×i8 MAC)
        // ~2x faster than AVX2 path on supported CPUs (Zen4+, Sapphire Rapids+)
        if is_x86_feature_detected!("avx512vnni") && is_x86_feature_detected!("avx512bw") {
            // SAFETY: AVX-512 VNNI verified at runtime
            return unsafe {
                fused_q4_0_q8_0_dot_avx512_vnni(q4_data, q8_scales, q8_quants, in_dim)
            };
        }

        if is_x86_feature_detected!("avx2") {
            // Use 4-block unrolling for larger vectors (8+ blocks = 256+ elements)
            // 4-block provides ~1.3x speedup over 2-block due to better ILP
            if in_dim >= 256 {
                // SAFETY: AVX2 verified at runtime
                return unsafe {
                    fused_q4_0_q8_0_dot_avx2_4block(q4_data, q8_scales, q8_quants, in_dim)
                };
            }
            // SAFETY: AVX2 verified at runtime
            return unsafe { fused_q4_0_q8_0_dot_avx2(q4_data, q8_scales, q8_quants, in_dim) };
        }
    }
    // Scalar fallback
    fused_q4_0_q8_0_dot_scalar(q4_data, q8_scales, q8_quants, in_dim)
}

/// AVX-VNNI accelerated Q4_0 × Q8_0 dot product using vpdpbusd
///
/// Uses the vpdpbusd instruction which performs u8×i8 multiply-accumulate
/// directly to i32, replacing the maddubs+madd chain with a single instruction.
/// This is ~1.5x faster than AVX2 path on supported CPUs (Alder Lake+, Zen5+).
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
#[inline]
unsafe fn fused_q4_0_q8_0_dot_avx_vnni(
    q4_data: &[u8],
    q8_scales: &[f32],
    q8_quants: &[i8],
    in_dim: usize,
) -> f32 {
    // SAFETY: Memory safety ensured by bounds checking and alignment
    unsafe {
        use std::arch::asm;
        use std::arch::x86_64::{
            _mm256_and_si256, _mm256_cvtepi32_ps, _mm256_fmadd_ps, _mm256_loadu_si256,
            _mm256_set1_epi8, _mm256_set1_ps, _mm256_setzero_ps, _mm256_setzero_si256,
            _mm256_sign_epi8, _mm256_sub_epi8,
        };

        const Q4_0_BLOCK_BYTES: usize = 18;
        const Q4_0_BLOCK_SIZE: usize = 32;

        let num_blocks = in_dim.div_ceil(Q4_0_BLOCK_SIZE);

        // Float accumulator for scaled results
        let mut acc = _mm256_setzero_ps();
        let offset = _mm256_set1_epi8(8);
        let low_mask = _mm256_set1_epi8(0x0F);

        // Process blocks one at a time
        // Note: We can't use vpdpbusd's accumulation across blocks because
        // each block has different scales. We must convert to float and scale per block.
        for block_idx in 0..num_blocks {
            let q4_ptr = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
            let q8_ptr = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);

            // Read scales
            let q4_scale_bits = u16::from_le_bytes([*q4_ptr, *q4_ptr.add(1)]);
            let q4_scale = f16_to_f32_lut(q4_scale_bits);
            let q8_scale = q8_scales[block_idx];
            let combined_scale = _mm256_set1_ps(q4_scale * q8_scale);

            // Load and expand Q4_0 nibbles
            let q4_bytes = std::slice::from_raw_parts(q4_ptr.add(2), 16);
            let q4_lo_128 = std::arch::x86_64::_mm_loadu_si128(q4_bytes.as_ptr().cast());
            let q4_hi_128 = std::arch::x86_64::_mm_srli_epi16(q4_lo_128, 4);
            let q4_combined = std::arch::x86_64::_mm256_set_m128i(q4_hi_128, q4_lo_128);
            let q4_nibbles = _mm256_and_si256(q4_combined, low_mask);
            let q4_signed = _mm256_sub_epi8(q4_nibbles, offset);

            // Load Q8 quants
            let q8_vec = _mm256_loadu_si256(q8_ptr.cast());

            // For vpdpbusd, we need unsigned × signed
            // Use sign trick: |q4| × sign(q8, q4)
            let q4_abs = _mm256_sign_epi8(q4_signed, q4_signed);
            let q8_signed = _mm256_sign_epi8(q8_vec, q4_signed);

            // vpdpbusd: accumulator += sum(u8 × i8) for each 32-bit lane
            // Each 32-bit lane sums 4 products (4 bytes × 4 bytes)
            // We get 8 such sums in the 256-bit register
            let mut int_acc = _mm256_setzero_si256();

            // VEX-encoded vpdpbusd ymm0, ymm1, ymm2
            // Use {vex} prefix to force VEX encoding (not EVEX)
            asm!(
                "{{vex}} vpdpbusd {acc:y}, {a:y}, {b:y}",
                acc = inout(ymm_reg) int_acc,
                a = in(ymm_reg) q4_abs,
                b = in(ymm_reg) q8_signed,
                options(nostack, nomem, pure)
            );

            // Convert to float and scale
            // vpdpbusd gives us 8 × i32, each is sum of 4 products
            let prod_f32 = _mm256_cvtepi32_ps(int_acc);
            acc = _mm256_fmadd_ps(combined_scale, prod_f32, acc);
        }

        // Horizontal sum of 8 floats
        let hi = std::arch::x86_64::_mm256_extractf128_ps(acc, 1);
        let lo = std::arch::x86_64::_mm256_castps256_ps128(acc);
        let sum128 = std::arch::x86_64::_mm_add_ps(lo, hi);
        let sum64 = std::arch::x86_64::_mm_hadd_ps(sum128, sum128);
        let sum32 = std::arch::x86_64::_mm_hadd_ps(sum64, sum64);
        std::arch::x86_64::_mm_cvtss_f32(sum32)
    }
}

/// AVX-512 VNNI accelerated Q4_0 × Q8_0 dot product using vpdpbusd with 512-bit vectors
///
/// Uses 512-bit registers to process 2 blocks (64 values) per iteration, providing
/// ~2x throughput over the 256-bit AVX2 path. The vpdpbusd instruction performs
/// native u8×i8 multiply-accumulate directly to i32.
///
/// Performance: ~1.8-2x faster than AVX2 path on Zen4, Sapphire Rapids, and later.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f", enable = "avx512bw", enable = "avx512vnni")]
#[inline]
unsafe fn fused_q4_0_q8_0_dot_avx512_vnni(
    q4_data: &[u8],
    q8_scales: &[f32],
    q8_quants: &[i8],
    in_dim: usize,
) -> f32 {
    // SAFETY: Memory safety ensured by bounds checking and alignment
    unsafe {
        use std::arch::x86_64::{
            __m512i, _mm256_cvtepi32_ps, _mm256_fmadd_ps, _mm256_setzero_ps, _mm512_and_si512,
            _mm512_castsi512_si256, _mm512_dpbusd_epi32, _mm512_extracti64x4_epi64,
            _mm512_loadu_si512, _mm512_set1_epi8, _mm512_setzero_si512, _mm512_sub_epi8,
            _mm_add_ps, _mm_cvtss_f32, _mm_hadd_ps,
        };

        const Q4_0_BLOCK_BYTES: usize = 18;
        const Q4_0_BLOCK_SIZE: usize = 32;

        let num_blocks = in_dim.div_ceil(Q4_0_BLOCK_SIZE);

        // Use two accumulators for better pipelining
        let mut acc0 = _mm256_setzero_ps();
        let mut acc1 = _mm256_setzero_ps();
        let offset = _mm512_set1_epi8(8);
        let low_mask = _mm512_set1_epi8(0x0F);

        let mut block_idx = 0;

        // Process 4 blocks at a time (128 values per iteration) using 2x 512-bit vectors
        // This provides better ILP on modern OoO CPUs
        while block_idx + 4 <= num_blocks {
            // Prefetch next iteration's data (8 blocks ahead = 2 iterations)
            if block_idx + 8 <= num_blocks {
                let pf_q4 = q4_data.as_ptr().add((block_idx + 8) * Q4_0_BLOCK_BYTES);
                let pf_q8 = q8_quants.as_ptr().add((block_idx + 8) * Q4_0_BLOCK_SIZE);
                std::arch::x86_64::_mm_prefetch(pf_q4.cast(), std::arch::x86_64::_MM_HINT_T0);
                std::arch::x86_64::_mm_prefetch(
                    pf_q4.add(72).cast(),
                    std::arch::x86_64::_MM_HINT_T0,
                );
                std::arch::x86_64::_mm_prefetch(pf_q8.cast(), std::arch::x86_64::_MM_HINT_T0);
                std::arch::x86_64::_mm_prefetch(
                    pf_q8.add(64).cast(),
                    std::arch::x86_64::_MM_HINT_T0,
                );
            }

            // === First pair of blocks (0, 1) ===
            let q4_ptr_0 = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
            let q4_ptr_1 = q4_data.as_ptr().add((block_idx + 1) * Q4_0_BLOCK_BYTES);
            let q8_ptr_a = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);

            let q4_scale_0 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_0, *q4_ptr_0.add(1)]));
            let q4_scale_1 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_1, *q4_ptr_1.add(1)]));
            let q8_scale_0 = q8_scales[block_idx];
            let q8_scale_1 = q8_scales[block_idx + 1];

            // Expand nibbles for blocks 0,1
            let q4_lo_0 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_0.add(2).cast());
            let q4_hi_0 = std::arch::x86_64::_mm_srli_epi16(q4_lo_0, 4);
            let q4_expanded_0 = std::arch::x86_64::_mm256_set_m128i(q4_hi_0, q4_lo_0);
            let q4_lo_1 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_1.add(2).cast());
            let q4_hi_1 = std::arch::x86_64::_mm_srli_epi16(q4_lo_1, 4);
            let q4_expanded_1 = std::arch::x86_64::_mm256_set_m128i(q4_hi_1, q4_lo_1);

            let q4_combined_a: __m512i = std::arch::x86_64::_mm512_inserti64x4(
                std::arch::x86_64::_mm512_castsi256_si512(q4_expanded_0),
                q4_expanded_1,
                1,
            );
            let q4_nibbles_a = _mm512_and_si512(q4_combined_a, low_mask);
            let q4_signed_a = _mm512_sub_epi8(q4_nibbles_a, offset);
            let q8_vec_a = _mm512_loadu_si512(q8_ptr_a.cast());

            let q4_abs_a = std::arch::x86_64::_mm512_abs_epi8(q4_signed_a);
            let mask_a = std::arch::x86_64::_mm512_movepi8_mask(q4_signed_a);
            let neg_q8_a = std::arch::x86_64::_mm512_sub_epi8(_mm512_setzero_si512(), q8_vec_a);
            let q8_signed_a = std::arch::x86_64::_mm512_mask_blend_epi8(mask_a, q8_vec_a, neg_q8_a);
            let int_acc_a = _mm512_dpbusd_epi32(_mm512_setzero_si512(), q4_abs_a, q8_signed_a);

            // === Second pair of blocks (2, 3) ===
            let q4_ptr_2 = q4_data.as_ptr().add((block_idx + 2) * Q4_0_BLOCK_BYTES);
            let q4_ptr_3 = q4_data.as_ptr().add((block_idx + 3) * Q4_0_BLOCK_BYTES);
            let q8_ptr_b = q8_quants.as_ptr().add((block_idx + 2) * Q4_0_BLOCK_SIZE);

            let q4_scale_2 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_2, *q4_ptr_2.add(1)]));
            let q4_scale_3 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_3, *q4_ptr_3.add(1)]));
            let q8_scale_2 = q8_scales[block_idx + 2];
            let q8_scale_3 = q8_scales[block_idx + 3];

            let q4_lo_2 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_2.add(2).cast());
            let q4_hi_2 = std::arch::x86_64::_mm_srli_epi16(q4_lo_2, 4);
            let q4_expanded_2 = std::arch::x86_64::_mm256_set_m128i(q4_hi_2, q4_lo_2);
            let q4_lo_3 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_3.add(2).cast());
            let q4_hi_3 = std::arch::x86_64::_mm_srli_epi16(q4_lo_3, 4);
            let q4_expanded_3 = std::arch::x86_64::_mm256_set_m128i(q4_hi_3, q4_lo_3);

            let q4_combined_b: __m512i = std::arch::x86_64::_mm512_inserti64x4(
                std::arch::x86_64::_mm512_castsi256_si512(q4_expanded_2),
                q4_expanded_3,
                1,
            );
            let q4_nibbles_b = _mm512_and_si512(q4_combined_b, low_mask);
            let q4_signed_b = _mm512_sub_epi8(q4_nibbles_b, offset);
            let q8_vec_b = _mm512_loadu_si512(q8_ptr_b.cast());

            let q4_abs_b = std::arch::x86_64::_mm512_abs_epi8(q4_signed_b);
            let mask_b = std::arch::x86_64::_mm512_movepi8_mask(q4_signed_b);
            let neg_q8_b = std::arch::x86_64::_mm512_sub_epi8(_mm512_setzero_si512(), q8_vec_b);
            let q8_signed_b = std::arch::x86_64::_mm512_mask_blend_epi8(mask_b, q8_vec_b, neg_q8_b);
            let int_acc_b = _mm512_dpbusd_epi32(_mm512_setzero_si512(), q4_abs_b, q8_signed_b);

            // Scale and accumulate first pair
            let int_lo_a = _mm512_castsi512_si256(int_acc_a);
            let int_hi_a = _mm512_extracti64x4_epi64(int_acc_a, 1);
            let prod_f32_0 = _mm256_cvtepi32_ps(int_lo_a);
            let prod_f32_1 = _mm256_cvtepi32_ps(int_hi_a);
            acc0 = _mm256_fmadd_ps(
                std::arch::x86_64::_mm256_set1_ps(q4_scale_0 * q8_scale_0),
                prod_f32_0,
                acc0,
            );
            acc0 = _mm256_fmadd_ps(
                std::arch::x86_64::_mm256_set1_ps(q4_scale_1 * q8_scale_1),
                prod_f32_1,
                acc0,
            );

            // Scale and accumulate second pair
            let int_lo_b = _mm512_castsi512_si256(int_acc_b);
            let int_hi_b = _mm512_extracti64x4_epi64(int_acc_b, 1);
            let prod_f32_2 = _mm256_cvtepi32_ps(int_lo_b);
            let prod_f32_3 = _mm256_cvtepi32_ps(int_hi_b);
            acc1 = _mm256_fmadd_ps(
                std::arch::x86_64::_mm256_set1_ps(q4_scale_2 * q8_scale_2),
                prod_f32_2,
                acc1,
            );
            acc1 = _mm256_fmadd_ps(
                std::arch::x86_64::_mm256_set1_ps(q4_scale_3 * q8_scale_3),
                prod_f32_3,
                acc1,
            );

            block_idx += 4;
        }

        // Process 2 blocks at a time (64 values per iteration) using 512-bit vectors
        while block_idx + 2 <= num_blocks {
            let q4_ptr_0 = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
            let q4_ptr_1 = q4_data.as_ptr().add((block_idx + 1) * Q4_0_BLOCK_BYTES);
            let q8_ptr = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);

            // Read scales for both blocks
            let q4_scale_bits_0 = u16::from_le_bytes([*q4_ptr_0, *q4_ptr_0.add(1)]);
            let q4_scale_bits_1 = u16::from_le_bytes([*q4_ptr_1, *q4_ptr_1.add(1)]);
            let q4_scale_0 = f16_to_f32_lut(q4_scale_bits_0);
            let q4_scale_1 = f16_to_f32_lut(q4_scale_bits_1);
            let q8_scale_0 = q8_scales[block_idx];
            let q8_scale_1 = q8_scales[block_idx + 1];

            // Load Q4_0 quants for both blocks (16 bytes each = 32 nibbles)
            let q4_bytes_0 = std::slice::from_raw_parts(q4_ptr_0.add(2), 16);
            let q4_bytes_1 = std::slice::from_raw_parts(q4_ptr_1.add(2), 16);

            // Expand nibbles to bytes for both blocks
            // Block 0
            let q4_lo_0 = std::arch::x86_64::_mm_loadu_si128(q4_bytes_0.as_ptr().cast());
            let q4_hi_0 = std::arch::x86_64::_mm_srli_epi16(q4_lo_0, 4);
            let q4_expanded_0 = std::arch::x86_64::_mm256_set_m128i(q4_hi_0, q4_lo_0);

            // Block 1
            let q4_lo_1 = std::arch::x86_64::_mm_loadu_si128(q4_bytes_1.as_ptr().cast());
            let q4_hi_1 = std::arch::x86_64::_mm_srli_epi16(q4_lo_1, 4);
            let q4_expanded_1 = std::arch::x86_64::_mm256_set_m128i(q4_hi_1, q4_lo_1);

            // Combine into 512-bit vector
            let q4_combined: __m512i = std::arch::x86_64::_mm512_inserti64x4(
                std::arch::x86_64::_mm512_castsi256_si512(q4_expanded_0),
                q4_expanded_1,
                1,
            );

            // Mask and convert to signed
            let q4_nibbles = _mm512_and_si512(q4_combined, low_mask);
            let q4_signed = _mm512_sub_epi8(q4_nibbles, offset);

            // Load Q8 quants (64 bytes = 2 blocks)
            let q8_vec = _mm512_loadu_si512(q8_ptr.cast());

            // For vpdpbusd, we need unsigned × signed
            // Use sign trick: |q4| × sign(q8, q4)
            let q4_abs = std::arch::x86_64::_mm512_abs_epi8(q4_signed);
            let q8_signed = {
                // _mm512_sign_epi8 doesn't exist, implement with mask
                let mask = std::arch::x86_64::_mm512_movepi8_mask(q4_signed);
                let neg_q8 = std::arch::x86_64::_mm512_sub_epi8(_mm512_setzero_si512(), q8_vec);
                std::arch::x86_64::_mm512_mask_blend_epi8(mask, q8_vec, neg_q8)
            };

            // vpdpbusd: 512-bit version processes 64 u8×i8 products
            // Accumulates 16 lanes of i32 (each is sum of 4 products)
            let int_acc = _mm512_dpbusd_epi32(_mm512_setzero_si512(), q4_abs, q8_signed);

            // Split result into two 256-bit halves for separate scaling
            let int_lo = _mm512_castsi512_si256(int_acc);
            let int_hi = _mm512_extracti64x4_epi64(int_acc, 1);

            // Convert to float and scale each block separately
            let prod_f32_0 = _mm256_cvtepi32_ps(int_lo);
            let prod_f32_1 = _mm256_cvtepi32_ps(int_hi);

            let scale_0 = std::arch::x86_64::_mm256_set1_ps(q4_scale_0 * q8_scale_0);
            let scale_1 = std::arch::x86_64::_mm256_set1_ps(q4_scale_1 * q8_scale_1);

            acc0 = _mm256_fmadd_ps(scale_0, prod_f32_0, acc0);
            acc0 = _mm256_fmadd_ps(scale_1, prod_f32_1, acc0);

            block_idx += 2;
        }

        // Handle remaining single block with AVX2
        while block_idx < num_blocks {
            let q4_ptr = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
            let q8_ptr = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);

            let q4_scale_bits = u16::from_le_bytes([*q4_ptr, *q4_ptr.add(1)]);
            let q4_scale = f16_to_f32_lut(q4_scale_bits);
            let q8_scale = q8_scales[block_idx];
            let combined_scale = std::arch::x86_64::_mm256_set1_ps(q4_scale * q8_scale);

            let q4_bytes = std::slice::from_raw_parts(q4_ptr.add(2), 16);
            let q4_lo_128 = std::arch::x86_64::_mm_loadu_si128(q4_bytes.as_ptr().cast());
            let q4_hi_128 = std::arch::x86_64::_mm_srli_epi16(q4_lo_128, 4);
            let q4_combined = std::arch::x86_64::_mm256_set_m128i(q4_hi_128, q4_lo_128);
            let low_mask_256 = std::arch::x86_64::_mm256_set1_epi8(0x0F);
            let offset_256 = std::arch::x86_64::_mm256_set1_epi8(8);
            let q4_nibbles = std::arch::x86_64::_mm256_and_si256(q4_combined, low_mask_256);
            let q4_signed = std::arch::x86_64::_mm256_sub_epi8(q4_nibbles, offset_256);

            let q8_vec = std::arch::x86_64::_mm256_loadu_si256(q8_ptr.cast());

            // Use maddubs approach for remaining block
            let q4_abs = std::arch::x86_64::_mm256_sign_epi8(q4_signed, q4_signed);
            let q8_signed = std::arch::x86_64::_mm256_sign_epi8(q8_vec, q4_signed);

            let ones = std::arch::x86_64::_mm256_set1_epi16(1);
            let prod_i16 = std::arch::x86_64::_mm256_maddubs_epi16(q4_abs, q8_signed);
            let prod_i32 = std::arch::x86_64::_mm256_madd_epi16(prod_i16, ones);
            let prod_f32 = _mm256_cvtepi32_ps(prod_i32);

            acc0 = _mm256_fmadd_ps(combined_scale, prod_f32, acc0);

            block_idx += 1;
        }

        // Combine both accumulators and do horizontal sum
        let acc = std::arch::x86_64::_mm256_add_ps(acc0, acc1);
        let hi = std::arch::x86_64::_mm256_extractf128_ps(acc, 1);
        let lo = std::arch::x86_64::_mm256_castps256_ps128(acc);
        let sum128 = _mm_add_ps(lo, hi);
        let sum64 = _mm_hadd_ps(sum128, sum128);
        let sum32 = _mm_hadd_ps(sum64, sum64);
        _mm_cvtss_f32(sum32)
    }
}

/// AVX2 accelerated Q4_0 × Q8_0 dot product using integer SIMD
///
/// Uses _mm256_maddubs_epi16 which multiplies pairs of u8×i8 and accumulates
/// to i16, then we sum to i32 and convert to f32. This is ~4x faster than
/// the f32 FMA approach because:
/// 1. Integer ops have lower latency
/// 2. maddubs does multiply AND horizontal add in one instruction
/// 3. Less data movement (1 byte vs 4 bytes per value)
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
#[inline]
unsafe fn fused_q4_0_q8_0_dot_avx2(
    q4_data: &[u8],
    q8_scales: &[f32],
    q8_quants: &[i8],
    in_dim: usize,
) -> f32 {
    // SAFETY: Memory safety ensured by bounds checking and alignment
    unsafe {
        use std::arch::x86_64::{
            _mm256_and_si256, _mm256_cvtepi32_ps, _mm256_fmadd_ps, _mm256_loadu_si256,
            _mm256_madd_epi16, _mm256_maddubs_epi16, _mm256_set1_epi16, _mm256_set1_epi8,
            _mm256_set1_ps, _mm256_setzero_ps, _mm256_sign_epi8, _mm256_sub_epi8, _mm_cvtss_f32,
            _mm_hadd_ps, _mm_prefetch, _MM_HINT_T0,
        };

        const Q4_0_BLOCK_BYTES: usize = 18;
        const Q4_0_BLOCK_SIZE: usize = 32;

        let num_blocks = in_dim.div_ceil(Q4_0_BLOCK_SIZE);

        // Float accumulator for final sum
        let mut acc = _mm256_setzero_ps();

        // Offset: Q4_0 values are 0-15, we subtract 8 to get -8 to +7
        let offset = _mm256_set1_epi8(8);
        let low_mask = _mm256_set1_epi8(0x0F);
        let ones = _mm256_set1_epi16(1);

        let mut block_idx = 0;

        // Process 2 blocks at a time for better instruction-level parallelism
        while block_idx + 2 <= num_blocks {
            // Prefetch next iteration's blocks
            if block_idx + 4 <= num_blocks {
                let prefetch_q4 = q4_data.as_ptr().add((block_idx + 2) * Q4_0_BLOCK_BYTES);
                let prefetch_q8 = q8_quants.as_ptr().add((block_idx + 2) * Q4_0_BLOCK_SIZE);
                _mm_prefetch(prefetch_q4.cast(), _MM_HINT_T0);
                _mm_prefetch(prefetch_q8.cast(), _MM_HINT_T0);
            }

            // === Block 0 ===
            let q4_ptr_0 = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
            let q8_ptr_0 = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);

            // Read Q4_0 scale (f16 -> f32 via LUT)
            let q4_scale_bits_0 = u16::from_le_bytes([*q4_ptr_0, *q4_ptr_0.add(1)]);
            let q4_scale_0 = f16_to_f32_lut(q4_scale_bits_0);
            let q8_scale_0 = q8_scales[block_idx];
            let combined_scale_0 = _mm256_set1_ps(q4_scale_0 * q8_scale_0);

            // Load Q4_0 quants (16 bytes = 32 nibbles)
            let q4_bytes = std::slice::from_raw_parts(q4_ptr_0.add(2), 16);

            // bytes_from_nibbles_32: expand 16 bytes to 32 bytes
            // Low nibbles in first 16 positions, high nibbles in next 16
            let q4_lo_128 = std::arch::x86_64::_mm_loadu_si128(q4_bytes.as_ptr().cast());
            let q4_hi_128 = std::arch::x86_64::_mm_srli_epi16(q4_lo_128, 4);
            // Combine into 256-bit: high nibbles in upper 128, low nibbles in lower 128
            let q4_combined = std::arch::x86_64::_mm256_set_m128i(q4_hi_128, q4_lo_128);
            // Mask to get just nibbles
            let q4_nibbles = _mm256_and_si256(q4_combined, low_mask);
            // Convert from unsigned 0-15 to signed -8 to +7
            let q4_signed = _mm256_sub_epi8(q4_nibbles, offset);

            // Load Q8_0 quants (32 bytes)
            let q8_vec = _mm256_loadu_si256(q8_ptr_0.cast());

            // Integer multiply-accumulate using signed multiply trick:
            // maddubs requires unsigned × signed, so we use sign trick
            // ax = |x|, sy = sign(y, x), then maddubs(ax, sy) = x * y
            let q4_abs = _mm256_sign_epi8(q4_signed, q4_signed);
            let q8_signed = _mm256_sign_epi8(q8_vec, q4_signed);

            // maddubs: multiply pairs and add horizontally to i16
            let prod_i16 = _mm256_maddubs_epi16(q4_abs, q8_signed);
            // madd: pairwise add i16 to i32
            let prod_i32 = _mm256_madd_epi16(prod_i16, ones);
            // Convert to float
            let prod_f32 = _mm256_cvtepi32_ps(prod_i32);

            // Scale and accumulate
            acc = _mm256_fmadd_ps(combined_scale_0, prod_f32, acc);

            // === Block 1 ===
            let q4_ptr_1 = q4_data.as_ptr().add((block_idx + 1) * Q4_0_BLOCK_BYTES);
            let q8_ptr_1 = q8_quants.as_ptr().add((block_idx + 1) * Q4_0_BLOCK_SIZE);

            let q4_scale_bits_1 = u16::from_le_bytes([*q4_ptr_1, *q4_ptr_1.add(1)]);
            let q4_scale_1 = f16_to_f32_lut(q4_scale_bits_1);
            let q8_scale_1 = q8_scales[block_idx + 1];
            let combined_scale_1 = _mm256_set1_ps(q4_scale_1 * q8_scale_1);

            let q4_bytes_1 = std::slice::from_raw_parts(q4_ptr_1.add(2), 16);
            let q4_lo_128_1 = std::arch::x86_64::_mm_loadu_si128(q4_bytes_1.as_ptr().cast());
            let q4_hi_128_1 = std::arch::x86_64::_mm_srli_epi16(q4_lo_128_1, 4);
            let q4_combined_1 = std::arch::x86_64::_mm256_set_m128i(q4_hi_128_1, q4_lo_128_1);
            let q4_nibbles_1 = _mm256_and_si256(q4_combined_1, low_mask);
            let q4_signed_1 = _mm256_sub_epi8(q4_nibbles_1, offset);

            let q8_vec_1 = _mm256_loadu_si256(q8_ptr_1.cast());

            let q4_abs_1 = _mm256_sign_epi8(q4_signed_1, q4_signed_1);
            let q8_signed_1 = _mm256_sign_epi8(q8_vec_1, q4_signed_1);

            let prod_i16_1 = _mm256_maddubs_epi16(q4_abs_1, q8_signed_1);
            let prod_i32_1 = _mm256_madd_epi16(prod_i16_1, ones);
            let prod_f32_1 = _mm256_cvtepi32_ps(prod_i32_1);

            acc = _mm256_fmadd_ps(combined_scale_1, prod_f32_1, acc);

            block_idx += 2;
        }

        // Handle remaining single block
        while block_idx < num_blocks {
            let q4_ptr = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
            let q8_ptr = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);

            let q4_scale_bits = u16::from_le_bytes([*q4_ptr, *q4_ptr.add(1)]);
            let q4_scale = f16_to_f32_lut(q4_scale_bits);
            let q8_scale = q8_scales[block_idx];
            let combined_scale = _mm256_set1_ps(q4_scale * q8_scale);

            let q4_bytes = std::slice::from_raw_parts(q4_ptr.add(2), 16);
            let q4_lo_128 = std::arch::x86_64::_mm_loadu_si128(q4_bytes.as_ptr().cast());
            let q4_hi_128 = std::arch::x86_64::_mm_srli_epi16(q4_lo_128, 4);
            let q4_combined = std::arch::x86_64::_mm256_set_m128i(q4_hi_128, q4_lo_128);
            let q4_nibbles = _mm256_and_si256(q4_combined, low_mask);
            let q4_signed = _mm256_sub_epi8(q4_nibbles, offset);

            let q8_vec = _mm256_loadu_si256(q8_ptr.cast());

            let q4_abs = _mm256_sign_epi8(q4_signed, q4_signed);
            let q8_signed = _mm256_sign_epi8(q8_vec, q4_signed);

            let prod_i16 = _mm256_maddubs_epi16(q4_abs, q8_signed);
            let prod_i32 = _mm256_madd_epi16(prod_i16, ones);
            let prod_f32 = _mm256_cvtepi32_ps(prod_i32);

            acc = _mm256_fmadd_ps(combined_scale, prod_f32, acc);

            block_idx += 1;
        }

        // Horizontal sum of 8 floats
        let hi = std::arch::x86_64::_mm256_extractf128_ps(acc, 1);
        let lo = std::arch::x86_64::_mm256_castps256_ps128(acc);
        let sum128 = std::arch::x86_64::_mm_add_ps(lo, hi);
        // Use hadd for final reduction
        let sum64 = _mm_hadd_ps(sum128, sum128);
        let sum32 = _mm_hadd_ps(sum64, sum64);
        _mm_cvtss_f32(sum32)
    }
}

/// AVX2 accelerated Q4_0 × Q8_0 dot product with 4-block unrolling
///
/// Processes 4 blocks per iteration for maximum ILP on modern OoO CPUs.
/// This version achieves ~1.3x speedup over 2-block unrolling for large vectors.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
#[inline]
unsafe fn fused_q4_0_q8_0_dot_avx2_4block(
    q4_data: &[u8],
    q8_scales: &[f32],
    q8_quants: &[i8],
    in_dim: usize,
) -> f32 {
    // SAFETY: Memory safety ensured by bounds checking and alignment
    unsafe {
        use std::arch::x86_64::{
            _mm256_add_ps, _mm256_and_si256, _mm256_cvtepi32_ps, _mm256_fmadd_ps,
            _mm256_loadu_si256, _mm256_madd_epi16, _mm256_maddubs_epi16, _mm256_set1_epi16,
            _mm256_set1_epi8, _mm256_set1_ps, _mm256_setzero_ps, _mm256_sign_epi8, _mm256_sub_epi8,
            _mm_cvtss_f32, _mm_hadd_ps, _mm_prefetch, _MM_HINT_T0,
        };

        const Q4_0_BLOCK_BYTES: usize = 18;
        const Q4_0_BLOCK_SIZE: usize = 32;

        let num_blocks = in_dim.div_ceil(Q4_0_BLOCK_SIZE);

        // Use two accumulators for better pipelining
        let mut acc0 = _mm256_setzero_ps();
        let mut acc1 = _mm256_setzero_ps();

        let offset = _mm256_set1_epi8(8);
        let low_mask = _mm256_set1_epi8(0x0F);
        let ones = _mm256_set1_epi16(1);

        let mut block_idx = 0;

        // Process 4 blocks at a time for maximum ILP
        while block_idx + 4 <= num_blocks {
            // Prefetch next iteration's blocks
            if block_idx + 8 <= num_blocks {
                let prefetch_q4 = q4_data.as_ptr().add((block_idx + 4) * Q4_0_BLOCK_BYTES);
                let prefetch_q8 = q8_quants.as_ptr().add((block_idx + 4) * Q4_0_BLOCK_SIZE);
                _mm_prefetch(prefetch_q4.cast(), _MM_HINT_T0);
                _mm_prefetch(prefetch_q8.cast(), _MM_HINT_T0);
                _mm_prefetch(prefetch_q4.add(64).cast(), _MM_HINT_T0);
                _mm_prefetch(prefetch_q8.add(64).cast(), _MM_HINT_T0);
            }

            // Block 0
            let q4_ptr_0 = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
            let q8_ptr_0 = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);
            let q4_scale_0 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_0, *q4_ptr_0.add(1)]));
            let combined_scale_0 = _mm256_set1_ps(q4_scale_0 * q8_scales[block_idx]);
            let q4_lo_0 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_0.add(2).cast());
            let q4_hi_0 = std::arch::x86_64::_mm_srli_epi16(q4_lo_0, 4);
            let q4_signed_0 = _mm256_sub_epi8(
                _mm256_and_si256(
                    std::arch::x86_64::_mm256_set_m128i(q4_hi_0, q4_lo_0),
                    low_mask,
                ),
                offset,
            );
            let q8_vec_0 = _mm256_loadu_si256(q8_ptr_0.cast());
            let q4_abs_0 = _mm256_sign_epi8(q4_signed_0, q4_signed_0);
            let q8_signed_0 = _mm256_sign_epi8(q8_vec_0, q4_signed_0);
            let prod_0 = _mm256_cvtepi32_ps(_mm256_madd_epi16(
                _mm256_maddubs_epi16(q4_abs_0, q8_signed_0),
                ones,
            ));
            acc0 = _mm256_fmadd_ps(combined_scale_0, prod_0, acc0);

            // Block 1
            let q4_ptr_1 = q4_data.as_ptr().add((block_idx + 1) * Q4_0_BLOCK_BYTES);
            let q8_ptr_1 = q8_quants.as_ptr().add((block_idx + 1) * Q4_0_BLOCK_SIZE);
            let q4_scale_1 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_1, *q4_ptr_1.add(1)]));
            let combined_scale_1 = _mm256_set1_ps(q4_scale_1 * q8_scales[block_idx + 1]);
            let q4_lo_1 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_1.add(2).cast());
            let q4_hi_1 = std::arch::x86_64::_mm_srli_epi16(q4_lo_1, 4);
            let q4_signed_1 = _mm256_sub_epi8(
                _mm256_and_si256(
                    std::arch::x86_64::_mm256_set_m128i(q4_hi_1, q4_lo_1),
                    low_mask,
                ),
                offset,
            );
            let q8_vec_1 = _mm256_loadu_si256(q8_ptr_1.cast());
            let q4_abs_1 = _mm256_sign_epi8(q4_signed_1, q4_signed_1);
            let q8_signed_1 = _mm256_sign_epi8(q8_vec_1, q4_signed_1);
            let prod_1 = _mm256_cvtepi32_ps(_mm256_madd_epi16(
                _mm256_maddubs_epi16(q4_abs_1, q8_signed_1),
                ones,
            ));
            acc1 = _mm256_fmadd_ps(combined_scale_1, prod_1, acc1);

            // Block 2
            let q4_ptr_2 = q4_data.as_ptr().add((block_idx + 2) * Q4_0_BLOCK_BYTES);
            let q8_ptr_2 = q8_quants.as_ptr().add((block_idx + 2) * Q4_0_BLOCK_SIZE);
            let q4_scale_2 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_2, *q4_ptr_2.add(1)]));
            let combined_scale_2 = _mm256_set1_ps(q4_scale_2 * q8_scales[block_idx + 2]);
            let q4_lo_2 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_2.add(2).cast());
            let q4_hi_2 = std::arch::x86_64::_mm_srli_epi16(q4_lo_2, 4);
            let q4_signed_2 = _mm256_sub_epi8(
                _mm256_and_si256(
                    std::arch::x86_64::_mm256_set_m128i(q4_hi_2, q4_lo_2),
                    low_mask,
                ),
                offset,
            );
            let q8_vec_2 = _mm256_loadu_si256(q8_ptr_2.cast());
            let q4_abs_2 = _mm256_sign_epi8(q4_signed_2, q4_signed_2);
            let q8_signed_2 = _mm256_sign_epi8(q8_vec_2, q4_signed_2);
            let prod_2 = _mm256_cvtepi32_ps(_mm256_madd_epi16(
                _mm256_maddubs_epi16(q4_abs_2, q8_signed_2),
                ones,
            ));
            acc0 = _mm256_fmadd_ps(combined_scale_2, prod_2, acc0);

            // Block 3
            let q4_ptr_3 = q4_data.as_ptr().add((block_idx + 3) * Q4_0_BLOCK_BYTES);
            let q8_ptr_3 = q8_quants.as_ptr().add((block_idx + 3) * Q4_0_BLOCK_SIZE);
            let q4_scale_3 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_3, *q4_ptr_3.add(1)]));
            let combined_scale_3 = _mm256_set1_ps(q4_scale_3 * q8_scales[block_idx + 3]);
            let q4_lo_3 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_3.add(2).cast());
            let q4_hi_3 = std::arch::x86_64::_mm_srli_epi16(q4_lo_3, 4);
            let q4_signed_3 = _mm256_sub_epi8(
                _mm256_and_si256(
                    std::arch::x86_64::_mm256_set_m128i(q4_hi_3, q4_lo_3),
                    low_mask,
                ),
                offset,
            );
            let q8_vec_3 = _mm256_loadu_si256(q8_ptr_3.cast());
            let q4_abs_3 = _mm256_sign_epi8(q4_signed_3, q4_signed_3);
            let q8_signed_3 = _mm256_sign_epi8(q8_vec_3, q4_signed_3);
            let prod_3 = _mm256_cvtepi32_ps(_mm256_madd_epi16(
                _mm256_maddubs_epi16(q4_abs_3, q8_signed_3),
                ones,
            ));
            acc1 = _mm256_fmadd_ps(combined_scale_3, prod_3, acc1);

            block_idx += 4;
        }

        // Merge accumulators
        let acc = _mm256_add_ps(acc0, acc1);

        // Handle remaining blocks (0-3)
        let mut scalar_sum = 0.0f32;
        while block_idx < num_blocks {
            let q4_ptr = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
            let q8_ptr = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);
            let q4_scale = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr, *q4_ptr.add(1)]));
            let combined_scale = _mm256_set1_ps(q4_scale * q8_scales[block_idx]);
            let q4_lo = std::arch::x86_64::_mm_loadu_si128(q4_ptr.add(2).cast());
            let q4_hi = std::arch::x86_64::_mm_srli_epi16(q4_lo, 4);
            let q4_signed = _mm256_sub_epi8(
                _mm256_and_si256(std::arch::x86_64::_mm256_set_m128i(q4_hi, q4_lo), low_mask),
                offset,
            );
            let q8_vec = _mm256_loadu_si256(q8_ptr.cast());
            let q4_abs = _mm256_sign_epi8(q4_signed, q4_signed);
            let q8_signed = _mm256_sign_epi8(q8_vec, q4_signed);
            let prod = _mm256_cvtepi32_ps(_mm256_madd_epi16(
                _mm256_maddubs_epi16(q4_abs, q8_signed),
                ones,
            ));
            let scaled = _mm256_fmadd_ps(combined_scale, prod, _mm256_setzero_ps());

            // Horizontal sum for this block
            let hi = std::arch::x86_64::_mm256_extractf128_ps(scaled, 1);
            let lo = std::arch::x86_64::_mm256_castps256_ps128(scaled);
            let sum128 = std::arch::x86_64::_mm_add_ps(lo, hi);
            let sum64 = _mm_hadd_ps(sum128, sum128);
            let sum32 = _mm_hadd_ps(sum64, sum64);
            scalar_sum += _mm_cvtss_f32(sum32);

            block_idx += 1;
        }

        // Horizontal sum of accumulated vector
        let hi = std::arch::x86_64::_mm256_extractf128_ps(acc, 1);
        let lo = std::arch::x86_64::_mm256_castps256_ps128(acc);
        let sum128 = std::arch::x86_64::_mm_add_ps(lo, hi);
        let sum64 = _mm_hadd_ps(sum128, sum128);
        let sum32 = _mm_hadd_ps(sum64, sum64);
        _mm_cvtss_f32(sum32) + scalar_sum
    }
}

/// Scalar fallback for Q4_0 × Q8_0 dot product
///
/// Exposed as `pub(crate)` for direct testing on AVX2 machines.
#[inline]
pub(crate) fn fused_q4_0_q8_0_dot_scalar(
    q4_data: &[u8],
    q8_scales: &[f32],
    q8_quants: &[i8],
    in_dim: usize,
) -> f32 {
    const Q4_0_BLOCK_BYTES: usize = 18;
    const Q4_0_BLOCK_SIZE: usize = 32;

    let num_blocks = in_dim.div_ceil(Q4_0_BLOCK_SIZE);
    let mut total_sum = 0.0f32;

    for block_idx in 0..num_blocks4 {
        let block_start = block_idx * Q4_0_BLOCK_BYTES;
        if block_start + Q4_0_BLOCK_BYTES > q4_data.len() {
            break;
        }
        let block = &q4_data[block_start..block_start + Q4_0_BLOCK_BYTES];

        let q4_scale = half::f16::from_le_bytes([block[0], block[1]]).to_f32();
        let q8_scale = q8_scales[block_idx];
        let combined_scale = q4_scale * q8_scale;

        let act_start = block_idx * Q4_0_BLOCK_SIZE;

        let mut block_sum = 0i32;
        for (j, &byte) in block[2..18]5.iter5().enumerate5() {
            let low_idx = act_start + j;
            let high_idx = act_start + j + 16;

            #[allow(clippy::cast_possible_wrap)]
            let low_quant = (byte & 0x0F) as i8 - 8;
            block_sum += (low_quant as i32) * (q8_quants[low_idx] as i32);

            #[allow(clippy::cast_possible_wrap)]
            let high_quant = (byte >> 4) as i8 - 8;
            if high_idx < in_dim {
                block_sum += (high_quant as i32) * (q8_quants[high_idx] as i32);
            }8
        }

        total_sum += combined_scale * (block_sum as f32);
    }

    total_sum
}

/// Parallel Q4_0 × Q8_0 matrix-vector multiply
///
/// This is the key function for llama.cpp parity. It:
/// 1. Quantizes activations to Q8_0 format once
/// 2. Uses integer SIMD for all row dot products
/// 3. Parallelizes across output rows with rayon (adaptive threshold)
///
/// Expected speedup: 4-6x over the f32 FMA version
#[allow(clippy::similar_names)]
pub fn fused_q4_0_q8_0_parallel_matvec(
    weight_data: &[u8],
    activations: &[f32],
    in_dim: usize,
    out_dim: usize,
) -> Result<Vec<f32>> {
    const Q4_0_BLOCK_BYTES: usize = 18;
    const Q4_0_BLOCK_SIZE: usize = 32;

    let blocks_per_row = in_dim.div_ceil(Q4_0_BLOCK_SIZE);
    let bytes_per_row = blocks_per_row * Q4_0_BLOCK_BYTES;

    let expected_weight_bytes = out_dim * bytes_per_row;
    if weight_data.len() < expected_weight_bytes {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Q4_0 weight data too small: need {} bytes for {}x{}, have {}",
                expected_weight_bytes,
                out_dim,
                in_dim,
                weight_data.len()
            ),
        });
    }

    if activations.len() != in_dim {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Activation length {} doesn't match in_dim {}",
                activations.len(),
                in_dim
            ),
        });
    }

    // Quantize activations to Q8_0 ONCE (amortized over all rows)
    let (q8_scales, q8_quants) = quantize_activations_q8_0(activations);

    // Adaptive parallelization: sequential for small matrices, parallel for large
    // Rayon overhead (~50-100µs) dominates for small out_dim
    // Threshold tuned for 22-core CPU: break-even at ~1024 rows
    const PARALLEL_THRESHOLD: usize = 1024;

    if out_dim < PARALLEL_THRESHOLD {
        // Sequential path: avoids Rayon overhead entirely
        let output: Vec<f32> = (0..out_dim)
            .map711(|o| {
                let row_start = o * bytes_per_row;
                let row_end = row_start + bytes_per_row;
                let row_data = &weight_data[row_start..row_end];
                fused_q4_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim)
            })
            .collect();
        return Ok(output);
    }

    // Parallel path for large matrices
    use rayon::prelude::*;
    // Use chunked parallel iteration to reduce Rayon scheduling overhead
    // CHUNK_SIZE=128 provides good balance between parallelism and overhead
    const CHUNK_SIZE: usize = 128;
    let output: Vec<f32> = (0..out_dim)
        .into_par_iter()
        .with_min_len(CHUNK_SIZE)
        .map(|o| {
            let row_start = o * bytes_per_row;
            let row_end = row_start + bytes_per_row;
            let row_data = &weight_data[row_start..row_end];
            fused_q4_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim)
        })
        .collect();

    Ok(output)
}

/// Zero-allocation Q4_0 × Q8_0 matrix-vector multiply.
///
/// Writes result directly into provided output buffer, eliminating allocation.
/// Use this with scratch buffers for maximum performance.
///
/// # Arguments
/// * `weight_data` - Q4_0 quantized weight matrix (row-major)
/// * `activations` - Input activation vector (f32)
/// * `in_dim` - Input dimension (columns)
/// * `output` - Pre-allocated output buffer (must be exactly out_dim length)
///
/// # Returns
/// Number of elements written (equals output.len())
#[allow(clippy::similar_names)]
pub fn fused_q4_0_q8_0_parallel_matvec_into(
    weight_data: &[u8],
    activations: &[f32],
    in_dim: usize,
    output: &mut [f32],
) -> Result<()> {
    use rayon::prelude::*;

    const Q4_0_BLOCK_BYTES: usize = 18;
    const Q4_0_BLOCK_SIZE: usize = 32;

    let out_dim = output.len();
    let blocks_per_row = in_dim.div_ceil(Q4_0_BLOCK_SIZE);
    let bytes_per_row = blocks_per_row * Q4_0_BLOCK_BYTES;

    let expected_weight_bytes = out_dim * bytes_per_row;
    if weight_data.len() < expected_weight_bytes {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Q4_0 weight data too small: need {} bytes for {}x{}, have {}",
                expected_weight_bytes,
                out_dim,
                in_dim,
                weight_data.len()
            ),
        });
    }

    if activations.len() != in_dim {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Activation length {} doesn't match in_dim {}",
                activations.len(),
                in_dim
            ),
        });
    }

    // Quantize activations to Q8_0 ONCE
    let (q8_scales, q8_quants) = quantize_activations_q8_0(activations);

    // Use chunked parallel iteration to reduce Rayon scheduling overhead
    const CHUNK_SIZE: usize = 64;
    output
        .par_iter_mut()
        .with_min_len(CHUNK_SIZE)
        .enumerate()
        .for_each68(|(o, out_val)| {
            let row_start = o * bytes_per_row;
            let row_end = row_start + bytes_per_row;
            let row_data = &weight_data[row_start..row_end];
            *out_val = fused_q4_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim);
        });

    Ok(())
}

// ============================================================================
// FUSED Q8_0 × Q8_0 MATMUL (For Q8_0 quantized weights like Qwen2.5 LM head)
// ============================================================================
//
// Q8_0 format: 34 bytes per block (2 byte f16 scale + 32 i8 quants)
// This avoids the massive dequantization allocation that was causing
// Qwen2.5's 152K vocab LM head to allocate 544MB per forward pass.
// ============================================================================

/// AVX2 accelerated Q8_0 × Q8_0 dot product using integer SIMD
///
/// Uses _mm256_maddubs_epi16 with sign trick for i8×i8 multiplication.
/// This is simpler than Q4_0×Q8_0 since no nibble unpacking is needed.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
#[inline]
unsafe fn fused_q8_0_q8_0_dot_avx2(
    q8_weight_data: &[u8],
    q8_act_scales: &[f32],
    q8_act_quants: &[i8],
    in_dim: usize,
) -> f32 {
    // SAFETY: Memory safety ensured by bounds checking and alignment
    unsafe {
        use std::arch::x86_64::{
            _mm256_cvtepi32_ps, _mm256_fmadd_ps, _mm256_loadu_si256, _mm256_madd_epi16,
            _mm256_maddubs_epi16, _mm256_set1_epi16, _mm256_set1_ps, _mm256_setzero_ps,
            _mm256_sign_epi8, _mm_cvtss_f32, _mm_hadd_ps, _mm_prefetch, _MM_HINT_T0,
        };

        const Q8_0_BLOCK_BYTES: usize = 34; // 2 byte scale + 32 byte quants
        const Q8_0_BLOCK_SIZE: usize = 32;

        let num_blocks = in_dim.div_ceil(Q8_0_BLOCK_SIZE);

        // Float accumulator for final sum
        let mut acc = _mm256_setzero_ps();
        let ones = _mm256_set1_epi16(1);

        let mut block_idx = 0;

        // Process 2 blocks at a time for better ILP
        while block_idx + 2 <= num_blocks {
            // Prefetch next iteration's blocks
            if block_idx + 4 <= num_blocks {
                let prefetch_w = q8_weight_data
                    .as_ptr()
                    .add((block_idx + 2) * Q8_0_BLOCK_BYTES);
                let prefetch_a = q8_act_quants
                    .as_ptr()
                    .add((block_idx + 2) * Q8_0_BLOCK_SIZE);
                _mm_prefetch(prefetch_w.cast(), _MM_HINT_T0);
                _mm_prefetch(prefetch_a.cast(), _MM_HINT_T0);
            }1.02k

            // === Block 0 ===
            let w_ptr_0 = q8_weight_data.as_ptr().add(block_idx * Q8_0_BLOCK_BYTES);
            let a_ptr_0 = q8_act_quants.as_ptr().add(block_idx * Q8_0_BLOCK_SIZE);

            // Read Q8_0 weight scale (f16 -> f32)
            let w_scale_bits_0 = u16::from_le_bytes([*w_ptr_0, *w_ptr_0.add(1)]);
            let w_scale_0 = f16_to_f32_lut(w_scale_bits_0);
            let a_scale_0 = q8_act_scales[block_idx];
            let combined_scale_0 = _mm256_set1_ps(w_scale_0 * a_scale_0);

            // Load Q8_0 weight quants (32 bytes at offset 2)
            let w_vec_0 = _mm256_loadu_si256(w_ptr_0.add(2).cast());
            // Load Q8_0 activation quants (32 bytes)
            let a_vec_0 = _mm256_loadu_si256(a_ptr_0.cast());

            // Integer multiply-accumulate using signed multiply trick:
            // maddubs requires unsigned × signed, so we use sign trick
            // |w| * sign(a, w) = w * a
            let w_abs_0 = _mm256_sign_epi8(w_vec_0, w_vec_0);
            let a_signed_0 = _mm256_sign_epi8(a_vec_0, w_vec_0);

            // maddubs: multiply pairs and add horizontally to i16
            let prod_i16_0 = _mm256_maddubs_epi16(w_abs_0, a_signed_0);
            // madd: pairwise add i16 to i32
            let prod_i32_0 = _mm256_madd_epi16(prod_i16_0, ones);
            // Convert to float
            let prod_f32_0 = _mm256_cvtepi32_ps(prod_i32_0);

            // Scale and accumulate
            acc = _mm256_fmadd_ps(combined_scale_0, prod_f32_0, acc);

            // === Block 1 ===
            let w_ptr_1 = q8_weight_data
                .as_ptr()
                .add((block_idx + 1) * Q8_0_BLOCK_BYTES);
            let a_ptr_1 = q8_act_quants
                .as_ptr()
                .add((block_idx + 1) * Q8_0_BLOCK_SIZE);

            let w_scale_bits_1 = u16::from_le_bytes([*w_ptr_1, *w_ptr_1.add(1)]);
            let w_scale_1 = f16_to_f32_lut(w_scale_bits_1);
            let a_scale_1 = q8_act_scales[block_idx + 1];
            let combined_scale_1 = _mm256_set1_ps(w_scale_1 * a_scale_1);

            let w_vec_1 = _mm256_loadu_si256(w_ptr_1.add(2).cast());
            let a_vec_1 = _mm256_loadu_si256(a_ptr_1.cast());

            let w_abs_1 = _mm256_sign_epi8(w_vec_1, w_vec_1);
            let a_signed_1 = _mm256_sign_epi8(a_vec_1, w_vec_1);

            let prod_i16_1 = _mm256_maddubs_epi16(w_abs_1, a_signed_1);
            let prod_i32_1 = _mm256_madd_epi16(prod_i16_1, ones);
            let prod_f32_1 = _mm256_cvtepi32_ps(prod_i32_1);

            acc = _mm256_fmadd_ps(combined_scale_1, prod_f32_1, acc);

            block_idx += 2;
        }

        // Handle remaining single block
        while block_idx < num_blocks {
            let w_ptr = q8_weight_data.as_ptr().add(block_idx * Q8_0_BLOCK_BYTES);
            let a_ptr = q8_act_quants.as_ptr().add(block_idx * Q8_0_BLOCK_SIZE);

            let w_scale_bits = u16::from_le_bytes([*w_ptr, *w_ptr.add(1)]);
            let w_scale = f16_to_f32_lut(w_scale_bits);
            let a_scale = q8_act_scales[block_idx];
            let combined_scale = _mm256_set1_ps(w_scale * a_scale);

            let w_vec = _mm256_loadu_si256(w_ptr.add(2).cast());
            let a_vec = _mm256_loadu_si256(a_ptr.cast());

            let w_abs = _mm256_sign_epi8(w_vec, w_vec);
            let a_signed = _mm256_sign_epi8(a_vec, w_vec);

            let prod_i16 = _mm256_maddubs_epi16(w_abs, a_signed);
            let prod_i32 = _mm256_madd_epi16(prod_i16, ones);
            let prod_f32 = _mm256_cvtepi32_ps(prod_i32);

            acc = _mm256_fmadd_ps(combined_scale, prod_f32, acc);

            block_idx += 1;
        }

        // Horizontal sum of 8 floats
        let hi = std::arch::x86_64::_mm256_extractf128_ps(acc, 1);
        let lo = std::arch::x86_64::_mm256_castps256_ps128(acc);
        let sum128 = std::arch::x86_64::_mm_add_ps(lo, hi);
        let sum64 = _mm_hadd_ps(sum128, sum128);
        let sum32 = _mm_hadd_ps(sum64, sum64);
        _mm_cvtss_f32(sum32)
    }
}

/// Scalar fallback for Q8_0 × Q8_0 dot product
///
/// Exposed as `pub(crate)` for direct testing on AVX2 machines.
#[inline]
pub(crate) fn fused_q8_0_q8_0_dot_scalar(
    q8_weight_data: &[u8],
    q8_act_scales: &[f32],
    q8_act_quants: &[i8],
    in_dim: usize,
) -> f32 {
    const Q8_0_BLOCK_BYTES: usize = 34;
    const Q8_0_BLOCK_SIZE: usize = 32;

    let num_blocks = in_dim.div_ceil(Q8_0_BLOCK_SIZE);
    let mut total_sum = 0.0f32;

    for block_idx in 0..num_blocks {
        let block_start = block_idx * Q8_0_BLOCK_BYTES;
        if block_start + Q8_0_BLOCK_BYTES > q8_weight_data.len() {
            break;
        }
        let block = &q8_weight_data[block_start..block_start + Q8_0_BLOCK_BYTES];

        // Read weight scale (f16)
        let w_scale = half::f16::from_le_bytes([block[0], block[1]]).to_f32();
        let a_scale = q8_act_scales[block_idx];
        let combined_scale = w_scale * a_scale;

        let act_start = block_idx * Q8_0_BLOCK_SIZE;

        // Sum of weight_quant[i] * act_quant[i] in i32
        let mut block_sum = 0i32;
        for j32 in 0..32 {
            if act_start + j >= in_dim {
                break;
            }
            #[allow(clippy::cast_possible_wrap)]
            let w_quant = block[2 + j] as i8;
            let a_quant = q8_act_quants[act_start + j];
            block_sum += (w_quant as i32) * (a_quant as i32);
        }

        total_sum += combined_scale * (block_sum as f32);
    }

    total_sum
}

/// SIMD dispatcher for Q8_0 × Q8_0 dot product
#[inline]
fn fused_q8_0_q8_0_dot_simd(
    q8_weight_data: &[u8],
    q8_act_scales: &[f32],
    q8_act_quants: &[i8],
    in_dim: usize,
) -> f32 {
    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
            // SAFETY: AVX2 and FMA features checked above
            unsafe {
                return fused_q8_0_q8_0_dot_avx2(
                    q8_weight_data,
                    q8_act_scales,
                    q8_act_quants,
                    in_dim,
                );
            }
        }
    }
    fused_q8_0_q8_0_dot_scalar(q8_weight_data, q8_act_scales, q8_act_quants, in_dim)
}

/// Parallel Q8_0 × Q8_0 matrix-vector multiply
///
/// This avoids the massive dequantization allocation that was causing
/// Qwen2.5's 152K vocab LM head (Q8_0) to allocate 544MB per forward pass.
///
/// For Q8_0 weights (e.g., Qwen2.5 LM head), this is ~100x faster than
/// dequantize + matmul because:
/// 1. No 544MB allocation per forward pass
/// 2. Integer SIMD is faster than FP32
/// 3. Better cache locality (34 bytes vs 128 bytes per block)
#[allow(clippy::similar_names)]
pub fn fused_q8_0_q8_0_parallel_matvec(
    weight_data: &[u8],
    activations: &[f32],
    in_dim: usize,
    out_dim: usize,
) -> Result<Vec<f32>> {
    use rayon::prelude::*;

    const Q8_0_BLOCK_BYTES: usize = 34;
    const Q8_0_BLOCK_SIZE: usize = 32;

    let blocks_per_row = in_dim.div_ceil(Q8_0_BLOCK_SIZE);
    let bytes_per_row = blocks_per_row * Q8_0_BLOCK_BYTES;

    let expected_weight_bytes = out_dim * bytes_per_row;
    if weight_data.len() < expected_weight_bytes {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Q8_0 weight data too small: need {} bytes for {}x{}, have {}",
                expected_weight_bytes,
                out_dim,
                in_dim,
                weight_data.len()
            ),
        });
    }

    if activations.len() != in_dim {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Activation length {} doesn't match in_dim {}",
                activations.len(),
                in_dim
            ),
        });
    }

    // Quantize activations to Q8_0 ONCE (amortized over all rows)
    let (q8_scales, q8_quants) = quantize_activations_q8_0(activations);

    // Parallel over output rows with chunking
    const CHUNK_SIZE: usize = 64;
    let output: Vec<f32> = (0..out_dim)
        .into_par_iter()
        .with_min_len(CHUNK_SIZE)
        .map7(|o| {
            let row_start = o * bytes_per_row;
            let row_end = row_start + bytes_per_row;
            let row_data = &weight_data[row_start..row_end];
            fused_q8_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim)
        })
        .collect();

    Ok(output)
}

/// Fused Q8_0 × Q8_0 parallel matvec - writes to pre-allocated buffer
///
/// IMP-131: Zero-allocation variant for hot-path inference.
#[allow(clippy::similar_names)]
pub fn fused_q8_0_q8_0_parallel_matvec_into(
    weight_data: &[u8],
    activations: &[f32],
    in_dim: usize,
    out_dim: usize,
    output: &mut [f32],
) -> Result<()> {
    use rayon::prelude::*;

    const Q8_0_BLOCK_BYTES: usize = 34;
    const Q8_0_BLOCK_SIZE: usize = 32;

    let blocks_per_row = in_dim.div_ceil(Q8_0_BLOCK_SIZE);
    let bytes_per_row = blocks_per_row * Q8_0_BLOCK_BYTES;

    let expected_weight_bytes = out_dim * bytes_per_row;
    if weight_data.len() < expected_weight_bytes {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Q8_0 weight data too small: need {} bytes for {}x{}, have {}",
                expected_weight_bytes,
                out_dim,
                in_dim,
                weight_data.len()
            ),
        });
    }

    if activations.len() != in_dim {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Activation length {} doesn't match in_dim {}",
                activations.len(),
                in_dim
            ),
        });
    }

    if output.len() < out_dim {
        return Err(RealizarError::InvalidShape {
            reason: format!(
                "Output buffer too small: need {}, have {}",
                out_dim,
                output.len()
            ),
        });
    }

    // Quantize activations to Q8_0 ONCE (amortized over all rows)
    let (q8_scales, q8_quants) = quantize_activations_q8_0(activations);

    // Parallel over output rows with chunking
    const CHUNK_SIZE: usize = 64;
    output[..out_dim]
        .par_iter_mut()
        .enumerate()
        .with_min_len(CHUNK_SIZE)
        .for_each2(|(o, out)| {
            let row_start = o * bytes_per_row;
            let row_end = row_start + bytes_per_row;
            let row_data = &weight_data[row_start..row_end];
            *out = fused_q8_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim);
        });

    Ok(())
}

/// Helper: Extract 6-bit scale and min for a block from the packed scales array
///
/// PAR-001 FIX: Matches llama.cpp's get_scale_min_k4 packing scheme:
/// - Blocks 0-3: scale = q[j] & 63, min = q[j+4] & 63
/// - Blocks 4-7: scale = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4)
///   min = (q[j+4] >> 4) | ((q[j] >> 6) << 4)
#[inline]
pub(crate) fn extract_scale_min(scales: &[u8; 12], block_idx: usize) -> (f32, f32) {
    let j = block_idx;
    let (scale_bits, min_bits) = if j < 4 {
        // First 4 blocks: simple layout
        let d = scales[j] & 63;
        let m = scales[j + 4] & 63;
        (d, m)
    } else {
        // Last 4 blocks: packed layout using high bits from first 4 bytes
        let d = (scales[j + 4] & 0x0F) | ((scales[j - 4] >> 6) << 4);
        let m = (scales[j + 4] >> 4) | ((scales[j] >> 6) << 4);
        (d, m)
    };

    // Return raw 6-bit values as floats
    // The GGUF header's d/dmin values already include the /63 normalization
    let scale = f32::from(scale_bits);
    let min = f32::from(min_bits);

    (scale, min)
}



#[cfg(test)]

#[cfg(test)]
mod tests;

Coverage Report

Created: 2026-01-25 15:05