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

    Line
    Count
    Source
    1
    //! Parallel and tiled matrix-vector operations for K-quantization (PMAT-802)
    2
    //!
    3
    //! Implements L2-aware tiled and parallel matvec operations:
    4
    //! - `fused_q4k_tiled_matvec` - L2-aware tiled matmul
    5
    //! - `fused_q4k_parallel_matvec`, `fused_q4k_parallel_matvec_into` - Parallel Q4_K
    6
    //! - `fused_q5k_parallel_matvec`, `fused_q5k_parallel_matvec_into` - Parallel Q5_K
    7
    //! - `fused_q6k_parallel_matvec`, `fused_q6k_parallel_matvec_into` - Parallel Q6_K
    8
    //!
    9
    //! Per Goto & Van Geijn "Anatomy of High-Performance Matrix Multiplication":
    10
    //! - GEBP (General Block Panel) tiling maximizes cache reuse
    11
    //! - Tile size should fit in L2 cache (~256KB-512KB typically)
    12
    13
    use crate::error::{RealizarError, Result};
    14
    use super::fused_k::{fused_q4k_dot_simd, fused_q4k_q8k_dot_simd, fused_q4k_q8k_dot_4rows_avx512vnni};
    15
    use super::fused_q5k_q6k::{fused_q5k_dot_simd, fused_q6k_dot_simd};
    16
    use super::types::QK_K;
    17
    18
    // ============================================================================
    19
    20
    /// Default tile size for L2-aware tiled matmul
    21
    ///
    22
    /// Chosen to fit in L2 cache while maximizing parallelism:
    23
    /// - Typical L2 size: 256KB-512KB
    24
    /// - Q4_K row size for hidden_dim=2560: ~1440 bytes
    25
    /// - 64 rows = ~92KB of weight data, plus activations
    26
    const DEFAULT_OUTPUT_TILE_SIZE: usize = 64;
    27
    28
    /// Fused Q4_K matrix-vector multiply with L2-aware tiling
    29
    ///
    30
    /// Processes outputs in tiles to maximize L2 cache reuse.
    31
    /// Each tile loads weight data once and computes multiple outputs.
    32
    ///
    33
    /// # Arguments
    34
    ///
    35
    /// * `weight_data` - Raw Q4_K quantized weight data
    36
    /// * `activations` - Input activations [in_dim]
    37
    /// * `in_dim` - Input dimension (must be multiple of 256 for Q4_K)
    38
    /// * `out_dim` - Output dimension
    39
    /// * `tile_size` - Number of outputs to process per tile (default: 64)
    40
    ///
    41
    /// # Returns
    42
    ///
    43
    /// Output vector [out_dim]
    44
    ///
    45
    /// # Errors
    46
    ///
    47
    /// Returns error if dimensions don't match weight data
    48
    ///
    49
    /// # Performance
    50
    ///
    51
    /// - **L2-aware**: Tiles fit in L2 cache, reducing DRAM traffic
    52
    /// - **Fused**: Dequantize inline with dot product (8x bandwidth reduction)
    53
    /// - **SIMD**: Uses AVX2 when available for 4-8x compute speedup
    54
    #[allow(clippy::similar_names)]
    55
    1.23k
    pub fn fused_q4k_tiled_matvec(
    56
    1.23k
        weight_data: &[u8],
    57
    1.23k
        activations: &[f32],
    58
    1.23k
        in_dim: usize,
    59
    1.23k
        out_dim: usize,
    60
    1.23k
        tile_size: Option<usize>,
    61
    1.23k
    ) -> Result<Vec<f32>> {
    62
    1.23k
        let tile_size = tile_size.unwrap_or(DEFAULT_OUTPUT_TILE_SIZE);
    63
    64
        // Calculate bytes per output row
    65
    1.23k
        let super_blocks_per_row = in_dim.div_ceil(QK_K);
    66
    1.23k
        let bytes_per_row = super_blocks_per_row * 144; // Q4_K: 144 bytes per super-block
    67
    68
        // Validate dimensions
    69
    1.23k
        let expected_weight_bytes = out_dim * bytes_per_row;
    70
    1.23k
        if weight_data.len() < expected_weight_bytes {
    71
    3
            return Err(RealizarError::InvalidShape {
    72
    3
                reason: format!(
    73
    3
                    "Q4_K weight data too small: need {} bytes for {}x{}, have {}",
    74
    3
                    expected_weight_bytes,
    75
    3
                    out_dim,
    76
    3
                    in_dim,
    77
    3
                    weight_data.len()
    78
    3
                ),
    79
    3
            });
    80
    1.23k
        }
    81
    82
    1.23k
        if activations.len() != in_dim {
    83
    2
            return Err(RealizarError::InvalidShape {
    84
    2
                reason: format!(
    85
    2
                    "Activation length {} doesn't match in_dim {}",
    86
    2
                    activations.len(),
    87
    2
                    in_dim
    88
    2
                ),
    89
    2
            });
    90
    1.22k
        }
    91
    92
    1.22k
        let mut output = vec![0.0f32; out_dim];
    93
    94
        // Process outputs in tiles for L2 cache efficiency
    95
    1.22k
        let num_tiles = out_dim.div_ceil(tile_size);
    96
    97
    16.3k
        for tile_idx in 0..num_tiles1.22k {
    98
    16.3k
            let tile_start = tile_idx * tile_size;
    99
    16.3k
            let tile_end = (tile_start + tile_size).min(out_dim);
    100
    101
            // Prefetch next tile's weight data (if available)
    102
            #[cfg(target_arch = "x86_64")]
    103
    16.3k
            if tile_idx + 1 < num_tiles {
    104
    15.0k
                let next_tile_start = (tile_idx + 1) * tile_size;
    105
    15.0k
                let next_row_start = next_tile_start * bytes_per_row;
    106
    15.0k
                if next_row_start < weight_data.len() {
    107
                    // SAFETY: Prefetch is a hint, no memory safety requirements
    108
                    unsafe {
    109
                        use std::arch::x86_64::_mm_prefetch;
    110
                        use std::arch::x86_64::_MM_HINT_T0;
    111
    15.0k
                        let ptr = weight_data.as_ptr().add(next_row_start);
    112
    15.0k
                        _mm_prefetch(ptr.cast::<i8>(), _MM_HINT_T0);
    113
                    }
    114
    0
                }
    115
    1.22k
            }
    116
    117
            // Process tile: compute dot products for tile_start..tile_end
    118
    1.03M
            for (idx, out_slot) in output16.3k[tile_start..tile_end].iter_mut16.3k().enumerate16.3k() {
    119
    1.03M
                let o = tile_start + idx;
    120
    1.03M
                let row_start = o * bytes_per_row;
    121
    1.03M
                let row_end = row_start + bytes_per_row;
    122
    1.03M
                let row_data = &weight_data[row_start..row_end];
    123
    124
                // Fused dequant + dot product
    125
    1.03M
                *out_slot = fused_q4k_dot_simd(row_data, activations)?0;
    126
            }
    127
        }
    128
    129
    1.22k
        Ok(output)
    130
    1.23k
    }
    131
    132
    // ============================================================================
    133
    // PARALLEL TILED MATRIX-VECTOR MULTIPLICATION (Phase 2 + 3)
    134
    // ============================================================================
    135
    //
    136
    // Per Blumofe & Leiserson [6] "Scheduling Multithreaded Computations by Work Stealing":
    137
    // - Work-stealing schedulers like rayon maximize CPU utilization
    138
    // - Each output row is independent → trivially parallelizable
    139
    // - Expected speedup: ~Nx on N-core systems for memory-bound workloads
    140
    // ============================================================================
    141
    142
    /// Parallel fused Q4_K matrix-vector multiply with L2-aware tiling
    143
    ///
    144
    /// Uses rayon parallel iterators for multi-core acceleration.
    145
    /// Per Valiant's BSP model [14], synchronization happens at tile boundaries.
    146
    ///
    147
    /// # Performance
    148
    ///
    149
    /// - **Multi-core**: Linear speedup up to memory bandwidth saturation
    150
    /// - **L2-aware**: Tiles fit in L2 cache
    151
    /// - **Fused**: 8x memory bandwidth reduction
    152
    /// - **SIMD**: AVX2 when available
    153
    /// - **Adaptive parallelism**: Sequential for small matrices, parallel for large (IMP-103)
    154
    ///
    155
    /// # Errors
    156
    ///
    157
    /// Returns error if:
    158
    /// - Weight data is too small for the given dimensions
    159
    /// - Activation length doesn't match input dimension
    160
    #[allow(clippy::similar_names)]
    161
    3.69k
    pub fn fused_q4k_parallel_matvec(
    162
    3.69k
        weight_data: &[u8],
    163
    3.69k
        activations: &[f32],
    164
    3.69k
        in_dim: usize,
    165
    3.69k
        out_dim: usize,
    166
    3.69k
    ) -> Result<Vec<f32>> {
    167
        // PAR-126: Five-Whys fix - parallel threshold was too high
    168
        // OLD: PARALLEL_THRESHOLD=4096 meant FFN down (out_dim=1536) used sequential path
    169
        // PROBLEM: 1.5B model was 11 tok/s instead of 200 tok/s due to single-threaded matmuls
    170
        //
    171
        // ANALYSIS (for 32-core system with in_dim=8960):
    172
        // - Per-row time: 8960/256 superblocks × ~50ns/superblock = ~1.75µs
    173
        // - Rayon overhead: ~10µs (reduced with work-stealing)
    174
        // - Break-even: 10µs / (1.75µs/32) = ~183 rows
    175
        // SOLUTION: Lower threshold to 256 to enable parallelism for all practical matmuls
    176
        const PARALLEL_THRESHOLD: usize = 256;
    177
    178
        // Calculate bytes per output row
    179
    3.69k
        let super_blocks_per_row = in_dim.div_ceil(QK_K);
    180
    3.69k
        let bytes_per_row = super_blocks_per_row * 144; // Q4_K: 144 bytes per super-block
    181
    182
        // Validate dimensions
    183
    3.69k
        let expected_weight_bytes = out_dim * bytes_per_row;
    184
    3.69k
        if weight_data.len() < expected_weight_bytes {
    185
    3
            return Err(RealizarError::InvalidShape {
    186
    3
                reason: format!(
    187
    3
                    "Q4_K weight data too small: need {} bytes for {}x{}, have {}",
    188
    3
                    expected_weight_bytes,
    189
    3
                    out_dim,
    190
    3
                    in_dim,
    191
    3
                    weight_data.len()
    192
    3
                ),
    193
    3
            });
    194
    3.69k
        }
    195
    196
    3.69k
        if activations.len() != in_dim {
    197
    1
            return Err(RealizarError::InvalidShape {
    198
    1
                reason: format!(
    199
    1
                    "Activation length {} doesn't match in_dim {}",
    200
    1
                    activations.len(),
    201
    1
                    in_dim
    202
    1
                ),
    203
    1
            });
    204
    3.69k
        }
    205
    206
    3.69k
        if out_dim < PARALLEL_THRESHOLD {
    207
            // Sequential path: avoids rayon overhead for small matrices
    208
    3.37k
            let output: Vec<f32> = (0..out_dim)
    209
    326k
                .map3.37k(|o| {
    210
    326k
                    let row_start = o * bytes_per_row;
    211
    326k
                    let row_end = row_start + bytes_per_row;
    212
    326k
                    let row_data = &weight_data[row_start..row_end];
    213
    326k
                    fused_q4k_dot_simd(row_data, activations).unwrap_or(0.0)
    214
    326k
                })
    215
    3.37k
                .collect();
    216
    217
    3.37k
            Ok(output)
    218
        } else {
    219
            // Parallel path: better for large matrices
    220
            use rayon::prelude::*;
    221
    222
            // Use chunked parallel iteration with optimal chunk size
    223
            // Chunk size tuned for L2 cache (~256KB): process ~64 rows per chunk
    224
            const CHUNK_SIZE: usize = 64;
    225
    226
    316
            let output: Vec<f32> = (0..out_dim)
    227
    316
                .into_par_iter()
    228
    316
                .with_min_len(CHUNK_SIZE)
    229
    134k
                .map316(|o| {
    230
    134k
                    let row_start = o * bytes_per_row;
    231
    134k
                    let row_end = row_start + bytes_per_row;
    232
    134k
                    let row_data = &weight_data[row_start..row_end];
    233
    134k
                    fused_q4k_dot_simd(row_data, activations).unwrap_or(0.0)
    234
    134k
                })
    235
    316
                .collect();
    236
    237
    316
            Ok(output)
    238
        }
    239
    3.69k
    }
    240
    241
    /// Parallel fused Q4_K matrix-vector multiply - writes to pre-allocated buffer
    242
    ///
    243
    /// IMP-131: Zero-allocation variant for hot-path inference.
    244
    /// This avoids Vec allocation overhead that causes 30-40% performance loss.
    245
    ///
    246
    /// # Arguments
    247
    /// * `weight_data` - Raw Q4_K quantized weights [out_dim, in_dim]
    248
    /// * `activations` - Input activations [in_dim]
    249
    /// * `in_dim` - Input dimension (must match activations length)
    250
    /// * `out_dim` - Output dimension (must match output buffer length)
    251
    /// * `output` - Pre-allocated output buffer [out_dim]
    252
    ///
    253
    /// # Errors
    254
    ///
    255
    /// Returns error if:
    256
    /// - Weight data is too small for the given dimensions
    257
    /// - Activation length doesn't match input dimension
    258
    /// - Output buffer length doesn't match out_dim
    259
    #[allow(clippy::similar_names)]
    260
    2
    pub fn fused_q4k_parallel_matvec_into(
    261
    2
        weight_data: &[u8],
    262
    2
        activations: &[f32],
    263
    2
        in_dim: usize,
    264
    2
        out_dim: usize,
    265
    2
        output: &mut [f32],
    266
    2
    ) -> Result<()> {
    267
        // PAR-126: Match threshold from allocating version (was 4096, caused 25% perf loss)
    268
        // Analysis: For 32-core system with in_dim=8960:
    269
        // - Per-row time: ~1.75µs, Rayon overhead: ~10µs
    270
        // - Break-even: ~183 rows, so 256 is safe threshold
    271
        const PARALLEL_THRESHOLD: usize = 256;
    272
    273
    2
        let super_blocks_per_row = in_dim.div_ceil(QK_K);
    274
    2
        let bytes_per_row = super_blocks_per_row * 144;
    275
    276
    2
        let expected_weight_bytes = out_dim * bytes_per_row;
    277
    2
        if weight_data.len() < expected_weight_bytes {
    278
    0
            return Err(RealizarError::InvalidShape {
    279
    0
                reason: format!(
    280
    0
                    "Q4_K weight data too small: need {} bytes for {}x{}, have {}",
    281
    0
                    expected_weight_bytes,
    282
    0
                    out_dim,
    283
    0
                    in_dim,
    284
    0
                    weight_data.len()
    285
    0
                ),
    286
    0
            });
    287
    2
        }
    288
    289
    2
        if activations.len() != in_dim {
    290
    1
            return Err(RealizarError::InvalidShape {
    291
    1
                reason: format!(
    292
    1
                    "Activation length {} doesn't match in_dim {}",
    293
    1
                    activations.len(),
    294
    1
                    in_dim
    295
    1
                ),
    296
    1
            });
    297
    1
        }
    298
    299
    1
        if output.len() < out_dim {
    300
    0
            return Err(RealizarError::InvalidShape {
    301
    0
                reason: format!(
    302
    0
                    "Output buffer too small: need {}, have {}",
    303
    0
                    out_dim,
    304
    0
                    output.len()
    305
    0
                ),
    306
    0
            });
    307
    1
        }
    308
    309
    1
        if out_dim < PARALLEL_THRESHOLD {
    310
            // Sequential path
    311
    32
            for o in 0..out_dim1 {
    312
    32
                let row_start = o * bytes_per_row;
    313
    32
                let row_end = row_start + bytes_per_row;
    314
    32
                let row_data = &weight_data[row_start..row_end];
    315
    32
                output[o] = fused_q4k_dot_simd(row_data, activations).unwrap_or(0.0);
    316
    32
            }
    317
        } else {
    318
            // Parallel path with TCB-style midi-tile chunking
    319
            // Process rows in 64-row chunks to maximize activation cache reuse
    320
            use rayon::prelude::*;
    321
            const MIDI_TILE_M: usize = 64;
    322
    323
    0
            output[..out_dim]
    324
    0
                .par_chunks_mut(MIDI_TILE_M)
    325
    0
                .enumerate()
    326
    0
                .for_each(|(midi_idx, midi_chunk)| {
    327
    0
                    let midi_start = midi_idx * MIDI_TILE_M;
    328
    329
                    // Process each row in this midi-tile
    330
                    // All rows share the same activation vector (kept in L2 cache)
    331
    0
                    for (local_idx, out) in midi_chunk.iter_mut().enumerate() {
    332
    0
                        let row = midi_start + local_idx;
    333
    0
                        let row_start = row * bytes_per_row;
    334
    0
                        let row_end = row_start + bytes_per_row;
    335
    0
                        let row_data = &weight_data[row_start..row_end];
    336
    0
                        *out = fused_q4k_dot_simd(row_data, activations).unwrap_or(0.0);
    337
    0
                    }
    338
    0
                });
    339
        }
    340
    341
    1
        Ok(())
    342
    2
    }
    343
    344
    /// Parallel fused Q5_K matrix-vector multiply
    345
    ///
    346
    /// # Errors
    347
    ///
    348
    /// Returns error if:
    349
    /// - Weight data is too small for the given dimensions
    350
    /// - Activation length doesn't match input dimension
    351
    #[allow(clippy::similar_names)]
    352
    7
    pub fn fused_q5k_parallel_matvec(
    353
    7
        weight_data: &[u8],
    354
    7
        activations: &[f32],
    355
    7
        in_dim: usize,
    356
    7
        out_dim: usize,
    357
    7
    ) -> Result<Vec<f32>> {
    358
        use rayon::prelude::*;
    359
    360
    7
        let super_blocks_per_row = in_dim.div_ceil(QK_K);
    361
    7
        let bytes_per_row = super_blocks_per_row * 176; // Q5_K: 176 bytes per super-block
    362
    363
    7
        let expected_weight_bytes = out_dim * bytes_per_row;
    364
    7
        if weight_data.len() < expected_weight_bytes {
    365
    0
            return Err(RealizarError::InvalidShape {
    366
    0
                reason: format!(
    367
    0
                    "Q5_K weight data too small: need {} bytes for {}x{}, have {}",
    368
    0
                    expected_weight_bytes,
    369
    0
                    out_dim,
    370
    0
                    in_dim,
    371
    0
                    weight_data.len()
    372
    0
                ),
    373
    0
            });
    374
    7
        }
    375
    376
    7
        if activations.len() != in_dim {
    377
    0
            return Err(RealizarError::InvalidShape {
    378
    0
                reason: format!(
    379
    0
                    "Activation length {} doesn't match in_dim {}",
    380
    0
                    activations.len(),
    381
    0
                    in_dim
    382
    0
                ),
    383
    0
            });
    384
    7
        }
    385
    386
    7
        let output: Vec<f32> = (0..out_dim)
    387
    7
            .into_par_iter()
    388
    2.11k
            .map7(|o| {
    389
    2.11k
                let row_start = o * bytes_per_row;
    390
    2.11k
                let row_end = row_start + bytes_per_row;
    391
    2.11k
                let row_data = &weight_data[row_start..row_end];
    392
    393
    2.11k
                fused_q5k_dot_simd(row_data, activations).unwrap_or(0.0)
    394
    2.11k
            })
    395
    7
            .collect();
    396
    397
    7
        Ok(output)
    398
    7
    }
    399
    400
    /// Parallel fused Q5_K matrix-vector multiply - writes to pre-allocated buffer
    401
    ///
    402
    /// IMP-131: Zero-allocation variant for hot-path inference.
    403
    #[allow(clippy::similar_names)]
    404
    3
    pub fn fused_q5k_parallel_matvec_into(
    405
    3
        weight_data: &[u8],
    406
    3
        activations: &[f32],
    407
    3
        in_dim: usize,
    408
    3
        out_dim: usize,
    409
    3
        output: &mut [f32],
    410
    3
    ) -> Result<()> {
    411
        use rayon::prelude::*;
    412
    413
    3
        let super_blocks_per_row = in_dim.div_ceil(QK_K);
    414
    3
        let bytes_per_row = super_blocks_per_row * 176;
    415
    416
    3
        let expected_weight_bytes = out_dim * bytes_per_row;
    417
    3
        if weight_data.len() < expected_weight_bytes {
    418
    2
            return Err(RealizarError::InvalidShape {
    419
    2
                reason: format!(
    420
    2
                    "Q5_K weight data too small: need {} bytes for {}x{}, have {}",
    421
    2
                    expected_weight_bytes,
    422
    2
                    out_dim,
    423
    2
                    in_dim,
    424
    2
                    weight_data.len()
    425
    2
                ),
    426
    2
            });
    427
    1
        }
    428
    429
    1
        if activations.len() != in_dim {
    430
    1
            return Err(RealizarError::InvalidShape {
    431
    1
                reason: format!(
    432
    1
                    "Activation length {} doesn't match in_dim {}",
    433
    1
                    activations.len(),
    434
    1
                    in_dim
    435
    1
                ),
    436
    1
            });
    437
    0
        }
    438
    439
    0
        if output.len() < out_dim {
    440
    0
            return Err(RealizarError::InvalidShape {
    441
    0
                reason: format!(
    442
    0
                    "Output buffer too small: need {}, have {}",
    443
    0
                    out_dim,
    444
    0
                    output.len()
    445
    0
                ),
    446
    0
            });
    447
    0
        }
    448
    449
    0
        output[..out_dim]
    450
    0
            .par_iter_mut()
    451
    0
            .enumerate()
    452
    0
            .for_each(|(o, out)| {
    453
    0
                let row_start = o * bytes_per_row;
    454
    0
                let row_end = row_start + bytes_per_row;
    455
    0
                let row_data = &weight_data[row_start..row_end];
    456
    0
                *out = fused_q5k_dot_simd(row_data, activations).unwrap_or(0.0);
    457
    0
            });
    458
    459
    0
        Ok(())
    460
    3
    }
    461
    462
    /// Parallel fused Q6_K matrix-vector multiply
    463
    ///
    464
    /// # Errors
    465
    ///
    466
    /// Returns error if:
    467
    /// - Weight data is too small for the given dimensions
    468
    /// - Activation length doesn't match input dimension
    469
    #[allow(clippy::similar_names)]
    470
    8
    pub fn fused_q6k_parallel_matvec(
    471
    8
        weight_data: &[u8],
    472
    8
        activations: &[f32],
    473
    8
        in_dim: usize,
    474
    8
        out_dim: usize,
    475
    8
    ) -> Result<Vec<f32>> {
    476
        use rayon::prelude::*;
    477
    478
    8
        let super_blocks_per_row = in_dim.div_ceil(QK_K);
    479
    8
        let bytes_per_row = super_blocks_per_row * 210; // Q6_K: 210 bytes per super-block
    480
    481
    8
        let expected_weight_bytes = out_dim * bytes_per_row;
    482
    8
        if weight_data.len() < expected_weight_bytes {
    483
    0
            return Err(RealizarError::InvalidShape {
    484
    0
                reason: format!(
    485
    0
                    "Q6_K weight data too small: need {} bytes for {}x{}, have {}",
    486
    0
                    expected_weight_bytes,
    487
    0
                    out_dim,
    488
    0
                    in_dim,
    489
    0
                    weight_data.len()
    490
    0
                ),
    491
    0
            });
    492
    8
        }
    493
    494
    8
        if activations.len() != in_dim {
    495
    0
            return Err(RealizarError::InvalidShape {
    496
    0
                reason: format!(
    497
    0
                    "Activation length {} doesn't match in_dim {}",
    498
    0
                    activations.len(),
    499
    0
                    in_dim
    500
    0
                ),
    501
    0
            });
    502
    8
        }
    503
    504
    8
        let output: Vec<f32> = (0..out_dim)
    505
    8
            .into_par_iter()
    506
    2.14k
            .map8(|o| {
    507
    2.14k
                let row_start = o * bytes_per_row;
    508
    2.14k
                let row_end = row_start + bytes_per_row;
    509
    2.14k
                let row_data = &weight_data[row_start..row_end];
    510
    511
    2.14k
                fused_q6k_dot_simd(row_data, activations).unwrap_or(0.0)
    512
    2.14k
            })
    513
    8
            .collect();
    514
    515
    8
        Ok(output)
    516
    8
    }
    517
    518
    /// Parallel fused Q6_K matrix-vector multiply - writes to pre-allocated buffer
    519
    ///
    520
    /// IMP-131: Zero-allocation variant for hot-path inference.
    521
    #[allow(clippy::similar_names)]
    522
    3
    pub fn fused_q6k_parallel_matvec_into(
    523
    3
        weight_data: &[u8],
    524
    3
        activations: &[f32],
    525
    3
        in_dim: usize,
    526
    3
        out_dim: usize,
    527
    3
        output: &mut [f32],
    528
    3
    ) -> Result<()> {
    529
        use rayon::prelude::*;
    530
    531
    3
        let super_blocks_per_row = in_dim.div_ceil(QK_K);
    532
    3
        let bytes_per_row = super_blocks_per_row * 210;
    533
    534
    3
        let expected_weight_bytes = out_dim * bytes_per_row;
    535
    3
        if weight_data.len() < expected_weight_bytes {
    536
    2
            return Err(RealizarError::InvalidShape {
    537
    2
                reason: format!(
    538
    2
                    "Q6_K weight data too small: need {} bytes for {}x{}, have {}",
    539
    2
                    expected_weight_bytes,
    540
    2
                    out_dim,
    541
    2
                    in_dim,
    542
    2
                    weight_data.len()
    543
    2
                ),
    544
    2
            });
    545
    1
        }
    546
    547
    1
        if activations.len() != in_dim {
    548
    1
            return Err(RealizarError::InvalidShape {
    549
    1
                reason: format!(
    550
    1
                    "Activation length {} doesn't match in_dim {}",
    551
    1
                    activations.len(),
    552
    1
                    in_dim
    553
    1
                ),
    554
    1
            });
    555
    0
        }
    556
    557
    0
        if output.len() < out_dim {
    558
    0
            return Err(RealizarError::InvalidShape {
    559
    0
                reason: format!(
    560
    0
                    "Output buffer too small: need {}, have {}",
    561
    0
                    out_dim,
    562
    0
                    output.len()
    563
    0
                ),
    564
    0
            });
    565
    0
        }
    566
    567
        // TCB Tiling: Process rows in midi-tiles (64 rows) to maximize activation cache reuse
    568
        // While Q6K×f32 doesn't have integer Q8K, the f32 activation vector still benefits
    569
        // from being kept in L2 cache while processing multiple output rows.
    570
        const MIDI_TILE_M: usize = 64;
    571
    572
    0
        output[..out_dim]
    573
    0
            .par_chunks_mut(MIDI_TILE_M)
    574
    0
            .enumerate()
    575
    0
            .for_each(|(midi_idx, midi_chunk)| {
    576
    0
                let midi_start = midi_idx * MIDI_TILE_M;
    577
    578
                // Process each row in this midi-tile
    579
                // All rows share the same activation vector (kept in L2 cache)
    580
    0
                for (local_idx, out) in midi_chunk.iter_mut().enumerate() {
    581
    0
                    let row = midi_start + local_idx;
    582
    0
                    let row_start = row * bytes_per_row;
    583
    0
                    let row_end = row_start + bytes_per_row;
    584
    0
                    let row_data = &weight_data[row_start..row_end];
    585
    0
                    *out = fused_q6k_dot_simd(row_data, activations).unwrap_or(0.0);
    586
    0
                }
    587
    0
            });
    588
    589
    0
        Ok(())
    590
    3
    }
    591
    592
    // ============================================================================
    593
    594
    /// Backwards-compatible alias for `fused_q6k_parallel_matvec`.
    595
    ///
    596
    /// The column-major layout is now the default for the parallel implementation.
    597
    #[inline]
    598
    1
    pub fn fused_q6k_colmajor_matvec(
    599
    1
        weight_data: &[u8],
    600
    1
        activations: &[f32],
    601
    1
        in_dim: usize,
    602
    1
        out_dim: usize,
    603
    1
    ) -> Result<Vec<f32>> {
    604
    1
        fused_q6k_parallel_matvec(weight_data, activations, in_dim, out_dim)
    605
    1
    }
    606
    607
    /// Backwards-compatible alias for `fused_q4k_parallel_matvec_into`.
    608
    ///
    609
    /// The "auto" naming referred to automatic thread dispatch which is now the default.
    610
    #[inline]
    611
    1
    pub fn fused_q4k_auto_matvec_into(
    612
    1
        weight_data: &[u8],
    613
    1
        activations: &[f32],
    614
    1
        in_dim: usize,
    615
    1
        out_dim: usize,
    616
    1
        output: &mut [f32],
    617
    1
    ) -> Result<()> {
    618
    1
        fused_q4k_parallel_matvec_into(weight_data, activations, in_dim, out_dim, output)
    619
    1
    }
    620
    621
    /// Parallel Q4_K × Q8_K matrix-vector multiply with TCB tiling
    622
    ///
    623
    /// Uses rayon parallel iterators for multi-core acceleration and TCB tiling
    624
    /// pattern for cache optimization.
    625
    5
    pub fn fused_q4k_q8k_parallel_matvec_into(
    626
    5
        weight_data: &[u8],
    627
    5
        q8k_scales: &[f32],
    628
    5
        q8k_quants: &[i8],
    629
    5
        in_dim: usize,
    630
    5
        out_dim: usize,
    631
    5
        output: &mut [f32],
    632
    5
    ) -> Result<()> {
    633
        use rayon::prelude::*;
    634
    635
        const SUPER_BLOCK_BYTES: usize = 144;
    636
    637
    5
        let super_blocks_per_row = in_dim.div_ceil(QK_K);
    638
    5
        let bytes_per_row = super_blocks_per_row * SUPER_BLOCK_BYTES;
    639
    640
    5
        let expected_weight_bytes = out_dim * bytes_per_row;
    641
    5
        if weight_data.len() < expected_weight_bytes {
    642
    3
            return Err(RealizarError::InvalidShape {
    643
    3
                reason: format!(
    644
    3
                    "Weight data too small: need {} bytes, have {}",
    645
    3
                    expected_weight_bytes,
    646
    3
                    weight_data.len()
    647
    3
                ),
    648
    3
            });
    649
    2
        }
    650
    651
    2
        if output.len() < out_dim {
    652
    0
            return Err(RealizarError::InvalidShape {
    653
    0
                reason: format!(
    654
    0
                    "Output buffer too small: need {}, have {}",
    655
    0
                    out_dim,
    656
    0
                    output.len()
    657
    0
                ),
    658
    0
            });
    659
    2
        }
    660
    661
        // TCB Tiling Parameters (from trueno::tiling::TilingConfig::cpu_avx512_vnni_q4k_q8k)
    662
        // - Midi-tile: 64 rows (fits in L2 cache with Q8K input)
    663
        // - Micro-tile: 4 rows (processed simultaneously sharing Q8K loads)
    664
        const MIDI_TILE_M: usize = 64;
    665
        const MICRO_TILE_M: usize = 4;
    666
    667
        // Check if we can use the optimized 4-row micro-kernel
    668
        #[cfg(target_arch = "x86_64")]
    669
    2
        let use_4row_kernel =
    670
    2
            is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512vnni");
    671
        #[cfg(not(target_arch = "x86_64"))]
    672
        let use_4row_kernel = false;
    673
    674
    2
        if use_4row_kernel && out_dim >= MICRO_TILE_M {
    675
            // TCB-style tiled execution with 4-row micro-kernel
    676
            // Process MIDI_TILE_M rows per parallel chunk to maximize Q8K sharing
    677
    678
            // Split output into midi-tiles for rayon parallelism
    679
    0
            output[..out_dim]
    680
    0
                .par_chunks_mut(MIDI_TILE_M)
    681
    0
                .enumerate()
    682
    0
                .for_each(|(midi_idx, midi_chunk)| {
    683
    0
                    let midi_start = midi_idx * MIDI_TILE_M;
    684
    0
                    let midi_rows = midi_chunk.len();
    685
    686
                    // Process micro-tiles (4 rows at a time) within this midi-tile
    687
    0
                    let full_micro_tiles = midi_rows / MICRO_TILE_M;
    688
    0
                    let remainder = midi_rows % MICRO_TILE_M;
    689
    690
    0
                    for micro_idx in 0..full_micro_tiles {
    691
    0
                        let row_base = midi_start + micro_idx * MICRO_TILE_M;
    692
    0
    693
    0
                        // Build row pointers for 4-row kernel
    694
    0
                        let row_ptrs: [*const u8; 4] = [
    695
    0
                            weight_data.as_ptr().wrapping_add(row_base * bytes_per_row),
    696
    0
                            weight_data
    697
    0
                                .as_ptr()
    698
    0
                                .wrapping_add((row_base + 1) * bytes_per_row),
    699
    0
                            weight_data
    700
    0
                                .as_ptr()
    701
    0
                                .wrapping_add((row_base + 2) * bytes_per_row),
    702
    0
                            weight_data
    703
    0
                                .as_ptr()
    704
    0
                                .wrapping_add((row_base + 3) * bytes_per_row),
    705
    0
                        ];
    706
    0
    707
    0
                        // SAFETY: AVX-512 VNNI detected, pointers are within weight_data bounds
    708
    0
                        #[cfg(target_arch = "x86_64")]
    709
    0
                        let outputs = unsafe {
    710
    0
                            fused_q4k_q8k_dot_4rows_avx512vnni(
    711
    0
                                row_ptrs,
    712
    0
                                bytes_per_row,
    713
    0
                                q8k_scales,
    714
    0
                                q8k_quants,
    715
    0
                            )
    716
    0
                        };
    717
    0
    718
    0
                        #[cfg(not(target_arch = "x86_64"))]
    719
    0
                        let outputs = [0.0f32; 4];
    720
    0
    721
    0
                        let local_base = micro_idx * MICRO_TILE_M;
    722
    0
                        midi_chunk[local_base] = outputs[0];
    723
    0
                        midi_chunk[local_base + 1] = outputs[1];
    724
    0
                        midi_chunk[local_base + 2] = outputs[2];
    725
    0
                        midi_chunk[local_base + 3] = outputs[3];
    726
    0
                    }
    727
    728
                    // Handle remainder rows (< 4) with single-row kernel
    729
    0
                    for r in 0..remainder {
    730
    0
                        let row = midi_start + full_micro_tiles * MICRO_TILE_M + r;
    731
    0
                        let row_start = row * bytes_per_row;
    732
    0
                        let row_data = &weight_data[row_start..row_start + bytes_per_row];
    733
    0
                        let local_idx = full_micro_tiles * MICRO_TILE_M + r;
    734
    0
                        midi_chunk[local_idx] =
    735
    0
                            fused_q4k_q8k_dot_simd(row_data, q8k_scales, q8k_quants).unwrap_or(0.0);
    736
    0
                    }
    737
    0
                });
    738
        } else {
    739
            // Fallback: per-row execution (no TCB optimization)
    740
    2
            output[..out_dim]
    741
    2
                .par_iter_mut()
    742
    2
                .enumerate()
    743
    4
                .for_each2(|(o, out)| {
    744
    4
                    let row_start = o * bytes_per_row;
    745
    4
                    let row_data = &weight_data[row_start..row_start + bytes_per_row];
    746
    4
                    *out = fused_q4k_q8k_dot_simd(row_data, q8k_scales, q8k_quants).unwrap_or(0.0);
    747
    4
                });
    748
        }
    749
    750
    2
        Ok(())
    751
    5
    }
    752
    753
    /// Fused FFN up+gate projection in single parallel region
    754
    ///
    755
    /// Eliminates rayon::join overhead by processing both up and gate weights
    756
    /// in a single par_chunks_mut call. Both projections share the same Q8K
    757
    /// quantized input, so we only load it once per midi-tile.
    758
    ///
    759
    /// # Performance
    760
    ///
    761
    /// Reduces parallel region spawns from 2 to 1 per FFN layer, saving ~10-50µs
    762
    /// per layer. For 28 layers, this is 280-1400µs per token.
    763
    ///
    764
    /// # Arguments
    765
    ///
    766
    /// * `up_weight` - Q4K weight data for FFN up projection
    767
    /// * `gate_weight` - Q4K weight data for FFN gate projection
    768
    /// * `q8k_scales` - Pre-quantized activation scales
    769
    /// * `q8k_quants` - Pre-quantized activation values
    770
    /// * `in_dim` - Input dimension (hidden_dim)
    771
    /// * `out_dim` - Output dimension (intermediate_dim)
    772
    /// * `up_output` - Output buffer for up projection
    773
    /// * `gate_output` - Output buffer for gate projection
    774
    #[allow(clippy::too_many_arguments)]
    775
    6
    pub fn fused_q4k_q8k_ffn_up_gate_into(
    776
    6
        up_weight: &[u8],
    777
    6
        gate_weight: &[u8],
    778
    6
        q8k_scales: &[f32],
    779
    6
        q8k_quants: &[i8],
    780
    6
        in_dim: usize,
    781
    6
        out_dim: usize,
    782
    6
        up_output: &mut [f32],
    783
    6
        gate_output: &mut [f32],
    784
    6
    ) -> Result<()> {
    785
        use rayon::prelude::*;
    786
    787
        const SUPER_BLOCK_BYTES: usize = 144;
    788
        const MIDI_TILE_M: usize = 64;
    789
    790
    6
        let super_blocks_per_row = in_dim.div_ceil(QK_K);
    791
    6
        let bytes_per_row = super_blocks_per_row * SUPER_BLOCK_BYTES;
    792
    793
    6
        let expected_weight_bytes = out_dim * bytes_per_row;
    794
    6
        if up_weight.len() < expected_weight_bytes || gate_weight3.len() < expected_weight_bytes {
    795
    4
            return Err(RealizarError::InvalidShape {
    796
    4
                reason: format!(
    797
    4
                    "Weight data too small: need {} bytes",
    798
    4
                    expected_weight_bytes
    799
    4
                ),
    800
    4
            });
    801
    2
        }
    802
    803
    2
        if up_output.len() < out_dim || gate_output.len() < out_dim {
    804
    0
            return Err(RealizarError::InvalidShape {
    805
    0
                reason: format!("Output buffers too small: need {}", out_dim),
    806
    0
            });
    807
    2
        }
    808
    809
        // Process both up and gate in a single parallel region
    810
        // Each thread handles a midi-tile of rows for BOTH projections
    811
        // We use zip + par_chunks_mut to ensure thread-safe non-overlapping access
    812
    2
        up_output[..out_dim]
    813
    2
            .par_chunks_mut(MIDI_TILE_M)
    814
    2
            .zip(gate_output[..out_dim].par_chunks_mut(MIDI_TILE_M))
    815
    2
            .enumerate()
    816
    5
            .for_each2(|(midi_idx, (up_chunk, gate_chunk))| {
    817
    5
                let midi_start = midi_idx * MIDI_TILE_M;
    818
    819
    257
                for (local_row, (up_out, gate_out)) in
    820
    5
                    up_chunk.iter_mut().zip(gate_chunk.iter_mut()).enumerate()
    821
    257
                {
    822
    257
                    let row = midi_start + local_row;
    823
    257
                    let row_start = row * bytes_per_row;
    824
    257
    825
    257
                    // Compute up projection for this row
    826
    257
                    let up_row = &up_weight[row_start..row_start + bytes_per_row];
    827
    257
                    *up_out = fused_q4k_q8k_dot_simd(up_row, q8k_scales, q8k_quants).unwrap_or(0.0);
    828
    257
    829
    257
                    // Compute gate projection for this row
    830
    257
                    let gate_row = &gate_weight[row_start..row_start + bytes_per_row];
    831
    257
                    *gate_out = fused_q4k_q8k_dot_simd(gate_row, q8k_scales, q8k_quants).unwrap_or(0.0);
    832
    257
                }
    833
    5
            });
    834
    835
    2
        Ok(())
    836
    6
    }