## Coverage Report

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

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

    /home/noah/src/trueno/src/backends/q4k/gemv.rs

    Line
    Count
    Source
    1
    //! Row-major Q4_K matrix-vector multiplication.
    2
    //!
    3
    //! This module implements row-major GEMV where weights are stored row-first.
    4
    //! Includes scalar, AVX2-optimized, and parallel dispatch implementations.
    5
    6
    use super::{parse_q4k_header, SUPER_BLOCK_BYTES, SUPER_BLOCK_SIZE};
    7
    8
    0
    pub fn matmul_q4k_f32_scalar(
    9
    0
        q4k_data: &[u8],
    10
    0
        input: &[f32],
    11
    0
        out_dim: usize,
    12
    0
        in_dim: usize,
    13
    0
    ) -> Vec<f32> {
    14
    0
        assert_eq!(input.len(), in_dim, "Input length mismatch");
    15
    0
        assert!(
    16
    0
            in_dim % SUPER_BLOCK_SIZE == 0 || in_dim < SUPER_BLOCK_SIZE,
    17
    0
            "in_dim must be multiple of 256 (or smaller for padding)"
    18
        );
    19
    20
    0
        let num_blocks_per_row = (in_dim + SUPER_BLOCK_SIZE - 1) / SUPER_BLOCK_SIZE;
    21
    0
        let row_bytes = num_blocks_per_row * SUPER_BLOCK_BYTES;
    22
    0
        let expected_size = out_dim * row_bytes;
    23
    24
    0
        assert!(
    25
    0
            q4k_data.len() >= expected_size,
    26
    0
            "Q4K data too small: {} < {}",
    27
    0
            q4k_data.len(),
    28
            expected_size
    29
        );
    30
    31
    0
        let mut output = vec![0.0f32; out_dim];
    32
    33
    0
        for out_idx in 0..out_dim {
    34
    0
            let row_start = out_idx * row_bytes;
    35
    0
            let mut sum = 0.0f32;
    36
    37
    0
            for sb_idx in 0..num_blocks_per_row {
    38
    0
                let sb_start = row_start + sb_idx * SUPER_BLOCK_BYTES;
    39
    0
                let sb_data = &q4k_data[sb_start..sb_start + SUPER_BLOCK_BYTES];
    40
    41
                // Parse header
    42
    0
                let (d, dmin, scales, mins) = parse_q4k_header(sb_data);
    43
    0
                let qs = &sb_data[16..144];
    44
    45
                // Input offset for this super-block
    46
    0
                let input_offset = sb_idx * SUPER_BLOCK_SIZE;
    47
    48
                // Process 4 chunks of 64 values each
    49
    0
                for chunk in 0..4 {
    50
    0
                    let chunk_start = chunk * 64;
    51
    0
                    let q_start = chunk * 32;
    52
    53
                    // Scale indices for this chunk
    54
    0
                    let scale_idx_low = chunk * 2;
    55
    0
                    let scale_idx_high = chunk * 2 + 1;
    56
    57
    0
                    let d1 = d * f32::from(scales[scale_idx_low]);
    58
    0
                    let dm1 = dmin * f32::from(mins[scale_idx_low]);
    59
    0
                    let d2 = d * f32::from(scales[scale_idx_high]);
    60
    0
                    let dm2 = dmin * f32::from(mins[scale_idx_high]);
    61
    62
                    // First 32 values: low nibbles
    63
    0
                    for i in 0..32 {
    64
    0
                        let q_val = (qs[q_start + i] & 0x0F) as f32;
    65
    0
                        let dequant = d1 * q_val - dm1;
    66
    0
                        let input_idx = input_offset + chunk_start + i;
    67
    0
                        if input_idx < in_dim {
    68
    0
                            sum += dequant * input[input_idx];
    69
    0
                        }
    70
                    }
    71
    72
                    // Next 32 values: high nibbles
    73
    0
                    for i in 0..32 {
    74
    0
                        let q_val = (qs[q_start + i] >> 4) as f32;
    75
    0
                        let dequant = d2 * q_val - dm2;
    76
    0
                        let input_idx = input_offset + chunk_start + 32 + i;
    77
    0
                        if input_idx < in_dim {
    78
    0
                            sum += dequant * input[input_idx];
    79
    0
                        }
    80
                    }
    81
                }
    82
            }
    83
    84
    0
            output[out_idx] = sum;
    85
        }
    86
    87
    0
        output
    88
    0
    }
    89
    90
    /// Fused Q4_K matrix-vector multiply (optimized with 4-way unrolling)
    91
    ///
    92
    /// This version uses 4 independent accumulators to improve instruction-level
    93
    /// parallelism while maintaining scalar correctness.
    94
    ///
    95
    /// # Arguments
    96
    /// Same as `matmul_q4k_f32_scalar`
    97
    0
    pub fn matmul_q4k_f32(
    98
    0
        q4k_data: &[u8],
    99
    0
        input: &[f32],
    100
    0
        out_dim: usize,
    101
    0
        in_dim: usize,
    102
    0
    ) -> Vec<f32> {
    103
    0
        assert_eq!(input.len(), in_dim, "Input length mismatch");
    104
    105
    0
        let num_blocks_per_row = (in_dim + SUPER_BLOCK_SIZE - 1) / SUPER_BLOCK_SIZE;
    106
    0
        let row_bytes = num_blocks_per_row * SUPER_BLOCK_BYTES;
    107
    108
    0
        let mut output = vec![0.0f32; out_dim];
    109
    110
    0
        for out_idx in 0..out_dim {
    111
    0
            let row_start = out_idx * row_bytes;
    112
    113
            // 4 independent accumulators for better ILP
    114
    0
            let mut acc0 = 0.0f32;
    115
    0
            let mut acc1 = 0.0f32;
    116
    0
            let mut acc2 = 0.0f32;
    117
    0
            let mut acc3 = 0.0f32;
    118
    119
    0
            for sb_idx in 0..num_blocks_per_row {
    120
    0
                let sb_start = row_start + sb_idx * SUPER_BLOCK_BYTES;
    121
    0
                let sb_data = &q4k_data[sb_start..sb_start + SUPER_BLOCK_BYTES];
    122
    123
                // Parse header
    124
    0
                let (d, dmin, scales, mins) = parse_q4k_header(sb_data);
    125
    0
                let qs = &sb_data[16..144];
    126
    127
                // Input offset for this super-block
    128
    0
                let input_offset = sb_idx * SUPER_BLOCK_SIZE;
    129
    130
                // Process 4 chunks, accumulating to different registers
    131
    0
                for chunk in 0..4 {
    132
    0
                    let chunk_start = chunk * 64;
    133
    0
                    let q_start = chunk * 32;
    134
    135
    0
                    let scale_idx_low = chunk * 2;
    136
    0
                    let scale_idx_high = chunk * 2 + 1;
    137
    138
    0
                    let d1 = d * f32::from(scales[scale_idx_low]);
    139
    0
                    let dm1 = dmin * f32::from(mins[scale_idx_low]);
    140
    0
                    let d2 = d * f32::from(scales[scale_idx_high]);
    141
    0
                    let dm2 = dmin * f32::from(mins[scale_idx_high]);
    142
    143
                    // Process low nibbles (first 32) with 4-way unroll
    144
    0
                    let mut i = 0;
    145
    0
                    while i + 3 < 32 {
    146
    0
                        let input_base = input_offset + chunk_start + i;
    147
    0
                        if input_base + 3 < in_dim {
    148
    0
                            let q0 = (qs[q_start + i] & 0x0F) as f32;
    149
    0
                            let q1 = (qs[q_start + i + 1] & 0x0F) as f32;
    150
    0
                            let q2 = (qs[q_start + i + 2] & 0x0F) as f32;
    151
    0
                            let q3 = (qs[q_start + i + 3] & 0x0F) as f32;
    152
    0
    153
    0
                            acc0 = (d1 * q0 - dm1).mul_add(input[input_base], acc0);
    154
    0
                            acc1 = (d1 * q1 - dm1).mul_add(input[input_base + 1], acc1);
    155
    0
                            acc2 = (d1 * q2 - dm1).mul_add(input[input_base + 2], acc2);
    156
    0
                            acc3 = (d1 * q3 - dm1).mul_add(input[input_base + 3], acc3);
    157
    0
                        }
    158
    0
                        i += 4;
    159
                    }
    160
                    // Handle remainder
    161
    0
                    while i < 32 {
    162
    0
                        let input_idx = input_offset + chunk_start + i;
    163
    0
                        if input_idx < in_dim {
    164
    0
                            let q_val = (qs[q_start + i] & 0x0F) as f32;
    165
    0
                            acc0 = (d1 * q_val - dm1).mul_add(input[input_idx], acc0);
    166
    0
                        }
    167
    0
                        i += 1;
    168
                    }
    169
    170
                    // Process high nibbles (next 32) with 4-way unroll
    171
    0
                    let mut i = 0;
    172
    0
                    while i + 3 < 32 {
    173
    0
                        let input_base = input_offset + chunk_start + 32 + i;
    174
    0
                        if input_base + 3 < in_dim {
    175
    0
                            let q0 = (qs[q_start + i] >> 4) as f32;
    176
    0
                            let q1 = (qs[q_start + i + 1] >> 4) as f32;
    177
    0
                            let q2 = (qs[q_start + i + 2] >> 4) as f32;
    178
    0
                            let q3 = (qs[q_start + i + 3] >> 4) as f32;
    179
    0
    180
    0
                            acc0 = (d2 * q0 - dm2).mul_add(input[input_base], acc0);
    181
    0
                            acc1 = (d2 * q1 - dm2).mul_add(input[input_base + 1], acc1);
    182
    0
                            acc2 = (d2 * q2 - dm2).mul_add(input[input_base + 2], acc2);
    183
    0
                            acc3 = (d2 * q3 - dm2).mul_add(input[input_base + 3], acc3);
    184
    0
                        }
    185
    0
                        i += 4;
    186
                    }
    187
                    // Handle remainder
    188
    0
                    while i < 32 {
    189
    0
                        let input_idx = input_offset + chunk_start + 32 + i;
    190
    0
                        if input_idx < in_dim {
    191
    0
                            let q_val = (qs[q_start + i] >> 4) as f32;
    192
    0
                            acc0 = (d2 * q_val - dm2).mul_add(input[input_idx], acc0);
    193
    0
                        }
    194
    0
                        i += 1;
    195
                    }
    196
                }
    197
            }
    198
    199
            // Combine all accumulators
    200
    0
            output[out_idx] = (acc0 + acc1) + (acc2 + acc3);
    201
        }
    202
    203
    0
        output
    204
    0
    }
    205
    206
    /// Fused Q4_K matrix-vector multiply with AVX2 SIMD (8-wide)
    207
    ///
    208
    /// Processes 8 elements at a time using AVX2 intrinsics.
    209
    /// Falls back to scalar for remainder elements.
    210
    #[cfg(target_arch = "x86_64")]
    211
    #[target_feature(enable = "avx2", enable = "fma")]
    212
    0
    unsafe fn matmul_q4k_f32_avx2(
    213
    0
        q4k_data: &[u8],
    214
    0
        input: &[f32],
    215
    0
        out_dim: usize,
    216
    0
        in_dim: usize,
    217
    0
    ) -> Vec<f32> {
    218
        #[cfg(target_arch = "x86_64")]
    219
        use std::arch::x86_64::*;
    220
    221
    0
        let num_blocks_per_row = (in_dim + SUPER_BLOCK_SIZE - 1) / SUPER_BLOCK_SIZE;
    222
    0
        let row_bytes = num_blocks_per_row * SUPER_BLOCK_BYTES;
    223
    224
    0
        let mut output = vec![0.0f32; out_dim];
    225
    226
        // Mask for extracting low 4 bits
    227
    0
        let low_mask = _mm256_set1_epi32(0x0F);
    228
    229
    0
        for out_idx in 0..out_dim {
    230
    0
            let row_start = out_idx * row_bytes;
    231
    232
            // 8-wide accumulator
    233
    0
            let mut acc = _mm256_setzero_ps();
    234
    235
    0
            for sb_idx in 0..num_blocks_per_row {
    236
    0
                let sb_start = row_start + sb_idx * SUPER_BLOCK_BYTES;
    237
    0
                let sb_data = &q4k_data[sb_start..sb_start + SUPER_BLOCK_BYTES];
    238
    239
                // Parse header
    240
    0
                let (d, dmin, scales, mins) = parse_q4k_header(sb_data);
    241
    0
                let qs = &sb_data[16..144];
    242
    243
    0
                let input_offset = sb_idx * SUPER_BLOCK_SIZE;
    244
    245
                // Process 4 chunks of 64 values each
    246
    0
                for chunk in 0..4 {
    247
    0
                    let chunk_start = chunk * 64;
    248
    0
                    let q_start = chunk * 32;
    249
    250
    0
                    let scale_idx_low = chunk * 2;
    251
    0
                    let scale_idx_high = chunk * 2 + 1;
    252
    253
    0
                    let d1 = d * f32::from(scales[scale_idx_low]);
    254
    0
                    let dm1 = dmin * f32::from(mins[scale_idx_low]);
    255
    0
                    let d2 = d * f32::from(scales[scale_idx_high]);
    256
    0
                    let dm2 = dmin * f32::from(mins[scale_idx_high]);
    257
    258
                    // Broadcast scales
    259
    0
                    let d1_vec = _mm256_set1_ps(d1);
    260
    0
                    let dm1_vec = _mm256_set1_ps(dm1);
    261
    0
                    let d2_vec = _mm256_set1_ps(d2);
    262
    0
                    let dm2_vec = _mm256_set1_ps(dm2);
    263
    264
                    // Process low nibbles (32 values) in groups of 8
    265
    0
                    let mut i = 0;
    266
    0
                    while i + 8 <= 32 {
    267
    0
                        let input_base = input_offset + chunk_start + i;
    268
    0
                        if input_base + 8 <= in_dim {
    269
    0
                            // Load 8 bytes of quantized values
    270
    0
                            let q_bytes = _mm_loadl_epi64(
    271
    0
                                qs.as_ptr().add(q_start + i) as *const __m128i
    272
    0
                            );
    273
    0
    274
    0
                            // Zero-extend u8 to i32: [b0, b1, ..., b7, 0, 0, ...] -> [b0, b1, ..., b7] as i32
    275
    0
                            let q_i32 = _mm256_cvtepu8_epi32(q_bytes);
    276
    0
    277
    0
                            // Mask low nibbles
    278
    0
                            let q_low = _mm256_and_si256(q_i32, low_mask);
    279
    0
    280
    0
                            // Convert to f32
    281
    0
                            let q_f32 = _mm256_cvtepi32_ps(q_low);
    282
    0
    283
    0
                            // Load 8 input values
    284
    0
                            let x = _mm256_loadu_ps(input.as_ptr().add(input_base));
    285
    0
    286
    0
                            // dequant = d1 * q - dm1
    287
    0
                            let dequant = _mm256_fmsub_ps(d1_vec, q_f32, dm1_vec);
    288
    0
    289
    0
                            // acc += dequant * x
    290
    0
                            acc = _mm256_fmadd_ps(dequant, x, acc);
    291
    0
                        }
    292
    0
                        i += 8;
    293
                    }
    294
    295
                    // Process high nibbles (32 values) in groups of 8
    296
    0
                    let mut i = 0;
    297
    0
                    while i + 8 <= 32 {
    298
    0
                        let input_base = input_offset + chunk_start + 32 + i;
    299
    0
                        if input_base + 8 <= in_dim {
    300
    0
                            // Load 8 bytes of quantized values
    301
    0
                            let q_bytes = _mm_loadl_epi64(
    302
    0
                                qs.as_ptr().add(q_start + i) as *const __m128i
    303
    0
                            );
    304
    0
    305
    0
                            // Zero-extend u8 to i32
    306
    0
                            let q_i32 = _mm256_cvtepu8_epi32(q_bytes);
    307
    0
    308
    0
                            // Shift right 4 bits to get high nibbles
    309
    0
                            let q_high = _mm256_srli_epi32(q_i32, 4);
    310
    0
    311
    0
                            // Convert to f32
    312
    0
                            let q_f32 = _mm256_cvtepi32_ps(q_high);
    313
    0
    314
    0
                            // Load 8 input values
    315
    0
                            let x = _mm256_loadu_ps(input.as_ptr().add(input_base));
    316
    0
    317
    0
                            // dequant = d2 * q - dm2
    318
    0
                            let dequant = _mm256_fmsub_ps(d2_vec, q_f32, dm2_vec);
    319
    0
    320
    0
                            // acc += dequant * x
    321
    0
                            acc = _mm256_fmadd_ps(dequant, x, acc);
    322
    0
                        }
    323
    0
                        i += 8;
    324
                    }
    325
                }
    326
            }
    327
    328
            // Horizontal sum of 8-wide accumulator
    329
            // acc = [a0, a1, a2, a3, a4, a5, a6, a7]
    330
    0
            let hi128 = _mm256_extractf128_ps(acc, 1);
    331
    0
            let lo128 = _mm256_castps256_ps128(acc);
    332
    0
            let sum128 = _mm_add_ps(lo128, hi128);
    333
            // sum128 = [a0+a4, a1+a5, a2+a6, a3+a7]
    334
    0
            let hi64 = _mm_movehl_ps(sum128, sum128);
    335
    0
            let sum64 = _mm_add_ps(sum128, hi64);
    336
            // sum64 = [a0+a2+a4+a6, a1+a3+a5+a7, ...]
    337
    0
            let hi32 = _mm_shuffle_ps(sum64, sum64, 1);
    338
    0
            let sum32 = _mm_add_ss(sum64, hi32);
    339
    340
    0
            output[out_idx] = _mm_cvtss_f32(sum32);
    341
        }
    342
    343
    0
        output
    344
    0
    }
    345
    346
    /// Runtime dispatch for Q4K matmul - uses AVX2 if available, otherwise scalar
    347
    #[inline]
    348
    0
    pub fn matmul_q4k_f32_dispatch(
    349
    0
        q4k_data: &[u8],
    350
    0
        input: &[f32],
    351
    0
        out_dim: usize,
    352
    0
        in_dim: usize,
    353
    0
    ) -> Vec<f32> {
    354
        #[cfg(target_arch = "x86_64")]
    355
        {
    356
            // For large matmuls (total work >= ~8M ops), use parallel execution
    357
            // This catches FFN layers (8960x1536) and lm_head (151936x1536)
    358
            // Also catches ffn_down (1536x8960) where out_dim is small but in_dim is large
    359
    0
            let total_work = out_dim * in_dim;
    360
    0
            if total_work >= 8_000_000 {
    361
    0
                return matmul_q4k_f32_parallel(q4k_data, input, out_dim, in_dim);
    362
    0
            }
    363
    364
    0
            if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
    365
                // SAFETY: We just verified AVX2 + FMA are available
    366
    0
                return unsafe { matmul_q4k_f32_avx2(q4k_data, input, out_dim, in_dim) };
    367
    0
            }
    368
        }
    369
    370
        // Fallback to scalar with 4-way unroll
    371
    0
        matmul_q4k_f32(q4k_data, input, out_dim, in_dim)
    372
    0
    }
    373
    374
    /// Fused Q4_K matrix-vector multiply for GGML column-major layout
    375
    ///
    376
    /// Computes: output = input @ Q4K_weight (GGML convention: y = x @ W)
    377
    /// where weight is stored in Q4_K format with GGML column-major super-block organization.
    378
    ///
    379
    /// # GGML Column-Major Layout (PMAT-103)
    380
    ///
    381
    /// For a weight tensor with shape [ne0, ne1] in GGML notation:
    382
    /// - ne0 is the output dimension (rows)
    383
    /// - ne1 is the input/reduction dimension (columns)
    384
    /// - Elements are stored column-major: W[i,j] at offset i + j*ne0
    385
    /// - Each column j (length ne0) contains weights from input[j] to all outputs
    386
    /// - Super-blocks are organized by columns: column j uses super-blocks [j*blocks_per_col, (j+1)*blocks_per_col)
    387
    ///
    388
    /// This matches GGUF tensor storage and enables fused kernel execution without transposition.
    389
    ///
    390
    /// # Arguments
    391
    /// * `q4k_data` - Raw Q4K bytes in GGML column-major layout [ne0, ne1]
    392
    /// * `input` - F32 input vector [ne1] (input/reduction dimension)
    393
    /// * `ne0` - Size of output dimension (rows in GGML, output size)
    394
    /// * `ne1` - Size of input/reduction dimension (columns in GGML, input size)
    395
    ///
    396
    /// # Returns
    397
    /// F32 output vector [ne0]
    398
    ///
    399
    /// # Example
    400
    /// ```rust,ignore
    401
    /// // GGUF ffn_gate: shape [intermediate_dim, hidden_dim] = [8960, 1536]
    402
    /// // Computes: intermediate = hidden @ ffn_gate
    403
    /// let output = matmul_q4k_f32_colmajor(&q4k_bytes, &hidden, 8960, 1536);
    404
    /// // output has 8960 elements
    405
    /// ```
    406
    407
    // ============================================================================
    408
    // Parallel Execution Helpers
    409
    // ============================================================================
    410
    411
    0
    fn matmul_q4k_f32_parallel(
    412
    0
        q4k_data: &[u8],
    413
    0
        input: &[f32],
    414
    0
        out_dim: usize,
    415
    0
        in_dim: usize,
    416
    0
    ) -> Vec<f32> {
    417
        use std::thread;
    418
    419
        // Use fewer threads with larger chunks for better cache efficiency
    420
    0
        let num_threads = thread::available_parallelism()
    421
    0
            .map(|p| p.get())
    422
    0
            .unwrap_or(4)
    423
    0
            .min(12);
    424
    425
    0
        let chunk_size = (out_dim + num_threads - 1) / num_threads;
    426
    0
        let num_blocks_per_row = (in_dim + SUPER_BLOCK_SIZE - 1) / SUPER_BLOCK_SIZE;
    427
    0
        let row_bytes = num_blocks_per_row * SUPER_BLOCK_BYTES;
    428
    429
    0
        let mut output = vec![0.0f32; out_dim];
    430
    0
        let has_avx2 = is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma");
    431
    432
    0
        thread::scope(|s| {
    433
    0
            let input_ref = input;
    434
    0
            let q4k_ref = q4k_data;
    435
    0
            let chunks: Vec<_> = output.chunks_mut(chunk_size).enumerate().collect();
    436
    437
    0
            for (chunk_idx, chunk) in chunks {
    438
    0
                let start_row = chunk_idx * chunk_size;
    439
    440
    0
                s.spawn(move || {
    441
    0
                    if has_avx2 {
    442
    0
                        unsafe {
    443
    0
                            compute_chunk_q4k_avx2(
    444
    0
                                q4k_ref,
    445
    0
                                input_ref,
    446
    0
                                chunk,
    447
    0
                                start_row,
    448
    0
                                out_dim,
    449
    0
                                in_dim,
    450
    0
                                num_blocks_per_row,
    451
    0
                                row_bytes,
    452
    0
                            );
    453
    0
                        }
    454
    0
                    } else {
    455
    0
                        compute_chunk_q4k_scalar(
    456
    0
                            q4k_ref,
    457
    0
                            input_ref,
    458
    0
                            chunk,
    459
    0
                            start_row,
    460
    0
                            out_dim,
    461
    0
                            in_dim,
    462
    0
                            num_blocks_per_row,
    463
    0
                            row_bytes,
    464
    0
                        );
    465
    0
                    }
    466
    0
                });
    467
            }
    468
    0
        });
    469
    470
    0
        output
    471
    0
    }
    472
    473
    /// Fallback for non-x86_64
    474
    #[cfg(not(target_arch = "x86_64"))]
    475
    fn matmul_q4k_f32_parallel(
    476
        q4k_data: &[u8],
    477
        input: &[f32],
    478
        out_dim: usize,
    479
        in_dim: usize,
    480
    ) -> Vec<f32> {
    481
        matmul_q4k_f32(q4k_data, input, out_dim, in_dim)
    482
    }
    483
    484
    #[cfg(target_arch = "x86_64")]
    485
    #[target_feature(enable = "avx2", enable = "fma")]
    486
    0
    unsafe fn compute_chunk_q4k_avx2(
    487
    0
        q4k_data: &[u8],
    488
    0
        input: &[f32],
    489
    0
        chunk: &mut [f32],
    490
    0
        start_row: usize,
    491
    0
        out_dim: usize,
    492
    0
        in_dim: usize,
    493
    0
        num_blocks_per_row: usize,
    494
    0
        row_bytes: usize,
    495
    0
    ) {
    496
        use std::arch::x86_64::*;
    497
    498
    0
        let low_mask = _mm256_set1_epi32(0x0F);
    499
    500
    0
        for (local_idx, out_val) in chunk.iter_mut().enumerate() {
    501
    0
            let out_idx = start_row + local_idx;
    502
    0
            if out_idx >= out_dim {
    503
    0
                break;
    504
    0
            }
    505
    506
    0
            let row_start = out_idx * row_bytes;
    507
    0
            let mut acc = _mm256_setzero_ps();
    508
    509
    0
            for sb_idx in 0..num_blocks_per_row {
    510
    0
                let sb_start = row_start + sb_idx * SUPER_BLOCK_BYTES;
    511
    0
                if sb_start + SUPER_BLOCK_BYTES > q4k_data.len() {
    512
    0
                    break;
    513
    0
                }
    514
    0
                let sb_data = &q4k_data[sb_start..sb_start + SUPER_BLOCK_BYTES];
    515
    516
    0
                let (d, dmin, scales, mins) = parse_q4k_header(sb_data);
    517
    0
                let qs = &sb_data[16..144];
    518
    519
    0
                let input_offset = sb_idx * SUPER_BLOCK_SIZE;
    520
    521
                // Process 4 chunks of 64 values each
    522
    0
                for chunk_i in 0..4 {
    523
    0
                    let chunk_start = chunk_i * 64;
    524
    0
                    let q_start = chunk_i * 32;
    525
    526
    0
                    let scale_idx_low = chunk_i * 2;
    527
    0
                    let scale_idx_high = chunk_i * 2 + 1;
    528
    529
    0
                    let d1 = d * f32::from(scales[scale_idx_low]);
    530
    0
                    let dm1 = dmin * f32::from(mins[scale_idx_low]);
    531
    0
                    let d2 = d * f32::from(scales[scale_idx_high]);
    532
    0
                    let dm2 = dmin * f32::from(mins[scale_idx_high]);
    533
    534
    0
                    let d1_vec = _mm256_set1_ps(d1);
    535
    0
                    let dm1_vec = _mm256_set1_ps(dm1);
    536
    0
                    let d2_vec = _mm256_set1_ps(d2);
    537
    0
                    let dm2_vec = _mm256_set1_ps(dm2);
    538
    539
                    // Process low nibbles (32 values) in groups of 8
    540
    0
                    let mut i = 0;
    541
    0
                    while i + 8 <= 32 {
    542
    0
                        let input_base = input_offset + chunk_start + i;
    543
    0
                        if input_base + 8 <= in_dim {
    544
    0
                            let q_bytes = _mm_loadl_epi64(
    545
    0
                                qs.as_ptr().add(q_start + i) as *const __m128i
    546
    0
                            );
    547
    0
                            let q_i32 = _mm256_cvtepu8_epi32(q_bytes);
    548
    0
                            let q_low = _mm256_and_si256(q_i32, low_mask);
    549
    0
                            let q_f32 = _mm256_cvtepi32_ps(q_low);
    550
    0
                            let x = _mm256_loadu_ps(input.as_ptr().add(input_base));
    551
    0
                            let dequant = _mm256_fmsub_ps(d1_vec, q_f32, dm1_vec);
    552
    0
                            acc = _mm256_fmadd_ps(dequant, x, acc);
    553
    0
                        }
    554
    0
                        i += 8;
    555
                    }
    556
    557
                    // Process high nibbles (32 values) in groups of 8
    558
    0
                    let mut i = 0;
    559
    0
                    while i + 8 <= 32 {
    560
    0
                        let input_base = input_offset + chunk_start + 32 + i;
    561
    0
                        if input_base + 8 <= in_dim {
    562
    0
                            let q_bytes = _mm_loadl_epi64(
    563
    0
                                qs.as_ptr().add(q_start + i) as *const __m128i
    564
    0
                            );
    565
    0
                            let q_i32 = _mm256_cvtepu8_epi32(q_bytes);
    566
    0
                            let q_high = _mm256_srli_epi32(q_i32, 4);
    567
    0
                            let q_f32 = _mm256_cvtepi32_ps(q_high);
    568
    0
                            let x = _mm256_loadu_ps(input.as_ptr().add(input_base));
    569
    0
                            let dequant = _mm256_fmsub_ps(d2_vec, q_f32, dm2_vec);
    570
    0
                            acc = _mm256_fmadd_ps(dequant, x, acc);
    571
    0
                        }
    572
    0
                        i += 8;
    573
                    }
    574
                }
    575
            }
    576
    577
            // Horizontal sum
    578
    0
            let hi128 = _mm256_extractf128_ps(acc, 1);
    579
    0
            let lo128 = _mm256_castps256_ps128(acc);
    580
    0
            let sum128 = _mm_add_ps(lo128, hi128);
    581
    0
            let hi64 = _mm_movehl_ps(sum128, sum128);
    582
    0
            let sum64 = _mm_add_ps(sum128, hi64);
    583
    0
            let hi32 = _mm_shuffle_ps(sum64, sum64, 1);
    584
    0
            let sum32 = _mm_add_ss(sum64, hi32);
    585
    586
    0
            *out_val = _mm_cvtss_f32(sum32);
    587
        }
    588
    0
    }
    589
    590
    #[allow(dead_code)]
    591
    0
    pub(crate) fn compute_chunk_q4k_scalar(
    592
    0
        q4k_data: &[u8],
    593
    0
        input: &[f32],
    594
    0
        chunk: &mut [f32],
    595
    0
        start_row: usize,
    596
    0
        out_dim: usize,
    597
    0
        in_dim: usize,
    598
    0
        num_blocks_per_row: usize,
    599
    0
        row_bytes: usize,
    600
    0
    ) {
    601
    0
        for (local_idx, out_val) in chunk.iter_mut().enumerate() {
    602
    0
            let out_idx = start_row + local_idx;
    603
    0
            if out_idx >= out_dim {
    604
    0
                break;
    605
    0
            }
    606
    607
    0
            let row_start = out_idx * row_bytes;
    608
    0
            let mut sum = 0.0f32;
    609
    610
    0
            for sb_idx in 0..num_blocks_per_row {
    611
    0
                let sb_start = row_start + sb_idx * SUPER_BLOCK_BYTES;
    612
    0
                if sb_start + SUPER_BLOCK_BYTES > q4k_data.len() {
    613
    0
                    break;
    614
    0
                }
    615
    0
                let sb_data = &q4k_data[sb_start..sb_start + SUPER_BLOCK_BYTES];
    616
    617
    0
                let (d, dmin, scales, mins) = parse_q4k_header(sb_data);
    618
    0
                let qs = &sb_data[16..144];
    619
    620
    0
                let input_offset = sb_idx * SUPER_BLOCK_SIZE;
    621
    622
    0
                for chunk_i in 0..4 {
    623
    0
                    let chunk_start = chunk_i * 64;
    624
    0
                    let q_start = chunk_i * 32;
    625
    626
    0
                    let scale_idx_low = chunk_i * 2;
    627
    0
                    let scale_idx_high = chunk_i * 2 + 1;
    628
    629
    0
                    let d1 = d * f32::from(scales[scale_idx_low]);
    630
    0
                    let dm1 = dmin * f32::from(mins[scale_idx_low]);
    631
    0
                    let d2 = d * f32::from(scales[scale_idx_high]);
    632
    0
                    let dm2 = dmin * f32::from(mins[scale_idx_high]);
    633
    634
                    // Low nibbles
    635
    0
                    for i in 0..32 {
    636
    0
                        let input_idx = input_offset + chunk_start + i;
    637
    0
                        if input_idx < in_dim {
    638
    0
                            let q_val = (qs[q_start + i] & 0x0F) as f32;
    639
    0
                            sum += (d1 * q_val - dm1) * input[input_idx];
    640
    0
                        }
    641
                    }
    642
    643
                    // High nibbles
    644
    0
                    for i in 0..32 {
    645
    0
                        let input_idx = input_offset + chunk_start + 32 + i;
    646
    0
                        if input_idx < in_dim {
    647
    0
                            let q_val = (qs[q_start + i] >> 4) as f32;
    648
    0
                            sum += (d2 * q_val - dm2) * input[input_idx];
    649
    0
                        }
    650
                    }
    651
                }
    652
            }
    653
    654
    0
            *out_val = sum;
    655
        }
    656
    0
    }
    657