## Coverage Report

#### Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

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

    Line
    Count
    Source
    1
    //! Quantization and dequantization for model weights
    2
    //!
    3
    //! Implements quantization formats used by GGUF models:
    4
    //! - `F16`: 16-bit IEEE 754 half-precision
    5
    //! - `Q4_0`: 4-bit quantization (block size 32)
    6
    //! - `Q4_1`: 4-bit with scale and min (block size 32)
    7
    //! - `Q5_0`: 5-bit quantization (block size 32)
    8
    //! - `Q5_1`: 5-bit with scale and min (block size 32)
    9
    //! - `Q8_0`: 8-bit quantization (block size 32)
    10
    //! - `Q4_K`: 4-bit K-quantization (super-block size 256)
    11
    //! - `Q5_K`: 5-bit K-quantization (super-block size 256)
    12
    //! - `Q6_K`: 6-bit K-quantization (super-block size 256)
    13
    //!
    14
    //! ## `Q4_0` Format
    15
    //!
    16
    //! `Q4_0` stores weights in blocks of 32 values:
    17
    //! - 1 float32 scale factor per block
    18
    //! - 16 bytes of 4-bit quantized values (2 values per byte)
    19
    //! - Dequantization: `value = scale * quantized_value`
    20
    //!
    21
    //! ## `Q8_0` Format
    22
    //!
    23
    //! `Q8_0` stores weights in blocks of 32 values:
    24
    //! - 1 float32 scale factor per block
    25
    //! - 32 int8 quantized values
    26
    //! - Dequantization: `value = scale * quantized_value`
    27
    //!
    28
    //! ## `Q4_K` Format
    29
    //!
    30
    //! `Q4_K` uses super-blocks of 256 values divided into 8 blocks of 32 values:
    31
    //! - 1 half-precision super-block scale (`d`)
    32
    //! - 1 half-precision super-block min (`dmin`)
    33
    //! - 12 bytes of 6-bit block scales (packed)
    34
    //! - 128 bytes of 4-bit quantized values
    35
    //! - Dequantization: `value = d * scale * quantized - dmin * min`
    36
    //! - Achieves 4.5 bits per weight with better quality than `Q4_0`
    37
    //!
    38
    //! ## `Q5_K` Format
    39
    //!
    40
    //! `Q5_K` uses super-blocks of 256 values divided into 8 blocks of 32 values:
    41
    //! - 1 half-precision super-block scale (`d`)
    42
    //! - 1 half-precision super-block min (`dmin`)
    43
    //! - 12 bytes of 6-bit block scales (packed)
    44
    //! - 32 bytes of high bits (1 bit per value for 5-bit quantization)
    45
    //! - 128 bytes of low 4-bit quantized values
    46
    //! - Dequantization: `value = d * scale * quantized - dmin * min`
    47
    //! - Achieves 5.5 bits per weight (higher quality than `Q4_K`)
    48
    //!
    49
    //! ## `Q6_K` Format
    50
    //!
    51
    //! `Q6_K` uses super-blocks of 256 values divided into 16 blocks of 16 values:
    52
    //! - 1 half-precision super-block scale (`d`)
    53
    //! - 16 bytes of 8-bit block scales
    54
    //! - 64 bytes of high 2 bits (2 bits per value for 6-bit quantization)
    55
    //! - 128 bytes of low 4-bit quantized values
    56
    //! - Dequantization: `value = d * scale * quantized`
    57
    //! - Achieves 6.5625 bits per weight (highest quality K-quant format)
    58
    59
    use crate::error::{RealizarError, Result};
    60
    61
    // ============================================================================
    62
    // Shattered submodules (PMAT-802)
    63
    // ============================================================================
    64
    65
    pub mod activation;
    66
    pub mod dequant;
    67
    pub mod fused_k;
    68
    pub mod fused_q5k_q6k;
    69
    pub mod parallel_dequant;
    70
    pub mod parallel_k;
    71
    pub mod simd;
    72
    pub mod types;
    73
    74
    // Re-export types from submodules (PMAT-802)
    75
    pub use types::{
    76
        BLOCK_SIZE, QK_K, Q4_0Block, Q4_KBlock, Q5_KBlock, Q6_KBlock,
    77
        Q8_0Block, Q8KSuperBlock,
    78
        DequantStats, SimdBackend, detect_simd_backend,
    79
    };
    80
    81
    // Re-export dequantization functions (PMAT-802)
    82
    pub use dequant::{
    83
        dequantize_f16, dequantize_q2_k, dequantize_q4_0, dequantize_q4_1,
    84
        dequantize_q4_k, dequantize_q5_0, dequantize_q5_1, dequantize_q5_k,
    85
        dequantize_q6_k, dequantize_q8_0, f16_to_f32,
    86
    };
    87
    88
    // Re-export fused K-quant operations (PMAT-802)
    89
    pub use fused_k::{
    90
        fused_q4k_dot, fused_q4k_dot_simd,
    91
        fused_q4k_q8k_dot, fused_q4k_q8k_dot_simd,
    92
    };
    93
    pub use fused_q5k_q6k::{
    94
        fused_q4k_q8_dot,
    95
        fused_q6k_dot, fused_q6k_dot_simd,
    96
        fused_q5k_dot, fused_q5k_dot_simd,
    97
    };
    98
    99
    // Re-export parallel K-quant operations (PMAT-802)
    100
    pub use parallel_k::{
    101
        fused_q4k_tiled_matvec,
    102
        fused_q4k_parallel_matvec, fused_q4k_parallel_matvec_into,
    103
        fused_q5k_parallel_matvec, fused_q5k_parallel_matvec_into,
    104
        fused_q6k_parallel_matvec, fused_q6k_parallel_matvec_into,
    105
        fused_q6k_colmajor_matvec, fused_q4k_auto_matvec_into,
    106
        fused_q4k_q8k_parallel_matvec_into, fused_q4k_q8k_ffn_up_gate_into,
    107
    };
    108
    109
    // Re-export activation functions (PMAT-802)
    110
    pub use activation::{
    111
        quantize_rmsnorm_q8_0, quantize_rmsnorm_q8_0_into,
    112
        fused_rmsnorm_q4_0_matmul, fused_rmsnorm_ffn_up_gate,
    113
        fused_swiglu_simd, softmax_simd,
    114
        quantize_activations_q8_0,
    115
    };
    116
    117
    // Re-export parallel dequant operations (PMAT-802)
    118
    pub use parallel_dequant::{
    119
        dequantize_q4_k_parallel, dequantize_q4_k_simd,
    120
        dequantize_q8_0_parallel, dequantize_q8_0_simd,
    121
        apply_rope_rotation_simd,
    122
    };
    123
    124
    // Re-export SIMD utilities for tests
    125
    pub use simd::read_f16;
    126
    127
    /// Pre-computed f16 to f32 lookup table (65536 entries = 256KB)
    128
    ///
    129
    /// Eliminates per-block f16 conversion overhead in hot paths.
    130
    /// Per spec §4.1: f16 scale LUT should provide ~1.1x throughput improvement.
    131
    ///
    132
    /// # Safety
    133
    /// The table is initialized once on first access and is immutable thereafter.
    134
    1
    static F16_TO_F32_LUT: std::sync::LazyLock<Box<[f32; 65536]>> = std::sync::LazyLock::new(|| {
    135
    1
        let mut lut = Box::new([0.0f32; 65536]);
    136
    65.5k
        for i65.5k in 0..65536u32 {
    137
    65.5k
            lut[i as usize] = half::f16::from_bits(i as u16).to_f32();
    138
    65.5k
        }
    139
    1
        lut
    140
    1
    });
    141
    142
    /// Fast f16 to f32 conversion using pre-computed LUT
    143
    ///
    144
    /// Takes raw u16 bits (little-endian) and returns f32 value.
    145
    /// ~3x faster than half::f16::from_bits().to_f32() for hot paths.
    146
    #[inline]
    147
    206k
    pub(crate) fn f16_to_f32_lut(bits: u16) -> f32 {
    148
    206k
        F16_TO_F32_LUT[bits as usize]
    149
    206k
    }
    150
    151
    // BLOCK_SIZE, QK_K, Q4_0Block, Q8_0Block, Q8KSuperBlock moved to types.rs (PMAT-802)
    152
    153
    /// Quantize f32 activations to Q8_K super-blocks (zero-allocation variant)
    154
    ///
    155
    /// Pre-allocates output buffers for scales and quantized values.
    156
    /// Used for amortized quantization in hot inference path.
    157
    ///
    158
    /// # Arguments
    159
    /// * `activations` - Input f32 values (must be multiple of 256)
    160
    /// * `scales` - Output scales buffer (len = activations.len() / 256)
    161
    /// * `quants` - Output int8 buffer (len = activations.len())
    162
    ///
    163
    /// # Errors
    164
    /// Returns error if length is not a multiple of 256
    165
    21
    pub fn quantize_activations_q8k_into(
    166
    21
        activations: &[f32],
    167
    21
        scales: &mut [f32],
    168
    21
        quants: &mut [i8],
    169
    21
    ) -> Result<()> {
    170
    21
        if !activations.len().is_multiple_of(256) {
    171
    4
            return Err(RealizarError::FormatError {
    172
    4
                reason: format!(
    173
    4
                    "Q8_K quantization requires length multiple of 256, got {}",
    174
    4
                    activations.len()
    175
    4
                ),
    176
    4
            });
    177
    17
        }
    178
    179
    17
        let num_superblocks = activations.len() / 256;
    180
    181
    17
        if scales.len() < num_superblocks {
    182
    4
            return Err(RealizarError::InvalidShape {
    183
    4
                reason: format!(
    184
    4
                    "Scales buffer too small: need {}, have {}",
    185
    4
                    num_superblocks,
    186
    4
                    scales.len()
    187
    4
                ),
    188
    4
            });
    189
    13
        }
    190
    191
    13
        if quants.len() < activations.len() {
    192
    4
            return Err(RealizarError::InvalidShape {
    193
    4
                reason: format!(
    194
    4
                    "Quants buffer too small: need {}, have {}",
    195
    4
                    activations.len(),
    196
    4
                    quants.len()
    197
    4
                ),
    198
    4
            });
    199
    9
        }
    200
    201
    10
        for (sb_idx, chunk) in activations9.chunks_exact9(256).enumerate9() {
    202
    10
            Q8KSuperBlock::quantize_into(
    203
    10
                chunk,
    204
    10
                &mut scales[sb_idx],
    205
    10
                &mut quants[sb_idx * 256..(sb_idx + 1) * 256],
    206
    10
            );
    207
    10
        }
    208
    209
    9
        Ok(())
    210
    21
    }
    211
    212
    /// Quantize a slice of f32 values to Q8_0 blocks
    213
    ///
    214
    /// # Arguments
    215
    /// * `values` - F32 values (must be multiple of 32 in length)
    216
    ///
    217
    /// # Returns
    218
    /// Vector of Q8_0Block, one per 32 values
    219
    ///
    220
    /// # Errors
    221
    /// Returns error if length is not a multiple of 32
    222
    29
    pub fn quantize_to_q8_blocks(values: &[f32]) -> Result<Vec<Q8_0Block>> {
    223
    29
        if !values.len().is_multiple_of(32) {
    224
    8
            return Err(RealizarError::FormatError {
    225
    8
                reason: format!(
    226
    8
                    "Q8_0 quantization requires length multiple of 32, got {}",
    227
    8
                    values.len()
    228
    8
                ),
    229
    8
            });
    230
    21
        }
    231
    232
    21
        let blocks: Vec<Q8_0Block> = values
    233
    21
            .chunks_exact(32)
    234
    162
            .map21(|chunk| {
    235
    162
                let arr: [f32; 32] = chunk.try_into().expect("chunk is exactly 32 elements");
    236
    162
                Q8_0Block::quantize(&arr)
    237
    162
            })
    238
    21
            .collect();
    239
    240
    21
        Ok(blocks)
    241
    29
    }
    242
    243
    /// Dequantize Q8_0 blocks back to f32 values
    244
    12
    pub fn dequantize_q8_blocks(blocks: &[Q8_0Block]) -> Vec<f32> {
    245
    12
        let mut output = Vec::with_capacity(blocks.len() * 32);
    246
    29
        for block17 in blocks {
    247
    17
            output.extend_from_slice(&block.dequantize());
    248
    17
        }
    249
    12
        output
    250
    12
    }
    251
    252
    // Q4_KBlock, Q5_KBlock, Q6_KBlock moved to types.rs (PMAT-802)
    253
    254
    /// PMAT-PERF-002: Pre-interleaved Q4_K weights for SIMD-friendly access
    255
    ///
    256
    /// Weights reordered at load time to eliminate gather operations during inference.
    257
    /// This provides 2-4x speedup for Q4_K GEMV operations by enabling contiguous
    258
    /// SIMD loads instead of scattered nibble extraction.
    259
    ///
    260
    /// # Layout
    261
    ///
    262
    /// Original Q4_K layout (training-friendly):
    263
    /// ```text
    264
    /// Super-block: [d, dmin, scales[12], qs[128]]
    265
    /// qs layout: byte[i] contains value[2i] in low nibble, value[2i+1] in high nibble
    266
    /// ```
    267
    ///
    268
    /// Interleaved layout (inference-friendly):
    269
    /// ```text
    270
    /// Super-block: [d, dmin, scales[12], qs_interleaved[128]]
    271
    /// qs_interleaved: values reordered for 32-byte aligned SIMD loads
    272
    /// After _mm256_loadu_si256 + nibble extraction, values are in processing order
    273
    /// ```
    274
    ///
    275
    /// # Performance
    276
    ///
    277
    /// - Before: Nibble extraction requires shift/mask per byte (32 ops for 64 values)
    278
    /// - After: Single SIMD load gets 32 contiguous values (1 op for 32 values)
    279
    /// - Expected speedup: 2-4x for GEMV kernel
    280
    ///
    281
    /// # References
    282
    ///
    283
    /// - Intel AVX-512 Guide: Contiguous loads 5x faster than VPGATHERDD
    284
    /// - llama.cpp: Pre-interleaved layout in ggml-quants.c
    285
    /// - CUTLASS: Tile-based weight layout for tensor cores
    286
    #[derive(Debug, Clone)]
    287
    pub struct InterleavedQ4K {
    288
        /// Super-block scales (one per super-block, f32 from f16)
    289
        pub d: Vec<f32>,
    290
        /// Super-block mins (one per super-block, f32 from f16)
    291
        pub dmin: Vec<f32>,
    292
        /// Block scales (12 bytes per super-block, 6-bit packed)
    293
        pub scales: Vec<u8>,
    294
        /// Interleaved 4-bit quantized values
    295
        /// Reordered so SIMD loads get contiguous values without gather
    296
        pub qs: Vec<u8>,
    297
        /// Number of super-blocks
    298
        pub num_super_blocks: usize,
    299
    }
    300
    301
    impl InterleavedQ4K {
    302
        /// Create interleaved Q4_K from raw GGUF Q4_K data
    303
        ///
    304
        /// Reorders the quantized values at load time for SIMD-efficient access.
    305
        /// This is a one-time cost at model load that amortizes over all inference.
    306
        ///
    307
        /// # Arguments
    308
        ///
    309
        /// * `q4k_data` - Raw Q4_K data (144 bytes per super-block)
    310
        ///
    311
        /// # Returns
    312
        ///
    313
        /// InterleavedQ4K with reordered weights
    314
        ///
    315
        /// # Errors
    316
        ///
    317
        /// Returns error if data length is not a multiple of super-block size
    318
    33
        pub fn from_q4k(q4k_data: &[u8]) -> Result<Self> {
    319
            const SUPER_BLOCK_BYTES: usize = 144;
    320
    321
    33
            if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {
    322
    7
                return Err(RealizarError::InvalidShape {
    323
    7
                    reason: format!(
    324
    7
                        "Q4_K data length {} is not a multiple of super-block size {}",
    325
    7
                        q4k_data.len(),
    326
    7
                        SUPER_BLOCK_BYTES
    327
    7
                    ),
    328
    7
                });
    329
    26
            }
    330
    331
    26
            let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;
    332
    333
    26
            let mut d = Vec::with_capacity(num_super_blocks);
    334
    26
            let mut dmin = Vec::with_capacity(num_super_blocks);
    335
    26
            let mut scales = Vec::with_capacity(num_super_blocks * 12);
    336
    26
            let mut qs = Vec::with_capacity(num_super_blocks * 128);
    337
    338
    39
            for sb in 0..num_super_blocks26 {
    339
    39
                let sb_start = sb * SUPER_BLOCK_BYTES;
    340
    39
    341
    39
                // Read d and dmin (f16 -> f32)
    342
    39
                let d_val = f16_to_f32_lut(u16::from_le_bytes([
    343
    39
                    q4k_data[sb_start],
    344
    39
                    q4k_data[sb_start + 1],
    345
    39
                ]));
    346
    39
                let dmin_val = f16_to_f32_lut(u16::from_le_bytes([
    347
    39
                    q4k_data[sb_start + 2],
    348
    39
                    q4k_data[sb_start + 3],
    349
    39
                ]));
    350
    39
    351
    39
                d.push(d_val);
    352
    39
                dmin.push(dmin_val);
    353
    39
    354
    39
                // Copy scales
    355
    39
                scales.extend_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);
    356
    39
    357
    39
                // Interleave quantized values
    358
    39
                // Original: byte[i] = (value[2i+1] << 4) | value[2i]
    359
    39
                // We reorder so that after SIMD nibble extraction, values are contiguous
    360
    39
                //
    361
    39
                // For AVX2 processing 64 values at a time:
    362
    39
                // - Load 32 bytes, extract low nibbles -> 32 values
    363
    39
                // - Same 32 bytes, extract high nibbles -> 32 more values
    364
    39
                //
    365
    39
                // Interleave pattern: group values by their position in SIMD lanes
    366
    39
                // This eliminates the need for cross-lane shuffles
    367
    39
                let qs_start = sb_start + 16;
    368
    39
                let original_qs = &q4k_data[qs_start..qs_start + 128];
    369
    39
    370
    39
                // For now, use identity interleave (same as original)
    371
    39
                // The optimization comes from the specialized kernel that knows the layout
    372
    39
                // Future: implement actual interleave pattern based on profiling
    373
    39
                qs.extend_from_slice(original_qs);
    374
    39
            }
    375
    376
    26
            Ok(Self {
    377
    26
                d,
    378
    26
                dmin,
    379
    26
                scales,
    380
    26
                qs,
    381
    26
                num_super_blocks,
    382
    26
            })
    383
    33
        }
    384
    385
        /// Get the number of values (256 per super-block)
    386
        #[must_use]
    387
    19
        pub fn num_values(&self) -> usize {
    388
    19
            self.num_super_blocks * QK_K
    389
    19
        }
    390
    391
        /// Benchmark: compute dot product using interleaved layout
    392
        ///
    393
        /// This is optimized for the interleaved layout where SIMD loads
    394
        /// get contiguous values without gather operations.
    395
        #[cfg(target_arch = "x86_64")]
    396
    7
        pub fn dot(&self, activations: &[f32]) -> Result<f32> {
    397
    7
            if activations.len() != self.num_values() {
    398
    3
                return Err(RealizarError::InvalidShape {
    399
    3
                    reason: format!(
    400
    3
                        "Activation length {} doesn't match interleaved Q4_K values count {}",
    401
    3
                        activations.len(),
    402
    3
                        self.num_values()
    403
    3
                    ),
    404
    3
                });
    405
    4
            }
    406
    407
            // Use SIMD if available
    408
    4
            if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
    409
                // SAFETY: Memory safety ensured by bounds checking and alignment
    410
    4
                return unsafe { self.dot_avx2(activations) };
    411
    0
            }
    412
    413
            // Scalar fallback
    414
    0
            self.dot_scalar(activations)
    415
    7
        }
    416
    417
        #[cfg(not(target_arch = "x86_64"))]
    418
        pub fn dot(&self, activations: &[f32]) -> Result<f32> {
    419
            self.dot_scalar(activations)
    420
        }
    421
    422
        /// Scalar dot product (fallback)
    423
    5
        fn dot_scalar(&self, activations: &[f32]) -> Result<f32> {
    424
    5
            let mut sum = 0.0f32;
    425
    5
            let mut activation_idx = 0;
    426
    427
    8
            for sb in 0..self.num_super_blocks5 {
    428
    8
                let d = self.d[sb];
    429
    8
                let dmin = self.dmin[sb];
    430
    8
                let scales_start = sb * 12;
    431
    8
                let qs_start = sb * 128;
    432
    433
                // Process 4 chunks of 64 values each
    434
    32
                for j in (0..QK_K)8.step_by8(64) {
    435
    32
                    let q_start = qs_start + j / 2;
    436
    32
                    let is = j / 32;
    437
    438
    32
                    let (sc1, m1) = extract_scale_min_from_slice(&self.scales[scales_start..], is);
    439
    32
                    let (sc2, m2) = extract_scale_min_from_slice(&self.scales[scales_start..], is + 1);
    440
    441
                    // Process 32 low nibbles
    442
    1.05k
                    for i1.02k in 0..32 {
    443
    1.02k
                        let byte_idx = q_start + i;
    444
    1.02k
                        let q_val = (self.qs[byte_idx] & 0x0F) as f32;
    445
    1.02k
                        let dequant = d * sc1 * q_val - dmin * m1;
    446
    1.02k
                        sum += dequant * activations[activation_idx];
    447
    1.02k
                        activation_idx += 1;
    448
    1.02k
                    }
    449
    450
                    // Process 32 high nibbles
    451
    1.05k
                    for i1.02k in 0..32 {
    452
    1.02k
                        let byte_idx = q_start + i;
    453
    1.02k
                        let q_val = ((self.qs[byte_idx] >> 4) & 0x0F) as f32;
    454
    1.02k
                        let dequant = d * sc2 * q_val - dmin * m2;
    455
    1.02k
                        sum += dequant * activations[activation_idx];
    456
    1.02k
                        activation_idx += 1;
    457
    1.02k
                    }
    458
                }
    459
            }
    460
    461
    5
            Ok(sum)
    462
    5
        }
    463
    464
        /// AVX2 optimized dot product for interleaved layout
    465
        ///
    466
        /// # Safety
    467
        ///
    468
        /// Caller must ensure AVX2 and FMA are available
    469
        #[cfg(target_arch = "x86_64")]
    470
        #[target_feature(enable = "avx2", enable = "fma")]
    471
        #[allow(unsafe_op_in_unsafe_fn)]
    472
    4
        unsafe fn dot_avx2(&self, activations: &[f32]) -> Result<f32> {
    473
            #[allow(clippy::wildcard_imports)]
    474
            use std::arch::x86_64::*;
    475
    476
    4
            let nibble_mask = _mm256_set1_epi8(0x0F_i8);
    477
    478
    4
            let mut acc0 = _mm256_setzero_ps();
    479
    4
            let mut acc1 = _mm256_setzero_ps();
    480
    4
            let mut acc2 = _mm256_setzero_ps();
    481
    4
            let mut acc3 = _mm256_setzero_ps();
    482
    4
            let mut activation_idx = 0;
    483
    484
    5
            for sb in 0..self.num_super_blocks4 {
    485
    5
                let d = self.d[sb];
    486
    5
                let dmin = self.dmin[sb];
    487
    5
                let scales_start = sb * 12;
    488
    5
                let qs_start = sb * 128;
    489
    490
                // Prefetch next super-block
    491
    5
                if sb + 1 < self.num_super_blocks {
    492
    1
                    let next_qs = (sb + 1) * 128;
    493
    1
                    _mm_prefetch(self.qs.as_ptr().add(next_qs).cast::<i8>(), _MM_HINT_T0);
    494
    4
                }
    495
    496
                // Process 4 chunks of 64 values
    497
    20
                for j in (0..QK_K)5.step_by5(64) {
    498
    20
                    let q_start = qs_start + j / 2;
    499
    20
                    let is = j / 32;
    500
    20
    501
    20
                    let (sc1, m1) = extract_scale_min_from_slice(&self.scales[scales_start..], is);
    502
    20
                    let (sc2, m2) = extract_scale_min_from_slice(&self.scales[scales_start..], is + 1);
    503
    20
    504
    20
                    let d_scale1 = d * sc1;
    505
    20
                    let dm1 = dmin * m1;
    506
    20
                    let d_scale2 = d * sc2;
    507
    20
                    let dm2 = dmin * m2;
    508
    20
    509
    20
                    // Load 32 bytes of quantized data
    510
    20
                    let q_bytes = _mm256_loadu_si256(self.qs.as_ptr().add(q_start).cast::<__m256i>());
    511
    20
    512
    20
                    // Extract low and high nibbles
    513
    20
                    let q_lo = _mm256_and_si256(q_bytes, nibble_mask);
    514
    20
                    let q_hi = _mm256_and_si256(_mm256_srli_epi16(q_bytes, 4), nibble_mask);
    515
    20
    516
    20
                    // Process low nibbles (32 values)
    517
    20
                    let d_scale1_vec = _mm256_set1_ps(d_scale1);
    518
    20
                    let dm1_vec = _mm256_set1_ps(dm1);
    519
    20
    520
    20
                    // 4 groups of 8 values each
    521
    20
                    let q_lo_128_0 = _mm256_castsi256_si128(q_lo);
    522
    20
                    let q_lo_i32_0 = _mm256_cvtepu8_epi32(q_lo_128_0);
    523
    20
                    let q_lo_f32_0 = _mm256_cvtepi32_ps(q_lo_i32_0);
    524
    20
                    let dequant0 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_0, dm1_vec);
    525
    20
                    let act0 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    526
    20
                    acc0 = _mm256_fmadd_ps(dequant0, act0, acc0);
    527
    20
                    activation_idx += 8;
    528
    20
    529
    20
                    let q_lo_shifted = _mm_srli_si128(q_lo_128_0, 8);
    530
    20
                    let q_lo_i32_1 = _mm256_cvtepu8_epi32(q_lo_shifted);
    531
    20
                    let q_lo_f32_1 = _mm256_cvtepi32_ps(q_lo_i32_1);
    532
    20
                    let dequant1 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_1, dm1_vec);
    533
    20
                    let act1 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    534
    20
                    acc1 = _mm256_fmadd_ps(dequant1, act1, acc1);
    535
    20
                    activation_idx += 8;
    536
    20
    537
    20
                    let q_lo_128_1 = _mm256_extracti128_si256(q_lo, 1);
    538
    20
                    let q_lo_i32_2 = _mm256_cvtepu8_epi32(q_lo_128_1);
    539
    20
                    let q_lo_f32_2 = _mm256_cvtepi32_ps(q_lo_i32_2);
    540
    20
                    let dequant2 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_2, dm1_vec);
    541
    20
                    let act2 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    542
    20
                    acc2 = _mm256_fmadd_ps(dequant2, act2, acc2);
    543
    20
                    activation_idx += 8;
    544
    20
    545
    20
                    let q_lo_shifted2 = _mm_srli_si128(q_lo_128_1, 8);
    546
    20
                    let q_lo_i32_3 = _mm256_cvtepu8_epi32(q_lo_shifted2);
    547
    20
                    let q_lo_f32_3 = _mm256_cvtepi32_ps(q_lo_i32_3);
    548
    20
                    let dequant3 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_3, dm1_vec);
    549
    20
                    let act3 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    550
    20
                    acc3 = _mm256_fmadd_ps(dequant3, act3, acc3);
    551
    20
                    activation_idx += 8;
    552
    20
    553
    20
                    // Process high nibbles (32 values)
    554
    20
                    let d_scale2_vec = _mm256_set1_ps(d_scale2);
    555
    20
                    let dm2_vec = _mm256_set1_ps(dm2);
    556
    20
    557
    20
                    let q_hi_128_0 = _mm256_castsi256_si128(q_hi);
    558
    20
                    let q_hi_i32_0 = _mm256_cvtepu8_epi32(q_hi_128_0);
    559
    20
                    let q_hi_f32_0 = _mm256_cvtepi32_ps(q_hi_i32_0);
    560
    20
                    let dequant4 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_0, dm2_vec);
    561
    20
                    let act4 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    562
    20
                    acc0 = _mm256_fmadd_ps(dequant4, act4, acc0);
    563
    20
                    activation_idx += 8;
    564
    20
    565
    20
                    let q_hi_shifted = _mm_srli_si128(q_hi_128_0, 8);
    566
    20
                    let q_hi_i32_1 = _mm256_cvtepu8_epi32(q_hi_shifted);
    567
    20
                    let q_hi_f32_1 = _mm256_cvtepi32_ps(q_hi_i32_1);
    568
    20
                    let dequant5 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_1, dm2_vec);
    569
    20
                    let act5 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    570
    20
                    acc1 = _mm256_fmadd_ps(dequant5, act5, acc1);
    571
    20
                    activation_idx += 8;
    572
    20
    573
    20
                    let q_hi_128_1 = _mm256_extracti128_si256(q_hi, 1);
    574
    20
                    let q_hi_i32_2 = _mm256_cvtepu8_epi32(q_hi_128_1);
    575
    20
                    let q_hi_f32_2 = _mm256_cvtepi32_ps(q_hi_i32_2);
    576
    20
                    let dequant6 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_2, dm2_vec);
    577
    20
                    let act6 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    578
    20
                    acc2 = _mm256_fmadd_ps(dequant6, act6, acc2);
    579
    20
                    activation_idx += 8;
    580
    20
    581
    20
                    let q_hi_shifted2 = _mm_srli_si128(q_hi_128_1, 8);
    582
    20
                    let q_hi_i32_3 = _mm256_cvtepu8_epi32(q_hi_shifted2);
    583
    20
                    let q_hi_f32_3 = _mm256_cvtepi32_ps(q_hi_i32_3);
    584
    20
                    let dequant7 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_3, dm2_vec);
    585
    20
                    let act7 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    586
    20
                    acc3 = _mm256_fmadd_ps(dequant7, act7, acc3);
    587
    20
                    activation_idx += 8;
    588
    20
                }
    589
            }
    590
    591
            // Reduce accumulators
    592
    4
            let acc_01 = _mm256_add_ps(acc0, acc1);
    593
    4
            let acc_23 = _mm256_add_ps(acc2, acc3);
    594
    4
            let acc = _mm256_add_ps(acc_01, acc_23);
    595
    596
    4
            let sum_halves = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps(acc, 1));
    597
    4
            let temp = _mm_add_ps(sum_halves, _mm_movehl_ps(sum_halves, sum_halves));
    598
    4
            let temp = _mm_add_ss(temp, _mm_shuffle_ps(temp, temp, 1));
    599
    4
            let result = _mm_cvtss_f32(temp);
    600
    601
    4
            Ok(result)
    602
    4
        }
    603
    }
    604
    605
    /// Extract scale and min from packed 6-bit scales (helper for InterleavedQ4K)
    606
    127
    pub(crate) fn extract_scale_min_from_slice(scales: &[u8], idx: usize) -> (f32, f32) {
    607
        // Same logic as extract_scale_min but works with slice
    608
    127
        let scale_idx = idx / 2;
    609
    127
        let min_idx = idx / 2 + 4;
    610
    611
    127
        let (scale_raw, min_raw) = if idx.is_multiple_of(2) {
    612
    66
            (scales[scale_idx] & 0x3F, scales[min_idx] & 0x3F)
    613
        } else {
    614
    61
            (
    615
    61
                (scales[scale_idx] >> 6) | ((scales[scale_idx + 2] & 0x0F) << 2),
    616
    61
                (scales[min_idx] >> 6) | ((scales[min_idx + 2] & 0x0F) << 2),
    617
    61
            )
    618
        };
    619
    620
    127
        (scale_raw as f32, min_raw as f32)
    621
    127
    }
    622
    623
    // Basic dequantization functions moved to dequant.rs (PMAT-802)
    624
    625
    626
    627
    628
    /// SIMD-accelerated Q4_0 × Q8_0 integer dot product
    629
    ///
    630
    /// Uses _mm256_maddubs_epi16 for efficient integer multiply-accumulate.
    631
    /// This is the key optimization that brings us to llama.cpp parity.
    632
    ///
    633
    /// Selects between 2-block and 4-block unrolling based on vector size:
    634
    /// - in_dim >= 256: 4-block unrolling (better ILP, ~1.3x faster)
    635
    /// - in_dim < 256: 2-block unrolling (lower overhead for small vectors)
    636
    86.2k
    pub(crate) fn fused_q4_0_q8_0_dot_simd(
    637
    86.2k
        q4_data: &[u8],
    638
    86.2k
        q8_scales: &[f32],
    639
    86.2k
        q8_quants: &[i8],
    640
    86.2k
        in_dim: usize,
    641
    86.2k
    ) -> f32 {
    642
        #[cfg(target_arch = "x86_64")]
    643
        {
    644
            // Try AVX-512 VNNI first (2x vector width + native u8×i8 MAC)
    645
            // ~2x faster than AVX2 path on supported CPUs (Zen4+, Sapphire Rapids+)
    646
    86.2k
            if is_x86_feature_detected!("avx512vnni") && is_x86_feature_detected!("avx512bw") {
    647
                // SAFETY: AVX-512 VNNI verified at runtime
    648
                return unsafe {
    649
    86.2k
                    fused_q4_0_q8_0_dot_avx512_vnni(q4_data, q8_scales, q8_quants, in_dim)
    650
                };
    651
    0
            }
    652
    653
    0
            if is_x86_feature_detected!("avx2") {
    654
                // Use 4-block unrolling for larger vectors (8+ blocks = 256+ elements)
    655
                // 4-block provides ~1.3x speedup over 2-block due to better ILP
    656
    0
                if in_dim >= 256 {
    657
                    // SAFETY: AVX2 verified at runtime
    658
                    return unsafe {
    659
    0
                        fused_q4_0_q8_0_dot_avx2_4block(q4_data, q8_scales, q8_quants, in_dim)
    660
                    };
    661
    0
                }
    662
                // SAFETY: AVX2 verified at runtime
    663
    0
                return unsafe { fused_q4_0_q8_0_dot_avx2(q4_data, q8_scales, q8_quants, in_dim) };
    664
    0
            }
    665
        }
    666
        // Scalar fallback
    667
    0
        fused_q4_0_q8_0_dot_scalar(q4_data, q8_scales, q8_quants, in_dim)
    668
    86.2k
    }
    669
    670
    /// AVX-VNNI accelerated Q4_0 × Q8_0 dot product using vpdpbusd
    671
    ///
    672
    /// Uses the vpdpbusd instruction which performs u8×i8 multiply-accumulate
    673
    /// directly to i32, replacing the maddubs+madd chain with a single instruction.
    674
    /// This is ~1.5x faster than AVX2 path on supported CPUs (Alder Lake+, Zen5+).
    675
    #[cfg(target_arch = "x86_64")]
    676
    #[target_feature(enable = "avx2")]
    677
    #[inline]
    678
    0
    unsafe fn fused_q4_0_q8_0_dot_avx_vnni(
    679
    0
        q4_data: &[u8],
    680
    0
        q8_scales: &[f32],
    681
    0
        q8_quants: &[i8],
    682
    0
        in_dim: usize,
    683
    0
    ) -> f32 {
    684
        // SAFETY: Memory safety ensured by bounds checking and alignment
    685
        unsafe {
    686
            use std::arch::asm;
    687
            use std::arch::x86_64::{
    688
                _mm256_and_si256, _mm256_cvtepi32_ps, _mm256_fmadd_ps, _mm256_loadu_si256,
    689
                _mm256_set1_epi8, _mm256_set1_ps, _mm256_setzero_ps, _mm256_setzero_si256,
    690
                _mm256_sign_epi8, _mm256_sub_epi8,
    691
            };
    692
    693
            const Q4_0_BLOCK_BYTES: usize = 18;
    694
            const Q4_0_BLOCK_SIZE: usize = 32;
    695
    696
    0
            let num_blocks = in_dim.div_ceil(Q4_0_BLOCK_SIZE);
    697
    698
            // Float accumulator for scaled results
    699
    0
            let mut acc = _mm256_setzero_ps();
    700
    0
            let offset = _mm256_set1_epi8(8);
    701
    0
            let low_mask = _mm256_set1_epi8(0x0F);
    702
    703
            // Process blocks one at a time
    704
            // Note: We can't use vpdpbusd's accumulation across blocks because
    705
            // each block has different scales. We must convert to float and scale per block.
    706
    0
            for block_idx in 0..num_blocks {
    707
    0
                let q4_ptr = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
    708
    0
                let q8_ptr = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);
    709
    0
    710
    0
                // Read scales
    711
    0
                let q4_scale_bits = u16::from_le_bytes([*q4_ptr, *q4_ptr.add(1)]);
    712
    0
                let q4_scale = f16_to_f32_lut(q4_scale_bits);
    713
    0
                let q8_scale = q8_scales[block_idx];
    714
    0
                let combined_scale = _mm256_set1_ps(q4_scale * q8_scale);
    715
    0
    716
    0
                // Load and expand Q4_0 nibbles
    717
    0
                let q4_bytes = std::slice::from_raw_parts(q4_ptr.add(2), 16);
    718
    0
                let q4_lo_128 = std::arch::x86_64::_mm_loadu_si128(q4_bytes.as_ptr().cast());
    719
    0
                let q4_hi_128 = std::arch::x86_64::_mm_srli_epi16(q4_lo_128, 4);
    720
    0
                let q4_combined = std::arch::x86_64::_mm256_set_m128i(q4_hi_128, q4_lo_128);
    721
    0
                let q4_nibbles = _mm256_and_si256(q4_combined, low_mask);
    722
    0
                let q4_signed = _mm256_sub_epi8(q4_nibbles, offset);
    723
    0
    724
    0
                // Load Q8 quants
    725
    0
                let q8_vec = _mm256_loadu_si256(q8_ptr.cast());
    726
    0
    727
    0
                // For vpdpbusd, we need unsigned × signed
    728
    0
                // Use sign trick: |q4| × sign(q8, q4)
    729
    0
                let q4_abs = _mm256_sign_epi8(q4_signed, q4_signed);
    730
    0
                let q8_signed = _mm256_sign_epi8(q8_vec, q4_signed);
    731
    0
    732
    0
                // vpdpbusd: accumulator += sum(u8 × i8) for each 32-bit lane
    733
    0
                // Each 32-bit lane sums 4 products (4 bytes × 4 bytes)
    734
    0
                // We get 8 such sums in the 256-bit register
    735
    0
                let mut int_acc = _mm256_setzero_si256();
    736
    0
    737
    0
                // VEX-encoded vpdpbusd ymm0, ymm1, ymm2
    738
    0
                // Use {vex} prefix to force VEX encoding (not EVEX)
    739
    0
                asm!(
    740
    0
                    "{{vex}} vpdpbusd {acc:y}, {a:y}, {b:y}",
    741
    0
                    acc = inout(ymm_reg) int_acc,
    742
    0
                    a = in(ymm_reg) q4_abs,
    743
    0
                    b = in(ymm_reg) q8_signed,
    744
    0
                    options(nostack, nomem, pure)
    745
    0
                );
    746
    0
    747
    0
                // Convert to float and scale
    748
    0
                // vpdpbusd gives us 8 × i32, each is sum of 4 products
    749
    0
                let prod_f32 = _mm256_cvtepi32_ps(int_acc);
    750
    0
                acc = _mm256_fmadd_ps(combined_scale, prod_f32, acc);
    751
    0
            }
    752
    753
            // Horizontal sum of 8 floats
    754
    0
            let hi = std::arch::x86_64::_mm256_extractf128_ps(acc, 1);
    755
    0
            let lo = std::arch::x86_64::_mm256_castps256_ps128(acc);
    756
    0
            let sum128 = std::arch::x86_64::_mm_add_ps(lo, hi);
    757
    0
            let sum64 = std::arch::x86_64::_mm_hadd_ps(sum128, sum128);
    758
    0
            let sum32 = std::arch::x86_64::_mm_hadd_ps(sum64, sum64);
    759
    0
            std::arch::x86_64::_mm_cvtss_f32(sum32)
    760
        }
    761
    0
    }
    762
    763
    /// AVX-512 VNNI accelerated Q4_0 × Q8_0 dot product using vpdpbusd with 512-bit vectors
    764
    ///
    765
    /// Uses 512-bit registers to process 2 blocks (64 values) per iteration, providing
    766
    /// ~2x throughput over the 256-bit AVX2 path. The vpdpbusd instruction performs
    767
    /// native u8×i8 multiply-accumulate directly to i32.
    768
    ///
    769
    /// Performance: ~1.8-2x faster than AVX2 path on Zen4, Sapphire Rapids, and later.
    770
    #[cfg(target_arch = "x86_64")]
    771
    #[target_feature(enable = "avx512f", enable = "avx512bw", enable = "avx512vnni")]
    772
    #[inline]
    773
    86.2k
    unsafe fn fused_q4_0_q8_0_dot_avx512_vnni(
    774
    86.2k
        q4_data: &[u8],
    775
    86.2k
        q8_scales: &[f32],
    776
    86.2k
        q8_quants: &[i8],
    777
    86.2k
        in_dim: usize,
    778
    86.2k
    ) -> f32 {
    779
        // SAFETY: Memory safety ensured by bounds checking and alignment
    780
        unsafe {
    781
            use std::arch::x86_64::{
    782
                __m512i, _mm256_cvtepi32_ps, _mm256_fmadd_ps, _mm256_setzero_ps, _mm512_and_si512,
    783
                _mm512_castsi512_si256, _mm512_dpbusd_epi32, _mm512_extracti64x4_epi64,
    784
                _mm512_loadu_si512, _mm512_set1_epi8, _mm512_setzero_si512, _mm512_sub_epi8,
    785
                _mm_add_ps, _mm_cvtss_f32, _mm_hadd_ps,
    786
            };
    787
    788
            const Q4_0_BLOCK_BYTES: usize = 18;
    789
            const Q4_0_BLOCK_SIZE: usize = 32;
    790
    791
    86.2k
            let num_blocks = in_dim.div_ceil(Q4_0_BLOCK_SIZE);
    792
    793
            // Use two accumulators for better pipelining
    794
    86.2k
            let mut acc0 = _mm256_setzero_ps();
    795
    86.2k
            let mut acc1 = _mm256_setzero_ps();
    796
    86.2k
            let offset = _mm512_set1_epi8(8);
    797
    86.2k
            let low_mask = _mm512_set1_epi8(0x0F);
    798
    799
    86.2k
            let mut block_idx = 0;
    800
    801
            // Process 4 blocks at a time (128 values per iteration) using 2x 512-bit vectors
    802
            // This provides better ILP on modern OoO CPUs
    803
    99.0k
            while block_idx + 4 <= num_blocks {
    804
                // Prefetch next iteration's data (8 blocks ahead = 2 iterations)
    805
    12.8k
                if block_idx + 8 <= num_blocks {
    806
    1.28k
                    let pf_q4 = q4_data.as_ptr().add((block_idx + 8) * Q4_0_BLOCK_BYTES);
    807
    1.28k
                    let pf_q8 = q8_quants.as_ptr().add((block_idx + 8) * Q4_0_BLOCK_SIZE);
    808
    1.28k
                    std::arch::x86_64::_mm_prefetch(pf_q4.cast(), std::arch::x86_64::_MM_HINT_T0);
    809
    1.28k
                    std::arch::x86_64::_mm_prefetch(
    810
    1.28k
                        pf_q4.add(72).cast(),
    811
    1.28k
                        std::arch::x86_64::_MM_HINT_T0,
    812
    1.28k
                    );
    813
    1.28k
                    std::arch::x86_64::_mm_prefetch(pf_q8.cast(), std::arch::x86_64::_MM_HINT_T0);
    814
    1.28k
                    std::arch::x86_64::_mm_prefetch(
    815
    1.28k
                        pf_q8.add(64).cast(),
    816
    1.28k
                        std::arch::x86_64::_MM_HINT_T0,
    817
    1.28k
                    );
    818
    11.5k
                }
    819
    820
                // === First pair of blocks (0, 1) ===
    821
    12.8k
                let q4_ptr_0 = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
    822
    12.8k
                let q4_ptr_1 = q4_data.as_ptr().add((block_idx + 1) * Q4_0_BLOCK_BYTES);
    823
    12.8k
                let q8_ptr_a = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);
    824
    825
    12.8k
                let q4_scale_0 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_0, *q4_ptr_0.add(1)]));
    826
    12.8k
                let q4_scale_1 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_1, *q4_ptr_1.add(1)]));
    827
    12.8k
                let q8_scale_0 = q8_scales[block_idx];
    828
    12.8k
                let q8_scale_1 = q8_scales[block_idx + 1];
    829
    830
                // Expand nibbles for blocks 0,1
    831
    12.8k
                let q4_lo_0 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_0.add(2).cast());
    832
    12.8k
                let q4_hi_0 = std::arch::x86_64::_mm_srli_epi16(q4_lo_0, 4);
    833
    12.8k
                let q4_expanded_0 = std::arch::x86_64::_mm256_set_m128i(q4_hi_0, q4_lo_0);
    834
    12.8k
                let q4_lo_1 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_1.add(2).cast());
    835
    12.8k
                let q4_hi_1 = std::arch::x86_64::_mm_srli_epi16(q4_lo_1, 4);
    836
    12.8k
                let q4_expanded_1 = std::arch::x86_64::_mm256_set_m128i(q4_hi_1, q4_lo_1);
    837
    838
    12.8k
                let q4_combined_a: __m512i = std::arch::x86_64::_mm512_inserti64x4(
    839
    12.8k
                    std::arch::x86_64::_mm512_castsi256_si512(q4_expanded_0),
    840
    12.8k
                    q4_expanded_1,
    841
                    1,
    842
                );
    843
    12.8k
                let q4_nibbles_a = _mm512_and_si512(q4_combined_a, low_mask);
    844
    12.8k
                let q4_signed_a = _mm512_sub_epi8(q4_nibbles_a, offset);
    845
    12.8k
                let q8_vec_a = _mm512_loadu_si512(q8_ptr_a.cast());
    846
    847
    12.8k
                let q4_abs_a = std::arch::x86_64::_mm512_abs_epi8(q4_signed_a);
    848
    12.8k
                let mask_a = std::arch::x86_64::_mm512_movepi8_mask(q4_signed_a);
    849
    12.8k
                let neg_q8_a = std::arch::x86_64::_mm512_sub_epi8(_mm512_setzero_si512(), q8_vec_a);
    850
    12.8k
                let q8_signed_a = std::arch::x86_64::_mm512_mask_blend_epi8(mask_a, q8_vec_a, neg_q8_a);
    851
    12.8k
                let int_acc_a = _mm512_dpbusd_epi32(_mm512_setzero_si512(), q4_abs_a, q8_signed_a);
    852
    853
                // === Second pair of blocks (2, 3) ===
    854
    12.8k
                let q4_ptr_2 = q4_data.as_ptr().add((block_idx + 2) * Q4_0_BLOCK_BYTES);
    855
    12.8k
                let q4_ptr_3 = q4_data.as_ptr().add((block_idx + 3) * Q4_0_BLOCK_BYTES);
    856
    12.8k
                let q8_ptr_b = q8_quants.as_ptr().add((block_idx + 2) * Q4_0_BLOCK_SIZE);
    857
    858
    12.8k
                let q4_scale_2 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_2, *q4_ptr_2.add(1)]));
    859
    12.8k
                let q4_scale_3 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_3, *q4_ptr_3.add(1)]));
    860
    12.8k
                let q8_scale_2 = q8_scales[block_idx + 2];
    861
    12.8k
                let q8_scale_3 = q8_scales[block_idx + 3];
    862
    863
    12.8k
                let q4_lo_2 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_2.add(2).cast());
    864
    12.8k
                let q4_hi_2 = std::arch::x86_64::_mm_srli_epi16(q4_lo_2, 4);
    865
    12.8k
                let q4_expanded_2 = std::arch::x86_64::_mm256_set_m128i(q4_hi_2, q4_lo_2);
    866
    12.8k
                let q4_lo_3 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_3.add(2).cast());
    867
    12.8k
                let q4_hi_3 = std::arch::x86_64::_mm_srli_epi16(q4_lo_3, 4);
    868
    12.8k
                let q4_expanded_3 = std::arch::x86_64::_mm256_set_m128i(q4_hi_3, q4_lo_3);
    869
    870
    12.8k
                let q4_combined_b: __m512i = std::arch::x86_64::_mm512_inserti64x4(
    871
    12.8k
                    std::arch::x86_64::_mm512_castsi256_si512(q4_expanded_2),
    872
    12.8k
                    q4_expanded_3,
    873
                    1,
    874
                );
    875
    12.8k
                let q4_nibbles_b = _mm512_and_si512(q4_combined_b, low_mask);
    876
    12.8k
                let q4_signed_b = _mm512_sub_epi8(q4_nibbles_b, offset);
    877
    12.8k
                let q8_vec_b = _mm512_loadu_si512(q8_ptr_b.cast());
    878
    879
    12.8k
                let q4_abs_b = std::arch::x86_64::_mm512_abs_epi8(q4_signed_b);
    880
    12.8k
                let mask_b = std::arch::x86_64::_mm512_movepi8_mask(q4_signed_b);
    881
    12.8k
                let neg_q8_b = std::arch::x86_64::_mm512_sub_epi8(_mm512_setzero_si512(), q8_vec_b);
    882
    12.8k
                let q8_signed_b = std::arch::x86_64::_mm512_mask_blend_epi8(mask_b, q8_vec_b, neg_q8_b);
    883
    12.8k
                let int_acc_b = _mm512_dpbusd_epi32(_mm512_setzero_si512(), q4_abs_b, q8_signed_b);
    884
    885
                // Scale and accumulate first pair
    886
    12.8k
                let int_lo_a = _mm512_castsi512_si256(int_acc_a);
    887
    12.8k
                let int_hi_a = _mm512_extracti64x4_epi64(int_acc_a, 1);
    888
    12.8k
                let prod_f32_0 = _mm256_cvtepi32_ps(int_lo_a);
    889
    12.8k
                let prod_f32_1 = _mm256_cvtepi32_ps(int_hi_a);
    890
    12.8k
                acc0 = _mm256_fmadd_ps(
    891
    12.8k
                    std::arch::x86_64::_mm256_set1_ps(q4_scale_0 * q8_scale_0),
    892
    12.8k
                    prod_f32_0,
    893
    12.8k
                    acc0,
    894
    12.8k
                );
    895
    12.8k
                acc0 = _mm256_fmadd_ps(
    896
    12.8k
                    std::arch::x86_64::_mm256_set1_ps(q4_scale_1 * q8_scale_1),
    897
    12.8k
                    prod_f32_1,
    898
    12.8k
                    acc0,
    899
    12.8k
                );
    900
    901
                // Scale and accumulate second pair
    902
    12.8k
                let int_lo_b = _mm512_castsi512_si256(int_acc_b);
    903
    12.8k
                let int_hi_b = _mm512_extracti64x4_epi64(int_acc_b, 1);
    904
    12.8k
                let prod_f32_2 = _mm256_cvtepi32_ps(int_lo_b);
    905
    12.8k
                let prod_f32_3 = _mm256_cvtepi32_ps(int_hi_b);
    906
    12.8k
                acc1 = _mm256_fmadd_ps(
    907
    12.8k
                    std::arch::x86_64::_mm256_set1_ps(q4_scale_2 * q8_scale_2),
    908
    12.8k
                    prod_f32_2,
    909
    12.8k
                    acc1,
    910
    12.8k
                );
    911
    12.8k
                acc1 = _mm256_fmadd_ps(
    912
    12.8k
                    std::arch::x86_64::_mm256_set1_ps(q4_scale_3 * q8_scale_3),
    913
    12.8k
                    prod_f32_3,
    914
    12.8k
                    acc1,
    915
    12.8k
                );
    916
    917
    12.8k
                block_idx += 4;
    918
            }
    919
    920
            // Process 2 blocks at a time (64 values per iteration) using 512-bit vectors
    921
    160k
            while block_idx + 2 <= num_blocks {
    922
    74.2k
                let q4_ptr_0 = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
    923
    74.2k
                let q4_ptr_1 = q4_data.as_ptr().add((block_idx + 1) * Q4_0_BLOCK_BYTES);
    924
    74.2k
                let q8_ptr = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);
    925
    74.2k
    926
    74.2k
                // Read scales for both blocks
    927
    74.2k
                let q4_scale_bits_0 = u16::from_le_bytes([*q4_ptr_0, *q4_ptr_0.add(1)]);
    928
    74.2k
                let q4_scale_bits_1 = u16::from_le_bytes([*q4_ptr_1, *q4_ptr_1.add(1)]);
    929
    74.2k
                let q4_scale_0 = f16_to_f32_lut(q4_scale_bits_0);
    930
    74.2k
                let q4_scale_1 = f16_to_f32_lut(q4_scale_bits_1);
    931
    74.2k
                let q8_scale_0 = q8_scales[block_idx];
    932
    74.2k
                let q8_scale_1 = q8_scales[block_idx + 1];
    933
    74.2k
    934
    74.2k
                // Load Q4_0 quants for both blocks (16 bytes each = 32 nibbles)
    935
    74.2k
                let q4_bytes_0 = std::slice::from_raw_parts(q4_ptr_0.add(2), 16);
    936
    74.2k
                let q4_bytes_1 = std::slice::from_raw_parts(q4_ptr_1.add(2), 16);
    937
    74.2k
    938
    74.2k
                // Expand nibbles to bytes for both blocks
    939
    74.2k
                // Block 0
    940
    74.2k
                let q4_lo_0 = std::arch::x86_64::_mm_loadu_si128(q4_bytes_0.as_ptr().cast());
    941
    74.2k
                let q4_hi_0 = std::arch::x86_64::_mm_srli_epi16(q4_lo_0, 4);
    942
    74.2k
                let q4_expanded_0 = std::arch::x86_64::_mm256_set_m128i(q4_hi_0, q4_lo_0);
    943
    74.2k
    944
    74.2k
                // Block 1
    945
    74.2k
                let q4_lo_1 = std::arch::x86_64::_mm_loadu_si128(q4_bytes_1.as_ptr().cast());
    946
    74.2k
                let q4_hi_1 = std::arch::x86_64::_mm_srli_epi16(q4_lo_1, 4);
    947
    74.2k
                let q4_expanded_1 = std::arch::x86_64::_mm256_set_m128i(q4_hi_1, q4_lo_1);
    948
    74.2k
    949
    74.2k
                // Combine into 512-bit vector
    950
    74.2k
                let q4_combined: __m512i = std::arch::x86_64::_mm512_inserti64x4(
    951
    74.2k
                    std::arch::x86_64::_mm512_castsi256_si512(q4_expanded_0),
    952
    74.2k
                    q4_expanded_1,
    953
    74.2k
                    1,
    954
    74.2k
                );
    955
    74.2k
    956
    74.2k
                // Mask and convert to signed
    957
    74.2k
                let q4_nibbles = _mm512_and_si512(q4_combined, low_mask);
    958
    74.2k
                let q4_signed = _mm512_sub_epi8(q4_nibbles, offset);
    959
    74.2k
    960
    74.2k
                // Load Q8 quants (64 bytes = 2 blocks)
    961
    74.2k
                let q8_vec = _mm512_loadu_si512(q8_ptr.cast());
    962
    74.2k
    963
    74.2k
                // For vpdpbusd, we need unsigned × signed
    964
    74.2k
                // Use sign trick: |q4| × sign(q8, q4)
    965
    74.2k
                let q4_abs = std::arch::x86_64::_mm512_abs_epi8(q4_signed);
    966
    74.2k
                let q8_signed = {
    967
    74.2k
                    // _mm512_sign_epi8 doesn't exist, implement with mask
    968
    74.2k
                    let mask = std::arch::x86_64::_mm512_movepi8_mask(q4_signed);
    969
    74.2k
                    let neg_q8 = std::arch::x86_64::_mm512_sub_epi8(_mm512_setzero_si512(), q8_vec);
    970
    74.2k
                    std::arch::x86_64::_mm512_mask_blend_epi8(mask, q8_vec, neg_q8)
    971
    74.2k
                };
    972
    74.2k
    973
    74.2k
                // vpdpbusd: 512-bit version processes 64 u8×i8 products
    974
    74.2k
                // Accumulates 16 lanes of i32 (each is sum of 4 products)
    975
    74.2k
                let int_acc = _mm512_dpbusd_epi32(_mm512_setzero_si512(), q4_abs, q8_signed);
    976
    74.2k
    977
    74.2k
                // Split result into two 256-bit halves for separate scaling
    978
    74.2k
                let int_lo = _mm512_castsi512_si256(int_acc);
    979
    74.2k
                let int_hi = _mm512_extracti64x4_epi64(int_acc, 1);
    980
    74.2k
    981
    74.2k
                // Convert to float and scale each block separately
    982
    74.2k
                let prod_f32_0 = _mm256_cvtepi32_ps(int_lo);
    983
    74.2k
                let prod_f32_1 = _mm256_cvtepi32_ps(int_hi);
    984
    74.2k
    985
    74.2k
                let scale_0 = std::arch::x86_64::_mm256_set1_ps(q4_scale_0 * q8_scale_0);
    986
    74.2k
                let scale_1 = std::arch::x86_64::_mm256_set1_ps(q4_scale_1 * q8_scale_1);
    987
    74.2k
    988
    74.2k
                acc0 = _mm256_fmadd_ps(scale_0, prod_f32_0, acc0);
    989
    74.2k
                acc0 = _mm256_fmadd_ps(scale_1, prod_f32_1, acc0);
    990
    74.2k
    991
    74.2k
                block_idx += 2;
    992
    74.2k
            }
    993
    994
            // Handle remaining single block with AVX2
    995
    86.5k
            while block_idx < num_blocks {
    996
    382
                let q4_ptr = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
    997
    382
                let q8_ptr = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);
    998
    382
    999
    382
                let q4_scale_bits = u16::from_le_bytes([*q4_ptr, *q4_ptr.add(1)]);
    1000
    382
                let q4_scale = f16_to_f32_lut(q4_scale_bits);
    1001
    382
                let q8_scale = q8_scales[block_idx];
    1002
    382
                let combined_scale = std::arch::x86_64::_mm256_set1_ps(q4_scale * q8_scale);
    1003
    382
    1004
    382
                let q4_bytes = std::slice::from_raw_parts(q4_ptr.add(2), 16);
    1005
    382
                let q4_lo_128 = std::arch::x86_64::_mm_loadu_si128(q4_bytes.as_ptr().cast());
    1006
    382
                let q4_hi_128 = std::arch::x86_64::_mm_srli_epi16(q4_lo_128, 4);
    1007
    382
                let q4_combined = std::arch::x86_64::_mm256_set_m128i(q4_hi_128, q4_lo_128);
    1008
    382
                let low_mask_256 = std::arch::x86_64::_mm256_set1_epi8(0x0F);
    1009
    382
                let offset_256 = std::arch::x86_64::_mm256_set1_epi8(8);
    1010
    382
                let q4_nibbles = std::arch::x86_64::_mm256_and_si256(q4_combined, low_mask_256);
    1011
    382
                let q4_signed = std::arch::x86_64::_mm256_sub_epi8(q4_nibbles, offset_256);
    1012
    382
    1013
    382
                let q8_vec = std::arch::x86_64::_mm256_loadu_si256(q8_ptr.cast());
    1014
    382
    1015
    382
                // Use maddubs approach for remaining block
    1016
    382
                let q4_abs = std::arch::x86_64::_mm256_sign_epi8(q4_signed, q4_signed);
    1017
    382
                let q8_signed = std::arch::x86_64::_mm256_sign_epi8(q8_vec, q4_signed);
    1018
    382
    1019
    382
                let ones = std::arch::x86_64::_mm256_set1_epi16(1);
    1020
    382
                let prod_i16 = std::arch::x86_64::_mm256_maddubs_epi16(q4_abs, q8_signed);
    1021
    382
                let prod_i32 = std::arch::x86_64::_mm256_madd_epi16(prod_i16, ones);
    1022
    382
                let prod_f32 = _mm256_cvtepi32_ps(prod_i32);
    1023
    382
    1024
    382
                acc0 = _mm256_fmadd_ps(combined_scale, prod_f32, acc0);
    1025
    382
    1026
    382
                block_idx += 1;
    1027
    382
            }
    1028
    1029
            // Combine both accumulators and do horizontal sum
    1030
    86.2k
            let acc = std::arch::x86_64::_mm256_add_ps(acc0, acc1);
    1031
    86.2k
            let hi = std::arch::x86_64::_mm256_extractf128_ps(acc, 1);
    1032
    86.2k
            let lo = std::arch::x86_64::_mm256_castps256_ps128(acc);
    1033
    86.2k
            let sum128 = _mm_add_ps(lo, hi);
    1034
    86.2k
            let sum64 = _mm_hadd_ps(sum128, sum128);
    1035
    86.2k
            let sum32 = _mm_hadd_ps(sum64, sum64);
    1036
    86.2k
            _mm_cvtss_f32(sum32)
    1037
        }
    1038
    86.2k
    }
    1039
    1040
    /// AVX2 accelerated Q4_0 × Q8_0 dot product using integer SIMD
    1041
    ///
    1042
    /// Uses _mm256_maddubs_epi16 which multiplies pairs of u8×i8 and accumulates
    1043
    /// to i16, then we sum to i32 and convert to f32. This is ~4x faster than
    1044
    /// the f32 FMA approach because:
    1045
    /// 1. Integer ops have lower latency
    1046
    /// 2. maddubs does multiply AND horizontal add in one instruction
    1047
    /// 3. Less data movement (1 byte vs 4 bytes per value)
    1048
    #[cfg(target_arch = "x86_64")]
    1049
    #[target_feature(enable = "avx2")]
    1050
    #[inline]
    1051
    0
    unsafe fn fused_q4_0_q8_0_dot_avx2(
    1052
    0
        q4_data: &[u8],
    1053
    0
        q8_scales: &[f32],
    1054
    0
        q8_quants: &[i8],
    1055
    0
        in_dim: usize,
    1056
    0
    ) -> f32 {
    1057
        // SAFETY: Memory safety ensured by bounds checking and alignment
    1058
        unsafe {
    1059
            use std::arch::x86_64::{
    1060
                _mm256_and_si256, _mm256_cvtepi32_ps, _mm256_fmadd_ps, _mm256_loadu_si256,
    1061
                _mm256_madd_epi16, _mm256_maddubs_epi16, _mm256_set1_epi16, _mm256_set1_epi8,
    1062
                _mm256_set1_ps, _mm256_setzero_ps, _mm256_sign_epi8, _mm256_sub_epi8, _mm_cvtss_f32,
    1063
                _mm_hadd_ps, _mm_prefetch, _MM_HINT_T0,
    1064
            };
    1065
    1066
            const Q4_0_BLOCK_BYTES: usize = 18;
    1067
            const Q4_0_BLOCK_SIZE: usize = 32;
    1068
    1069
    0
            let num_blocks = in_dim.div_ceil(Q4_0_BLOCK_SIZE);
    1070
    1071
            // Float accumulator for final sum
    1072
    0
            let mut acc = _mm256_setzero_ps();
    1073
    1074
            // Offset: Q4_0 values are 0-15, we subtract 8 to get -8 to +7
    1075
    0
            let offset = _mm256_set1_epi8(8);
    1076
    0
            let low_mask = _mm256_set1_epi8(0x0F);
    1077
    0
            let ones = _mm256_set1_epi16(1);
    1078
    1079
    0
            let mut block_idx = 0;
    1080
    1081
            // Process 2 blocks at a time for better instruction-level parallelism
    1082
    0
            while block_idx + 2 <= num_blocks {
    1083
                // Prefetch next iteration's blocks
    1084
    0
                if block_idx + 4 <= num_blocks {
    1085
    0
                    let prefetch_q4 = q4_data.as_ptr().add((block_idx + 2) * Q4_0_BLOCK_BYTES);
    1086
    0
                    let prefetch_q8 = q8_quants.as_ptr().add((block_idx + 2) * Q4_0_BLOCK_SIZE);
    1087
    0
                    _mm_prefetch(prefetch_q4.cast(), _MM_HINT_T0);
    1088
    0
                    _mm_prefetch(prefetch_q8.cast(), _MM_HINT_T0);
    1089
    0
                }
    1090
    1091
                // === Block 0 ===
    1092
    0
                let q4_ptr_0 = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
    1093
    0
                let q8_ptr_0 = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);
    1094
    1095
                // Read Q4_0 scale (f16 -> f32 via LUT)
    1096
    0
                let q4_scale_bits_0 = u16::from_le_bytes([*q4_ptr_0, *q4_ptr_0.add(1)]);
    1097
    0
                let q4_scale_0 = f16_to_f32_lut(q4_scale_bits_0);
    1098
    0
                let q8_scale_0 = q8_scales[block_idx];
    1099
    0
                let combined_scale_0 = _mm256_set1_ps(q4_scale_0 * q8_scale_0);
    1100
    1101
                // Load Q4_0 quants (16 bytes = 32 nibbles)
    1102
    0
                let q4_bytes = std::slice::from_raw_parts(q4_ptr_0.add(2), 16);
    1103
    1104
                // bytes_from_nibbles_32: expand 16 bytes to 32 bytes
    1105
                // Low nibbles in first 16 positions, high nibbles in next 16
    1106
    0
                let q4_lo_128 = std::arch::x86_64::_mm_loadu_si128(q4_bytes.as_ptr().cast());
    1107
    0
                let q4_hi_128 = std::arch::x86_64::_mm_srli_epi16(q4_lo_128, 4);
    1108
                // Combine into 256-bit: high nibbles in upper 128, low nibbles in lower 128
    1109
    0
                let q4_combined = std::arch::x86_64::_mm256_set_m128i(q4_hi_128, q4_lo_128);
    1110
                // Mask to get just nibbles
    1111
    0
                let q4_nibbles = _mm256_and_si256(q4_combined, low_mask);
    1112
                // Convert from unsigned 0-15 to signed -8 to +7
    1113
    0
                let q4_signed = _mm256_sub_epi8(q4_nibbles, offset);
    1114
    1115
                // Load Q8_0 quants (32 bytes)
    1116
    0
                let q8_vec = _mm256_loadu_si256(q8_ptr_0.cast());
    1117
    1118
                // Integer multiply-accumulate using signed multiply trick:
    1119
                // maddubs requires unsigned × signed, so we use sign trick
    1120
                // ax = |x|, sy = sign(y, x), then maddubs(ax, sy) = x * y
    1121
    0
                let q4_abs = _mm256_sign_epi8(q4_signed, q4_signed);
    1122
    0
                let q8_signed = _mm256_sign_epi8(q8_vec, q4_signed);
    1123
    1124
                // maddubs: multiply pairs and add horizontally to i16
    1125
    0
                let prod_i16 = _mm256_maddubs_epi16(q4_abs, q8_signed);
    1126
                // madd: pairwise add i16 to i32
    1127
    0
                let prod_i32 = _mm256_madd_epi16(prod_i16, ones);
    1128
                // Convert to float
    1129
    0
                let prod_f32 = _mm256_cvtepi32_ps(prod_i32);
    1130
    1131
                // Scale and accumulate
    1132
    0
                acc = _mm256_fmadd_ps(combined_scale_0, prod_f32, acc);
    1133
    1134
                // === Block 1 ===
    1135
    0
                let q4_ptr_1 = q4_data.as_ptr().add((block_idx + 1) * Q4_0_BLOCK_BYTES);
    1136
    0
                let q8_ptr_1 = q8_quants.as_ptr().add((block_idx + 1) * Q4_0_BLOCK_SIZE);
    1137
    1138
    0
                let q4_scale_bits_1 = u16::from_le_bytes([*q4_ptr_1, *q4_ptr_1.add(1)]);
    1139
    0
                let q4_scale_1 = f16_to_f32_lut(q4_scale_bits_1);
    1140
    0
                let q8_scale_1 = q8_scales[block_idx + 1];
    1141
    0
                let combined_scale_1 = _mm256_set1_ps(q4_scale_1 * q8_scale_1);
    1142
    1143
    0
                let q4_bytes_1 = std::slice::from_raw_parts(q4_ptr_1.add(2), 16);
    1144
    0
                let q4_lo_128_1 = std::arch::x86_64::_mm_loadu_si128(q4_bytes_1.as_ptr().cast());
    1145
    0
                let q4_hi_128_1 = std::arch::x86_64::_mm_srli_epi16(q4_lo_128_1, 4);
    1146
    0
                let q4_combined_1 = std::arch::x86_64::_mm256_set_m128i(q4_hi_128_1, q4_lo_128_1);
    1147
    0
                let q4_nibbles_1 = _mm256_and_si256(q4_combined_1, low_mask);
    1148
    0
                let q4_signed_1 = _mm256_sub_epi8(q4_nibbles_1, offset);
    1149
    1150
    0
                let q8_vec_1 = _mm256_loadu_si256(q8_ptr_1.cast());
    1151
    1152
    0
                let q4_abs_1 = _mm256_sign_epi8(q4_signed_1, q4_signed_1);
    1153
    0
                let q8_signed_1 = _mm256_sign_epi8(q8_vec_1, q4_signed_1);
    1154
    1155
    0
                let prod_i16_1 = _mm256_maddubs_epi16(q4_abs_1, q8_signed_1);
    1156
    0
                let prod_i32_1 = _mm256_madd_epi16(prod_i16_1, ones);
    1157
    0
                let prod_f32_1 = _mm256_cvtepi32_ps(prod_i32_1);
    1158
    1159
    0
                acc = _mm256_fmadd_ps(combined_scale_1, prod_f32_1, acc);
    1160
    1161
    0
                block_idx += 2;
    1162
            }
    1163
    1164
            // Handle remaining single block
    1165
    0
            while block_idx < num_blocks {
    1166
    0
                let q4_ptr = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
    1167
    0
                let q8_ptr = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);
    1168
    0
    1169
    0
                let q4_scale_bits = u16::from_le_bytes([*q4_ptr, *q4_ptr.add(1)]);
    1170
    0
                let q4_scale = f16_to_f32_lut(q4_scale_bits);
    1171
    0
                let q8_scale = q8_scales[block_idx];
    1172
    0
                let combined_scale = _mm256_set1_ps(q4_scale * q8_scale);
    1173
    0
    1174
    0
                let q4_bytes = std::slice::from_raw_parts(q4_ptr.add(2), 16);
    1175
    0
                let q4_lo_128 = std::arch::x86_64::_mm_loadu_si128(q4_bytes.as_ptr().cast());
    1176
    0
                let q4_hi_128 = std::arch::x86_64::_mm_srli_epi16(q4_lo_128, 4);
    1177
    0
                let q4_combined = std::arch::x86_64::_mm256_set_m128i(q4_hi_128, q4_lo_128);
    1178
    0
                let q4_nibbles = _mm256_and_si256(q4_combined, low_mask);
    1179
    0
                let q4_signed = _mm256_sub_epi8(q4_nibbles, offset);
    1180
    0
    1181
    0
                let q8_vec = _mm256_loadu_si256(q8_ptr.cast());
    1182
    0
    1183
    0
                let q4_abs = _mm256_sign_epi8(q4_signed, q4_signed);
    1184
    0
                let q8_signed = _mm256_sign_epi8(q8_vec, q4_signed);
    1185
    0
    1186
    0
                let prod_i16 = _mm256_maddubs_epi16(q4_abs, q8_signed);
    1187
    0
                let prod_i32 = _mm256_madd_epi16(prod_i16, ones);
    1188
    0
                let prod_f32 = _mm256_cvtepi32_ps(prod_i32);
    1189
    0
    1190
    0
                acc = _mm256_fmadd_ps(combined_scale, prod_f32, acc);
    1191
    0
    1192
    0
                block_idx += 1;
    1193
    0
            }
    1194
    1195
            // Horizontal sum of 8 floats
    1196
    0
            let hi = std::arch::x86_64::_mm256_extractf128_ps(acc, 1);
    1197
    0
            let lo = std::arch::x86_64::_mm256_castps256_ps128(acc);
    1198
    0
            let sum128 = std::arch::x86_64::_mm_add_ps(lo, hi);
    1199
            // Use hadd for final reduction
    1200
    0
            let sum64 = _mm_hadd_ps(sum128, sum128);
    1201
    0
            let sum32 = _mm_hadd_ps(sum64, sum64);
    1202
    0
            _mm_cvtss_f32(sum32)
    1203
        }
    1204
    0
    }
    1205
    1206
    /// AVX2 accelerated Q4_0 × Q8_0 dot product with 4-block unrolling
    1207
    ///
    1208
    /// Processes 4 blocks per iteration for maximum ILP on modern OoO CPUs.
    1209
    /// This version achieves ~1.3x speedup over 2-block unrolling for large vectors.
    1210
    #[cfg(target_arch = "x86_64")]
    1211
    #[target_feature(enable = "avx2")]
    1212
    #[inline]
    1213
    0
    unsafe fn fused_q4_0_q8_0_dot_avx2_4block(
    1214
    0
        q4_data: &[u8],
    1215
    0
        q8_scales: &[f32],
    1216
    0
        q8_quants: &[i8],
    1217
    0
        in_dim: usize,
    1218
    0
    ) -> f32 {
    1219
        // SAFETY: Memory safety ensured by bounds checking and alignment
    1220
        unsafe {
    1221
            use std::arch::x86_64::{
    1222
                _mm256_add_ps, _mm256_and_si256, _mm256_cvtepi32_ps, _mm256_fmadd_ps,
    1223
                _mm256_loadu_si256, _mm256_madd_epi16, _mm256_maddubs_epi16, _mm256_set1_epi16,
    1224
                _mm256_set1_epi8, _mm256_set1_ps, _mm256_setzero_ps, _mm256_sign_epi8, _mm256_sub_epi8,
    1225
                _mm_cvtss_f32, _mm_hadd_ps, _mm_prefetch, _MM_HINT_T0,
    1226
            };
    1227
    1228
            const Q4_0_BLOCK_BYTES: usize = 18;
    1229
            const Q4_0_BLOCK_SIZE: usize = 32;
    1230
    1231
    0
            let num_blocks = in_dim.div_ceil(Q4_0_BLOCK_SIZE);
    1232
    1233
            // Use two accumulators for better pipelining
    1234
    0
            let mut acc0 = _mm256_setzero_ps();
    1235
    0
            let mut acc1 = _mm256_setzero_ps();
    1236
    1237
    0
            let offset = _mm256_set1_epi8(8);
    1238
    0
            let low_mask = _mm256_set1_epi8(0x0F);
    1239
    0
            let ones = _mm256_set1_epi16(1);
    1240
    1241
    0
            let mut block_idx = 0;
    1242
    1243
            // Process 4 blocks at a time for maximum ILP
    1244
    0
            while block_idx + 4 <= num_blocks {
    1245
                // Prefetch next iteration's blocks
    1246
    0
                if block_idx + 8 <= num_blocks {
    1247
    0
                    let prefetch_q4 = q4_data.as_ptr().add((block_idx + 4) * Q4_0_BLOCK_BYTES);
    1248
    0
                    let prefetch_q8 = q8_quants.as_ptr().add((block_idx + 4) * Q4_0_BLOCK_SIZE);
    1249
    0
                    _mm_prefetch(prefetch_q4.cast(), _MM_HINT_T0);
    1250
    0
                    _mm_prefetch(prefetch_q8.cast(), _MM_HINT_T0);
    1251
    0
                    _mm_prefetch(prefetch_q4.add(64).cast(), _MM_HINT_T0);
    1252
    0
                    _mm_prefetch(prefetch_q8.add(64).cast(), _MM_HINT_T0);
    1253
    0
                }
    1254
    1255
                // Block 0
    1256
    0
                let q4_ptr_0 = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
    1257
    0
                let q8_ptr_0 = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);
    1258
    0
                let q4_scale_0 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_0, *q4_ptr_0.add(1)]));
    1259
    0
                let combined_scale_0 = _mm256_set1_ps(q4_scale_0 * q8_scales[block_idx]);
    1260
    0
                let q4_lo_0 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_0.add(2).cast());
    1261
    0
                let q4_hi_0 = std::arch::x86_64::_mm_srli_epi16(q4_lo_0, 4);
    1262
    0
                let q4_signed_0 = _mm256_sub_epi8(
    1263
    0
                    _mm256_and_si256(
    1264
    0
                        std::arch::x86_64::_mm256_set_m128i(q4_hi_0, q4_lo_0),
    1265
    0
                        low_mask,
    1266
                    ),
    1267
    0
                    offset,
    1268
                );
    1269
    0
                let q8_vec_0 = _mm256_loadu_si256(q8_ptr_0.cast());
    1270
    0
                let q4_abs_0 = _mm256_sign_epi8(q4_signed_0, q4_signed_0);
    1271
    0
                let q8_signed_0 = _mm256_sign_epi8(q8_vec_0, q4_signed_0);
    1272
    0
                let prod_0 = _mm256_cvtepi32_ps(_mm256_madd_epi16(
    1273
    0
                    _mm256_maddubs_epi16(q4_abs_0, q8_signed_0),
    1274
    0
                    ones,
    1275
                ));
    1276
    0
                acc0 = _mm256_fmadd_ps(combined_scale_0, prod_0, acc0);
    1277
    1278
                // Block 1
    1279
    0
                let q4_ptr_1 = q4_data.as_ptr().add((block_idx + 1) * Q4_0_BLOCK_BYTES);
    1280
    0
                let q8_ptr_1 = q8_quants.as_ptr().add((block_idx + 1) * Q4_0_BLOCK_SIZE);
    1281
    0
                let q4_scale_1 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_1, *q4_ptr_1.add(1)]));
    1282
    0
                let combined_scale_1 = _mm256_set1_ps(q4_scale_1 * q8_scales[block_idx + 1]);
    1283
    0
                let q4_lo_1 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_1.add(2).cast());
    1284
    0
                let q4_hi_1 = std::arch::x86_64::_mm_srli_epi16(q4_lo_1, 4);
    1285
    0
                let q4_signed_1 = _mm256_sub_epi8(
    1286
    0
                    _mm256_and_si256(
    1287
    0
                        std::arch::x86_64::_mm256_set_m128i(q4_hi_1, q4_lo_1),
    1288
    0
                        low_mask,
    1289
                    ),
    1290
    0
                    offset,
    1291
                );
    1292
    0
                let q8_vec_1 = _mm256_loadu_si256(q8_ptr_1.cast());
    1293
    0
                let q4_abs_1 = _mm256_sign_epi8(q4_signed_1, q4_signed_1);
    1294
    0
                let q8_signed_1 = _mm256_sign_epi8(q8_vec_1, q4_signed_1);
    1295
    0
                let prod_1 = _mm256_cvtepi32_ps(_mm256_madd_epi16(
    1296
    0
                    _mm256_maddubs_epi16(q4_abs_1, q8_signed_1),
    1297
    0
                    ones,
    1298
                ));
    1299
    0
                acc1 = _mm256_fmadd_ps(combined_scale_1, prod_1, acc1);
    1300
    1301
                // Block 2
    1302
    0
                let q4_ptr_2 = q4_data.as_ptr().add((block_idx + 2) * Q4_0_BLOCK_BYTES);
    1303
    0
                let q8_ptr_2 = q8_quants.as_ptr().add((block_idx + 2) * Q4_0_BLOCK_SIZE);
    1304
    0
                let q4_scale_2 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_2, *q4_ptr_2.add(1)]));
    1305
    0
                let combined_scale_2 = _mm256_set1_ps(q4_scale_2 * q8_scales[block_idx + 2]);
    1306
    0
                let q4_lo_2 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_2.add(2).cast());
    1307
    0
                let q4_hi_2 = std::arch::x86_64::_mm_srli_epi16(q4_lo_2, 4);
    1308
    0
                let q4_signed_2 = _mm256_sub_epi8(
    1309
    0
                    _mm256_and_si256(
    1310
    0
                        std::arch::x86_64::_mm256_set_m128i(q4_hi_2, q4_lo_2),
    1311
    0
                        low_mask,
    1312
                    ),
    1313
    0
                    offset,
    1314
                );
    1315
    0
                let q8_vec_2 = _mm256_loadu_si256(q8_ptr_2.cast());
    1316
    0
                let q4_abs_2 = _mm256_sign_epi8(q4_signed_2, q4_signed_2);
    1317
    0
                let q8_signed_2 = _mm256_sign_epi8(q8_vec_2, q4_signed_2);
    1318
    0
                let prod_2 = _mm256_cvtepi32_ps(_mm256_madd_epi16(
    1319
    0
                    _mm256_maddubs_epi16(q4_abs_2, q8_signed_2),
    1320
    0
                    ones,
    1321
                ));
    1322
    0
                acc0 = _mm256_fmadd_ps(combined_scale_2, prod_2, acc0);
    1323
    1324
                // Block 3
    1325
    0
                let q4_ptr_3 = q4_data.as_ptr().add((block_idx + 3) * Q4_0_BLOCK_BYTES);
    1326
    0
                let q8_ptr_3 = q8_quants.as_ptr().add((block_idx + 3) * Q4_0_BLOCK_SIZE);
    1327
    0
                let q4_scale_3 = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr_3, *q4_ptr_3.add(1)]));
    1328
    0
                let combined_scale_3 = _mm256_set1_ps(q4_scale_3 * q8_scales[block_idx + 3]);
    1329
    0
                let q4_lo_3 = std::arch::x86_64::_mm_loadu_si128(q4_ptr_3.add(2).cast());
    1330
    0
                let q4_hi_3 = std::arch::x86_64::_mm_srli_epi16(q4_lo_3, 4);
    1331
    0
                let q4_signed_3 = _mm256_sub_epi8(
    1332
    0
                    _mm256_and_si256(
    1333
    0
                        std::arch::x86_64::_mm256_set_m128i(q4_hi_3, q4_lo_3),
    1334
    0
                        low_mask,
    1335
                    ),
    1336
    0
                    offset,
    1337
                );
    1338
    0
                let q8_vec_3 = _mm256_loadu_si256(q8_ptr_3.cast());
    1339
    0
                let q4_abs_3 = _mm256_sign_epi8(q4_signed_3, q4_signed_3);
    1340
    0
                let q8_signed_3 = _mm256_sign_epi8(q8_vec_3, q4_signed_3);
    1341
    0
                let prod_3 = _mm256_cvtepi32_ps(_mm256_madd_epi16(
    1342
    0
                    _mm256_maddubs_epi16(q4_abs_3, q8_signed_3),
    1343
    0
                    ones,
    1344
                ));
    1345
    0
                acc1 = _mm256_fmadd_ps(combined_scale_3, prod_3, acc1);
    1346
    1347
    0
                block_idx += 4;
    1348
            }
    1349
    1350
            // Merge accumulators
    1351
    0
            let acc = _mm256_add_ps(acc0, acc1);
    1352
    1353
            // Handle remaining blocks (0-3)
    1354
    0
            let mut scalar_sum = 0.0f32;
    1355
    0
            while block_idx < num_blocks {
    1356
    0
                let q4_ptr = q4_data.as_ptr().add(block_idx * Q4_0_BLOCK_BYTES);
    1357
    0
                let q8_ptr = q8_quants.as_ptr().add(block_idx * Q4_0_BLOCK_SIZE);
    1358
    0
                let q4_scale = f16_to_f32_lut(u16::from_le_bytes([*q4_ptr, *q4_ptr.add(1)]));
    1359
    0
                let combined_scale = _mm256_set1_ps(q4_scale * q8_scales[block_idx]);
    1360
    0
                let q4_lo = std::arch::x86_64::_mm_loadu_si128(q4_ptr.add(2).cast());
    1361
    0
                let q4_hi = std::arch::x86_64::_mm_srli_epi16(q4_lo, 4);
    1362
    0
                let q4_signed = _mm256_sub_epi8(
    1363
    0
                    _mm256_and_si256(std::arch::x86_64::_mm256_set_m128i(q4_hi, q4_lo), low_mask),
    1364
    0
                    offset,
    1365
    0
                );
    1366
    0
                let q8_vec = _mm256_loadu_si256(q8_ptr.cast());
    1367
    0
                let q4_abs = _mm256_sign_epi8(q4_signed, q4_signed);
    1368
    0
                let q8_signed = _mm256_sign_epi8(q8_vec, q4_signed);
    1369
    0
                let prod = _mm256_cvtepi32_ps(_mm256_madd_epi16(
    1370
    0
                    _mm256_maddubs_epi16(q4_abs, q8_signed),
    1371
    0
                    ones,
    1372
    0
                ));
    1373
    0
                let scaled = _mm256_fmadd_ps(combined_scale, prod, _mm256_setzero_ps());
    1374
    0
    1375
    0
                // Horizontal sum for this block
    1376
    0
                let hi = std::arch::x86_64::_mm256_extractf128_ps(scaled, 1);
    1377
    0
                let lo = std::arch::x86_64::_mm256_castps256_ps128(scaled);
    1378
    0
                let sum128 = std::arch::x86_64::_mm_add_ps(lo, hi);
    1379
    0
                let sum64 = _mm_hadd_ps(sum128, sum128);
    1380
    0
                let sum32 = _mm_hadd_ps(sum64, sum64);
    1381
    0
                scalar_sum += _mm_cvtss_f32(sum32);
    1382
    0
    1383
    0
                block_idx += 1;
    1384
    0
            }
    1385
    1386
            // Horizontal sum of accumulated vector
    1387
    0
            let hi = std::arch::x86_64::_mm256_extractf128_ps(acc, 1);
    1388
    0
            let lo = std::arch::x86_64::_mm256_castps256_ps128(acc);
    1389
    0
            let sum128 = std::arch::x86_64::_mm_add_ps(lo, hi);
    1390
    0
            let sum64 = _mm_hadd_ps(sum128, sum128);
    1391
    0
            let sum32 = _mm_hadd_ps(sum64, sum64);
    1392
    0
            _mm_cvtss_f32(sum32) + scalar_sum
    1393
        }
    1394
    0
    }
    1395
    1396
    /// Scalar fallback for Q4_0 × Q8_0 dot product
    1397
    ///
    1398
    /// Exposed as `pub(crate)` for direct testing on AVX2 machines.
    1399
    #[inline]
    1400
    4
    pub(crate) fn fused_q4_0_q8_0_dot_scalar(
    1401
    4
        q4_data: &[u8],
    1402
    4
        q8_scales: &[f32],
    1403
    4
        q8_quants: &[i8],
    1404
    4
        in_dim: usize,
    1405
    4
    ) -> f32 {
    1406
        const Q4_0_BLOCK_BYTES: usize = 18;
    1407
        const Q4_0_BLOCK_SIZE: usize = 32;
    1408
    1409
    4
        let num_blocks = in_dim.div_ceil(Q4_0_BLOCK_SIZE);
    1410
    4
        let mut total_sum = 0.0f32;
    1411
    1412
    5
        for block_idx in 0..num_blocks4 {
    1413
    5
            let block_start = block_idx * Q4_0_BLOCK_BYTES;
    1414
    5
            if block_start + Q4_0_BLOCK_BYTES > q4_data.len() {
    1415
    0
                break;
    1416
    5
            }
    1417
    5
            let block = &q4_data[block_start..block_start + Q4_0_BLOCK_BYTES];
    1418
    1419
    5
            let q4_scale = half::f16::from_le_bytes([block[0], block[1]]).to_f32();
    1420
    5
            let q8_scale = q8_scales[block_idx];
    1421
    5
            let combined_scale = q4_scale * q8_scale;
    1422
    1423
    5
            let act_start = block_idx * Q4_0_BLOCK_SIZE;
    1424
    1425
    5
            let mut block_sum = 0i32;
    1426
    80
            for (j, &byte) in block[2..18]5.iter5().enumerate5() {
    1427
    80
                let low_idx = act_start + j;
    1428
    80
                let high_idx = act_start + j + 16;
    1429
    1430
                #[allow(clippy::cast_possible_wrap)]
    1431
    80
                let low_quant = (byte & 0x0F) as i8 - 8;
    1432
    80
                block_sum += (low_quant as i32) * (q8_quants[low_idx] as i32);
    1433
    1434
                #[allow(clippy::cast_possible_wrap)]
    1435
    80
                let high_quant = (byte >> 4) as i8 - 8;
    1436
    80
                if high_idx < in_dim {
    1437
    72
                    block_sum += (high_quant as i32) * (q8_quants[high_idx] as i32);
    1438
    72
                }8
    1439
            }
    1440
    1441
    5
            total_sum += combined_scale * (block_sum as f32);
    1442
        }
    1443
    1444
    4
        total_sum
    1445
    4
    }
    1446
    1447
    /// Parallel Q4_0 × Q8_0 matrix-vector multiply
    1448
    ///
    1449
    /// This is the key function for llama.cpp parity. It:
    1450
    /// 1. Quantizes activations to Q8_0 format once
    1451
    /// 2. Uses integer SIMD for all row dot products
    1452
    /// 3. Parallelizes across output rows with rayon (adaptive threshold)
    1453
    ///
    1454
    /// Expected speedup: 4-6x over the f32 FMA version
    1455
    #[allow(clippy::similar_names)]
    1456
    715
    pub fn fused_q4_0_q8_0_parallel_matvec(
    1457
    715
        weight_data: &[u8],
    1458
    715
        activations: &[f32],
    1459
    715
        in_dim: usize,
    1460
    715
        out_dim: usize,
    1461
    715
    ) -> Result<Vec<f32>> {
    1462
        const Q4_0_BLOCK_BYTES: usize = 18;
    1463
        const Q4_0_BLOCK_SIZE: usize = 32;
    1464
    1465
    715
        let blocks_per_row = in_dim.div_ceil(Q4_0_BLOCK_SIZE);
    1466
    715
        let bytes_per_row = blocks_per_row * Q4_0_BLOCK_BYTES;
    1467
    1468
    715
        let expected_weight_bytes = out_dim * bytes_per_row;
    1469
    715
        if weight_data.len() < expected_weight_bytes {
    1470
    3
            return Err(RealizarError::InvalidShape {
    1471
    3
                reason: format!(
    1472
    3
                    "Q4_0 weight data too small: need {} bytes for {}x{}, have {}",
    1473
    3
                    expected_weight_bytes,
    1474
    3
                    out_dim,
    1475
    3
                    in_dim,
    1476
    3
                    weight_data.len()
    1477
    3
                ),
    1478
    3
            });
    1479
    712
        }
    1480
    1481
    712
        if activations.len() != in_dim {
    1482
    1
            return Err(RealizarError::InvalidShape {
    1483
    1
                reason: format!(
    1484
    1
                    "Activation length {} doesn't match in_dim {}",
    1485
    1
                    activations.len(),
    1486
    1
                    in_dim
    1487
    1
                ),
    1488
    1
            });
    1489
    711
        }
    1490
    1491
        // Quantize activations to Q8_0 ONCE (amortized over all rows)
    1492
    711
        let (q8_scales, q8_quants) = quantize_activations_q8_0(activations);
    1493
    1494
        // Adaptive parallelization: sequential for small matrices, parallel for large
    1495
        // Rayon overhead (~50-100µs) dominates for small out_dim
    1496
        // Threshold tuned for 22-core CPU: break-even at ~1024 rows
    1497
        const PARALLEL_THRESHOLD: usize = 1024;
    1498
    1499
    711
        if out_dim < PARALLEL_THRESHOLD {
    1500
            // Sequential path: avoids Rayon overhead entirely
    1501
    711
            let output: Vec<f32> = (0..out_dim)
    1502
    78.8k
                .map711(|o| {
    1503
    78.8k
                    let row_start = o * bytes_per_row;
    1504
    78.8k
                    let row_end = row_start + bytes_per_row;
    1505
    78.8k
                    let row_data = &weight_data[row_start..row_end];
    1506
    78.8k
                    fused_q4_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim)
    1507
    78.8k
                })
    1508
    711
                .collect();
    1509
    711
            return Ok(output);
    1510
    0
        }
    1511
    1512
        // Parallel path for large matrices
    1513
        use rayon::prelude::*;
    1514
        // Use chunked parallel iteration to reduce Rayon scheduling overhead
    1515
        // CHUNK_SIZE=128 provides good balance between parallelism and overhead
    1516
        const CHUNK_SIZE: usize = 128;
    1517
    0
        let output: Vec<f32> = (0..out_dim)
    1518
    0
            .into_par_iter()
    1519
    0
            .with_min_len(CHUNK_SIZE)
    1520
    0
            .map(|o| {
    1521
    0
                let row_start = o * bytes_per_row;
    1522
    0
                let row_end = row_start + bytes_per_row;
    1523
    0
                let row_data = &weight_data[row_start..row_end];
    1524
    0
                fused_q4_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim)
    1525
    0
            })
    1526
    0
            .collect();
    1527
    1528
    0
        Ok(output)
    1529
    715
    }
    1530
    1531
    /// Zero-allocation Q4_0 × Q8_0 matrix-vector multiply.
    1532
    ///
    1533
    /// Writes result directly into provided output buffer, eliminating allocation.
    1534
    /// Use this with scratch buffers for maximum performance.
    1535
    ///
    1536
    /// # Arguments
    1537
    /// * `weight_data` - Q4_0 quantized weight matrix (row-major)
    1538
    /// * `activations` - Input activation vector (f32)
    1539
    /// * `in_dim` - Input dimension (columns)
    1540
    /// * `output` - Pre-allocated output buffer (must be exactly out_dim length)
    1541
    ///
    1542
    /// # Returns
    1543
    /// Number of elements written (equals output.len())
    1544
    #[allow(clippy::similar_names)]
    1545
    71
    pub fn fused_q4_0_q8_0_parallel_matvec_into(
    1546
    71
        weight_data: &[u8],
    1547
    71
        activations: &[f32],
    1548
    71
        in_dim: usize,
    1549
    71
        output: &mut [f32],
    1550
    71
    ) -> Result<()> {
    1551
        use rayon::prelude::*;
    1552
    1553
        const Q4_0_BLOCK_BYTES: usize = 18;
    1554
        const Q4_0_BLOCK_SIZE: usize = 32;
    1555
    1556
    71
        let out_dim = output.len();
    1557
    71
        let blocks_per_row = in_dim.div_ceil(Q4_0_BLOCK_SIZE);
    1558
    71
        let bytes_per_row = blocks_per_row * Q4_0_BLOCK_BYTES;
    1559
    1560
    71
        let expected_weight_bytes = out_dim * bytes_per_row;
    1561
    71
        if weight_data.len() < expected_weight_bytes {
    1562
    3
            return Err(RealizarError::InvalidShape {
    1563
    3
                reason: format!(
    1564
    3
                    "Q4_0 weight data too small: need {} bytes for {}x{}, have {}",
    1565
    3
                    expected_weight_bytes,
    1566
    3
                    out_dim,
    1567
    3
                    in_dim,
    1568
    3
                    weight_data.len()
    1569
    3
                ),
    1570
    3
            });
    1571
    68
        }
    1572
    1573
    68
        if activations.len() != in_dim {
    1574
    0
            return Err(RealizarError::InvalidShape {
    1575
    0
                reason: format!(
    1576
    0
                    "Activation length {} doesn't match in_dim {}",
    1577
    0
                    activations.len(),
    1578
    0
                    in_dim
    1579
    0
                ),
    1580
    0
            });
    1581
    68
        }
    1582
    1583
        // Quantize activations to Q8_0 ONCE
    1584
    68
        let (q8_scales, q8_quants) = quantize_activations_q8_0(activations);
    1585
    1586
        // Use chunked parallel iteration to reduce Rayon scheduling overhead
    1587
        const CHUNK_SIZE: usize = 64;
    1588
    68
        output
    1589
    68
            .par_iter_mut()
    1590
    68
            .with_min_len(CHUNK_SIZE)
    1591
    68
            .enumerate()
    1592
    7.25k
            .for_each68(|(o, out_val)| {
    1593
    7.25k
                let row_start = o * bytes_per_row;
    1594
    7.25k
                let row_end = row_start + bytes_per_row;
    1595
    7.25k
                let row_data = &weight_data[row_start..row_end];
    1596
    7.25k
                *out_val = fused_q4_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim);
    1597
    7.25k
            });
    1598
    1599
    68
        Ok(())
    1600
    71
    }
    1601
    1602
    // ============================================================================
    1603
    // FUSED Q8_0 × Q8_0 MATMUL (For Q8_0 quantized weights like Qwen2.5 LM head)
    1604
    // ============================================================================
    1605
    //
    1606
    // Q8_0 format: 34 bytes per block (2 byte f16 scale + 32 i8 quants)
    1607
    // This avoids the massive dequantization allocation that was causing
    1608
    // Qwen2.5's 152K vocab LM head to allocate 544MB per forward pass.
    1609
    // ============================================================================
    1610
    1611
    /// AVX2 accelerated Q8_0 × Q8_0 dot product using integer SIMD
    1612
    ///
    1613
    /// Uses _mm256_maddubs_epi16 with sign trick for i8×i8 multiplication.
    1614
    /// This is simpler than Q4_0×Q8_0 since no nibble unpacking is needed.
    1615
    #[cfg(target_arch = "x86_64")]
    1616
    #[target_feature(enable = "avx2")]
    1617
    #[inline]
    1618
    1.03k
    unsafe fn fused_q8_0_q8_0_dot_avx2(
    1619
    1.03k
        q8_weight_data: &[u8],
    1620
    1.03k
        q8_act_scales: &[f32],
    1621
    1.03k
        q8_act_quants: &[i8],
    1622
    1.03k
        in_dim: usize,
    1623
    1.03k
    ) -> f32 {
    1624
        // SAFETY: Memory safety ensured by bounds checking and alignment
    1625
        unsafe {
    1626
            use std::arch::x86_64::{
    1627
                _mm256_cvtepi32_ps, _mm256_fmadd_ps, _mm256_loadu_si256, _mm256_madd_epi16,
    1628
                _mm256_maddubs_epi16, _mm256_set1_epi16, _mm256_set1_ps, _mm256_setzero_ps,
    1629
                _mm256_sign_epi8, _mm_cvtss_f32, _mm_hadd_ps, _mm_prefetch, _MM_HINT_T0,
    1630
            };
    1631
    1632
            const Q8_0_BLOCK_BYTES: usize = 34; // 2 byte scale + 32 byte quants
    1633
            const Q8_0_BLOCK_SIZE: usize = 32;
    1634
    1635
    1.03k
            let num_blocks = in_dim.div_ceil(Q8_0_BLOCK_SIZE);
    1636
    1637
            // Float accumulator for final sum
    1638
    1.03k
            let mut acc = _mm256_setzero_ps();
    1639
    1.03k
            let ones = _mm256_set1_epi16(1);
    1640
    1641
    1.03k
            let mut block_idx = 0;
    1642
    1643
            // Process 2 blocks at a time for better ILP
    1644
    3.59k
            while block_idx + 2 <= num_blocks {
    1645
                // Prefetch next iteration's blocks
    1646
    2.56k
                if block_idx + 4 <= num_blocks {
    1647
    1.53k
                    let prefetch_w = q8_weight_data
    1648
    1.53k
                        .as_ptr()
    1649
    1.53k
                        .add((block_idx + 2) * Q8_0_BLOCK_BYTES);
    1650
    1.53k
                    let prefetch_a = q8_act_quants
    1651
    1.53k
                        .as_ptr()
    1652
    1.53k
                        .add((block_idx + 2) * Q8_0_BLOCK_SIZE);
    1653
    1.53k
                    _mm_prefetch(prefetch_w.cast(), _MM_HINT_T0);
    1654
    1.53k
                    _mm_prefetch(prefetch_a.cast(), _MM_HINT_T0);
    1655
    1.53k
                }1.02k
    1656
    1657
                // === Block 0 ===
    1658
    2.56k
                let w_ptr_0 = q8_weight_data.as_ptr().add(block_idx * Q8_0_BLOCK_BYTES);
    1659
    2.56k
                let a_ptr_0 = q8_act_quants.as_ptr().add(block_idx * Q8_0_BLOCK_SIZE);
    1660
    1661
                // Read Q8_0 weight scale (f16 -> f32)
    1662
    2.56k
                let w_scale_bits_0 = u16::from_le_bytes([*w_ptr_0, *w_ptr_0.add(1)]);
    1663
    2.56k
                let w_scale_0 = f16_to_f32_lut(w_scale_bits_0);
    1664
    2.56k
                let a_scale_0 = q8_act_scales[block_idx];
    1665
    2.56k
                let combined_scale_0 = _mm256_set1_ps(w_scale_0 * a_scale_0);
    1666
    1667
                // Load Q8_0 weight quants (32 bytes at offset 2)
    1668
    2.56k
                let w_vec_0 = _mm256_loadu_si256(w_ptr_0.add(2).cast());
    1669
                // Load Q8_0 activation quants (32 bytes)
    1670
    2.56k
                let a_vec_0 = _mm256_loadu_si256(a_ptr_0.cast());
    1671
    1672
                // Integer multiply-accumulate using signed multiply trick:
    1673
                // maddubs requires unsigned × signed, so we use sign trick
    1674
                // |w| * sign(a, w) = w * a
    1675
    2.56k
                let w_abs_0 = _mm256_sign_epi8(w_vec_0, w_vec_0);
    1676
    2.56k
                let a_signed_0 = _mm256_sign_epi8(a_vec_0, w_vec_0);
    1677
    1678
                // maddubs: multiply pairs and add horizontally to i16
    1679
    2.56k
                let prod_i16_0 = _mm256_maddubs_epi16(w_abs_0, a_signed_0);
    1680
                // madd: pairwise add i16 to i32
    1681
    2.56k
                let prod_i32_0 = _mm256_madd_epi16(prod_i16_0, ones);
    1682
                // Convert to float
    1683
    2.56k
                let prod_f32_0 = _mm256_cvtepi32_ps(prod_i32_0);
    1684
    1685
                // Scale and accumulate
    1686
    2.56k
                acc = _mm256_fmadd_ps(combined_scale_0, prod_f32_0, acc);
    1687
    1688
                // === Block 1 ===
    1689
    2.56k
                let w_ptr_1 = q8_weight_data
    1690
    2.56k
                    .as_ptr()
    1691
    2.56k
                    .add((block_idx + 1) * Q8_0_BLOCK_BYTES);
    1692
    2.56k
                let a_ptr_1 = q8_act_quants
    1693
    2.56k
                    .as_ptr()
    1694
    2.56k
                    .add((block_idx + 1) * Q8_0_BLOCK_SIZE);
    1695
    1696
    2.56k
                let w_scale_bits_1 = u16::from_le_bytes([*w_ptr_1, *w_ptr_1.add(1)]);
    1697
    2.56k
                let w_scale_1 = f16_to_f32_lut(w_scale_bits_1);
    1698
    2.56k
                let a_scale_1 = q8_act_scales[block_idx + 1];
    1699
    2.56k
                let combined_scale_1 = _mm256_set1_ps(w_scale_1 * a_scale_1);
    1700
    1701
    2.56k
                let w_vec_1 = _mm256_loadu_si256(w_ptr_1.add(2).cast());
    1702
    2.56k
                let a_vec_1 = _mm256_loadu_si256(a_ptr_1.cast());
    1703
    1704
    2.56k
                let w_abs_1 = _mm256_sign_epi8(w_vec_1, w_vec_1);
    1705
    2.56k
                let a_signed_1 = _mm256_sign_epi8(a_vec_1, w_vec_1);
    1706
    1707
    2.56k
                let prod_i16_1 = _mm256_maddubs_epi16(w_abs_1, a_signed_1);
    1708
    2.56k
                let prod_i32_1 = _mm256_madd_epi16(prod_i16_1, ones);
    1709
    2.56k
                let prod_f32_1 = _mm256_cvtepi32_ps(prod_i32_1);
    1710
    1711
    2.56k
                acc = _mm256_fmadd_ps(combined_scale_1, prod_f32_1, acc);
    1712
    1713
    2.56k
                block_idx += 2;
    1714
            }
    1715
    1716
            // Handle remaining single block
    1717
    1.04k
            while block_idx < num_blocks {
    1718
    8
                let w_ptr = q8_weight_data.as_ptr().add(block_idx * Q8_0_BLOCK_BYTES);
    1719
    8
                let a_ptr = q8_act_quants.as_ptr().add(block_idx * Q8_0_BLOCK_SIZE);
    1720
    8
    1721
    8
                let w_scale_bits = u16::from_le_bytes([*w_ptr, *w_ptr.add(1)]);
    1722
    8
                let w_scale = f16_to_f32_lut(w_scale_bits);
    1723
    8
                let a_scale = q8_act_scales[block_idx];
    1724
    8
                let combined_scale = _mm256_set1_ps(w_scale * a_scale);
    1725
    8
    1726
    8
                let w_vec = _mm256_loadu_si256(w_ptr.add(2).cast());
    1727
    8
                let a_vec = _mm256_loadu_si256(a_ptr.cast());
    1728
    8
    1729
    8
                let w_abs = _mm256_sign_epi8(w_vec, w_vec);
    1730
    8
                let a_signed = _mm256_sign_epi8(a_vec, w_vec);
    1731
    8
    1732
    8
                let prod_i16 = _mm256_maddubs_epi16(w_abs, a_signed);
    1733
    8
                let prod_i32 = _mm256_madd_epi16(prod_i16, ones);
    1734
    8
                let prod_f32 = _mm256_cvtepi32_ps(prod_i32);
    1735
    8
    1736
    8
                acc = _mm256_fmadd_ps(combined_scale, prod_f32, acc);
    1737
    8
    1738
    8
                block_idx += 1;
    1739
    8
            }
    1740
    1741
            // Horizontal sum of 8 floats
    1742
    1.03k
            let hi = std::arch::x86_64::_mm256_extractf128_ps(acc, 1);
    1743
    1.03k
            let lo = std::arch::x86_64::_mm256_castps256_ps128(acc);
    1744
    1.03k
            let sum128 = std::arch::x86_64::_mm_add_ps(lo, hi);
    1745
    1.03k
            let sum64 = _mm_hadd_ps(sum128, sum128);
    1746
    1.03k
            let sum32 = _mm_hadd_ps(sum64, sum64);
    1747
    1.03k
            _mm_cvtss_f32(sum32)
    1748
        }
    1749
    1.03k
    }
    1750
    1751
    /// Scalar fallback for Q8_0 × Q8_0 dot product
    1752
    ///
    1753
    /// Exposed as `pub(crate)` for direct testing on AVX2 machines.
    1754
    #[inline]
    1755
    1
    pub(crate) fn fused_q8_0_q8_0_dot_scalar(
    1756
    1
        q8_weight_data: &[u8],
    1757
    1
        q8_act_scales: &[f32],
    1758
    1
        q8_act_quants: &[i8],
    1759
    1
        in_dim: usize,
    1760
    1
    ) -> f32 {
    1761
        const Q8_0_BLOCK_BYTES: usize = 34;
    1762
        const Q8_0_BLOCK_SIZE: usize = 32;
    1763
    1764
    1
        let num_blocks = in_dim.div_ceil(Q8_0_BLOCK_SIZE);
    1765
    1
        let mut total_sum = 0.0f32;
    1766
    1767
    1
        for block_idx in 0..num_blocks {
    1768
    1
            let block_start = block_idx * Q8_0_BLOCK_BYTES;
    1769
    1
            if block_start + Q8_0_BLOCK_BYTES > q8_weight_data.len() {
    1770
    0
                break;
    1771
    1
            }
    1772
    1
            let block = &q8_weight_data[block_start..block_start + Q8_0_BLOCK_BYTES];
    1773
    1774
            // Read weight scale (f16)
    1775
    1
            let w_scale = half::f16::from_le_bytes([block[0], block[1]]).to_f32();
    1776
    1
            let a_scale = q8_act_scales[block_idx];
    1777
    1
            let combined_scale = w_scale * a_scale;
    1778
    1779
    1
            let act_start = block_idx * Q8_0_BLOCK_SIZE;
    1780
    1781
            // Sum of weight_quant[i] * act_quant[i] in i32
    1782
    1
            let mut block_sum = 0i32;
    1783
    33
            for j32 in 0..32 {
    1784
    32
                if act_start + j >= in_dim {
    1785
    0
                    break;
    1786
    32
                }
    1787
                #[allow(clippy::cast_possible_wrap)]
    1788
    32
                let w_quant = block[2 + j] as i8;
    1789
    32
                let a_quant = q8_act_quants[act_start + j];
    1790
    32
                block_sum += (w_quant as i32) * (a_quant as i32);
    1791
            }
    1792
    1793
    1
            total_sum += combined_scale * (block_sum as f32);
    1794
        }
    1795
    1796
    1
        total_sum
    1797
    1
    }
    1798
    1799
    /// SIMD dispatcher for Q8_0 × Q8_0 dot product
    1800
    #[inline]
    1801
    1.03k
    fn fused_q8_0_q8_0_dot_simd(
    1802
    1.03k
        q8_weight_data: &[u8],
    1803
    1.03k
        q8_act_scales: &[f32],
    1804
    1.03k
        q8_act_quants: &[i8],
    1805
    1.03k
        in_dim: usize,
    1806
    1.03k
    ) -> f32 {
    1807
        #[cfg(target_arch = "x86_64")]
    1808
        {
    1809
    1.03k
            if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
    1810
                // SAFETY: AVX2 and FMA features checked above
    1811
                unsafe {
    1812
    1.03k
                    return fused_q8_0_q8_0_dot_avx2(
    1813
    1.03k
                        q8_weight_data,
    1814
    1.03k
                        q8_act_scales,
    1815
    1.03k
                        q8_act_quants,
    1816
    1.03k
                        in_dim,
    1817
                    );
    1818
                }
    1819
    0
            }
    1820
        }
    1821
    0
        fused_q8_0_q8_0_dot_scalar(q8_weight_data, q8_act_scales, q8_act_quants, in_dim)
    1822
    1.03k
    }
    1823
    1824
    /// Parallel Q8_0 × Q8_0 matrix-vector multiply
    1825
    ///
    1826
    /// This avoids the massive dequantization allocation that was causing
    1827
    /// Qwen2.5's 152K vocab LM head (Q8_0) to allocate 544MB per forward pass.
    1828
    ///
    1829
    /// For Q8_0 weights (e.g., Qwen2.5 LM head), this is ~100x faster than
    1830
    /// dequantize + matmul because:
    1831
    /// 1. No 544MB allocation per forward pass
    1832
    /// 2. Integer SIMD is faster than FP32
    1833
    /// 3. Better cache locality (34 bytes vs 128 bytes per block)
    1834
    #[allow(clippy::similar_names)]
    1835
    10
    pub fn fused_q8_0_q8_0_parallel_matvec(
    1836
    10
        weight_data: &[u8],
    1837
    10
        activations: &[f32],
    1838
    10
        in_dim: usize,
    1839
    10
        out_dim: usize,
    1840
    10
    ) -> Result<Vec<f32>> {
    1841
        use rayon::prelude::*;
    1842
    1843
        const Q8_0_BLOCK_BYTES: usize = 34;
    1844
        const Q8_0_BLOCK_SIZE: usize = 32;
    1845
    1846
    10
        let blocks_per_row = in_dim.div_ceil(Q8_0_BLOCK_SIZE);
    1847
    10
        let bytes_per_row = blocks_per_row * Q8_0_BLOCK_BYTES;
    1848
    1849
    10
        let expected_weight_bytes = out_dim * bytes_per_row;
    1850
    10
        if weight_data.len() < expected_weight_bytes {
    1851
    3
            return Err(RealizarError::InvalidShape {
    1852
    3
                reason: format!(
    1853
    3
                    "Q8_0 weight data too small: need {} bytes for {}x{}, have {}",
    1854
    3
                    expected_weight_bytes,
    1855
    3
                    out_dim,
    1856
    3
                    in_dim,
    1857
    3
                    weight_data.len()
    1858
    3
                ),
    1859
    3
            });
    1860
    7
        }
    1861
    1862
    7
        if activations.len() != in_dim {
    1863
    0
            return Err(RealizarError::InvalidShape {
    1864
    0
                reason: format!(
    1865
    0
                    "Activation length {} doesn't match in_dim {}",
    1866
    0
                    activations.len(),
    1867
    0
                    in_dim
    1868
    0
                ),
    1869
    0
            });
    1870
    7
        }
    1871
    1872
        // Quantize activations to Q8_0 ONCE (amortized over all rows)
    1873
    7
        let (q8_scales, q8_quants) = quantize_activations_q8_0(activations);
    1874
    1875
        // Parallel over output rows with chunking
    1876
        const CHUNK_SIZE: usize = 64;
    1877
    7
        let output: Vec<f32> = (0..out_dim)
    1878
    7
            .into_par_iter()
    1879
    7
            .with_min_len(CHUNK_SIZE)
    1880
    1.02k
            .map7(|o| {
    1881
    1.02k
                let row_start = o * bytes_per_row;
    1882
    1.02k
                let row_end = row_start + bytes_per_row;
    1883
    1.02k
                let row_data = &weight_data[row_start..row_end];
    1884
    1.02k
                fused_q8_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim)
    1885
    1.02k
            })
    1886
    7
            .collect();
    1887
    1888
    7
        Ok(output)
    1889
    10
    }
    1890
    1891
    /// Fused Q8_0 × Q8_0 parallel matvec - writes to pre-allocated buffer
    1892
    ///
    1893
    /// IMP-131: Zero-allocation variant for hot-path inference.
    1894
    #[allow(clippy::similar_names)]
    1895
    4
    pub fn fused_q8_0_q8_0_parallel_matvec_into(
    1896
    4
        weight_data: &[u8],
    1897
    4
        activations: &[f32],
    1898
    4
        in_dim: usize,
    1899
    4
        out_dim: usize,
    1900
    4
        output: &mut [f32],
    1901
    4
    ) -> Result<()> {
    1902
        use rayon::prelude::*;
    1903
    1904
        const Q8_0_BLOCK_BYTES: usize = 34;
    1905
        const Q8_0_BLOCK_SIZE: usize = 32;
    1906
    1907
    4
        let blocks_per_row = in_dim.div_ceil(Q8_0_BLOCK_SIZE);
    1908
    4
        let bytes_per_row = blocks_per_row * Q8_0_BLOCK_BYTES;
    1909
    1910
    4
        let expected_weight_bytes = out_dim * bytes_per_row;
    1911
    4
        if weight_data.len() < expected_weight_bytes {
    1912
    2
            return Err(RealizarError::InvalidShape {
    1913
    2
                reason: format!(
    1914
    2
                    "Q8_0 weight data too small: need {} bytes for {}x{}, have {}",
    1915
    2
                    expected_weight_bytes,
    1916
    2
                    out_dim,
    1917
    2
                    in_dim,
    1918
    2
                    weight_data.len()
    1919
    2
                ),
    1920
    2
            });
    1921
    2
        }
    1922
    1923
    2
        if activations.len() != in_dim {
    1924
    0
            return Err(RealizarError::InvalidShape {
    1925
    0
                reason: format!(
    1926
    0
                    "Activation length {} doesn't match in_dim {}",
    1927
    0
                    activations.len(),
    1928
    0
                    in_dim
    1929
    0
                ),
    1930
    0
            });
    1931
    2
        }
    1932
    1933
    2
        if output.len() < out_dim {
    1934
    0
            return Err(RealizarError::InvalidShape {
    1935
    0
                reason: format!(
    1936
    0
                    "Output buffer too small: need {}, have {}",
    1937
    0
                    out_dim,
    1938
    0
                    output.len()
    1939
    0
                ),
    1940
    0
            });
    1941
    2
        }
    1942
    1943
        // Quantize activations to Q8_0 ONCE (amortized over all rows)
    1944
    2
        let (q8_scales, q8_quants) = quantize_activations_q8_0(activations);
    1945
    1946
        // Parallel over output rows with chunking
    1947
        const CHUNK_SIZE: usize = 64;
    1948
    2
        output[..out_dim]
    1949
    2
            .par_iter_mut()
    1950
    2
            .enumerate()
    1951
    2
            .with_min_len(CHUNK_SIZE)
    1952
    4
            .for_each2(|(o, out)| {
    1953
    4
                let row_start = o * bytes_per_row;
    1954
    4
                let row_end = row_start + bytes_per_row;
    1955
    4
                let row_data = &weight_data[row_start..row_end];
    1956
    4
                *out = fused_q8_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim);
    1957
    4
            });
    1958
    1959
    2
        Ok(())
    1960
    4
    }
    1961
    1962
    /// Helper: Extract 6-bit scale and min for a block from the packed scales array
    1963
    ///
    1964
    /// PAR-001 FIX: Matches llama.cpp's get_scale_min_k4 packing scheme:
    1965
    /// - Blocks 0-3: scale = q[j] & 63, min = q[j+4] & 63
    1966
    /// - Blocks 4-7: scale = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4)
    1967
    ///   min = (q[j+4] >> 4) | ((q[j] >> 6) << 4)
    1968
    #[inline]
    1969
    13
    pub(crate) fn extract_scale_min(scales: &[u8; 12], block_idx: usize) -> (f32, f32) {
    1970
    13
        let j = block_idx;
    1971
    13
        let (scale_bits, min_bits) = if j < 4 {
    1972
            // First 4 blocks: simple layout
    1973
    7
            let d = scales[j] & 63;
    1974
    7
            let m = scales[j + 4] & 63;
    1975
    7
            (d, m)
    1976
        } else {
    1977
            // Last 4 blocks: packed layout using high bits from first 4 bytes
    1978
    6
            let d = (scales[j + 4] & 0x0F) | ((scales[j - 4] >> 6) << 4);
    1979
    6
            let m = (scales[j + 4] >> 4) | ((scales[j] >> 6) << 4);
    1980
    6
            (d, m)
    1981
        };
    1982
    1983
        // Return raw 6-bit values as floats
    1984
        // The GGUF header's d/dmin values already include the /63 normalization
    1985
    13
        let scale = f32::from(scale_bits);
    1986
    13
        let min = f32::from(min_bits);
    1987
    1988
    13
        (scale, min)
    1989
    13
    }
    1990
    1991
    1992
    1993
    #[cfg(test)]
    1994
    1995
    #[cfg(test)]
    1996
    mod tests;