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

    Line
    Count
    Source
    1
    //! Row-major Q6_K matrix-vector multiplication.
    2
    //!
    3
    //! This module implements row-major GEMV for Q6_K format.
    4
    //! Includes scalar, AVX2-optimized, and parallel dispatch implementations.
    5
    6
    use super::{f16_to_f32, SUPER_BLOCK_BYTES, SUPER_BLOCK_SIZE};
    7
    8
    /// Fused Q6_K matrix-vector multiply (scalar reference)
    9
    0
    pub fn matmul_q6k_f32_scalar(
    10
    0
        q6k_data: &[u8],
    11
    0
        input: &[f32],
    12
    0
        out_dim: usize,
    13
    0
        in_dim: usize,
    14
    0
    ) -> Vec<f32> {
    15
    0
        assert_eq!(input.len(), in_dim, "Input length mismatch");
    16
    17
    0
        let num_blocks_per_row = (in_dim + SUPER_BLOCK_SIZE - 1) / SUPER_BLOCK_SIZE;
    18
    0
        let row_bytes = num_blocks_per_row * SUPER_BLOCK_BYTES;
    19
    20
    0
        let mut output = vec![0.0f32; out_dim];
    21
    22
    0
        for out_idx in 0..out_dim {
    23
    0
            let row_start = out_idx * row_bytes;
    24
    0
            let mut sum = 0.0f32;
    25
    26
    0
            for sb_idx in 0..num_blocks_per_row {
    27
    0
                let sb_start = row_start + sb_idx * SUPER_BLOCK_BYTES;
    28
    0
                if sb_start + SUPER_BLOCK_BYTES > q6k_data.len() {
    29
    0
                    break;
    30
    0
                }
    31
    0
                let sb_data = &q6k_data[sb_start..sb_start + SUPER_BLOCK_BYTES];
    32
    33
                // Parse Q6K block: ql[128], qh[64], scales[16], d[2]
    34
    0
                let ql = &sb_data[0..128];
    35
    0
                let qh = &sb_data[128..192];
    36
    0
                let scales = &sb_data[192..208];
    37
    0
                let d = f16_to_f32(u16::from_le_bytes([sb_data[208], sb_data[209]]));
    38
    39
    0
                let input_offset = sb_idx * SUPER_BLOCK_SIZE;
    40
    41
                // Process 16 groups of 16 values each (256 total)
    42
    0
                for group in 0..16 {
    43
    0
                    let scale = (scales[group] as i8) as f32;
    44
    0
                    let group_offset = group * 16;
    45
    46
    0
                    for j in 0..16 {
    47
    0
                        let idx = group_offset + j;
    48
    0
                        let input_idx = input_offset + idx;
    49
    0
                        if input_idx >= in_dim {
    50
    0
                            continue;
    51
    0
                        }
    52
    53
                        // Extract 6-bit value: 4 low bits from ql, 2 high bits from qh
    54
    0
                        let ql_byte = ql[idx / 2];
    55
    0
                        let low4 = if idx % 2 == 0 {
    56
    0
                            ql_byte & 0x0F
    57
                        } else {
    58
    0
                            ql_byte >> 4
    59
                        };
    60
    61
                        // qh is packed: 4 values per byte (2 bits each)
    62
    0
                        let qh_byte = qh[idx / 4];
    63
    0
                        let qh_shift = (idx % 4) * 2;
    64
    0
                        let high2 = (qh_byte >> qh_shift) & 0x03;
    65
    66
                        // Combine to 6-bit value (0-63) then center to signed (-32 to 31)
    67
    0
                        let q6 = (low4 | (high2 << 4)) as i8 - 32;
    68
    69
                        // Dequantize: d * scale * q6
    70
    0
                        let dequant = d * scale * q6 as f32;
    71
    0
                        sum += dequant * input[input_idx];
    72
                    }
    73
                }
    74
            }
    75
    76
    0
            output[out_idx] = sum;
    77
        }
    78
    79
    0
        output
    80
    0
    }
    81
    82
    /// Fused Q6_K matrix-vector multiply with AVX2 SIMD
    83
    ///
    84
    /// Optimized to process groups of 8 values at a time, computing
    85
    /// dequant and dot product in one pass without intermediate buffer.
    86
    #[cfg(target_arch = "x86_64")]
    87
    #[target_feature(enable = "avx2", enable = "fma")]
    88
    0
    unsafe fn matmul_q6k_f32_avx2(
    89
    0
        q6k_data: &[u8],
    90
    0
        input: &[f32],
    91
    0
        out_dim: usize,
    92
    0
        in_dim: usize,
    93
    0
    ) -> Vec<f32> {
    94
        #[cfg(target_arch = "x86_64")]
    95
        use std::arch::x86_64::*;
    96
    97
    0
        let num_blocks_per_row = (in_dim + SUPER_BLOCK_SIZE - 1) / SUPER_BLOCK_SIZE;
    98
    0
        let row_bytes = num_blocks_per_row * SUPER_BLOCK_BYTES;
    99
    100
    0
        let mut output = vec![0.0f32; out_dim];
    101
    102
    0
        for out_idx in 0..out_dim {
    103
    0
            let row_start = out_idx * row_bytes;
    104
    0
            let mut acc = _mm256_setzero_ps();
    105
    106
    0
            for sb_idx in 0..num_blocks_per_row {
    107
    0
                let sb_start = row_start + sb_idx * SUPER_BLOCK_BYTES;
    108
    0
                if sb_start + SUPER_BLOCK_BYTES > q6k_data.len() {
    109
    0
                    break;
    110
    0
                }
    111
    0
                let sb_data = &q6k_data[sb_start..sb_start + SUPER_BLOCK_BYTES];
    112
    113
                // Parse Q6K block
    114
    0
                let ql = &sb_data[0..128];
    115
    0
                let qh = &sb_data[128..192];
    116
    0
                let scales = &sb_data[192..208];
    117
    0
                let d = f16_to_f32(u16::from_le_bytes([sb_data[208], sb_data[209]]));
    118
    119
    0
                let input_offset = sb_idx * SUPER_BLOCK_SIZE;
    120
    0
                let d_vec = _mm256_set1_ps(d);
    121
    122
                // Process each group of 16 values (scale is constant per group)
    123
    0
                for group in 0..16 {
    124
    0
                    let scale = (scales[group] as i8) as f32;
    125
    0
                    let scale_vec = _mm256_set1_ps(scale);
    126
    0
                    let ds_vec = _mm256_mul_ps(d_vec, scale_vec);
    127
    0
                    let group_offset = group * 16;
    128
    0
                    let input_group = input_offset + group_offset;
    129
    130
                    // Process 8 values at a time (2 iterations per group of 16)
    131
    0
                    for half in 0..2 {
    132
    0
                        let half_offset = half * 8;
    133
    0
                        let idx_base = group_offset + half_offset;
    134
    0
                        let input_base = input_group + half_offset;
    135
    136
    0
                        if input_base + 8 > in_dim {
    137
    0
                            continue;
    138
    0
                        }
    139
    140
                        // Extract 8 quantized values
    141
                        // Q6 value = (ql_low4 | (qh_2bit << 4)) - 32
    142
    0
                        let mut q6_vals = [0i32; 8];
    143
    0
                        for i in 0..8 {
    144
    0
                            let idx = idx_base + i;
    145
    0
                            let ql_byte = ql[idx / 2];
    146
    0
                            let low4 = if idx % 2 == 0 {
    147
    0
                                ql_byte & 0x0F
    148
                            } else {
    149
    0
                                ql_byte >> 4
    150
                            };
    151
    0
                            let qh_byte = qh[idx / 4];
    152
    0
                            let qh_shift = (idx % 4) * 2;
    153
    0
                            let high2 = (qh_byte >> qh_shift) & 0x03;
    154
    0
                            q6_vals[i] = ((low4 | (high2 << 4)) as i32) - 32;
    155
                        }
    156
    157
                        // Load into SIMD
    158
    0
                        let q6_i32 = _mm256_loadu_si256(q6_vals.as_ptr() as *const __m256i);
    159
    0
                        let q6_f32 = _mm256_cvtepi32_ps(q6_i32);
    160
    161
                        // Load input
    162
    0
                        let x = _mm256_loadu_ps(input.as_ptr().add(input_base));
    163
    164
                        // Compute: acc += (d * scale * q6) * x
    165
    0
                        let dequant = _mm256_mul_ps(ds_vec, q6_f32);
    166
    0
                        acc = _mm256_fmadd_ps(dequant, x, acc);
    167
                    }
    168
                }
    169
            }
    170
    171
            // Horizontal sum
    172
    0
            let hi128 = _mm256_extractf128_ps(acc, 1);
    173
    0
            let lo128 = _mm256_castps256_ps128(acc);
    174
    0
            let sum128 = _mm_add_ps(lo128, hi128);
    175
    0
            let hi64 = _mm_movehl_ps(sum128, sum128);
    176
    0
            let sum64 = _mm_add_ps(sum128, hi64);
    177
    0
            let hi32 = _mm_shuffle_ps(sum64, sum64, 1);
    178
    0
            let sum32 = _mm_add_ss(sum64, hi32);
    179
    180
    0
            output[out_idx] = _mm_cvtss_f32(sum32);
    181
        }
    182
    183
    0
        output
    184
    0
    }
    185
    186
    /// Runtime dispatch for Q6K matmul - uses AVX2 if available
    187
    #[inline]
    188
    0
    pub fn matmul_q6k_f32_dispatch(
    189
    0
        q6k_data: &[u8],
    190
    0
        input: &[f32],
    191
    0
        out_dim: usize,
    192
    0
        in_dim: usize,
    193
    0
    ) -> Vec<f32> {
    194
        // For large matmuls (total work >= ~8M ops), use parallel execution
    195
        // This catches FFN layers (8960x1536) and lm_head (151936x1536)
    196
        // Also catches ffn_down (1536x8960) where out_dim is small but in_dim is large
    197
    0
        let total_work = out_dim * in_dim;
    198
    0
        if total_work >= 8_000_000 {
    199
    0
            return matmul_q6k_f32_parallel(q6k_data, input, out_dim, in_dim);
    200
    0
        }
    201
    202
        #[cfg(target_arch = "x86_64")]
    203
        {
    204
    0
            if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
    205
    0
                return unsafe { matmul_q6k_f32_avx2(q6k_data, input, out_dim, in_dim) };
    206
    0
            }
    207
        }
    208
    0
        matmul_q6k_f32_scalar(q6k_data, input, out_dim, in_dim)
    209
    0
    }
    210
    211
    /// Parallel Q6K matmul using multiple threads with AVX2
    212
    #[cfg(target_arch = "x86_64")]
    213
    0
    fn matmul_q6k_f32_parallel(
    214
    0
        q6k_data: &[u8],
    215
    0
        input: &[f32],
    216
    0
        out_dim: usize,
    217
    0
        in_dim: usize,
    218
    0
    ) -> Vec<f32> {
    219
        use std::thread;
    220
    221
        // Use fewer threads with larger chunks for better cache efficiency
    222
    0
        let num_threads = thread::available_parallelism()
    223
    0
            .map(|p| p.get())
    224
    0
            .unwrap_or(4)
    225
    0
            .min(12); // Use 12 threads max for better cache behavior
    226
    227
    0
        let chunk_size = (out_dim + num_threads - 1) / num_threads;
    228
    0
        let num_blocks_per_row = (in_dim + SUPER_BLOCK_SIZE - 1) / SUPER_BLOCK_SIZE;
    229
    0
        let row_bytes = num_blocks_per_row * SUPER_BLOCK_BYTES;
    230
    231
    0
        let mut output = vec![0.0f32; out_dim];
    232
    0
        let has_avx2 = is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma");
    233
    234
    0
        thread::scope(|s| {
    235
    0
            let input_ref = input;
    236
    0
            let q6k_ref = q6k_data;
    237
    0
            let chunks: Vec<_> = output.chunks_mut(chunk_size).enumerate().collect();
    238
    239
    0
            for (chunk_idx, chunk) in chunks {
    240
    0
                let start_row = chunk_idx * chunk_size;
    241
    242
    0
                s.spawn(move || {
    243
    0
                    if has_avx2 {
    244
                        // Call AVX2 kernel for this chunk
    245
    0
                        unsafe {
    246
    0
                            compute_chunk_avx2(
    247
    0
                                q6k_ref,
    248
    0
                                input_ref,
    249
    0
                                chunk,
    250
    0
                                start_row,
    251
    0
                                out_dim,
    252
    0
                                in_dim,
    253
    0
                                num_blocks_per_row,
    254
    0
                                row_bytes,
    255
    0
                            );
    256
    0
                        }
    257
    0
                    } else {
    258
    0
                        compute_chunk_scalar(
    259
    0
                            q6k_ref,
    260
    0
                            input_ref,
    261
    0
                            chunk,
    262
    0
                            start_row,
    263
    0
                            out_dim,
    264
    0
                            in_dim,
    265
    0
                            num_blocks_per_row,
    266
    0
                            row_bytes,
    267
    0
                        );
    268
    0
                    }
    269
    0
                });
    270
            }
    271
    0
        });
    272
    273
    0
        output
    274
    0
    }
    275
    276
    /// Fallback for non-x86_64
    277
    #[cfg(not(target_arch = "x86_64"))]
    278
    fn matmul_q6k_f32_parallel(
    279
        q6k_data: &[u8],
    280
        input: &[f32],
    281
        out_dim: usize,
    282
        in_dim: usize,
    283
    ) -> Vec<f32> {
    284
        matmul_q6k_f32_scalar(q6k_data, input, out_dim, in_dim)
    285
    }
    286
    287
    #[cfg(target_arch = "x86_64")]
    288
    #[target_feature(enable = "avx2", enable = "fma")]
    289
    0
    unsafe fn compute_chunk_avx2(
    290
    0
        q6k_data: &[u8],
    291
    0
        input: &[f32],
    292
    0
        chunk: &mut [f32],
    293
    0
        start_row: usize,
    294
    0
        out_dim: usize,
    295
    0
        in_dim: usize,
    296
    0
        num_blocks_per_row: usize,
    297
    0
        row_bytes: usize,
    298
    0
    ) {
    299
        use std::arch::x86_64::*;
    300
    301
    0
        for (local_idx, out_val) in chunk.iter_mut().enumerate() {
    302
    0
            let out_idx = start_row + local_idx;
    303
    0
            if out_idx >= out_dim {
    304
    0
                break;
    305
    0
            }
    306
    307
    0
            let row_start = out_idx * row_bytes;
    308
    0
            let mut acc = _mm256_setzero_ps();
    309
    310
    0
            for sb_idx in 0..num_blocks_per_row {
    311
    0
                let sb_start = row_start + sb_idx * SUPER_BLOCK_BYTES;
    312
    0
                if sb_start + SUPER_BLOCK_BYTES > q6k_data.len() {
    313
    0
                    break;
    314
    0
                }
    315
    0
                let sb_data = &q6k_data[sb_start..sb_start + SUPER_BLOCK_BYTES];
    316
    317
    0
                let ql = &sb_data[0..128];
    318
    0
                let qh = &sb_data[128..192];
    319
    0
                let scales = &sb_data[192..208];
    320
    0
                let d = f16_to_f32(u16::from_le_bytes([sb_data[208], sb_data[209]]));
    321
    322
    0
                let input_offset = sb_idx * SUPER_BLOCK_SIZE;
    323
    0
                let d_vec = _mm256_set1_ps(d);
    324
    325
    0
                for group in 0..16 {
    326
    0
                    let scale = (scales[group] as i8) as f32;
    327
    0
                    let scale_vec = _mm256_set1_ps(scale);
    328
    0
                    let ds_vec = _mm256_mul_ps(d_vec, scale_vec);
    329
    0
                    let group_offset = group * 16;
    330
    0
                    let input_group = input_offset + group_offset;
    331
    332
    0
                    for half in 0..2 {
    333
    0
                        let half_offset = half * 8;
    334
    0
                        let idx_base = group_offset + half_offset;
    335
    0
                        let input_base = input_group + half_offset;
    336
    337
    0
                        if input_base + 8 > in_dim {
    338
    0
                            continue;
    339
    0
                        }
    340
    341
                        // Extract 8 quantized values
    342
    0
                        let mut q6_vals = [0i32; 8];
    343
    0
                        for i in 0..8 {
    344
    0
                            let idx = idx_base + i;
    345
    0
                            let ql_byte = ql[idx / 2];
    346
    0
                            let low4 = if idx % 2 == 0 {
    347
    0
                                ql_byte & 0x0F
    348
                            } else {
    349
    0
                                ql_byte >> 4
    350
                            };
    351
    0
                            let qh_byte = qh[idx / 4];
    352
    0
                            let qh_shift = (idx % 4) * 2;
    353
    0
                            let high2 = (qh_byte >> qh_shift) & 0x03;
    354
    0
                            q6_vals[i] = ((low4 | (high2 << 4)) as i32) - 32;
    355
                        }
    356
    357
    0
                        let q6_i32 = _mm256_loadu_si256(q6_vals.as_ptr() as *const __m256i);
    358
    0
                        let q6_f32 = _mm256_cvtepi32_ps(q6_i32);
    359
    0
                        let x = _mm256_loadu_ps(input.as_ptr().add(input_base));
    360
    0
                        let dequant = _mm256_mul_ps(ds_vec, q6_f32);
    361
    0
                        acc = _mm256_fmadd_ps(dequant, x, acc);
    362
                    }
    363
                }
    364
            }
    365
    366
            // Horizontal sum
    367
    0
            let hi128 = _mm256_extractf128_ps(acc, 1);
    368
    0
            let lo128 = _mm256_castps256_ps128(acc);
    369
    0
            let sum128 = _mm_add_ps(lo128, hi128);
    370
    0
            let hi64 = _mm_movehl_ps(sum128, sum128);
    371
    0
            let sum64 = _mm_add_ps(sum128, hi64);
    372
    0
            let hi32 = _mm_shuffle_ps(sum64, sum64, 1);
    373
    0
            let sum32 = _mm_add_ss(sum64, hi32);
    374
    375
    0
            *out_val = _mm_cvtss_f32(sum32);
    376
        }
    377
    0
    }
    378
    379
    #[allow(dead_code)]
    380
    0
    pub(crate) fn compute_chunk_scalar(
    381
    0
        q6k_data: &[u8],
    382
    0
        input: &[f32],
    383
    0
        chunk: &mut [f32],
    384
    0
        start_row: usize,
    385
    0
        out_dim: usize,
    386
    0
        in_dim: usize,
    387
    0
        num_blocks_per_row: usize,
    388
    0
        row_bytes: usize,
    389
    0
    ) {
    390
    0
        for (local_idx, out_val) in chunk.iter_mut().enumerate() {
    391
    0
            let out_idx = start_row + local_idx;
    392
    0
            if out_idx >= out_dim {
    393
    0
                break;
    394
    0
            }
    395
    396
    0
            let row_start = out_idx * row_bytes;
    397
    0
            let mut sum = 0.0f32;
    398
    399
    0
            for sb_idx in 0..num_blocks_per_row {
    400
    0
                let sb_start = row_start + sb_idx * SUPER_BLOCK_BYTES;
    401
    0
                if sb_start + SUPER_BLOCK_BYTES > q6k_data.len() {
    402
    0
                    break;
    403
    0
                }
    404
    0
                let sb_data = &q6k_data[sb_start..sb_start + SUPER_BLOCK_BYTES];
    405
    406
    0
                let ql = &sb_data[0..128];
    407
    0
                let qh = &sb_data[128..192];
    408
    0
                let scales = &sb_data[192..208];
    409
    0
                let d = f16_to_f32(u16::from_le_bytes([sb_data[208], sb_data[209]]));
    410
    411
    0
                let input_offset = sb_idx * SUPER_BLOCK_SIZE;
    412
    413
    0
                for group in 0..16 {
    414
    0
                    let scale = (scales[group] as i8) as f32;
    415
    0
                    let group_offset = group * 16;
    416
    417
    0
                    for j in 0..16 {
    418
    0
                        let idx = group_offset + j;
    419
    0
                        let input_idx = input_offset + idx;
    420
    0
                        if input_idx >= in_dim {
    421
    0
                            continue;
    422
    0
                        }
    423
    424
    0
                        let ql_byte = ql[idx / 2];
    425
    0
                        let low4 = if idx % 2 == 0 {
    426
    0
                            ql_byte & 0x0F
    427
                        } else {
    428
    0
                            ql_byte >> 4
    429
                        };
    430
    431
    0
                        let qh_byte = qh[idx / 4];
    432
    0
                        let qh_shift = (idx % 4) * 2;
    433
    0
                        let high2 = (qh_byte >> qh_shift) & 0x03;
    434
    435
    0
                        let q6 = (low4 | (high2 << 4)) as i8 - 32;
    436
    0
                        let dequant = d * scale * q6 as f32;
    437
    0
                        sum += dequant * input[input_idx];
    438
                    }
    439
                }
    440
            }
    441
    442
    0
            *out_val = sum;
    443
        }
    444
    0
    }
    445
    446
    /// Public alias for the optimized Q6K matmul
    447
    0
    pub fn matmul_q6k_f32(
    448
    0
        q6k_data: &[u8],
    449
    0
        input: &[f32],
    450
    0
        out_dim: usize,
    451
    0
        in_dim: usize,
    452
    0
    ) -> Vec<f32> {
    453
    0
        matmul_q6k_f32_dispatch(q6k_data, input, out_dim, in_dim)
    454
    0
    }