## 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/fused_k.rs

    Line
    Count
    Source
    1
    //! Fused quantized operations for K-quantization formats (PMAT-802)
    2
    //!
    3
    //! Implements fused dequant+dot operations that eliminate intermediate f32 buffer
    4
    //! allocation for 8x memory bandwidth reduction.
    5
    //!
    6
    //! Per llama-cpp-style-performance-spec.md:
    7
    //! - Memory wall is the bottleneck (Wulf & McKee [10])
    8
    //! - Fused operations keep data in registers, avoid memory round-trips
    9
    //! - ULP tolerance of ≤4 for numerical equivalence (Goldberg [9])
    10
    //!
    11
    //! Functions:
    12
    //! - `fused_q4k_dot`, `fused_q4k_dot_simd` - Q4_K dot products
    13
    //! - `fused_q4k_q8_dot` - Q4_K with Q8 activations
    14
    //! - `fused_q4k_q8k_dot`, `fused_q4k_q8k_dot_simd` - Q4_K with Q8_K activations
    15
    //! - `fused_q4k_q8k_parallel_matvec_into` - Parallel matrix-vector multiply
    16
    //! - `fused_q4k_q8k_ffn_up_gate_into` - Fused FFN up/gate computation
    17
    //! - `fused_q6k_dot`, `fused_q6k_dot_simd` - Q6_K dot products
    18
    //! - `fused_q5k_dot`, `fused_q5k_dot_simd` - Q5_K dot products
    19
    20
    use crate::error::{RealizarError, Result};
    21
    use super::dequant::read_f16;
    22
    use super::simd::extract_scale_min;
    23
    use super::types::QK_K;
    24
    25
    /// Fused Q4_K dequantize + dot product
    26
    ///
    27
    /// Computes the dot product of Q4_K quantized weights with f32 activations
    28
    /// WITHOUT allocating an intermediate f32 buffer. Dequantization happens
    29
    /// inline, accumulating directly into a register.
    30
    ///
    31
    /// # Arguments
    32
    ///
    33
    /// * `q4k_data` - Raw Q4_K quantized data (super-blocks of 144 bytes)
    34
    /// * `activations` - f32 activation values (must match dequantized length)
    35
    ///
    36
    /// # Returns
    37
    ///
    38
    /// The dot product as f32
    39
    ///
    40
    /// # Errors
    41
    ///
    42
    /// Returns error if:
    43
    /// - `q4k_data` length is not a multiple of 144 bytes (super-block size)
    44
    /// - `activations` length doesn't match the number of quantized values
    45
    ///
    46
    /// # Performance
    47
    ///
    48
    /// This function reduces memory traffic by 8x compared to separate
    49
    /// dequantize-then-dot operations:
    50
    /// - Naive: Read Q4_K (4.5 bits) → Write f32 (32 bits) → Read f32 → Compute
    51
    /// - Fused: Read Q4_K (4.5 bits) → Compute in registers
    52
    ///
    53
    /// # Examples
    54
    ///
    55
    /// ```rust,ignore
    56
    /// let weights_q4k = load_q4k_weights();
    57
    /// let activations = get_layer_activations();
    58
    /// let result = fused_q4k_dot(&weights_q4k, &activations)?;
    59
    /// ```
    60
    1.17k
    pub fn fused_q4k_dot(q4k_data: &[u8], activations: &[f32]) -> Result<f32> {
    61
        const SUPER_BLOCK_BYTES: usize = 144;
    62
    63
        // Validate Q4_K data length
    64
    1.17k
        if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {
    65
    7
            return Err(RealizarError::InvalidShape {
    66
    7
                reason: format!(
    67
    7
                    "Q4_K data length {} is not a multiple of super-block size {}",
    68
    7
                    q4k_data.len(),
    69
    7
                    SUPER_BLOCK_BYTES
    70
    7
                ),
    71
    7
            });
    72
    1.17k
        }
    73
    74
    1.17k
        let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;
    75
    1.17k
        let expected_values = num_super_blocks * QK_K;
    76
    77
        // Validate activation length matches
    78
    1.17k
        if activations.len() != expected_values {
    79
    9
            return Err(RealizarError::InvalidShape {
    80
    9
                reason: format!(
    81
    9
                    "Activation length {} doesn't match Q4_K values count {}",
    82
    9
                    activations.len(),
    83
    9
                    expected_values
    84
    9
                ),
    85
    9
            });
    86
    1.16k
        }
    87
    88
        // Accumulator for dot product result
    89
    1.16k
        let mut acc = 0.0f32;
    90
    1.16k
        let mut activation_idx = 0;
    91
    92
    16.3k
        for sb_idx in 0..num_super_blocks1.16k {
    93
    16.3k
            let sb_start = sb_idx * SUPER_BLOCK_BYTES;
    94
    95
            // Read d (f16 -> f32)
    96
    16.3k
            let d = read_f16(&q4k_data[sb_start..sb_start + 2]);
    97
    98
            // Read dmin (f16 -> f32)
    99
    16.3k
            let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);
    100
    101
            // Read scales (12 bytes)
    102
    16.3k
            let mut scales = [0u8; 12];
    103
    16.3k
            scales.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);
    104
    105
            // Read qs (128 bytes)
    106
    16.3k
            let qs_start = sb_start + 16;
    107
    16.3k
            let qs = &q4k_data[qs_start..qs_start + 128];
    108
    109
            // PAR-001: Match dequantize_q4_k layout (llama.cpp/candle compatible)
    110
            // Process 4 chunks of 64 values each (0, 64, 128, 192)
    111
            // Each chunk: 32 low nibbles, then 32 high nibbles from 32 consecutive bytes
    112
    65.5k
            for j in (0..QK_K)16.3k.step_by16.3k(64) {
    113
    65.5k
                let q = &qs[j / 2..j / 2 + 32];
    114
    115
                // Get scales for the two 32-value halves
    116
    65.5k
                let is = j / 32;
    117
    65.5k
                let (sc1, m1) = extract_scale_min(&scales, is);
    118
    65.5k
                let d1 = d * sc1;
    119
    65.5k
                let dm1 = dmin * m1;
    120
    121
    65.5k
                let (sc2, m2) = extract_scale_min(&scales, is + 1);
    122
    65.5k
                let d2 = d * sc2;
    123
    65.5k
                let dm2 = dmin * m2;
    124
    125
                // First pass: 32 low nibbles (use sc1, m1)
    126
    2.16M
                for &byte2.09M in q {
    127
    2.09M
                    let q_val = (byte & 0x0F) as f32;
    128
    2.09M
                    let value = d1 * q_val - dm1;
    129
    2.09M
                    acc += value * activations[activation_idx];
    130
    2.09M
                    activation_idx += 1;
    131
    2.09M
                }
    132
    133
                // Second pass: 32 high nibbles (use sc2, m2)
    134
    2.16M
                for &byte2.09M in q {
    135
    2.09M
                    let q_val = (byte >> 4) as f32;
    136
    2.09M
                    let value = d2 * q_val - dm2;
    137
    2.09M
                    acc += value * activations[activation_idx];
    138
    2.09M
                    activation_idx += 1;
    139
    2.09M
                }
    140
            }
    141
        }
    142
    143
    1.16k
        Ok(acc)
    144
    1.17k
    }
    145
    146
    /// Fused Q4_K dequantize + dot product with SIMD acceleration
    147
    ///
    148
    /// This is the public, safe API that automatically dispatches to the best
    149
    /// available implementation (AVX2 when available, scalar fallback otherwise).
    150
    ///
    151
    /// # Arguments
    152
    ///
    153
    /// * `q4k_data` - Raw Q4_K quantized data (super-blocks of 144 bytes)
    154
    /// * `activations` - f32 activation values (must match dequantized length)
    155
    ///
    156
    /// # Returns
    157
    ///
    158
    /// The dot product as f32, matching `fused_q4k_dot` within 4 ULPs
    159
    ///
    160
    /// # Errors
    161
    ///
    162
    /// Returns error if:
    163
    /// - `q4k_data` length is not a multiple of 144 bytes (super-block size)
    164
    /// - `activations` length doesn't match the number of quantized values
    165
    ///
    166
    /// # Performance
    167
    ///
    168
    /// - AVX2: ~8x speedup over scalar via 256-bit SIMD + FMA
    169
    /// - Fused operation: 8x memory bandwidth reduction vs dequant-then-dot
    170
    /// - Combined potential: Up to 64x improvement for memory-bound operations
    171
    1.49M
    pub fn fused_q4k_dot_simd(q4k_data: &[u8], activations: &[f32]) -> Result<f32> {
    172
        // Runtime feature detection with fallback (per RustBelt pattern)
    173
        #[cfg(target_arch = "x86_64")]
    174
        {
    175
            // PAR-126: AVX-512 VNNI requires pre-quantized activations (Q4K×Q8K format)
    176
            // For now, use AVX2 which works with f32 activations directly.
    177
            // Future optimization: pre-quantize activations to Q8_0 format once per matmul.
    178
    1.49M
            if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
    179
                // SAFETY: We've verified AVX2 and FMA are available at runtime
    180
                // The unsafe function performs the same logical operation as scalar
    181
                // SAFETY: Memory safety ensured by bounds checking and alignment
    182
    1.49M
                return unsafe { fused_q4k_dot_avx2(q4k_data, activations) };
    183
    0
            }
    184
        }
    185
    186
        // Fallback to scalar implementation
    187
    0
        fused_q4k_dot(q4k_data, activations)
    188
    1.49M
    }
    189
    190
    /// AVX-512 VNNI accelerated Q4_K dot product (PAR-126)
    191
    ///
    192
    /// Uses vpdpbusd for int8×int8 accumulation which is 4x faster than FP32 FMA.
    193
    /// Key insight from llama.cpp/llamafile: integer dot products dominate FP on modern CPUs.
    194
    ///
    195
    /// # Algorithm
    196
    /// 1. Quantize f32 activations to int8 (on the fly, per super-block)
    197
    /// 2. Expand 4-bit weights to int8
    198
    /// 3. Use vpdpbusd for packed u8×i8 dot product
    199
    /// 4. Scale and accumulate at the end
    200
    ///
    201
    /// # Safety
    202
    /// Requires AVX-512F, AVX-512BW, and AVX-512VNNI
    203
    #[cfg(target_arch = "x86_64")]
    204
    #[target_feature(enable = "avx512f", enable = "avx512bw", enable = "avx512vnni")]
    205
    #[allow(unsafe_op_in_unsafe_fn)]
    206
    0
    unsafe fn fused_q4k_dot_avx512_vnni(q4k_data: &[u8], activations: &[f32]) -> Result<f32> {
    207
        #[allow(clippy::wildcard_imports)]
    208
        use std::arch::x86_64::*;
    209
    210
        const SUPER_BLOCK_BYTES: usize = 144;
    211
    212
    0
        if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {
    213
    0
            return Err(RealizarError::InvalidShape {
    214
    0
                reason: format!(
    215
    0
                    "Q4_K data length {} is not a multiple of super-block size {}",
    216
    0
                    q4k_data.len(),
    217
    0
                    SUPER_BLOCK_BYTES
    218
    0
                ),
    219
    0
            });
    220
    0
        }
    221
    222
    0
        let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;
    223
    0
        let expected_values = num_super_blocks * QK_K;
    224
    225
    0
        if activations.len() != expected_values {
    226
    0
            return Err(RealizarError::InvalidShape {
    227
    0
                reason: format!(
    228
    0
                    "Activation length {} doesn't match Q4_K values count {}",
    229
    0
                    activations.len(),
    230
    0
                    expected_values
    231
    0
                ),
    232
    0
            });
    233
    0
        }
    234
    235
        // Final accumulator (FP32)
    236
    0
        let mut total_sum = 0.0f32;
    237
    0
        let mut activation_idx = 0;
    238
    239
        // Nibble mask for 4-bit extraction
    240
    0
        let nibble_mask = _mm512_set1_epi8(0x0F_i8);
    241
    242
    0
        for sb_idx in 0..num_super_blocks {
    243
    0
            let sb_start = sb_idx * SUPER_BLOCK_BYTES;
    244
    245
            // Prefetch next super-block
    246
    0
            if sb_idx + 1 < num_super_blocks {
    247
    0
                let next_sb = (sb_idx + 1) * SUPER_BLOCK_BYTES;
    248
    0
                _mm_prefetch(q4k_data.as_ptr().add(next_sb).cast::<i8>(), _MM_HINT_T0);
    249
    0
            }
    250
    251
            // Read d and dmin (f16 → f32)
    252
    0
            let d = read_f16(&q4k_data[sb_start..sb_start + 2]);
    253
    0
            let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);
    254
    255
            // Read scales (12 bytes)
    256
    0
            let mut scales = [0u8; 12];
    257
    0
            scales.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);
    258
    259
            // Pointer to quantized data (128 bytes)
    260
    0
            let qs_ptr = q4k_data.as_ptr().add(sb_start + 16);
    261
    262
            // Process 64 values at a time (matches AVX-512 width of 64 bytes)
    263
    0
            for j in (0..QK_K).step_by(64) {
    264
    0
                let q_start = j / 2;
    265
    266
                // Get scales for the two 32-value halves
    267
    0
                let is = j / 32;
    268
    0
                let (sc1, m1) = extract_scale_min(&scales, is);
    269
    0
                let (sc2, m2) = extract_scale_min(&scales, is + 1);
    270
    271
                // Load 32 bytes of quantized data (64 nibbles)
    272
    0
                let q_bytes_256 = _mm256_loadu_si256(qs_ptr.add(q_start).cast::<__m256i>());
    273
    274
                // Expand to 512-bit for AVX-512 operations
    275
    0
                let q_bytes = _mm512_castsi256_si512(q_bytes_256);
    276
    0
                let q_bytes = _mm512_inserti64x4(q_bytes, q_bytes_256, 1);
    277
    278
                // Extract low and high nibbles
    279
    0
                let q_lo = _mm512_and_si512(q_bytes, nibble_mask);
    280
    0
                let q_hi = _mm512_and_si512(_mm512_srli_epi16(q_bytes, 4), nibble_mask);
    281
    282
                // Process with integer dot products
    283
                // We need to quantize activations and use vpdpbusd
    284
    285
                // Load 32 f32 activations for low nibbles, convert to int8
    286
    0
                let act_slice = &activations[activation_idx..activation_idx + 32];
    287
    288
                // Find min/max for quantization
    289
    0
                let act_max = act_slice.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    290
    0
                let act_min = act_slice.iter().copied().fold(f32::INFINITY, f32::min);
    291
    0
                let act_scale = if act_max > act_min {
    292
    0
                    127.0 / (act_max - act_min)
    293
                } else {
    294
    0
                    1.0
    295
                };
    296
    297
                // Quantize activations to int8
    298
    0
                let mut act_i8 = [0i8; 32];
    299
    0
                for (i, &a) in act_slice.iter().enumerate() {
    300
    0
                    act_i8[i] = ((a - act_min) * act_scale).round() as i8;
    301
    0
                }
    302
    303
                // Load quantized activations
    304
    0
                let act_vec = _mm256_loadu_si256(act_i8.as_ptr().cast::<__m256i>());
    305
    306
                // Expand q_lo from u8 to match
    307
    0
                let q_lo_256 = _mm512_castsi512_si256(q_lo);
    308
    309
                // Use vpdpbusd: u8 × i8 → i32 accumulation
    310
    0
                let _zero = _mm512_setzero_si512();
    311
    0
                let q_lo_512 = _mm512_cvtepu8_epi16(q_lo_256);
    312
    0
                let act_512 = _mm512_cvtepi8_epi16(act_vec);
    313
    0
                let prod = _mm512_mullo_epi16(q_lo_512, act_512);
    314
    315
                // Horizontal sum
    316
    0
                let sum_256 = _mm256_add_epi16(
    317
    0
                    _mm512_castsi512_si256(prod),
    318
    0
                    _mm512_extracti64x4_epi64(prod, 1),
    319
                );
    320
    0
                let sum_128 = _mm_add_epi16(
    321
    0
                    _mm256_castsi256_si128(sum_256),
    322
    0
                    _mm256_extracti128_si256(sum_256, 1),
    323
                );
    324
    325
                // Expand to i32 and sum
    326
    0
                let sum_32 = _mm256_cvtepi16_epi32(sum_128);
    327
    0
                let sum_arr: [i32; 8] = std::mem::transmute(sum_32);
    328
    0
                let int_sum: i32 = sum_arr.iter().sum();
    329
    330
                // Scale back to float
    331
    0
                let float_sum = int_sum as f32 * d * sc1 / act_scale - (32.0 * dmin * m1);
    332
    0
                total_sum += float_sum;
    333
    0
                activation_idx += 32;
    334
    335
                // Process high nibbles similarly
    336
    0
                let act_slice2 = &activations[activation_idx..activation_idx + 32];
    337
    0
                let act_max2 = act_slice2.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    338
    0
                let act_min2 = act_slice2.iter().copied().fold(f32::INFINITY, f32::min);
    339
    0
                let act_scale2 = if act_max2 > act_min2 {
    340
    0
                    127.0 / (act_max2 - act_min2)
    341
                } else {
    342
    0
                    1.0
    343
                };
    344
    345
    0
                let mut act_i8_2 = [0i8; 32];
    346
    0
                for (i, &a) in act_slice2.iter().enumerate() {
    347
    0
                    act_i8_2[i] = ((a - act_min2) * act_scale2).round() as i8;
    348
    0
                }
    349
    350
    0
                let act_vec2 = _mm256_loadu_si256(act_i8_2.as_ptr().cast::<__m256i>());
    351
    0
                let q_hi_256 = _mm512_castsi512_si256(q_hi);
    352
    0
                let q_hi_512 = _mm512_cvtepu8_epi16(q_hi_256);
    353
    0
                let act_512_2 = _mm512_cvtepi8_epi16(act_vec2);
    354
    0
                let prod2 = _mm512_mullo_epi16(q_hi_512, act_512_2);
    355
    356
    0
                let sum_256_2 = _mm256_add_epi16(
    357
    0
                    _mm512_castsi512_si256(prod2),
    358
    0
                    _mm512_extracti64x4_epi64(prod2, 1),
    359
                );
    360
    0
                let sum_128_2 = _mm_add_epi16(
    361
    0
                    _mm256_castsi256_si128(sum_256_2),
    362
    0
                    _mm256_extracti128_si256(sum_256_2, 1),
    363
                );
    364
    0
                let sum_32_2 = _mm256_cvtepi16_epi32(sum_128_2);
    365
    0
                let sum_arr2: [i32; 8] = std::mem::transmute(sum_32_2);
    366
    0
                let int_sum2: i32 = sum_arr2.iter().sum();
    367
    368
    0
                let float_sum2 = int_sum2 as f32 * d * sc2 / act_scale2 - (32.0 * dmin * m2);
    369
    0
                total_sum += float_sum2;
    370
    0
                activation_idx += 32;
    371
            }
    372
        }
    373
    374
    0
        Ok(total_sum)
    375
    0
    }
    376
    377
    /// AVX2-accelerated fused Q4_K dequant+dot kernel (PARITY-003: llama.cpp-style SIMD)
    378
    ///
    379
    /// # Safety
    380
    ///
    381
    /// Caller must ensure:
    382
    /// 1. AVX2 and FMA CPU features are available (use `is_x86_feature_detected!`)
    383
    /// 2. Input slices are valid (handled by Rust's slice guarantees)
    384
    ///
    385
    /// This function is marked unsafe due to SIMD intrinsics, but is logically
    386
    /// equivalent to the scalar `fused_q4k_dot` (within ULP tolerance).
    387
    ///
    388
    /// # Optimizations (PARITY-003)
    389
    /// - SIMD loads: `_mm256_loadu_si256` for 32-byte bulk loads (vs scalar byte loads)
    390
    /// - SIMD nibble extraction: `_mm256_and_si256` with 0x0F mask (vs scalar & 0x0F)
    391
    /// - 4 independent accumulators to hide FMA latency
    392
    /// - Software prefetching for next super-block
    393
    /// - Matches llama.cpp ggml_vec_dot_q4_K_q8_K pattern
    394
    #[cfg(target_arch = "x86_64")]
    395
    #[target_feature(enable = "avx2", enable = "fma")]
    396
    #[allow(unsafe_op_in_unsafe_fn)]
    397
    1.49M
    unsafe fn fused_q4k_dot_avx2(q4k_data: &[u8], activations: &[f32]) -> Result<f32> {
    398
        // Allow wildcard import for SIMD intrinsics (standard pattern for arch-specific code)
    399
        #[allow(clippy::wildcard_imports)]
    400
        use std::arch::x86_64::*;
    401
    402
        const SUPER_BLOCK_BYTES: usize = 144;
    403
    404
        // Validate inputs (same as scalar)
    405
    1.49M
        if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {
    406
    3
            return Err(RealizarError::InvalidShape {
    407
    3
                reason: format!(
    408
    3
                    "Q4_K data length {} is not a multiple of super-block size {}",
    409
    3
                    q4k_data.len(),
    410
    3
                    SUPER_BLOCK_BYTES
    411
    3
                ),
    412
    3
            });
    413
    1.49M
        }
    414
    415
    1.49M
        let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;
    416
    1.49M
        let expected_values = num_super_blocks * QK_K;
    417
    418
    1.49M
        if activations.len() != expected_values {
    419
    317k
            return Err(RealizarError::InvalidShape {
    420
    317k
                reason: format!(
    421
    317k
                    "Activation length {} doesn't match Q4_K values count {}",
    422
    317k
                    activations.len(),
    423
    317k
                    expected_values
    424
    317k
                ),
    425
    317k
            });
    426
    1.18M
        }
    427
    428
        // Nibble mask for extracting 4-bit values (llama.cpp pattern)
    429
    1.18M
        let nibble_mask = _mm256_set1_epi8(0x0F_i8);
    430
    431
        // PARITY-003: 4 independent accumulators to hide FMA latency
    432
        // FMA latency = 4 cycles, throughput = 2/cycle
    433
        // With 4 independent chains, we saturate the FMA throughput
    434
    1.18M
        let mut acc0 = _mm256_setzero_ps();
    435
    1.18M
        let mut acc1 = _mm256_setzero_ps();
    436
    1.18M
        let mut acc2 = _mm256_setzero_ps();
    437
    1.18M
        let mut acc3 = _mm256_setzero_ps();
    438
    1.18M
        let mut activation_idx = 0;
    439
    440
    1.20M
        for sb_idx in 0..num_super_blocks1.18M {
    441
    1.20M
            let sb_start = sb_idx * SUPER_BLOCK_BYTES;
    442
    443
            // Prefetch next super-block while processing current
    444
    1.20M
            if sb_idx + 1 < num_super_blocks {
    445
    27.4k
                let next_sb = (sb_idx + 1) * SUPER_BLOCK_BYTES;
    446
    27.4k
                // SAFETY: Prefetch is a hint, pointer arithmetic is in bounds (checked above)
    447
    27.4k
                _mm_prefetch(q4k_data.as_ptr().add(next_sb).cast::<i8>(), _MM_HINT_T0);
    448
    1.18M
            }
    449
    450
            // Read d and dmin (f16 -> f32)
    451
    1.20M
            let d = read_f16(&q4k_data[sb_start..sb_start + 2]);
    452
    1.20M
            let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);
    453
    454
            // Read scales (12 bytes)
    455
    1.20M
            let mut scales = [0u8; 12];
    456
    1.20M
            scales.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);
    457
    458
            // Pointer to quantized data (128 bytes = 256 nibbles = 256 values)
    459
    1.20M
            let qs_ptr = q4k_data.as_ptr().add(sb_start + 16);
    460
    461
            // PAR-001: Match dequantize_q4_k layout (llama.cpp/candle compatible)
    462
            // Process 4 chunks of 64 values each (j=0, 64, 128, 192)
    463
            // Each chunk: 32 low nibbles (sc1), then 32 high nibbles (sc2)
    464
    4.83M
            for j in (0..QK_K)1.20M.step_by1.20M(64) {
    465
    4.83M
                let q_start = j / 2; // 32 bytes per 64-value chunk
    466
    4.83M
    467
    4.83M
                // Get scales for the two 32-value halves
    468
    4.83M
                let is = j / 32;
    469
    4.83M
                let (sc1, m1) = extract_scale_min(&scales, is);
    470
    4.83M
                let (sc2, m2) = extract_scale_min(&scales, is + 1);
    471
    4.83M
    472
    4.83M
                // Precompute d*scale and dmin*min for both halves
    473
    4.83M
                let d_scale1 = d * sc1;
    474
    4.83M
                let dm1 = dmin * m1;
    475
    4.83M
                let d_scale2 = d * sc2;
    476
    4.83M
                let dm2 = dmin * m2;
    477
    4.83M
    478
    4.83M
                // SIMD OPTIMIZATION: Load 32 bytes at once (64 nibbles = 64 values)
    479
    4.83M
                // llama.cpp pattern: _mm256_loadu_si256 + AND/SHIFT for nibble extraction
    480
    4.83M
                let q_bytes = _mm256_loadu_si256(qs_ptr.add(q_start).cast::<__m256i>());
    481
    4.83M
    482
    4.83M
                // Extract low nibbles (first 32 values) and high nibbles (second 32 values)
    483
    4.83M
                let q_lo = _mm256_and_si256(q_bytes, nibble_mask);
    484
    4.83M
                let q_hi = _mm256_and_si256(_mm256_srli_epi16(q_bytes, 4), nibble_mask);
    485
    4.83M
    486
    4.83M
                // Process low nibbles (32 values with scale sc1/m1)
    487
    4.83M
                // Split into 4 groups of 8 for f32 conversion (AVX2 can convert 8 i32→f32)
    488
    4.83M
                let d_scale1_vec = _mm256_set1_ps(d_scale1);
    489
    4.83M
                let dm1_vec = _mm256_set1_ps(dm1);
    490
    4.83M
    491
    4.83M
                // Extract bytes 0-7 (low nibbles) → convert to f32
    492
    4.83M
                let q_lo_128_0 = _mm256_castsi256_si128(q_lo);
    493
    4.83M
                let q_lo_i32_0 = _mm256_cvtepu8_epi32(q_lo_128_0);
    494
    4.83M
                let q_lo_f32_0 = _mm256_cvtepi32_ps(q_lo_i32_0);
    495
    4.83M
                let dequant0 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_0, dm1_vec);
    496
    4.83M
                let act0 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    497
    4.83M
                acc0 = _mm256_fmadd_ps(dequant0, act0, acc0);
    498
    4.83M
                activation_idx += 8;
    499
    4.83M
    500
    4.83M
                // Extract bytes 8-15 (low nibbles)
    501
    4.83M
                let q_lo_shifted = _mm_srli_si128(q_lo_128_0, 8);
    502
    4.83M
                let q_lo_i32_1 = _mm256_cvtepu8_epi32(q_lo_shifted);
    503
    4.83M
                let q_lo_f32_1 = _mm256_cvtepi32_ps(q_lo_i32_1);
    504
    4.83M
                let dequant1 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_1, dm1_vec);
    505
    4.83M
                let act1 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    506
    4.83M
                acc1 = _mm256_fmadd_ps(dequant1, act1, acc1);
    507
    4.83M
                activation_idx += 8;
    508
    4.83M
    509
    4.83M
                // Extract bytes 16-23 (low nibbles from high 128 bits)
    510
    4.83M
                let q_lo_128_1 = _mm256_extracti128_si256(q_lo, 1);
    511
    4.83M
                let q_lo_i32_2 = _mm256_cvtepu8_epi32(q_lo_128_1);
    512
    4.83M
                let q_lo_f32_2 = _mm256_cvtepi32_ps(q_lo_i32_2);
    513
    4.83M
                let dequant2 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_2, dm1_vec);
    514
    4.83M
                let act2 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    515
    4.83M
                acc2 = _mm256_fmadd_ps(dequant2, act2, acc2);
    516
    4.83M
                activation_idx += 8;
    517
    4.83M
    518
    4.83M
                // Extract bytes 24-31 (low nibbles)
    519
    4.83M
                let q_lo_shifted2 = _mm_srli_si128(q_lo_128_1, 8);
    520
    4.83M
                let q_lo_i32_3 = _mm256_cvtepu8_epi32(q_lo_shifted2);
    521
    4.83M
                let q_lo_f32_3 = _mm256_cvtepi32_ps(q_lo_i32_3);
    522
    4.83M
                let dequant3 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_3, dm1_vec);
    523
    4.83M
                let act3 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    524
    4.83M
                acc3 = _mm256_fmadd_ps(dequant3, act3, acc3);
    525
    4.83M
                activation_idx += 8;
    526
    4.83M
    527
    4.83M
                // Process high nibbles (32 values with scale sc2/m2)
    528
    4.83M
                let d_scale2_vec = _mm256_set1_ps(d_scale2);
    529
    4.83M
                let dm2_vec = _mm256_set1_ps(dm2);
    530
    4.83M
    531
    4.83M
                // Extract bytes 0-7 (high nibbles)
    532
    4.83M
                let q_hi_128_0 = _mm256_castsi256_si128(q_hi);
    533
    4.83M
                let q_hi_i32_0 = _mm256_cvtepu8_epi32(q_hi_128_0);
    534
    4.83M
                let q_hi_f32_0 = _mm256_cvtepi32_ps(q_hi_i32_0);
    535
    4.83M
                let dequant4 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_0, dm2_vec);
    536
    4.83M
                let act4 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    537
    4.83M
                acc0 = _mm256_fmadd_ps(dequant4, act4, acc0);
    538
    4.83M
                activation_idx += 8;
    539
    4.83M
    540
    4.83M
                // Extract bytes 8-15 (high nibbles)
    541
    4.83M
                let q_hi_shifted = _mm_srli_si128(q_hi_128_0, 8);
    542
    4.83M
                let q_hi_i32_1 = _mm256_cvtepu8_epi32(q_hi_shifted);
    543
    4.83M
                let q_hi_f32_1 = _mm256_cvtepi32_ps(q_hi_i32_1);
    544
    4.83M
                let dequant5 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_1, dm2_vec);
    545
    4.83M
                let act5 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    546
    4.83M
                acc1 = _mm256_fmadd_ps(dequant5, act5, acc1);
    547
    4.83M
                activation_idx += 8;
    548
    4.83M
    549
    4.83M
                // Extract bytes 16-23 (high nibbles from high 128 bits)
    550
    4.83M
                let q_hi_128_1 = _mm256_extracti128_si256(q_hi, 1);
    551
    4.83M
                let q_hi_i32_2 = _mm256_cvtepu8_epi32(q_hi_128_1);
    552
    4.83M
                let q_hi_f32_2 = _mm256_cvtepi32_ps(q_hi_i32_2);
    553
    4.83M
                let dequant6 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_2, dm2_vec);
    554
    4.83M
                let act6 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    555
    4.83M
                acc2 = _mm256_fmadd_ps(dequant6, act6, acc2);
    556
    4.83M
                activation_idx += 8;
    557
    4.83M
    558
    4.83M
                // Extract bytes 24-31 (high nibbles)
    559
    4.83M
                let q_hi_shifted2 = _mm_srli_si128(q_hi_128_1, 8);
    560
    4.83M
                let q_hi_i32_3 = _mm256_cvtepu8_epi32(q_hi_shifted2);
    561
    4.83M
                let q_hi_f32_3 = _mm256_cvtepi32_ps(q_hi_i32_3);
    562
    4.83M
                let dequant7 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_3, dm2_vec);
    563
    4.83M
                let act7 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));
    564
    4.83M
                acc3 = _mm256_fmadd_ps(dequant7, act7, acc3);
    565
    4.83M
                activation_idx += 8;
    566
    4.83M
            }
    567
        }
    568
    569
        // Combine 4 accumulators → single accumulator
    570
    1.18M
        let acc_01 = _mm256_add_ps(acc0, acc1);
    571
    1.18M
        let acc_23 = _mm256_add_ps(acc2, acc3);
    572
    1.18M
        let acc = _mm256_add_ps(acc_01, acc_23);
    573
    574
        // Horizontal sum: reduce 8 lanes to single value
    575
    1.18M
        let sum_halves = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps(acc, 1));
    576
    1.18M
        let temp = _mm_add_ps(sum_halves, _mm_movehl_ps(sum_halves, sum_halves));
    577
    1.18M
        let temp = _mm_add_ss(temp, _mm_shuffle_ps(temp, temp, 1));
    578
    1.18M
        let result = _mm_cvtss_f32(temp);
    579
    580
    1.18M
        Ok(result)
    581
    1.49M
    }
    582
    583
    // ============================================================================
    584
    // Q4_K × Q8_K KERNELS (Super-block aligned integer-only arithmetic)
    585
    // ============================================================================
    586
    587
    /// Fused Q4_K × Q8_K dot product (super-block aligned, llama.cpp-style)
    588
    ///
    589
    /// Uses Q8_K format (256 values per super-block, single scale) for maximum
    590
    /// SIMD efficiency. This matches llama.cpp's `ggml_vec_dot_q4_K_q8_K`.
    591
    ///
    592
    /// # Arguments
    593
    /// * `q4k_data` - Raw Q4_K quantized data (144 bytes per super-block)
    594
    /// * `q8k_scales` - Q8_K scales (one per super-block)
    595
    /// * `q8k_quants` - Q8_K quantized int8 values (256 per super-block)
    596
    ///
    597
    /// # Performance
    598
    ///
    599
    /// Compared to Q4_K × f32:
    600
    /// - 8x fewer memory reads for activations
    601
    /// - Integer-only inner loop (no f32 conversion until end)
    602
    /// - Single scale multiplication per super-block (vs 8 for Q8_0)
    603
    40
    pub fn fused_q4k_q8k_dot(q4k_data: &[u8], q8k_scales: &[f32], q8k_quants: &[i8]) -> Result<f32> {
    604
        const SUPER_BLOCK_BYTES: usize = 144;
    605
    606
    40
        if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {
    607
    7
            return Err(RealizarError::InvalidShape {
    608
    7
                reason: format!(
    609
    7
                    "Q4_K data length {} is not a multiple of {}",
    610
    7
                    q4k_data.len(),
    611
    7
                    SUPER_BLOCK_BYTES
    612
    7
                ),
    613
    7
            });
    614
    33
        }
    615
    616
    33
        let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;
    617
    33
        let expected_values = num_super_blocks * QK_K;
    618
    619
    33
        if q8k_scales.len() < num_super_blocks {
    620
    3
            return Err(RealizarError::InvalidShape {
    621
    3
                reason: format!(
    622
    3
                    "Q8_K scales count {} < expected {}",
    623
    3
                    q8k_scales.len(),
    624
    3
                    num_super_blocks
    625
    3
                ),
    626
    3
            });
    627
    30
        }
    628
    629
    30
        if q8k_quants.len() < expected_values {
    630
    4
            return Err(RealizarError::InvalidShape {
    631
    4
                reason: format!(
    632
    4
                    "Q8_K quants count {} < expected {}",
    633
    4
                    q8k_quants.len(),
    634
    4
                    expected_values
    635
    4
                ),
    636
    4
            });
    637
    26
        }
    638
    639
    26
        let mut total_acc = 0.0f32;
    640
    641
    45
        for sb_idx in 0..num_super_blocks26 {
    642
    45
            let sb_start = sb_idx * SUPER_BLOCK_BYTES;
    643
    45
            let q8_start = sb_idx * QK_K;
    644
    645
            // Read Q4_K super-block header
    646
    45
            let d = read_f16(&q4k_data[sb_start..sb_start + 2]);
    647
    45
            let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);
    648
    649
            // Read scales (12 bytes for 8 blocks)
    650
    45
            let mut scales = [0u8; 12];
    651
    45
            scales.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);
    652
    653
            // Q8_K scale for this super-block
    654
    45
            let q8_scale = q8k_scales[sb_idx];
    655
    656
            // Process 4 chunks of 64 values (matching dequantize_q4_k layout)
    657
            // The dequantized output order is: 32 low nibbles, then 32 high nibbles
    658
            // So activations[j..j+32] correspond to low nibbles, activations[j+32..j+64] to high
    659
    180
            for j in (0..QK_K)45.step_by45(64) {
    660
    180
                let q_offset = sb_start + 16 + j / 2; // 32 bytes per 64-value chunk
    661
    180
                let q8_offset = q8_start + j;
    662
    663
                // Get scales for low and high nibbles
    664
    180
                let is = j / 32;
    665
    180
                let (sc1, m1) = extract_scale_min(&scales, is);
    666
    180
                let (sc2, m2) = extract_scale_min(&scales, is + 1);
    667
    668
                // Combined scale factors
    669
    180
                let d_sc1_q8 = d * sc1 * q8_scale;
    670
    180
                let dm1_q8 = dmin * m1 * q8_scale;
    671
    180
                let d_sc2_q8 = d * sc2 * q8_scale;
    672
    180
                let dm2_q8 = dmin * m2 * q8_scale;
    673
    674
                // Accumulators for low and high nibbles
    675
    180
                let mut sum_lo: i32 = 0; // q4_lo × q8 (for activations[j..j+32])
    676
    180
                let mut sum_hi: i32 = 0; // q4_hi × q8 (for activations[j+32..j+64])
    677
    180
                let mut q8_sum_lo: i32 = 0;
    678
    180
                let mut q8_sum_hi: i32 = 0;
    679
    680
    5.94k
                for b5.76k in 0..32 {
    681
    5.76k
                    let q4_byte = q4k_data[q_offset + b];
    682
    5.76k
    683
    5.76k
                    // Low nibble × activation[j + b] (first 32 positions in dequant order)
    684
    5.76k
                    let q4_lo = (q4_byte & 0x0F) as i32;
    685
    5.76k
                    let q8_lo = q8k_quants[q8_offset + b] as i32;
    686
    5.76k
                    sum_lo += q4_lo * q8_lo;
    687
    5.76k
                    q8_sum_lo += q8_lo;
    688
    5.76k
    689
    5.76k
                    // High nibble × activation[j + 32 + b] (second 32 positions in dequant order)
    690
    5.76k
                    let q4_hi = ((q4_byte >> 4) & 0x0F) as i32;
    691
    5.76k
                    let q8_hi = q8k_quants[q8_offset + 32 + b] as i32;
    692
    5.76k
                    sum_hi += q4_hi * q8_hi;
    693
    5.76k
                    q8_sum_hi += q8_hi;
    694
    5.76k
                }
    695
    696
                // Apply formula: (d * scale * sum_q4_q8 - dmin * min * sum_q8) * q8_scale
    697
    180
                total_acc += d_sc1_q8 * (sum_lo as f32) - dm1_q8 * (q8_sum_lo as f32);
    698
    180
                total_acc += d_sc2_q8 * (sum_hi as f32) - dm2_q8 * (q8_sum_hi as f32);
    699
            }
    700
        }
    701
    702
    26
        Ok(total_acc)
    703
    40
    }
    704
    705
    /// SIMD-accelerated Q4_K × Q8_K dot product
    706
    ///
    707
    /// Uses AVX-512 VNNI (vpdpbusd) for maximum throughput, falls back to AVX2 or scalar.
    708
    /// Single scale per super-block eliminates per-block overhead.
    709
    536
    pub fn fused_q4k_q8k_dot_simd(
    710
    536
        q4k_data: &[u8],
    711
    536
        q8k_scales: &[f32],
    712
    536
        q8k_quants: &[i8],
    713
    536
    ) -> Result<f32> {
    714
        #[cfg(target_arch = "x86_64")]
    715
        {
    716
            // PAR-126: Use V2 optimized AVX-512 VNNI kernel (deferred horizontal sums)
    717
    536
            if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512vnni") {
    718
                // SAFETY: Memory safety ensured by bounds checking and alignment
    719
    536
                return unsafe { fused_q4k_q8k_dot_avx512vnni_v2(q4k_data, q8k_scales, q8k_quants) };
    720
    0
            }
    721
            // Fallback to AVX2 (layout issue resolved)
    722
    0
            if is_x86_feature_detected!("avx2") {
    723
                // SAFETY: Memory safety ensured by bounds checking and alignment
    724
    0
                return unsafe { fused_q4k_q8k_dot_avx2(q4k_data, q8k_scales, q8k_quants) };
    725
    0
            }
    726
        }
    727
    728
    0
        fused_q4k_q8k_dot(q4k_data, q8k_scales, q8k_quants)
    729
    536
    }
    730
    731
    /// PAR-126 V2: AVX-512 VNNI kernel with deferred horizontal sums
    732
    ///
    733
    /// Key optimization: Instead of extracting block sums to scalars in the inner loop,
    734
    /// keeps all 8 block accumulators in __m256i vectors and only reduces at the end.
    735
    /// This eliminates 24 horizontal sums per row (from 8 to 1).
    736
    #[cfg(target_arch = "x86_64")]
    737
    #[target_feature(enable = "avx512f", enable = "avx512vnni", enable = "avx512bw")]
    738
    #[allow(unsafe_op_in_unsafe_fn)]
    739
    #[allow(clippy::similar_names)]
    740
    #[allow(clippy::too_many_lines)]
    741
    536
    unsafe fn fused_q4k_q8k_dot_avx512vnni_v2(
    742
    536
        q4k_data: &[u8],
    743
    536
        q8k_scales: &[f32],
    744
    536
        q8k_quants: &[i8],
    745
    536
    ) -> Result<f32> {
    746
        #[allow(clippy::wildcard_imports)]
    747
        use std::arch::x86_64::*;
    748
    749
        const SUPER_BLOCK_BYTES: usize = 144;
    750
    751
    536
        if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {
    752
    3
            return Err(RealizarError::InvalidShape {
    753
    3
                reason: format!(
    754
    3
                    "Q4_K data length {} is not a multiple of {}",
    755
    3
                    q4k_data.len(),
    756
    3
                    SUPER_BLOCK_BYTES
    757
    3
                ),
    758
    3
            });
    759
    533
        }
    760
    761
    533
        let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;
    762
    533
        let expected_values = num_super_blocks * QK_K;
    763
    764
    533
        if q8k_scales.len() < num_super_blocks || q8k_quants.len() < expected_values {
    765
    1
            return Err(RealizarError::InvalidShape {
    766
    1
                reason: "Q8_K buffer too small".to_string(),
    767
    1
            });
    768
    532
        }
    769
    770
    532
        let nibble_mask = _mm256_set1_epi8(0x0F_i8);
    771
    532
        let ones_16 = _mm256_set1_epi16(1);
    772
    773
        // Global float accumulator
    774
    532
        let mut total_acc = _mm256_setzero_ps();
    775
    776
    553
        for sb_idx in 0..num_super_blocks532 {
    777
    553
            let sb_start = sb_idx * SUPER_BLOCK_BYTES;
    778
    553
            let q8_start = sb_idx * QK_K;
    779
    780
            // Prefetch next super-block
    781
    553
            if sb_idx + 1 < num_super_blocks {
    782
    22
                _mm_prefetch(
    783
    22
                    q4k_data
    784
    22
                        .as_ptr()
    785
    22
                        .add((sb_idx + 1) * SUPER_BLOCK_BYTES)
    786
    22
                        .cast::<i8>(),
    787
    22
                    _MM_HINT_T0,
    788
    22
                );
    789
    22
                _mm_prefetch(
    790
    22
                    q8k_quants.as_ptr().add((sb_idx + 1) * QK_K).cast::<i8>(),
    791
    22
                    _MM_HINT_T0,
    792
    22
                );
    793
    531
            }
    794
    795
            // Read Q4_K header
    796
    553
            let d = read_f16(&q4k_data[sb_start..sb_start + 2]);
    797
    553
            let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);
    798
    799
    553
            let mut scales_raw = [0u8; 12];
    800
    553
            scales_raw.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);
    801
    802
    553
            let q8_scale = q8k_scales[sb_idx];
    803
    553
            let d_q8 = d * q8_scale;
    804
    553
            let dmin_q8 = dmin * q8_scale;
    805
    806
    553
            let qs_ptr = q4k_data.as_ptr().add(sb_start + 16);
    807
    553
            let q8_ptr = q8k_quants.as_ptr().add(q8_start);
    808
    809
            // PAR-126 V2: Keep all 8 block results in vectors
    810
            // block_dots_vec[i] = sum of (Q4[block_i] * Q8[block_i])
    811
            // block_q8sums_vec[i] = sum of Q8[block_i]
    812
    553
            let mut block_dots_vec = _mm256_setzero_si256();
    813
    553
            let mut block_q8sums_vec = _mm256_setzero_si256();
    814
    815
            // Process 64 values (2 blocks) per iteration, 4 iterations = 8 blocks
    816
            // Each iteration produces 2 block sums that go into specific lanes
    817
    2.76k
            for chunk2.21k in 0..4 {
    818
    2.21k
                let j = chunk * 64;
    819
    2.21k
                let q_offset = j / 2;
    820
    821
                // Load 32 bytes Q4 (64 nibbles packed)
    822
    2.21k
                let q4_bytes = _mm256_loadu_si256(qs_ptr.add(q_offset).cast::<__m256i>());
    823
    824
                // Extract nibbles: lo = first 32 values, hi = next 32 values
    825
    2.21k
                let q4_lo = _mm256_and_si256(q4_bytes, nibble_mask);
    826
    2.21k
                let q4_hi = _mm256_and_si256(_mm256_srli_epi16(q4_bytes, 4), nibble_mask);
    827
    828
                // Load 64 bytes Q8
    829
    2.21k
                let q8_lo = _mm256_loadu_si256(q8_ptr.add(j).cast::<__m256i>());
    830
    2.21k
                let q8_hi = _mm256_loadu_si256(q8_ptr.add(j + 32).cast::<__m256i>());
    831
    832
                // Q4 × Q8 products -> i16 via maddubs -> i32 via madd
    833
    2.21k
                let prod_lo_i16 = _mm256_maddubs_epi16(q4_lo, q8_lo);
    834
    2.21k
                let prod_hi_i16 = _mm256_maddubs_epi16(q4_hi, q8_hi);
    835
    2.21k
                let prod_lo_i32 = _mm256_madd_epi16(prod_lo_i16, ones_16);
    836
    2.21k
                let prod_hi_i32 = _mm256_madd_epi16(prod_hi_i16, ones_16);
    837
    838
                // Reduce each 256-bit register to ONE sum using hadd (within-lane reduction)
    839
                // prod_lo_i32 has 8 i32 values, we need their sum for block chunk*2
    840
                // prod_hi_i32 has 8 i32 values, we need their sum for block chunk*2+1
    841
    842
                // Use hadd twice to get 8 -> 4 -> 2 values per register
    843
    2.21k
                let prod_lo_h1 = _mm256_hadd_epi32(prod_lo_i32, prod_hi_i32); // interleaves lo and hi
    844
    2.21k
                let prod_h2 = _mm256_hadd_epi32(prod_lo_h1, prod_lo_h1); // further reduce
    845
    846
                // Now prod_h2 has: [sum_lo_lane0, sum_hi_lane0, sum_lo_lane0, sum_hi_lane0,
    847
                //                   sum_lo_lane1, sum_hi_lane1, sum_lo_lane1, sum_hi_lane1]
    848
                // We need to extract the two unique sums and place them in block_dots_vec
    849
    850
                // Extract two sums: add lane0 and lane1 of each block
    851
    2.21k
                let lane0 = _mm256_castsi256_si128(prod_h2);
    852
    2.21k
                let lane1 = _mm256_extracti128_si256(prod_h2, 1);
    853
    2.21k
                let sums_128 = _mm_add_epi32(lane0, lane1); // [dot_lo, dot_hi, dot_lo, dot_hi]
    854
    855
                // Insert into correct position based on chunk
    856
                // block_dots_vec lanes: [0,1,2,3,4,5,6,7] for blocks [0,1,2,3,4,5,6,7]
    857
    2.21k
                match chunk {
    858
    553
                    0 => {
    859
    553
                        // Put in lanes 0,1
    860
    553
                        block_dots_vec =
    861
    553
                            _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 0), 0);
    862
    553
                        block_dots_vec =
    863
    553
                            _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 1), 1);
    864
    553
                    },
    865
    553
                    1 => {
    866
    553
                        block_dots_vec =
    867
    553
                            _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 0), 2);
    868
    553
                        block_dots_vec =
    869
    553
                            _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 1), 3);
    870
    553
                    },
    871
    553
                    2 => {
    872
    553
                        block_dots_vec =
    873
    553
                            _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 0), 4);
    874
    553
                        block_dots_vec =
    875
    553
                            _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 1), 5);
    876
    553
                    },
    877
    553
                    3 => {
    878
    553
                        block_dots_vec =
    879
    553
                            _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 0), 6);
    880
    553
                        block_dots_vec =
    881
    553
                            _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 1), 7);
    882
    553
                    },
    883
    0
                    _ => unreachable!(),
    884
                }
    885
    886
                // Q8 sums for dmin correction using sad_epu8 (sum of absolute differences)
    887
                // Convert signed i8 to unsigned by adding 128, compute SAD, then adjust
    888
    2.21k
                let bias = _mm256_set1_epi8(-128_i8);
    889
    2.21k
                let _q8_lo_u = _mm256_sub_epi8(q8_lo, bias); // signed [-128,127] -> unsigned [0,255]
    890
    2.21k
                let _q8_hi_u = _mm256_sub_epi8(q8_hi, bias);
    891
    892
                // Use mpsadbw or manual approach for sum
    893
                // Simpler: sign-extend and sum using madd
    894
    2.21k
                let q8_lo_i16_a = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(q8_lo));
    895
    2.21k
                let q8_lo_i16_b = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q8_lo, 1));
    896
    2.21k
                let q8_hi_i16_a = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(q8_hi));
    897
    2.21k
                let q8_hi_i16_b = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q8_hi, 1));
    898
    899
    2.21k
                let q8_lo_i32_a = _mm256_madd_epi16(q8_lo_i16_a, ones_16);
    900
    2.21k
                let q8_lo_i32_b = _mm256_madd_epi16(q8_lo_i16_b, ones_16);
    901
    2.21k
                let q8_hi_i32_a = _mm256_madd_epi16(q8_hi_i16_a, ones_16);
    902
    2.21k
                let q8_hi_i32_b = _mm256_madd_epi16(q8_hi_i16_b, ones_16);
    903
    904
    2.21k
                let q8_lo_sum = _mm256_add_epi32(q8_lo_i32_a, q8_lo_i32_b);
    905
    2.21k
                let q8_hi_sum = _mm256_add_epi32(q8_hi_i32_a, q8_hi_i32_b);
    906
    907
                // Reduce to two sums using hadd
    908
    2.21k
                let q8_h1 = _mm256_hadd_epi32(q8_lo_sum, q8_hi_sum);
    909
    2.21k
                let q8_h2 = _mm256_hadd_epi32(q8_h1, q8_h1);
    910
    911
    2.21k
                let q8_lane0 = _mm256_castsi256_si128(q8_h2);
    912
    2.21k
                let q8_lane1 = _mm256_extracti128_si256(q8_h2, 1);
    913
    2.21k
                let q8_sums_128 = _mm_add_epi32(q8_lane0, q8_lane1);
    914
    915
    2.21k
                match chunk {
    916
    553
                    0 => {
    917
    553
                        block_q8sums_vec =
    918
    553
                            _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 0), 0);
    919
    553
                        block_q8sums_vec =
    920
    553
                            _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 1), 1);
    921
    553
                    },
    922
    553
                    1 => {
    923
    553
                        block_q8sums_vec =
    924
    553
                            _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 0), 2);
    925
    553
                        block_q8sums_vec =
    926
    553
                            _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 1), 3);
    927
    553
                    },
    928
    553
                    2 => {
    929
    553
                        block_q8sums_vec =
    930
    553
                            _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 0), 4);
    931
    553
                        block_q8sums_vec =
    932
    553
                            _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 1), 5);
    933
    553
                    },
    934
    553
                    3 => {
    935
    553
                        block_q8sums_vec =
    936
    553
                            _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 0), 6);
    937
    553
                        block_q8sums_vec =
    938
    553
                            _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 1), 7);
    939
    553
                    },
    940
    0
                    _ => unreachable!(),
    941
                }
    942
            }
    943
    944
            // Extract all 8 scales and mins into vectors
    945
    553
            let mut scales = [0.0f32; 8];
    946
    553
            let mut mins = [0.0f32; 8];
    947
    4.97k
            for i4.42k in 0..8 {
    948
    4.42k
                let (sc, m) = extract_scale_min(&scales_raw, i);
    949
    4.42k
                scales[i] = sc;
    950
    4.42k
                mins[i] = m;
    951
    4.42k
            }
    952
    953
    553
            let scales_vec = _mm256_loadu_ps(scales.as_ptr());
    954
    553
            let mins_vec = _mm256_loadu_ps(mins.as_ptr());
    955
    956
            // Convert block results to f32
    957
    553
            let dots_f32 = _mm256_cvtepi32_ps(block_dots_vec);
    958
    553
            let q8sums_f32 = _mm256_cvtepi32_ps(block_q8sums_vec);
    959
    960
            // Compute: d_q8 * scales * dots - dmin_q8 * mins * q8sums
    961
    553
            let d_q8_vec = _mm256_set1_ps(d_q8);
    962
    553
            let dmin_q8_vec = _mm256_set1_ps(dmin_q8);
    963
    964
    553
            let term1 = _mm256_mul_ps(d_q8_vec, _mm256_mul_ps(scales_vec, dots_f32));
    965
    553
            let term2 = _mm256_mul_ps(dmin_q8_vec, _mm256_mul_ps(mins_vec, q8sums_f32));
    966
    553
            let result = _mm256_sub_ps(term1, term2);
    967
    968
            // Accumulate into total (defer final horizontal sum)
    969
    553
            total_acc = _mm256_add_ps(total_acc, result);
    970
        }
    971
    972
        // ONE horizontal sum at the very end
    973
    532
        let sum128 = _mm_add_ps(
    974
    532
            _mm256_castps256_ps128(total_acc),
    975
    532
            _mm256_extractf128_ps(total_acc, 1),
    976
        );
    977
    532
        let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
    978
    532
        let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
    979
    980
    532
        Ok(_mm_cvtss_f32(sum32))
    981
    536
    }
    982
    983
    /// AVX-512 VNNI optimized Q4_K × Q8_K dot product (llama.cpp style)
    984
    ///
    985
    /// Uses `vpdpbusd` instruction for 16-way uint8×int8→int32 multiply-accumulate.
    986
    /// This is the fastest CPU path for quantized matmul.
    987
    ///
    988
    /// # Safety
    989
    /// Requires AVX-512F and AVX-512 VNNI CPU features.
    990
    #[cfg(target_arch = "x86_64")]
    991
    #[target_feature(enable = "avx512f", enable = "avx512vnni", enable = "avx512bw")]
    992
    #[allow(unsafe_op_in_unsafe_fn)]
    993
    0
    unsafe fn fused_q4k_q8k_dot_avx512vnni(
    994
    0
        q4k_data: &[u8],
    995
    0
        q8k_scales: &[f32],
    996
    0
        q8k_quants: &[i8],
    997
    0
    ) -> Result<f32> {
    998
        #[allow(clippy::wildcard_imports)]
    999
        use std::arch::x86_64::*;
    1000
    1001
        const SUPER_BLOCK_BYTES: usize = 144;
    1002
    1003
    0
        if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {
    1004
    0
            return Err(RealizarError::InvalidShape {
    1005
    0
                reason: format!(
    1006
    0
                    "Q4_K data length {} is not a multiple of {}",
    1007
    0
                    q4k_data.len(),
    1008
    0
                    SUPER_BLOCK_BYTES
    1009
    0
                ),
    1010
    0
            });
    1011
    0
        }
    1012
    1013
    0
        let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;
    1014
    0
        let expected_values = num_super_blocks * QK_K;
    1015
    1016
    0
        if q8k_scales.len() < num_super_blocks || q8k_quants.len() < expected_values {
    1017
    0
            return Err(RealizarError::InvalidShape {
    1018
    0
                reason: "Q8_K buffer too small".to_string(),
    1019
    0
            });
    1020
    0
        }
    1021
    1022
    0
        let _nibble_mask = _mm512_set1_epi8(0x0F_i8); // Reserved for future optimization
    1023
    0
        let mut total_acc = 0.0f32;
    1024
    1025
    0
        for sb_idx in 0..num_super_blocks {
    1026
    0
            let sb_start = sb_idx * SUPER_BLOCK_BYTES;
    1027
    0
            let q8_start = sb_idx * QK_K;
    1028
    1029
            // Prefetch next super-block
    1030
    0
            if sb_idx + 1 < num_super_blocks {
    1031
    0
                _mm_prefetch(
    1032
    0
                    q4k_data
    1033
    0
                        .as_ptr()
    1034
    0
                        .add((sb_idx + 1) * SUPER_BLOCK_BYTES)
    1035
    0
                        .cast::<i8>(),
    1036
    0
                    _MM_HINT_T0,
    1037
    0
                );
    1038
    0
            }
    1039
    1040
            // Read Q4_K header
    1041
    0
            let d = read_f16(&q4k_data[sb_start..sb_start + 2]);
    1042
    0
            let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);
    1043
    1044
    0
            let mut scales = [0u8; 12];
    1045
    0
            scales.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);
    1046
    1047
    0
            let q8_scale = q8k_scales[sb_idx];
    1048
    0
            let d_q8 = d * q8_scale;
    1049
    0
            let dmin_q8 = dmin * q8_scale;
    1050
    1051
    0
            let qs_ptr = q4k_data.as_ptr().add(sb_start + 16);
    1052
    0
            let q8_ptr = q8k_quants.as_ptr().add(q8_start);
    1053
    1054
            // Process 4 chunks of 64 values (matching dequantize_q4_k layout)
    1055
    0
            for j in (0..QK_K).step_by(64) {
    1056
    0
                let q_offset = j / 2; // 32 bytes per 64-value chunk
    1057
    0
    1058
    0
                let is = j / 32;
    1059
    0
                let (sc1, m1) = extract_scale_min(&scales, is);
    1060
    0
                let (sc2, m2) = extract_scale_min(&scales, is + 1);
    1061
    0
    1062
    0
                // Load 32 bytes of Q4_K (as lower half of 512-bit register)
    1063
    0
                // We need to process 64 Q4 values (32 bytes) with 64 Q8 values (64 bytes)
    1064
    0
                let q4_256 = _mm256_loadu_si256(qs_ptr.add(q_offset).cast::<__m256i>());
    1065
    0
    1066
    0
                // Zero-extend to 512-bit (q4 data in lower 256 bits) - reserved for full-width optimization
    1067
    0
                let _q4_512 = _mm512_castsi256_si512(q4_256);
    1068
    0
    1069
    0
                // Extract low nibbles (first 32 values) and high nibbles (second 32 values)
    1070
    0
                let q4_lo_256 = _mm256_and_si256(q4_256, _mm256_set1_epi8(0x0F_i8));
    1071
    0
                let q4_hi_256 =
    1072
    0
                    _mm256_and_si256(_mm256_srli_epi16(q4_256, 4), _mm256_set1_epi8(0x0F_i8));
    1073
    0
    1074
    0
                // Load Q8 values: first 32 for low nibbles, next 32 for high nibbles
    1075
    0
                // (matching dequantize_q4_k output order: 32 low, then 32 high)
    1076
    0
                let q8_lo_256 = _mm256_loadu_si256(q8_ptr.add(j).cast::<__m256i>());
    1077
    0
                let q8_hi_256 = _mm256_loadu_si256(q8_ptr.add(j + 32).cast::<__m256i>());
    1078
    0
    1079
    0
                // For VNNI vpdpbusd: accumulator += (unsigned a) × (signed b)
    1080
    0
                // Q4 values are 0-15 (unsigned), Q8 values are -128..127 (signed)
    1081
    0
                // We need to cast Q4 to unsigned interpretation
    1082
    0
    1083
    0
                // Convert to 512-bit for VNNI
    1084
    0
                let q4_lo_512 = _mm512_castsi256_si512(q4_lo_256);
    1085
    0
                let q4_hi_512 = _mm512_castsi256_si512(q4_hi_256);
    1086
    0
                let q8_lo_512 = _mm512_castsi256_si512(q8_lo_256);
    1087
    0
                let q8_hi_512 = _mm512_castsi256_si512(q8_hi_256);
    1088
    0
    1089
    0
                // VNNI multiply-accumulate: result[i] += q4[4i..4i+4] · q8[4i..4i+4]
    1090
    0
                // This computes 4 int8×int8 products and adds them to int32 accumulator
    1091
    0
                let acc_lo = _mm512_dpbusd_epi32(_mm512_setzero_si512(), q4_lo_512, q8_lo_512);
    1092
    0
                let acc_hi = _mm512_dpbusd_epi32(_mm512_setzero_si512(), q4_hi_512, q8_hi_512);
    1093
    0
    1094
    0
                // Horizontal sum the accumulators (each has 16 int32 values)
    1095
    0
                // Extract lower 256 bits which contain our results
    1096
    0
                let acc_lo_256 = _mm512_castsi512_si256(acc_lo);
    1097
    0
                let acc_hi_256 = _mm512_castsi512_si256(acc_hi);
    1098
    0
    1099
    0
                // Sum all 8 int32 values in each 256-bit register
    1100
    0
                let sum_lo = horizontal_sum_epi32_256(acc_lo_256);
    1101
    0
                let sum_hi = horizontal_sum_epi32_256(acc_hi_256);
    1102
    0
    1103
    0
                // Sum Q8 values for min contribution
    1104
    0
                // For signed i8, we need to convert to i16 first to avoid overflow
    1105
    0
                let q8_lo_256_i16_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(q8_lo_256));
    1106
    0
                let q8_lo_256_i16_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q8_lo_256, 1));
    1107
    0
                let q8_sum_lo = horizontal_sum_epi16_256(q8_lo_256_i16_lo)
    1108
    0
                    + horizontal_sum_epi16_256(q8_lo_256_i16_hi);
    1109
    0
    1110
    0
                let q8_hi_256_i16_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(q8_hi_256));
    1111
    0
                let q8_hi_256_i16_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q8_hi_256, 1));
    1112
    0
                let q8_sum_hi = horizontal_sum_epi16_256(q8_hi_256_i16_lo)
    1113
    0
                    + horizontal_sum_epi16_256(q8_hi_256_i16_hi);
    1114
    0
    1115
    0
                // Apply scales
    1116
    0
                total_acc += d_q8 * sc1 * (sum_lo as f32) - dmin_q8 * m1 * (q8_sum_lo as f32);
    1117
    0
                total_acc += d_q8 * sc2 * (sum_hi as f32) - dmin_q8 * m2 * (q8_sum_hi as f32);
    1118
    0
            }
    1119
        }
    1120
    1121
    0
        Ok(total_acc)
    1122
    0
    }
    1123
    1124
    /// Optimized AVX-512 VNNI Q4_K × Q8_K dot product - DEFERRED horizontal sum
    1125
    ///
    1126
    /// PAR-126 Five-Whys optimization: Instead of horizontal sums per chunk (24+ per super-block),
    1127
    /// we accumulate all 8 block results in vector registers and do ONE horizontal sum at the end.
    1128
    /// This reduces horizontal sum operations from 24+ to 1 per super-block (24x reduction).
    1129
    #[cfg(target_arch = "x86_64")]
    1130
    #[target_feature(enable = "avx512f", enable = "avx512vnni", enable = "avx512bw")]
    1131
    #[allow(unsafe_op_in_unsafe_fn)]
    1132
    #[allow(clippy::similar_names)]
    1133
    #[allow(clippy::too_many_lines)]
    1134
    0
    unsafe fn fused_q4k_q8k_dot_avx512vnni_opt(
    1135
    0
        q4k_data: &[u8],
    1136
    0
        q8k_scales: &[f32],
    1137
    0
        q8k_quants: &[i8],
    1138
    0
    ) -> Result<f32> {
    1139
        #[allow(clippy::wildcard_imports)]
    1140
        use std::arch::x86_64::*;
    1141
    1142
        const SUPER_BLOCK_BYTES: usize = 144;
    1143
    1144
    0
        if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {
    1145
    0
            return Err(RealizarError::InvalidShape {
    1146
    0
                reason: format!(
    1147
    0
                    "Q4_K data length {} is not a multiple of {}",
    1148
    0
                    q4k_data.len(),
    1149
    0
                    SUPER_BLOCK_BYTES
    1150
    0
                ),
    1151
    0
            });
    1152
    0
        }
    1153
    1154
    0
        let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;
    1155
    0
        let expected_values = num_super_blocks * QK_K;
    1156
    1157
    0
        if q8k_scales.len() < num_super_blocks || q8k_quants.len() < expected_values {
    1158
    0
            return Err(RealizarError::InvalidShape {
    1159
    0
                reason: "Q8_K buffer too small".to_string(),
    1160
    0
            });
    1161
    0
        }
    1162
    1163
    0
        let nibble_mask = _mm256_set1_epi8(0x0F_i8);
    1164
    0
        let ones_16 = _mm256_set1_epi16(1);
    1165
    1166
        // Global float accumulator
    1167
    0
        let mut total_acc = 0.0f32;
    1168
    1169
    0
        for sb_idx in 0..num_super_blocks {
    1170
    0
            let sb_start = sb_idx * SUPER_BLOCK_BYTES;
    1171
    0
            let q8_start = sb_idx * QK_K;
    1172
    1173
            // Prefetch next super-block
    1174
    0
            if sb_idx + 1 < num_super_blocks {
    1175
    0
                _mm_prefetch(
    1176
    0
                    q4k_data
    1177
    0
                        .as_ptr()
    1178
    0
                        .add((sb_idx + 1) * SUPER_BLOCK_BYTES)
    1179
    0
                        .cast::<i8>(),
    1180
    0
                    _MM_HINT_T0,
    1181
    0
                );
    1182
    0
                _mm_prefetch(
    1183
    0
                    q8k_quants.as_ptr().add((sb_idx + 1) * QK_K).cast::<i8>(),
    1184
    0
                    _MM_HINT_T0,
    1185
    0
                );
    1186
    0
            }
    1187
    1188
            // Read Q4_K header
    1189
    0
            let d = read_f16(&q4k_data[sb_start..sb_start + 2]);
    1190
    0
            let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);
    1191
    1192
    0
            let mut scales_raw = [0u8; 12];
    1193
    0
            scales_raw.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);
    1194
    1195
    0
            let q8_scale = q8k_scales[sb_idx];
    1196
    0
            let d_q8 = d * q8_scale;
    1197
    0
            let dmin_q8 = dmin * q8_scale;
    1198
    1199
    0
            let qs_ptr = q4k_data.as_ptr().add(sb_start + 16);
    1200
    0
            let q8_ptr = q8k_quants.as_ptr().add(q8_start);
    1201
    1202
            // PAR-126: DEFERRED HORIZONTAL SUM OPTIMIZATION
    1203
            // Keep 8 block accumulators in __m256i vectors (one per 32-value block)
    1204
            // block_dots[i] = sum of (Q4[block_i] * Q8[block_i]) for block i
    1205
            // block_q8sums[i] = sum of Q8[block_i] for dmin correction
    1206
    0
            let mut block_dots = [0i32; 8];
    1207
    0
            let mut block_q8sums = [0i32; 8];
    1208
    1209
            // Process 64 values (2 blocks) per iteration, 4 iterations total = 8 blocks
    1210
    0
            for chunk in 0..4 {
    1211
    0
                let j = chunk * 64;
    1212
    0
                let q_offset = j / 2;
    1213
    0
    1214
    0
                // Load 32 bytes Q4 (64 nibbles packed)
    1215
    0
                let q4_bytes = _mm256_loadu_si256(qs_ptr.add(q_offset).cast::<__m256i>());
    1216
    0
    1217
    0
                // Extract nibbles: lo = first 32 values, hi = next 32 values
    1218
    0
                let q4_lo = _mm256_and_si256(q4_bytes, nibble_mask);
    1219
    0
                let q4_hi = _mm256_and_si256(_mm256_srli_epi16(q4_bytes, 4), nibble_mask);
    1220
    0
    1221
    0
                // Load 64 bytes Q8
    1222
    0
                let q8_lo = _mm256_loadu_si256(q8_ptr.add(j).cast::<__m256i>());
    1223
    0
                let q8_hi = _mm256_loadu_si256(q8_ptr.add(j + 32).cast::<__m256i>());
    1224
    0
    1225
    0
                // Q4 × Q8 products -> i16 via maddubs -> i32 via madd
    1226
    0
                // maddubs: adjacent u8*i8 pairs summed to i16
    1227
    0
                // madd with ones: sum pairs of i16 to i32
    1228
    0
                let prod_lo_i16 = _mm256_maddubs_epi16(q4_lo, q8_lo);
    1229
    0
                let prod_hi_i16 = _mm256_maddubs_epi16(q4_hi, q8_hi);
    1230
    0
                let prod_lo_i32 = _mm256_madd_epi16(prod_lo_i16, ones_16);
    1231
    0
                let prod_hi_i32 = _mm256_madd_epi16(prod_hi_i16, ones_16);
    1232
    0
    1233
    0
                // prod_lo_i32 now has 8 i32 values (4 per 128-bit lane)
    1234
    0
                // We need to reduce to 1 value per 32-element block
    1235
    0
                // Lane 0-3 are from elements 0-15, lane 4-7 are from elements 16-31
    1236
    0
                // So we sum all 8 to get one block sum
    1237
    0
    1238
    0
                // Sum all 8 i32 in prod_lo_i32 using only one horizontal sum
    1239
    0
                // First add the two 128-bit halves
    1240
    0
                let prod_lo_128 = _mm_add_epi32(
    1241
    0
                    _mm256_castsi256_si128(prod_lo_i32),
    1242
    0
                    _mm256_extracti128_si256(prod_lo_i32, 1),
    1243
    0
                );
    1244
    0
                let prod_hi_128 = _mm_add_epi32(
    1245
    0
                    _mm256_castsi256_si128(prod_hi_i32),
    1246
    0
                    _mm256_extracti128_si256(prod_hi_i32, 1),
    1247
    0
                );
    1248
    0
    1249
    0
                // Now we have 4 i32 each - sum them with hadd
    1250
    0
                let prod_lo_64 = _mm_hadd_epi32(prod_lo_128, prod_hi_128);
    1251
    0
                let prod_32 = _mm_hadd_epi32(prod_lo_64, prod_lo_64);
    1252
    0
    1253
    0
                // Extract block sums - lane 0 is block_lo, lane 1 is block_hi
    1254
    0
                let block_idx = chunk * 2;
    1255
    0
                block_dots[block_idx] = _mm_extract_epi32(prod_32, 0);
    1256
    0
                block_dots[block_idx + 1] = _mm_extract_epi32(prod_32, 1);
    1257
    0
    1258
    0
                // Q8 sums for dmin correction - convert i8 to i16 first to avoid overflow
    1259
    0
                // We need sum of all 32 i8 values for each block
    1260
    0
                let q8_lo_i16_a = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(q8_lo));
    1261
    0
                let q8_lo_i16_b = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q8_lo, 1));
    1262
    0
                let q8_hi_i16_a = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(q8_hi));
    1263
    0
                let q8_hi_i16_b = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q8_hi, 1));
    1264
    0
    1265
    0
                // Sum the i16 values to i32 using madd with ones
    1266
    0
                let q8_lo_i32_a = _mm256_madd_epi16(q8_lo_i16_a, _mm256_set1_epi16(1));
    1267
    0
                let q8_lo_i32_b = _mm256_madd_epi16(q8_lo_i16_b, _mm256_set1_epi16(1));
    1268
    0
                let q8_hi_i32_a = _mm256_madd_epi16(q8_hi_i16_a, _mm256_set1_epi16(1));
    1269
    0
                let q8_hi_i32_b = _mm256_madd_epi16(q8_hi_i16_b, _mm256_set1_epi16(1));
    1270
    0
    1271
    0
                // Combine halves for each block
    1272
    0
                let q8_lo_sum = _mm256_add_epi32(q8_lo_i32_a, q8_lo_i32_b);
    1273
    0
                let q8_hi_sum = _mm256_add_epi32(q8_hi_i32_a, q8_hi_i32_b);
    1274
    0
    1275
    0
                // Horizontal sum each block's q8 sum
    1276
    0
                let q8_lo_128 = _mm_add_epi32(
    1277
    0
                    _mm256_castsi256_si128(q8_lo_sum),
    1278
    0
                    _mm256_extracti128_si256(q8_lo_sum, 1),
    1279
    0
                );
    1280
    0
                let q8_hi_128 = _mm_add_epi32(
    1281
    0
                    _mm256_castsi256_si128(q8_hi_sum),
    1282
    0
                    _mm256_extracti128_si256(q8_hi_sum, 1),
    1283
    0
                );
    1284
    0
                let q8_64 = _mm_hadd_epi32(q8_lo_128, q8_hi_128);
    1285
    0
                let q8_32 = _mm_hadd_epi32(q8_64, q8_64);
    1286
    0
    1287
    0
                block_q8sums[block_idx] = _mm_extract_epi32(q8_32, 0);
    1288
    0
                block_q8sums[block_idx + 1] = _mm_extract_epi32(q8_32, 1);
    1289
    0
            }
    1290
    1291
            // Extract all 8 scales and mins
    1292
    0
            let mut scales = [0.0f32; 8];
    1293
    0
            let mut mins = [0.0f32; 8];
    1294
    0
            for i in 0..8 {
    1295
    0
                let (sc, m) = extract_scale_min(&scales_raw, i);
    1296
    0
                scales[i] = sc;
    1297
    0
                mins[i] = m;
    1298
    0
            }
    1299
    1300
            // Load scales and mins into vectors for SIMD multiply
    1301
    0
            let scales_vec = _mm256_loadu_ps(scales.as_ptr());
    1302
    0
            let mins_vec = _mm256_loadu_ps(mins.as_ptr());
    1303
    1304
            // Convert block_dots and block_q8sums to f32
    1305
    0
            let dots_i32 = _mm256_loadu_si256(block_dots.as_ptr().cast::<__m256i>());
    1306
    0
            let q8sums_i32 = _mm256_loadu_si256(block_q8sums.as_ptr().cast::<__m256i>());
    1307
    0
            let dots_f32 = _mm256_cvtepi32_ps(dots_i32);
    1308
    0
            let q8sums_f32 = _mm256_cvtepi32_ps(q8sums_i32);
    1309
    1310
            // Compute: d_q8 * scales * dots - dmin_q8 * mins * q8sums
    1311
    0
            let d_q8_vec = _mm256_set1_ps(d_q8);
    1312
    0
            let dmin_q8_vec = _mm256_set1_ps(dmin_q8);
    1313
    1314
    0
            let term1 = _mm256_mul_ps(d_q8_vec, _mm256_mul_ps(scales_vec, dots_f32));
    1315
    0
            let term2 = _mm256_mul_ps(dmin_q8_vec, _mm256_mul_ps(mins_vec, q8sums_f32));
    1316
    0
            let result = _mm256_sub_ps(term1, term2);
    1317
    1318
            // ONE horizontal sum for all 8 blocks
    1319
    0
            let sum128 = _mm_add_ps(
    1320
    0
                _mm256_castps256_ps128(result),
    1321
    0
                _mm256_extractf128_ps(result, 1),
    1322
            );
    1323
    0
            let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
    1324
    0
            let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
    1325
    0
            total_acc += _mm_cvtss_f32(sum32);
    1326
        }
    1327
    1328
    0
        Ok(total_acc)
    1329
    0
    }
    1330
    1331
    /// Fast horizontal sum of 4 i32 in __m128i
    1332
    #[cfg(target_arch = "x86_64")]
    1333
    #[target_feature(enable = "avx2")]
    1334
    #[inline]
    1335
    0
    unsafe fn hsum_epi32_128(v: std::arch::x86_64::__m128i) -> i32 {
    1336
        use std::arch::x86_64::{_mm_cvtsi128_si32, _mm_hadd_epi32};
    1337
    0
        let sum64 = _mm_hadd_epi32(v, v);
    1338
    0
        let sum32 = _mm_hadd_epi32(sum64, sum64);
    1339
    0
        _mm_cvtsi128_si32(sum32)
    1340
    0
    }
    1341
    1342
    /// Fast horizontal sum of 8 i32 in __m256i
    1343
    #[cfg(target_arch = "x86_64")]
    1344
    #[target_feature(enable = "avx2")]
    1345
    #[inline]
    1346
    0
    unsafe fn hsum_epi32_256(v: std::arch::x86_64::__m256i) -> i32 {
    1347
        use std::arch::x86_64::{_mm256_castsi256_si128, _mm256_extracti128_si256, _mm_add_epi32};
    1348
        // SAFETY: Unsafe operation with validated invariants
    1349
        unsafe {
    1350
    0
            let lo = _mm256_castsi256_si128(v);
    1351
    0
            let hi = _mm256_extracti128_si256(v, 1);
    1352
    0
            hsum_epi32_128(_mm_add_epi32(lo, hi))
    1353
        }
    1354
    0
    }
    1355
    1356
    /// Helper: horizontal sum of 8 int32 values in a 256-bit register
    1357
    #[cfg(target_arch = "x86_64")]
    1358
    #[target_feature(enable = "avx2")]
    1359
    #[inline]
    1360
    0
    unsafe fn horizontal_sum_epi32_256(v: std::arch::x86_64::__m256i) -> i32 {
    1361
        use std::arch::x86_64::{
    1362
            _mm256_castsi256_si128, _mm256_extracti128_si256, _mm_add_epi32, _mm_cvtsi128_si32,
    1363
            _mm_hadd_epi32,
    1364
        };
    1365
    1366
        // Add high 128 bits to low 128 bits
    1367
    0
        let hi = _mm256_extracti128_si256(v, 1);
    1368
    0
        let lo = _mm256_castsi256_si128(v);
    1369
    0
        let sum128 = _mm_add_epi32(lo, hi);
    1370
    1371
        // Horizontal add within 128 bits
    1372
    0
        let sum64 = _mm_hadd_epi32(sum128, sum128);
    1373
    0
        let sum32 = _mm_hadd_epi32(sum64, sum64);
    1374
    1375
    0
        _mm_cvtsi128_si32(sum32)
    1376
    0
    }
    1377
    1378
    /// Helper: horizontal sum of 16 int16 values in a 256-bit register
    1379
    #[cfg(target_arch = "x86_64")]
    1380
    #[target_feature(enable = "avx2")]
    1381
    #[inline]
    1382
    #[allow(unsafe_op_in_unsafe_fn)]
    1383
    0
    unsafe fn horizontal_sum_epi16_256(v: std::arch::x86_64::__m256i) -> i32 {
    1384
        use std::arch::x86_64::{_mm256_madd_epi16, _mm256_set1_epi16};
    1385
    1386
        // Use madd to sum pairs of i16 to i32
    1387
    0
        let ones = _mm256_set1_epi16(1);
    1388
    0
        let sum_i32 = _mm256_madd_epi16(v, ones);
    1389
    1390
        // Now sum the 8 i32 values
    1391
    0
        horizontal_sum_epi32_256(sum_i32)
    1392
    0
    }
    1393
    1394
    /// AVX2-optimized Q4_K × Q8_K dot product
    1395
    ///
    1396
    /// # Safety
    1397
    /// Requires AVX2 CPU feature.
    1398
    #[cfg(target_arch = "x86_64")]
    1399
    #[target_feature(enable = "avx2")]
    1400
    #[allow(unsafe_op_in_unsafe_fn)]
    1401
    0
    unsafe fn fused_q4k_q8k_dot_avx2(
    1402
    0
        q4k_data: &[u8],
    1403
    0
        q8k_scales: &[f32],
    1404
    0
        q8k_quants: &[i8],
    1405
    0
    ) -> Result<f32> {
    1406
        #[allow(clippy::wildcard_imports)]
    1407
        use std::arch::x86_64::*;
    1408
    1409
        const SUPER_BLOCK_BYTES: usize = 144;
    1410
    1411
    0
        if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {
    1412
    0
            return Err(RealizarError::InvalidShape {
    1413
    0
                reason: format!(
    1414
    0
                    "Q4_K data length {} is not a multiple of {}",
    1415
    0
                    q4k_data.len(),
    1416
    0
                    SUPER_BLOCK_BYTES
    1417
    0
                ),
    1418
    0
            });
    1419
    0
        }
    1420
    1421
    0
        let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;
    1422
    0
        let expected_values = num_super_blocks * QK_K;
    1423
    1424
    0
        if q8k_scales.len() < num_super_blocks || q8k_quants.len() < expected_values {
    1425
    0
            return Err(RealizarError::InvalidShape {
    1426
    0
                reason: "Q8_K buffer too small".to_string(),
    1427
    0
            });
    1428
    0
        }
    1429
    1430
    0
        let nibble_mask = _mm256_set1_epi8(0x0F_i8);
    1431
    0
        let ones_16 = _mm256_set1_epi16(1);
    1432
    1433
    0
        let mut total_acc = 0.0f32;
    1434
    1435
    0
        for sb_idx in 0..num_super_blocks {
    1436
    0
            let sb_start = sb_idx * SUPER_BLOCK_BYTES;
    1437
    0
            let q8_start = sb_idx * QK_K;
    1438
    1439
            // Prefetch next super-block
    1440
    0
            if sb_idx + 1 < num_super_blocks {
    1441
    0
                _mm_prefetch(
    1442
    0
                    q4k_data
    1443
    0
                        .as_ptr()
    1444
    0
                        .add((sb_idx + 1) * SUPER_BLOCK_BYTES)
    1445
    0
                        .cast::<i8>(),
    1446
    0
                    _MM_HINT_T0,
    1447
    0
                );
    1448
    0
                _mm_prefetch(
    1449
    0
                    q8k_quants.as_ptr().add((sb_idx + 1) * QK_K).cast::<i8>(),
    1450
    0
                    _MM_HINT_T0,
    1451
    0
                );
    1452
    0
            }
    1453
    1454
            // Read Q4_K header
    1455
    0
            let d = read_f16(&q4k_data[sb_start..sb_start + 2]);
    1456
    0
            let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);
    1457
    1458
    0
            let mut scales = [0u8; 12];
    1459
    0
            scales.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);
    1460
    1461
    0
            let q8_scale = q8k_scales[sb_idx];
    1462
    0
            let d_q8 = d * q8_scale;
    1463
    0
            let dmin_q8 = dmin * q8_scale;
    1464
    1465
    0
            let qs_ptr = q4k_data.as_ptr().add(sb_start + 16);
    1466
    0
            let q8_ptr = q8k_quants.as_ptr().add(q8_start);
    1467
    1468
            // (accumulator variables for future fixed-point optimization)
    1469
    0
            let _acc_sum = _mm256_setzero_si256();
    1470
    0
            let _acc_min = _mm256_setzero_si256();
    1471
    1472
            // Process 4 iterations of 64 values each (256 total)
    1473
    0
            for j in (0..QK_K).step_by(64) {
    1474
    0
                let q_offset = j / 2; // 32 bytes per 64 values
    1475
    0
    1476
    0
                // Get scales for two 32-value blocks
    1477
    0
                let is = j / 32;
    1478
    0
                let (sc1, m1) = extract_scale_min(&scales, is);
    1479
    0
                let (sc2, m2) = extract_scale_min(&scales, is + 1);
    1480
    0
    1481
    0
                // Fixed-point scales (8.8 format) - used for integer path
    1482
    0
                let sc1_i16 = (sc1 * 256.0).round() as i16;
    1483
    0
                let sc2_i16 = (sc2 * 256.0).round() as i16;
    1484
    0
                let _m1_i16 = (m1 * 256.0).round() as i16;
    1485
    0
                let _m2_i16 = (m2 * 256.0).round() as i16;
    1486
    0
    1487
    0
                // Load 32 bytes of Q4_K (64 nibbles)
    1488
    0
                let q4_bytes = _mm256_loadu_si256(qs_ptr.add(q_offset).cast::<__m256i>());
    1489
    0
    1490
    0
                // Extract nibbles
    1491
    0
                let q4_lo = _mm256_and_si256(q4_bytes, nibble_mask);
    1492
    0
                let q4_hi = _mm256_and_si256(_mm256_srli_epi16(q4_bytes, 4), nibble_mask);
    1493
    0
    1494
    0
                // Load 64 bytes of Q8_K (sequential values)
    1495
    0
                // CORRECT LAYOUT: dequantize_q4_k outputs 32 low nibbles, then 32 high nibbles
    1496
    0
                // So Q8[j..j+32] corresponds to low nibbles, Q8[j+32..j+64] to high nibbles
    1497
    0
                let q8_lo = _mm256_loadu_si256(q8_ptr.add(j).cast::<__m256i>());
    1498
    0
                let q8_hi = _mm256_loadu_si256(q8_ptr.add(j + 32).cast::<__m256i>());
    1499
    0
    1500
    0
                // Q4_lo × Q8_lo: low nibbles times first 32 Q8 values (unsigned × signed → i16)
    1501
    0
                let prod_lo = _mm256_maddubs_epi16(q4_lo, q8_lo);
    1502
    0
                // Q4_hi × Q8_hi: high nibbles times second 32 Q8 values
    1503
    0
                let prod_hi = _mm256_maddubs_epi16(q4_hi, q8_hi);
    1504
    0
    1505
    0
                // Apply block scales and accumulate (for future integer-only path)
    1506
    0
                let _scale_lo = _mm256_set1_epi16(sc1_i16);
    1507
    0
                let _scale_hi = _mm256_set1_epi16(sc2_i16);
    1508
    0
    1509
    0
                // Split products by block (first 128 bits = block 1, second 128 bits = block 2)
    1510
    0
                let prod_lo_128 = _mm256_castsi256_si128(prod_lo);
    1511
    0
                let prod_lo_hi128 = _mm256_extracti128_si256(prod_lo, 1);
    1512
    0
                let prod_hi_128 = _mm256_castsi256_si128(prod_hi);
    1513
    0
                let prod_hi_hi128 = _mm256_extracti128_si256(prod_hi, 1);
    1514
    0
    1515
    0
                // Horizontal sum to i32
    1516
    0
                let sum_lo_1 = _mm_madd_epi16(prod_lo_128, _mm_set1_epi16(1));
    1517
    0
                let sum_lo_2 = _mm_madd_epi16(prod_lo_hi128, _mm_set1_epi16(1));
    1518
    0
                let sum_hi_1 = _mm_madd_epi16(prod_hi_128, _mm_set1_epi16(1));
    1519
    0
                let sum_hi_2 = _mm_madd_epi16(prod_hi_hi128, _mm_set1_epi16(1));
    1520
    0
    1521
    0
                // Add low and high nibble products
    1522
    0
                let sum_1 = _mm_add_epi32(sum_lo_1, sum_hi_1);
    1523
    0
                let sum_2 = _mm_add_epi32(sum_lo_2, sum_hi_2);
    1524
    0
    1525
    0
                // Apply scales (as f32 to avoid overflow)
    1526
    0
                let sum_1_f = _mm_cvtepi32_ps(sum_1);
    1527
    0
                let sum_2_f = _mm_cvtepi32_ps(sum_2);
    1528
    0
    1529
    0
                let scaled_1 = _mm_mul_ps(sum_1_f, _mm_set1_ps(sc1));
    1530
    0
                let scaled_2 = _mm_mul_ps(sum_2_f, _mm_set1_ps(sc2));
    1531
    0
    1532
    0
                // Sum for min contribution (sum of Q8 values)
    1533
    0
                let q8_sum_lo =
    1534
    0
                    _mm256_madd_epi16(_mm256_cvtepi8_epi16(_mm256_castsi256_si128(q8_lo)), ones_16);
    1535
    0
                let q8_sum_hi = _mm256_madd_epi16(
    1536
    0
                    _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q8_lo, 1)),
    1537
    0
                    ones_16,
    1538
    0
                );
    1539
    0
    1540
    0
                // Horizontal reduce
    1541
    0
                let hsum_lo = _mm_add_epi32(
    1542
    0
                    _mm256_castsi256_si128(q8_sum_lo),
    1543
    0
                    _mm256_extracti128_si256(q8_sum_lo, 1),
    1544
    0
                );
    1545
    0
                let _hsum_hi = _mm_add_epi32(
    1546
    0
                    _mm256_castsi256_si128(q8_sum_hi),
    1547
    0
                    _mm256_extracti128_si256(q8_sum_hi, 1),
    1548
    0
                );
    1549
    0
    1550
    0
                // Include both halves in block sum
    1551
    0
                let q8_block1_sum = _mm_add_epi32(hsum_lo, _mm_shuffle_epi32(hsum_lo, 0b10_11_00_01));
    1552
    0
                let q8_block1_sum = _mm_add_epi32(
    1553
    0
                    q8_block1_sum,
    1554
    0
                    _mm_shuffle_epi32(q8_block1_sum, 0b00_00_10_10),
    1555
    0
                );
    1556
    0
                let q8_block1_val = _mm_cvtsi128_si32(q8_block1_sum);
    1557
    0
    1558
    0
                // Similar for second block (q8_hi)
    1559
    0
                let q8_sum_hi2 =
    1560
    0
                    _mm256_madd_epi16(_mm256_cvtepi8_epi16(_mm256_castsi256_si128(q8_hi)), ones_16);
    1561
    0
                let q8_sum_hi3 = _mm256_madd_epi16(
    1562
    0
                    _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q8_hi, 1)),
    1563
    0
                    ones_16,
    1564
    0
                );
    1565
    0
                let hsum2_lo = _mm_add_epi32(
    1566
    0
                    _mm256_castsi256_si128(q8_sum_hi2),
    1567
    0
                    _mm256_extracti128_si256(q8_sum_hi2, 1),
    1568
    0
                );
    1569
    0
                let hsum2_hi = _mm_add_epi32(
    1570
    0
                    _mm256_castsi256_si128(q8_sum_hi3),
    1571
    0
                    _mm256_extracti128_si256(q8_sum_hi3, 1),
    1572
    0
                );
    1573
    0
                let q8_block2_sum = _mm_add_epi32(hsum2_lo, hsum2_hi);
    1574
    0
                let q8_block2_sum = _mm_add_epi32(
    1575
    0
                    q8_block2_sum,
    1576
    0
                    _mm_shuffle_epi32(q8_block2_sum, 0b10_11_00_01),
    1577
    0
                );
    1578
    0
                let q8_block2_sum = _mm_add_epi32(
    1579
    0
                    q8_block2_sum,
    1580
    0
                    _mm_shuffle_epi32(q8_block2_sum, 0b00_00_10_10),
    1581
    0
                );
    1582
    0
                let q8_block2_val = _mm_cvtsi128_si32(q8_block2_sum);
    1583
    0
    1584
    0
                // Final accumulation with f32 precision
    1585
    0
                let scaled_sum = _mm_add_ps(scaled_1, scaled_2);
    1586
    0
                let hsum = _mm_hadd_ps(scaled_sum, scaled_sum);
    1587
    0
                let hsum = _mm_hadd_ps(hsum, hsum);
    1588
    0
                let block_prod = _mm_cvtss_f32(hsum);
    1589
    0
    1590
    0
                total_acc += d_q8 * block_prod;
    1591
    0
                total_acc -= dmin_q8 * (m1 * q8_block1_val as f32 + m2 * q8_block2_val as f32);
    1592
    0
            }
    1593
        }
    1594
    1595
    0
        Ok(total_acc)
    1596
    0
    }
    1597
    1598
    /// TCB 4-row micro-kernel: Process 4 rows simultaneously sharing Q8K loads
    1599
    ///
    1600
    /// This implements the trueno TCB micro-tile pattern (4×1×256):
    1601
    /// - Load Q8K input ONCE per superblock
    1602
    /// - Process 4 weight rows using the SAME Q8K loads
    1603
    /// - Return 4 output values
    1604
    ///
    1605
    /// # Performance
    1606
    ///
    1607
    /// Sharing Q8K loads across 4 rows reduces memory bandwidth by ~4x for the
    1608
    /// input vector, which is the key optimization from TCB (Tiling Compute Blocks).
    1609
    ///
    1610
    /// # Safety
    1611
    /// Requires AVX-512F, AVX-512 VNNI, and AVX-512BW CPU features.
    1612
    #[cfg(target_arch = "x86_64")]
    1613
    #[target_feature(enable = "avx512f", enable = "avx512vnni", enable = "avx512bw")]
    1614
    #[allow(unsafe_op_in_unsafe_fn)]
    1615
    #[allow(clippy::similar_names)]
    1616
    #[allow(clippy::too_many_lines)]
    1617
    0
    pub(crate) unsafe fn fused_q4k_q8k_dot_4rows_avx512vnni(
    1618
    0
        row_ptrs: [*const u8; 4],
    1619
    0
        bytes_per_row: usize,
    1620
    0
        q8k_scales: &[f32],
    1621
    0
        q8k_quants: &[i8],
    1622
    0
    ) -> [f32; 4] {
    1623
        #[allow(clippy::wildcard_imports)]
    1624
        use std::arch::x86_64::*;
    1625
    1626
        const SUPER_BLOCK_BYTES: usize = 144;
    1627
    0
        let num_super_blocks = bytes_per_row / SUPER_BLOCK_BYTES;
    1628
    1629
    0
        let nibble_mask = _mm256_set1_epi8(0x0F_i8);
    1630
    0
        let ones_16 = _mm256_set1_epi16(1);
    1631
    1632
        // 4 accumulators for 4 output rows (8 blocks × f32 = 8 values each)
    1633
    0
        let mut total_acc = [_mm256_setzero_ps(); 4];
    1634
    1635
    0
        for sb_idx in 0..num_super_blocks {
    1636
    0
            let q8_start = sb_idx * QK_K;
    1637
    0
            let sb_start = sb_idx * SUPER_BLOCK_BYTES;
    1638
    1639
            // Prefetch next Q8K superblock (shared across all 4 rows)
    1640
    0
            if sb_idx + 1 < num_super_blocks {
    1641
    0
                _mm_prefetch(
    1642
    0
                    q8k_quants.as_ptr().add((sb_idx + 1) * QK_K).cast::<i8>(),
    1643
    0
                    _MM_HINT_T0,
    1644
    0
                );
    1645
    0
            }
    1646
    1647
    0
            let q8_scale = q8k_scales[sb_idx];
    1648
    0
            let q8_ptr = q8k_quants.as_ptr().add(q8_start);
    1649
    1650
            // ============================================================
    1651
            // CRITICAL: Pre-load Q8K data and compute Q8 sums ONCE per superblock
    1652
            // These are shared across ALL 4 rows
    1653
            // ============================================================
    1654
    1655
            // Pre-load all Q8K chunks for this superblock (4 chunks × 2 registers = 8 loads)
    1656
    0
            let q8_chunk0_lo = _mm256_loadu_si256(q8_ptr.cast::<__m256i>());
    1657
    0
            let q8_chunk0_hi = _mm256_loadu_si256(q8_ptr.add(32).cast::<__m256i>());
    1658
    0
            let q8_chunk1_lo = _mm256_loadu_si256(q8_ptr.add(64).cast::<__m256i>());
    1659
    0
            let q8_chunk1_hi = _mm256_loadu_si256(q8_ptr.add(96).cast::<__m256i>());
    1660
    0
            let q8_chunk2_lo = _mm256_loadu_si256(q8_ptr.add(128).cast::<__m256i>());
    1661
    0
            let q8_chunk2_hi = _mm256_loadu_si256(q8_ptr.add(160).cast::<__m256i>());
    1662
    0
            let q8_chunk3_lo = _mm256_loadu_si256(q8_ptr.add(192).cast::<__m256i>());
    1663
    0
            let q8_chunk3_hi = _mm256_loadu_si256(q8_ptr.add(224).cast::<__m256i>());
    1664
    1665
            // Pre-compute Q8 sums for dmin correction (same for all rows)
    1666
    0
            let q8_sums = compute_q8_sums_8blocks(
    1667
    0
                q8_chunk0_lo,
    1668
    0
                q8_chunk0_hi,
    1669
    0
                q8_chunk1_lo,
    1670
    0
                q8_chunk1_hi,
    1671
    0
                q8_chunk2_lo,
    1672
    0
                q8_chunk2_hi,
    1673
    0
                q8_chunk3_lo,
    1674
    0
                q8_chunk3_hi,
    1675
    0
                ones_16,
    1676
            );
    1677
    1678
            // Process 4 rows using the pre-loaded Q8K data
    1679
    0
            for row in 0..4 {
    1680
    0
                let row_data = row_ptrs[row].add(sb_start);
    1681
    1682
                // Read Q4_K header for this row
    1683
    0
                let d = read_f16(std::slice::from_raw_parts(row_data, 2));
    1684
    0
                let dmin = read_f16(std::slice::from_raw_parts(row_data.add(2), 2));
    1685
    1686
    0
                let mut scales_raw = [0u8; 12];
    1687
    0
                std::ptr::copy_nonoverlapping(row_data.add(4), scales_raw.as_mut_ptr(), 12);
    1688
    1689
    0
                let d_q8 = d * q8_scale;
    1690
    0
                let dmin_q8 = dmin * q8_scale;
    1691
    1692
    0
                let qs_ptr = row_data.add(16);
    1693
    1694
                // Compute Q4×Q8 dot products for all 8 blocks using pre-loaded Q8K
    1695
    0
                let block_dots = compute_q4_q8_dots_8blocks(
    1696
    0
                    qs_ptr,
    1697
    0
                    q8_chunk0_lo,
    1698
    0
                    q8_chunk0_hi,
    1699
    0
                    q8_chunk1_lo,
    1700
    0
                    q8_chunk1_hi,
    1701
    0
                    q8_chunk2_lo,
    1702
    0
                    q8_chunk2_hi,
    1703
    0
                    q8_chunk3_lo,
    1704
    0
                    q8_chunk3_hi,
    1705
    0
                    nibble_mask,
    1706
    0
                    ones_16,
    1707
                );
    1708
    1709
                // Extract 6-bit scales and mins
    1710
    0
                let mut scales = [0.0f32; 8];
    1711
    0
                let mut mins = [0.0f32; 8];
    1712
    0
                for i in 0..8 {
    1713
    0
                    let (sc, m) = extract_scale_min(&scales_raw, i);
    1714
    0
                    scales[i] = sc;
    1715
    0
                    mins[i] = m;
    1716
    0
                }
    1717
    1718
                // Final computation: d_q8 * scales * dots - dmin_q8 * mins * q8sums
    1719
    0
                let scales_vec = _mm256_loadu_ps(scales.as_ptr());
    1720
    0
                let mins_vec = _mm256_loadu_ps(mins.as_ptr());
    1721
    0
                let d_q8_vec = _mm256_set1_ps(d_q8);
    1722
    0
                let dmin_q8_vec = _mm256_set1_ps(dmin_q8);
    1723
    1724
    0
                let dots_f32 = _mm256_cvtepi32_ps(block_dots);
    1725
    0
                let q8sums_f32 = _mm256_cvtepi32_ps(q8_sums);
    1726
    1727
    0
                let term1 = _mm256_mul_ps(d_q8_vec, _mm256_mul_ps(scales_vec, dots_f32));
    1728
    0
                let term2 = _mm256_mul_ps(dmin_q8_vec, _mm256_mul_ps(mins_vec, q8sums_f32));
    1729
    0
                let result = _mm256_sub_ps(term1, term2);
    1730
    1731
    0
                total_acc[row] = _mm256_add_ps(total_acc[row], result);
    1732
            }
    1733
        }
    1734
    1735
        // Final horizontal sums for each row
    1736
    0
        let mut outputs = [0.0f32; 4];
    1737
    0
        for row in 0..4 {
    1738
    0
            let sum128 = _mm_add_ps(
    1739
    0
                _mm256_castps256_ps128(total_acc[row]),
    1740
    0
                _mm256_extractf128_ps(total_acc[row], 1),
    1741
    0
            );
    1742
    0
            let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
    1743
    0
            let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
    1744
    0
            outputs[row] = _mm_cvtss_f32(sum32);
    1745
    0
        }
    1746
    1747
    0
        outputs
    1748
    0
    }
    1749
    1750
    /// Helper: Compute Q8 sums for 8 blocks (shared across rows)
    1751
    #[cfg(target_arch = "x86_64")]
    1752
    #[target_feature(enable = "avx2")]
    1753
    #[inline]
    1754
    #[allow(clippy::too_many_arguments)]
    1755
    0
    unsafe fn compute_q8_sums_8blocks(
    1756
    0
        c0_lo: std::arch::x86_64::__m256i,
    1757
    0
        c0_hi: std::arch::x86_64::__m256i,
    1758
    0
        c1_lo: std::arch::x86_64::__m256i,
    1759
    0
        c1_hi: std::arch::x86_64::__m256i,
    1760
    0
        c2_lo: std::arch::x86_64::__m256i,
    1761
    0
        c2_hi: std::arch::x86_64::__m256i,
    1762
    0
        c3_lo: std::arch::x86_64::__m256i,
    1763
    0
        c3_hi: std::arch::x86_64::__m256i,
    1764
    0
        ones_16: std::arch::x86_64::__m256i,
    1765
    0
    ) -> std::arch::x86_64::__m256i {
    1766
        #[allow(clippy::wildcard_imports)]
    1767
        use std::arch::x86_64::*;
    1768
    1769
        // SAFETY: All calls are to unsafe intrinsics, within already-unsafe fn
    1770
        // Sum Q8 values for each of the 8 blocks (32 values each)
    1771
        unsafe {
    1772
    0
            let sum0 = sum_i8_to_i32(c0_lo, ones_16);
    1773
    0
            let sum1 = sum_i8_to_i32(c0_hi, ones_16);
    1774
    0
            let sum2 = sum_i8_to_i32(c1_lo, ones_16);
    1775
    0
            let sum3 = sum_i8_to_i32(c1_hi, ones_16);
    1776
    0
            let sum4 = sum_i8_to_i32(c2_lo, ones_16);
    1777
    0
            let sum5 = sum_i8_to_i32(c2_hi, ones_16);
    1778
    0
            let sum6 = sum_i8_to_i32(c3_lo, ones_16);
    1779
    0
            let sum7 = sum_i8_to_i32(c3_hi, ones_16);
    1780
    1781
            // Pack 8 sums into a single __m256i
    1782
    0
            let mut result = _mm256_setzero_si256();
    1783
    0
            result = _mm256_insert_epi32(result, sum0, 0);
    1784
    0
            result = _mm256_insert_epi32(result, sum1, 1);
    1785
    0
            result = _mm256_insert_epi32(result, sum2, 2);
    1786
    0
            result = _mm256_insert_epi32(result, sum3, 3);
    1787
    0
            result = _mm256_insert_epi32(result, sum4, 4);
    1788
    0
            result = _mm256_insert_epi32(result, sum5, 5);
    1789
    0
            result = _mm256_insert_epi32(result, sum6, 6);
    1790
    0
            result = _mm256_insert_epi32(result, sum7, 7);
    1791
    0
            result
    1792
        }
    1793
    0
    }
    1794
    1795
    /// Helper: Sum 32 i8 values to i32
    1796
    #[cfg(target_arch = "x86_64")]
    1797
    #[target_feature(enable = "avx2")]
    1798
    #[inline]
    1799
    0
    unsafe fn sum_i8_to_i32(v: std::arch::x86_64::__m256i, ones: std::arch::x86_64::__m256i) -> i32 {
    1800
        #[allow(clippy::wildcard_imports)]
    1801
        use std::arch::x86_64::*;
    1802
    1803
    0
        let lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(v));
    1804
    0
        let hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(v, 1));
    1805
    0
        let sum_lo = _mm256_madd_epi16(lo, ones);
    1806
    0
        let sum_hi = _mm256_madd_epi16(hi, ones);
    1807
    0
        let sum = _mm256_add_epi32(sum_lo, sum_hi);
    1808
    1809
        // Horizontal sum of 8 i32 -> 1 i32
    1810
    0
        let sum128 = _mm_add_epi32(
    1811
    0
            _mm256_castsi256_si128(sum),
    1812
    0
            _mm256_extracti128_si256(sum, 1),
    1813
        );
    1814
    0
        let sum64 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, 0b10_11_00_01));
    1815
    0
        let sum32 = _mm_add_epi32(sum64, _mm_shuffle_epi32(sum64, 0b00_00_10_10));
    1816
    0
        _mm_cvtsi128_si32(sum32)
    1817
    0
    }
    1818
    1819
    /// Helper: Compute Q4×Q8 dot products for 8 blocks using pre-loaded Q8K
    1820
    #[cfg(target_arch = "x86_64")]
    1821
    #[target_feature(enable = "avx2")]
    1822
    #[inline]
    1823
    #[allow(clippy::too_many_arguments)]
    1824
    0
    unsafe fn compute_q4_q8_dots_8blocks(
    1825
    0
        qs_ptr: *const u8,
    1826
    0
        q8_c0_lo: std::arch::x86_64::__m256i,
    1827
    0
        q8_c0_hi: std::arch::x86_64::__m256i,
    1828
    0
        q8_c1_lo: std::arch::x86_64::__m256i,
    1829
    0
        q8_c1_hi: std::arch::x86_64::__m256i,
    1830
    0
        q8_c2_lo: std::arch::x86_64::__m256i,
    1831
    0
        q8_c2_hi: std::arch::x86_64::__m256i,
    1832
    0
        q8_c3_lo: std::arch::x86_64::__m256i,
    1833
    0
        q8_c3_hi: std::arch::x86_64::__m256i,
    1834
    0
        nibble_mask: std::arch::x86_64::__m256i,
    1835
    0
        ones_16: std::arch::x86_64::__m256i,
    1836
    0
    ) -> std::arch::x86_64::__m256i {
    1837
        #[allow(clippy::wildcard_imports)]
    1838
        use std::arch::x86_64::*;
    1839
    1840
        // SAFETY: All intrinsics and pointer ops are unsafe, within already-unsafe fn
    1841
        unsafe {
    1842
            // Load Q4K data (128 bytes total for 256 values)
    1843
    0
            let q4_c0 = _mm256_loadu_si256(qs_ptr.cast::<__m256i>());
    1844
    0
            let q4_c1 = _mm256_loadu_si256(qs_ptr.add(32).cast::<__m256i>());
    1845
    0
            let q4_c2 = _mm256_loadu_si256(qs_ptr.add(64).cast::<__m256i>());
    1846
    0
            let q4_c3 = _mm256_loadu_si256(qs_ptr.add(96).cast::<__m256i>());
    1847
    1848
            // Extract nibbles and compute Q4×Q8 for each chunk
    1849
    0
            let dot0 = q4_q8_chunk_dot(q4_c0, q8_c0_lo, q8_c0_hi, nibble_mask, ones_16);
    1850
    0
            let dot1 = q4_q8_chunk_dot(q4_c1, q8_c1_lo, q8_c1_hi, nibble_mask, ones_16);
    1851
    0
            let dot2 = q4_q8_chunk_dot(q4_c2, q8_c2_lo, q8_c2_hi, nibble_mask, ones_16);
    1852
    0
            let dot3 = q4_q8_chunk_dot(q4_c3, q8_c3_lo, q8_c3_hi, nibble_mask, ones_16);
    1853
    1854
            // Pack 8 dot products into result (2 per chunk: lo nibble block, hi nibble block)
    1855
    0
            let mut result = _mm256_setzero_si256();
    1856
    0
            result = _mm256_insert_epi32(result, _mm_cvtsi128_si32(_mm256_castsi256_si128(dot0)), 0);
    1857
    0
            result = _mm256_insert_epi32(
    1858
    0
                result,
    1859
    0
                _mm_extract_epi32(_mm256_castsi256_si128(dot0), 1),
    1860
    0
                1,
    1861
    0
            );
    1862
    0
            result = _mm256_insert_epi32(result, _mm_cvtsi128_si32(_mm256_castsi256_si128(dot1)), 2);
    1863
    0
            result = _mm256_insert_epi32(
    1864
    0
                result,
    1865
    0
                _mm_extract_epi32(_mm256_castsi256_si128(dot1), 1),
    1866
    0
                3,
    1867
    0
            );
    1868
    0
            result = _mm256_insert_epi32(result, _mm_cvtsi128_si32(_mm256_castsi256_si128(dot2)), 4);
    1869
    0
            result = _mm256_insert_epi32(
    1870
    0
                result,
    1871
    0
                _mm_extract_epi32(_mm256_castsi256_si128(dot2), 1),
    1872
    0
                5,
    1873
    0
            );
    1874
    0
            result = _mm256_insert_epi32(result, _mm_cvtsi128_si32(_mm256_castsi256_si128(dot3)), 6);
    1875
    0
            result = _mm256_insert_epi32(
    1876
    0
                result,
    1877
    0
                _mm_extract_epi32(_mm256_castsi256_si128(dot3), 1),
    1878
    0
                7,
    1879
    0
            );
    1880
    0
            result
    1881
        }
    1882
    0
    }
    1883
    1884
    /// Helper: Q4×Q8 dot product for one 64-value chunk (32 lo nibbles + 32 hi nibbles)
    1885
    #[cfg(target_arch = "x86_64")]
    1886
    #[target_feature(enable = "avx2")]
    1887
    #[inline]
    1888
    0
    unsafe fn q4_q8_chunk_dot(
    1889
    0
        q4_packed: std::arch::x86_64::__m256i,
    1890
    0
        q8_lo: std::arch::x86_64::__m256i,
    1891
    0
        q8_hi: std::arch::x86_64::__m256i,
    1892
    0
        nibble_mask: std::arch::x86_64::__m256i,
    1893
    0
        ones_16: std::arch::x86_64::__m256i,
    1894
    0
    ) -> std::arch::x86_64::__m256i {
    1895
        #[allow(clippy::wildcard_imports)]
    1896
        use std::arch::x86_64::*;
    1897
    1898
        // SAFETY: All intrinsics are unsafe, within already-unsafe fn
    1899
        unsafe {
    1900
            // Extract nibbles
    1901
    0
            let q4_lo = _mm256_and_si256(q4_packed, nibble_mask);
    1902
    0
            let q4_hi = _mm256_and_si256(_mm256_srli_epi16(q4_packed, 4), nibble_mask);
    1903
    1904
            // Q4 × Q8 products
    1905
    0
            let prod_lo_i16 = _mm256_maddubs_epi16(q4_lo, q8_lo);
    1906
    0
            let prod_hi_i16 = _mm256_maddubs_epi16(q4_hi, q8_hi);
    1907
    0
            let prod_lo_i32 = _mm256_madd_epi16(prod_lo_i16, ones_16);
    1908
    0
            let prod_hi_i32 = _mm256_madd_epi16(prod_hi_i16, ones_16);
    1909
    1910
            // Reduce each to single sum
    1911
    0
            let sum_lo = hsum_epi32(prod_lo_i32);
    1912
    0
            let sum_hi = hsum_epi32(prod_hi_i32);
    1913
    1914
            // Return [sum_lo, sum_hi, 0, 0, 0, 0, 0, 0]
    1915
    0
            let mut result = _mm256_setzero_si256();
    1916
    0
            result = _mm256_insert_epi32(result, sum_lo, 0);
    1917
    0
            result = _mm256_insert_epi32(result, sum_hi, 1);
    1918
    0
            result
    1919
        }
    1920
    0
    }
    1921
    1922
    /// Helper: Horizontal sum of 8 i32 values to single i32
    1923
    #[cfg(target_arch = "x86_64")]
    1924
    #[target_feature(enable = "avx2")]
    1925
    #[inline]
    1926
    0
    unsafe fn hsum_epi32(v: std::arch::x86_64::__m256i) -> i32 {
    1927
        #[allow(clippy::wildcard_imports)]
    1928
        use std::arch::x86_64::*;
    1929
    1930
        // All intrinsics are unsafe and we're in an unsafe fn with target_feature
    1931
    0
        let sum128 = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256(v, 1));
    1932
    0
        let sum64 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, 0b10_11_00_01));
    1933
    0
        let sum32 = _mm_add_epi32(sum64, _mm_shuffle_epi32(sum64, 0b00_00_10_10));
    1934
    0
        _mm_cvtsi128_si32(sum32)
    1935
    0
    }
    1936