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

    Line
    Count
    Source
    1
    //! Fused activation functions for quantized inference (PMAT-802)
    2
    //!
    3
    //! Implements fused operations combining normalization with quantization:
    4
    //! - `quantize_rmsnorm_q8_0` - RMSNorm with Q8_0 quantization
    5
    //! - `quantize_rmsnorm_q8_0_into` - Zero-allocation variant
    6
    //! - `fused_rmsnorm_q4_0_matmul` - RMSNorm + matmul fusion
    7
    //! - `fused_rmsnorm_ffn_up_gate` - RMSNorm + FFN up/gate fusion
    8
    //! - `fused_swiglu_simd` - SIMD-accelerated SwiGLU activation
    9
    //! - `softmax_simd` - SIMD-accelerated softmax
    10
    11
    use crate::error::{RealizarError, Result};
    12
    use super::fused_q4_0_q8_0_dot_simd;
    13
    14
    // ============================================================================
    15
    // Key insight: llama.cpp quantizes activations to Q8_0 and uses integer
    16
    // multiply-accumulate (maddubs_epi16), which is 4-5x faster than f32 FMA.
    17
    18
    /// Fused RMSNorm + Q8_0 quantization
    19
    ///
    20
    /// Computes RMSNorm and quantizes in a single pass:
    21
    /// normalized[i] = x[i] / sqrt(mean(x^2) + eps) * weight[i]
    22
    /// Then quantizes to Q8_0 format.
    23
    ///
    24
    /// This avoids allocating an intermediate normalized vector.
    25
    #[inline]
    26
    22
    pub fn quantize_rmsnorm_q8_0(input: &[f32], norm_weight: &[f32], eps: f32) -> (Vec<f32>, Vec<i8>) {
    27
        #[cfg(target_arch = "x86_64")]
    28
        {
    29
    22
            if is_x86_feature_detected!("avx2") {
    30
                // SAFETY: Memory safety ensured by bounds checking and alignment
    31
    22
                return unsafe { quantize_rmsnorm_q8_0_avx2(input, norm_weight, eps) };
    32
    0
            }
    33
        }
    34
    0
        quantize_rmsnorm_q8_0_scalar(input, norm_weight, eps)
    35
    22
    }
    36
    37
    /// Scalar implementation of fused RMSNorm + Q8_0 quantization
    38
    ///
    39
    /// This is exposed as `pub(crate)` for direct testing. The production code
    40
    /// uses the public `quantize_rmsnorm_q8_0` wrapper which dispatches to AVX2
    41
    /// when available.
    42
    0
    pub(crate) fn quantize_rmsnorm_q8_0_scalar(
    43
    0
        input: &[f32],
    44
    0
        norm_weight: &[f32],
    45
    0
        eps: f32,
    46
    0
    ) -> (Vec<f32>, Vec<i8>) {
    47
    0
        let hidden_dim = input.len();
    48
    0
        debug_assert_eq!(hidden_dim, norm_weight.len());
    49
    50
        // Compute sum of squares for RMSNorm
    51
    0
        let sum_sq: f32 = input.iter().map(|x| x * x).sum();
    52
    0
        let mean_sq = sum_sq / hidden_dim as f32;
    53
    0
        let inv_rms = 1.0 / (mean_sq + eps).sqrt();
    54
    55
        // Now quantize the normalized values directly
    56
    0
        let num_blocks = hidden_dim.div_ceil(32);
    57
    0
        let mut scales = Vec::with_capacity(num_blocks);
    58
    0
        let mut quants = Vec::with_capacity(num_blocks * 32);
    59
    60
    0
        for block_idx in 0..num_blocks {
    61
    0
            let start = block_idx * 32;
    62
    0
            let end = (start + 32).min(hidden_dim);
    63
    64
            // Find max absolute value of normalized values for this block
    65
    0
            let mut max_abs = 0.0f32;
    66
    0
            for i in start..end {
    67
                // Fused: x[i] * inv_rms * weight[i]
    68
    0
                let normalized = input[i] * inv_rms * norm_weight[i];
    69
    0
                let abs = normalized.abs();
    70
    0
                if abs > max_abs {
    71
    0
                    max_abs = abs;
    72
    0
                }
    73
            }
    74
    75
            // Compute scale
    76
    0
            let scale = if max_abs > 1e-10 {
    77
    0
                max_abs / 127.0
    78
            } else {
    79
    0
                1.0 / 127.0
    80
            };
    81
    0
            let inv_scale = 1.0 / scale;
    82
    0
            scales.push(scale);
    83
    84
            // Quantize normalized values
    85
    0
            for i in start..end {
    86
    0
                let normalized = input[i] * inv_rms * norm_weight[i];
    87
    0
                let q = (normalized * inv_scale).round();
    88
    0
                quants.push(q.clamp(-128.0, 127.0) as i8);
    89
    0
            }
    90
            // Pad to 32 if partial block
    91
    0
            for _ in end..(start + 32) {
    92
    0
                quants.push(0i8);
    93
    0
            }
    94
        }
    95
    96
    0
        (scales, quants)
    97
    0
    }
    98
    99
    /// AVX2-accelerated fused RMSNorm + Q8_0 quantization
    100
    ///
    101
    /// Processes 8 floats at a time using SIMD for:
    102
    /// - Sum of squares computation
    103
    /// - Max abs finding per block
    104
    /// - Normalization and quantization
    105
    #[cfg(target_arch = "x86_64")]
    106
    #[target_feature(enable = "avx2")]
    107
    #[inline]
    108
    22
    unsafe fn quantize_rmsnorm_q8_0_avx2(
    109
    22
        input: &[f32],
    110
    22
        norm_weight: &[f32],
    111
    22
        eps: f32,
    112
    22
    ) -> (Vec<f32>, Vec<i8>) {
    113
        // SAFETY: Memory safety ensured by bounds checking and alignment
    114
        unsafe {
    115
            use std::arch::x86_64::{
    116
                _mm256_add_ps, _mm256_and_ps, _mm256_andnot_ps, _mm256_castps256_ps128,
    117
                _mm256_castsi256_ps, _mm256_castsi256_si128, _mm256_cvtps_epi32, _mm256_extractf128_ps,
    118
                _mm256_extracti128_si256, _mm256_floor_ps, _mm256_fmadd_ps, _mm256_loadu_ps,
    119
                _mm256_max_ps, _mm256_min_ps, _mm256_mul_ps, _mm256_or_ps, _mm256_set1_epi32,
    120
                _mm256_set1_ps, _mm256_setzero_ps, _mm_add_ps, _mm_cvtss_f32, _mm_hadd_ps, _mm_max_ps,
    121
                _mm_movehl_ps, _mm_packs_epi16, _mm_packs_epi32, _mm_shuffle_ps, _mm_storel_epi64,
    122
            };
    123
    124
    22
            let hidden_dim = input.len();
    125
    22
            debug_assert_eq!(hidden_dim, norm_weight.len());
    126
    127
            // SIMD sum of squares
    128
    22
            let mut sum_sq_vec = _mm256_setzero_ps();
    129
    22
            let mut i = 0;
    130
    131
            // Process 8 floats at a time
    132
    354
            while i + 8 <= hidden_dim {
    133
    332
                let v = _mm256_loadu_ps(input.as_ptr().add(i));
    134
    332
                sum_sq_vec = _mm256_fmadd_ps(v, v, sum_sq_vec);
    135
    332
                i += 8;
    136
    332
            }
    137
    138
            // Horizontal sum
    139
    22
            let hi = _mm256_extractf128_ps(sum_sq_vec, 1);
    140
    22
            let lo = _mm256_castps256_ps128(sum_sq_vec);
    141
    22
            let sum128 = _mm_add_ps(lo, hi);
    142
    22
            let sum64 = _mm_hadd_ps(sum128, sum128);
    143
    22
            let sum32 = _mm_hadd_ps(sum64, sum64);
    144
    22
            let mut sum_sq = _mm_cvtss_f32(sum32);
    145
    146
            // Handle remaining elements
    147
    22
            while i < hidden_dim {
    148
    0
                sum_sq += input[i] * input[i];
    149
    0
                i += 1;
    150
    0
            }
    151
    152
    22
            let mean_sq = sum_sq / hidden_dim as f32;
    153
    22
            let inv_rms = 1.0 / (mean_sq + eps).sqrt();
    154
    22
            let inv_rms_vec = _mm256_set1_ps(inv_rms);
    155
    156
            // Quantize with SIMD
    157
    22
            let num_blocks = hidden_dim.div_ceil(32);
    158
    22
            let mut scales = Vec::with_capacity(num_blocks);
    159
    22
            let mut quants = vec![0i8; num_blocks * 32];
    160
    161
    22
            let abs_mask = _mm256_castsi256_ps(_mm256_set1_epi32(0x7FFFFFFF_u32 as i32));
    162
    22
            let round_const = _mm256_set1_ps(0.5);
    163
    22
            let clamp_min = _mm256_set1_ps(-128.0);
    164
    22
            let clamp_max = _mm256_set1_ps(127.0);
    165
    166
    84
            for block_idx in 0..num_blocks22 {
    167
    84
                let start = block_idx * 32;
    168
    84
                let block_end = (start + 32).min(hidden_dim);
    169
    84
                let valid_len = block_end - start;
    170
    171
                // Find max abs in this block using SIMD
    172
    84
                let mut max_vec = _mm256_setzero_ps();
    173
    84
                let mut j = 0;
    174
    416
                while j + 8 <= valid_len {
    175
    332
                    let idx = start + j;
    176
    332
                    let inp = _mm256_loadu_ps(input.as_ptr().add(idx));
    177
    332
                    let wgt = _mm256_loadu_ps(norm_weight.as_ptr().add(idx));
    178
    332
                    let normalized = _mm256_mul_ps(_mm256_mul_ps(inp, inv_rms_vec), wgt);
    179
    332
                    let abs_val = _mm256_and_ps(normalized, abs_mask);
    180
    332
                    max_vec = _mm256_max_ps(max_vec, abs_val);
    181
    332
                    j += 8;
    182
    332
                }
    183
    184
                // Horizontal max
    185
    84
                let max_hi = _mm256_extractf128_ps(max_vec, 1);
    186
    84
                let max_lo = _mm256_castps256_ps128(max_vec);
    187
    84
                let max_128 = _mm_max_ps(max_lo, max_hi);
    188
    84
                let max_64 = _mm_max_ps(max_128, _mm_movehl_ps(max_128, max_128));
    189
    84
                let max_32 = _mm_max_ps(max_64, _mm_shuffle_ps(max_64, max_64, 1));
    190
    84
                let mut max_abs = _mm_cvtss_f32(max_32);
    191
    192
                // Handle remaining elements in block
    193
    84
                while j < valid_len {
    194
    0
                    let normalized = input[start + j] * inv_rms * norm_weight[start + j];
    195
    0
                    let abs = normalized.abs();
    196
    0
                    if abs > max_abs {
    197
    0
                        max_abs = abs;
    198
    0
                    }
    199
    0
                    j += 1;
    200
                }
    201
    202
                // Compute scale
    203
    84
                let scale = if max_abs > 1e-10 {
    204
    84
                    max_abs / 127.0
    205
                } else {
    206
    0
                    1.0 / 127.0
    207
                };
    208
    84
                let inv_scale = 1.0 / scale;
    209
    84
                let inv_scale_vec = _mm256_set1_ps(inv_scale);
    210
    84
                scales.push(scale);
    211
    212
                // Quantize with SIMD
    213
    84
                let quant_ptr = quants.as_mut_ptr().add(block_idx * 32);
    214
    84
                let mut k = 0;
    215
    416
                while k + 8 <= valid_len {
    216
    332
                    let idx = start + k;
    217
    332
                    let inp = _mm256_loadu_ps(input.as_ptr().add(idx));
    218
    332
                    let wgt = _mm256_loadu_ps(norm_weight.as_ptr().add(idx));
    219
    332
                    let normalized = _mm256_mul_ps(_mm256_mul_ps(inp, inv_rms_vec), wgt);
    220
    332
                    let scaled = _mm256_mul_ps(normalized, inv_scale_vec);
    221
    332
    222
    332
                    // Round to nearest (add 0.5 and truncate, handle sign)
    223
    332
                    let sign = _mm256_and_ps(
    224
    332
                        scaled,
    225
    332
                        _mm256_castsi256_ps(_mm256_set1_epi32(0x80000000_u32 as i32)),
    226
    332
                    );
    227
    332
                    let abs_scaled = _mm256_andnot_ps(sign, scaled);
    228
    332
                    let rounded = _mm256_or_ps(
    229
    332
                        _mm256_floor_ps(_mm256_add_ps(abs_scaled, round_const)),
    230
    332
                        sign,
    231
    332
                    );
    232
    332
    233
    332
                    // Clamp to [-128, 127]
    234
    332
                    let clamped = _mm256_max_ps(clamp_min, _mm256_min_ps(clamp_max, rounded));
    235
    332
    236
    332
                    // Convert to int32 and extract to i8
    237
    332
                    let int32 = _mm256_cvtps_epi32(clamped);
    238
    332
    239
    332
                    // Pack i32 -> i16 -> i8 (only need lower 8 values)
    240
    332
                    let lo128 = _mm256_castsi256_si128(int32);
    241
    332
                    let hi128 = _mm256_extracti128_si256(int32, 1);
    242
    332
                    let packed16 = _mm_packs_epi32(lo128, hi128);
    243
    332
                    let packed8 = _mm_packs_epi16(packed16, packed16);
    244
    332
    245
    332
                    // Store 8 i8 values
    246
    332
                    _mm_storel_epi64(quant_ptr.add(k).cast(), packed8);
    247
    332
                    k += 8;
    248
    332
                }
    249
    250
                // Handle remaining elements
    251
    84
                while k < valid_len {
    252
    0
                    let normalized = input[start + k] * inv_rms * norm_weight[start + k];
    253
    0
                    let q = (normalized * inv_scale).round();
    254
    0
                    *quant_ptr.add(k) = q.clamp(-128.0, 127.0) as i8;
    255
    0
                    k += 1;
    256
    0
                }
    257
            }
    258
    259
    22
            (scales, quants)
    260
        }
    261
    22
    }
    262
    263
    /// Fused RMSNorm + Q4_0 matmul
    264
    ///
    265
    /// Combines RMSNorm normalization with quantized matmul in one operation:
    266
    /// 1. Computes inv_rms = 1 / sqrt(mean(x^2) + eps)
    267
    /// 2. Quantizes (x * inv_rms * norm_weight) to Q8_0
    268
    /// 3. Performs Q4_0 × Q8_0 integer matmul
    269
    ///
    270
    /// This eliminates the intermediate normalized vector allocation.
    271
    #[allow(clippy::similar_names)]
    272
    4
    pub fn fused_rmsnorm_q4_0_matmul(
    273
    4
        input: &[f32],
    274
    4
        norm_weight: &[f32],
    275
    4
        eps: f32,
    276
    4
        weight_data: &[u8],
    277
    4
        in_dim: usize,
    278
    4
        out_dim: usize,
    279
    4
    ) -> Result<Vec<f32>> {
    280
        use rayon::prelude::*;
    281
    282
        const Q4_0_BLOCK_BYTES: usize = 18;
    283
        const Q4_0_BLOCK_SIZE: usize = 32;
    284
    285
    4
        if input.len() != in_dim {
    286
    1
            return Err(RealizarError::InvalidShape {
    287
    1
                reason: format!(
    288
    1
                    "Input length {} doesn't match in_dim {}",
    289
    1
                    input.len(),
    290
    1
                    in_dim
    291
    1
                ),
    292
    1
            });
    293
    3
        }
    294
    295
    3
        let blocks_per_row = in_dim.div_ceil(Q4_0_BLOCK_SIZE);
    296
    3
        let bytes_per_row = blocks_per_row * Q4_0_BLOCK_BYTES;
    297
    298
    3
        let expected_weight_bytes = out_dim * bytes_per_row;
    299
    3
        if weight_data.len() < expected_weight_bytes {
    300
    1
            return Err(RealizarError::InvalidShape {
    301
    1
                reason: format!(
    302
    1
                    "Q4_0 weight data too small: need {} bytes for {}x{}, have {}",
    303
    1
                    expected_weight_bytes,
    304
    1
                    out_dim,
    305
    1
                    in_dim,
    306
    1
                    weight_data.len()
    307
    1
                ),
    308
    1
            });
    309
    2
        }
    310
    311
        // Fused RMSNorm + Q8_0 quantization (single pass, no intermediate allocation)
    312
    2
        let (q8_scales, q8_quants) = quantize_rmsnorm_q8_0(input, norm_weight, eps);
    313
    314
        // Parallel matmul with chunking
    315
        const CHUNK_SIZE: usize = 64;
    316
    2
        let output: Vec<f32> = (0..out_dim)
    317
    2
            .into_par_iter()
    318
    2
            .with_min_len(CHUNK_SIZE)
    319
    4
            .map2(|o| {
    320
    4
                let row_start = o * bytes_per_row;
    321
    4
                let row_end = row_start + bytes_per_row;
    322
    4
                let row_data = &weight_data[row_start..row_end];
    323
    4
                fused_q4_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim)
    324
    4
            })
    325
    2
            .collect();
    326
    327
    2
        Ok(output)
    328
    4
    }
    329
    330
    /// Fused RMSNorm + parallel FFN up/gate projections
    331
    ///
    332
    /// For SwiGLU models, FFN has two parallel matmuls (up and gate) that share
    333
    /// the same normalized input. This function:
    334
    /// 1. Computes inv_rms once
    335
    /// 2. Quantizes normalized input to Q8_0 once
    336
    /// 3. Runs both up and gate matmuls in parallel
    337
    ///
    338
    /// Eliminates: 1 RMSNorm pass, 1 intermediate allocation, 1 Q8_0 quantization
    339
    #[allow(clippy::similar_names)]
    340
    #[allow(clippy::too_many_arguments)]
    341
    6
    pub fn fused_rmsnorm_ffn_up_gate(
    342
    6
        input: &[f32],
    343
    6
        norm_weight: &[f32],
    344
    6
        eps: f32,
    345
    6
        up_weight_data: &[u8],
    346
    6
        gate_weight_data: &[u8],
    347
    6
        in_dim: usize,
    348
    6
        out_dim: usize,
    349
    6
    ) -> Result<(Vec<f32>, Vec<f32>)> {
    350
        use rayon::prelude::*;
    351
    352
        const Q4_0_BLOCK_BYTES: usize = 18;
    353
        const Q4_0_BLOCK_SIZE: usize = 32;
    354
    355
    6
        if input.len() != in_dim {
    356
    1
            return Err(RealizarError::InvalidShape {
    357
    1
                reason: format!(
    358
    1
                    "Input length {} doesn't match in_dim {}",
    359
    1
                    input.len(),
    360
    1
                    in_dim
    361
    1
                ),
    362
    1
            });
    363
    5
        }
    364
    365
    5
        let blocks_per_row = in_dim.div_ceil(Q4_0_BLOCK_SIZE);
    366
    5
        let bytes_per_row = blocks_per_row * Q4_0_BLOCK_BYTES;
    367
    5
        let expected_weight_bytes = out_dim * bytes_per_row;
    368
    369
    5
        if up_weight_data.len() < expected_weight_bytes {
    370
    2
            return Err(RealizarError::InvalidShape {
    371
    2
                reason: format!(
    372
    2
                    "FFN up weight data too small: need {} bytes, have {}",
    373
    2
                    expected_weight_bytes,
    374
    2
                    up_weight_data.len()
    375
    2
                ),
    376
    2
            });
    377
    3
        }
    378
    3
        if gate_weight_data.len() < expected_weight_bytes {
    379
    1
            return Err(RealizarError::InvalidShape {
    380
    1
                reason: format!(
    381
    1
                    "FFN gate weight data too small: need {} bytes, have {}",
    382
    1
                    expected_weight_bytes,
    383
    1
                    gate_weight_data.len()
    384
    1
                ),
    385
    1
            });
    386
    2
        }
    387
    388
        // Fused RMSNorm + Q8_0 quantization - computed ONCE for both matmuls
    389
    2
        let (q8_scales, q8_quants) = quantize_rmsnorm_q8_0(input, norm_weight, eps);
    390
    391
        // Run both matmuls in parallel using rayon::join
    392
        // Each matmul uses parallel iteration with chunking to reduce overhead
    393
    2
        let (up_output, gate_output) = rayon::join(
    394
    2
            || {
    395
                const CHUNK_SIZE: usize = 64;
    396
    2
                (0..out_dim)
    397
    2
                    .into_par_iter()
    398
    2
                    .with_min_len(CHUNK_SIZE)
    399
    65
                    .map2(|o| {
    400
    65
                        let row_start = o * bytes_per_row;
    401
    65
                        let row_end = row_start + bytes_per_row;
    402
    65
                        let row_data = &up_weight_data[row_start..row_end];
    403
    65
                        fused_q4_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim)
    404
    65
                    })
    405
    2
                    .collect::<Vec<f32>>()
    406
    2
            },
    407
    2
            || {
    408
                const CHUNK_SIZE: usize = 64;
    409
    2
                (0..out_dim)
    410
    2
                    .into_par_iter()
    411
    2
                    .with_min_len(CHUNK_SIZE)
    412
    65
                    .map2(|o| {
    413
    65
                        let row_start = o * bytes_per_row;
    414
    65
                        let row_end = row_start + bytes_per_row;
    415
    65
                        let row_data = &gate_weight_data[row_start..row_end];
    416
    65
                        fused_q4_0_q8_0_dot_simd(row_data, &q8_scales, &q8_quants, in_dim)
    417
    65
                    })
    418
    2
                    .collect::<Vec<f32>>()
    419
    2
            },
    420
        );
    421
    422
    2
        Ok((up_output, gate_output))
    423
    6
    }
    424
    425
    /// Zero-allocation variant of quantize_rmsnorm_q8_0
    426
    ///
    427
    /// Writes results directly into pre-allocated output buffers.
    428
    3
    pub fn quantize_rmsnorm_q8_0_into(
    429
    3
        input: &[f32],
    430
    3
        norm_weight: &[f32],
    431
    3
        eps: f32,
    432
    3
        scales: &mut [f32],
    433
    3
        quants: &mut [i8],
    434
    3
    ) {
    435
    3
        let hidden_dim = input.len();
    436
    3
        debug_assert_eq!(hidden_dim, norm_weight.len());
    437
    438
        // Compute sum of squares for RMSNorm
    439
    320
        let sum_sq3: f323 = input3.iter3().map3(|x| x * x).sum3();
    440
    3
        let mean_sq = sum_sq / hidden_dim as f32;
    441
    3
        let inv_rms = 1.0 / (mean_sq + eps).sqrt();
    442
    443
    3
        let num_blocks = hidden_dim.div_ceil(32);
    444
    445
    10
        for block_idx in 0..num_blocks3 {
    446
    10
            let start = block_idx * 32;
    447
    10
            let end = (start + 32).min(hidden_dim);
    448
    449
            // Find max absolute value of normalized values for this block
    450
    10
            let mut max_abs = 0.0f32;
    451
    320
            for i in start10..end10 {
    452
    320
                let normalized = input[i] * inv_rms * norm_weight[i];
    453
    320
                let abs = normalized.abs();
    454
    320
                if abs > max_abs {
    455
    40
                    max_abs = abs;
    456
    280
                }
    457
            }
    458
    459
            // Compute scale
    460
    10
            let scale = if max_abs > 1e-10 {
    461
    10
                max_abs / 127.0
    462
            } else {
    463
    0
                1.0 / 127.0
    464
            };
    465
    10
            let inv_scale = 1.0 / scale;
    466
    10
            scales[block_idx] = scale;
    467
    468
            // Quantize normalized values
    469
    10
            let quant_start = block_idx * 32;
    470
    320
            for i in start10..end10 {
    471
    320
                let normalized = input[i] * inv_rms * norm_weight[i];
    472
    320
                let q = (normalized * inv_scale).round();
    473
    320
                quants[quant_start + (i - start)] = q.clamp(-128.0, 127.0) as i8;
    474
    320
            }
    475
            // Pad to 32 if partial block
    476
    10
            for j0 in (end - start)..32 {
    477
    0
                quants[quant_start + j] = 0i8;
    478
    0
            }
    479
        }
    480
    3
    }
    481
    482
    /// SIMD-accelerated fused SwiGLU activation: silu(gate) * up
    483
    ///
    484
    /// Combines silu activation and element-wise multiply in a single pass
    485
    /// for better cache locality. Uses AVX2/AVX-512 SIMD where available.
    486
    ///
    487
    /// # Arguments
    488
    /// * `gate` - Gate values, modified in-place to contain result
    489
    /// * `up` - Up projection values
    490
    15
    pub fn fused_swiglu_simd(gate: &mut [f32], up: &[f32]) {
    491
    15
        debug_assert_eq!(gate.len(), up.len());
    492
    493
        #[cfg(target_arch = "x86_64")]
    494
        {
    495
    15
            if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
    496
                // SAFETY: AVX2 and FMA verified at runtime
    497
    15
                unsafe {
    498
    15
                    fused_swiglu_avx2(gate, up);
    499
    15
                }
    500
    15
                return;
    501
    0
            }
    502
        }
    503
    504
        // Scalar fallback
    505
    0
        fused_swiglu_scalar(gate, up);
    506
    15
    }
    507
    508
    /// Scalar fused SwiGLU: silu(gate) * up
    509
    ///
    510
    /// Exposed as `pub(crate)` for direct testing on AVX2 machines.
    511
    #[inline]
    512
    0
    pub(crate) fn fused_swiglu_scalar(gate: &mut [f32], up: &[f32]) {
    513
    0
        for (g, &u) in gate.iter_mut().zip(up.iter()) {
    514
    0
            // silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
    515
    0
            let silu_g = *g / (1.0 + (-*g).exp());
    516
    0
            *g = silu_g * u;
    517
    0
        }
    518
    0
    }
    519
    520
    /// AVX2 SIMD fused SwiGLU with FMA
    521
    ///
    522
    /// Computes silu(gate) * up using:
    523
    /// - Polynomial approximation for exp(-x)
    524
    /// - FMA for efficient multiply-add
    525
    /// - 8-wide AVX2 vectors
    526
    #[cfg(target_arch = "x86_64")]
    527
    #[target_feature(enable = "avx2", enable = "fma")]
    528
    #[inline]
    529
    #[allow(clippy::many_single_char_names)]
    530
    15
    unsafe fn fused_swiglu_avx2(gate: &mut [f32], up: &[f32]) {
    531
        use std::arch::x86_64::{
    532
            _mm256_add_epi32, _mm256_add_ps, _mm256_castsi256_ps, _mm256_cvtps_epi32, _mm256_floor_ps,
    533
            _mm256_fmadd_ps, _mm256_fnmadd_ps, _mm256_loadu_ps, _mm256_max_ps, _mm256_mul_ps,
    534
            _mm256_rcp_ps, _mm256_set1_epi32, _mm256_set1_ps, _mm256_setzero_ps, _mm256_slli_epi32,
    535
            _mm256_storeu_ps, _mm256_sub_ps,
    536
        };
    537
    538
        // SAFETY: Memory safety ensured by bounds checking and alignment
    539
        unsafe {
    540
    15
            let n = gate.len();
    541
    15
            let mut i = 0;
    542
    543
            // Constants for exp approximation (polynomial coefficients)
    544
            // Using 5th-degree polynomial approximation for exp(x) on [-87, 0]
    545
    15
            let one = _mm256_set1_ps(1.0);
    546
    15
            let ln2_inv = _mm256_set1_ps(1.442_695); // 1/ln(2)
    547
    15
            let ln2 = _mm256_set1_ps(0.693_147_2);
    548
    15
            let c0 = _mm256_set1_ps(1.0);
    549
    15
            let c1 = _mm256_set1_ps(0.693_147_2); // ln(2)
    550
    15
            let c2 = _mm256_set1_ps(0.240_226_5); // ln(2)^2 / 2!
    551
    15
            let c3 = _mm256_set1_ps(0.055_504_11); // ln(2)^3 / 3!
    552
    15
            let c4 = _mm256_set1_ps(0.009_618_13); // ln(2)^4 / 4!
    553
    15
            let c5 = _mm256_set1_ps(0.001_333_36); // ln(2)^5 / 5!
    554
    15
            let min_exp = _mm256_set1_ps(-87.0); // Minimum input to avoid underflow
    555
    15
            let two = _mm256_set1_ps(2.0); // For Newton-Raphson
    556
    557
            // Process 8 elements at a time
    558
    93
            while i + 8 <= n {
    559
    78
                // Load gate and up values
    560
    78
                let g = _mm256_loadu_ps(gate.as_ptr().add(i));
    561
    78
                let u = _mm256_loadu_ps(up.as_ptr().add(i));
    562
    78
    563
    78
                // Compute -g for sigmoid
    564
    78
                let neg_g = _mm256_sub_ps(_mm256_setzero_ps(), g);
    565
    78
    566
    78
                // Clamp to avoid exp underflow
    567
    78
                let neg_g_clamped = _mm256_max_ps(neg_g, min_exp);
    568
    78
    569
    78
                // Fast exp approximation using 2^(x/ln2) = 2^n * 2^f where n=floor, f=frac
    570
    78
                // n = floor(x * 1/ln2)
    571
    78
                let xln2 = _mm256_mul_ps(neg_g_clamped, ln2_inv);
    572
    78
                let n_f = _mm256_floor_ps(xln2);
    573
    78
                let n_i = _mm256_cvtps_epi32(n_f);
    574
    78
    575
    78
                // f = x - n * ln2 (fractional part scaled back)
    576
    78
                let f = _mm256_fnmadd_ps(n_f, ln2, neg_g_clamped);
    577
    78
    578
    78
                // Horner's method: c0 + f*(c1 + f*(c2 + f*(c3 + f*(c4 + f*c5))))
    579
    78
                let p = _mm256_fmadd_ps(f, c5, c4);
    580
    78
                let p = _mm256_fmadd_ps(f, p, c3);
    581
    78
                let p = _mm256_fmadd_ps(f, p, c2);
    582
    78
                let p = _mm256_fmadd_ps(f, p, c1);
    583
    78
                let p = _mm256_fmadd_ps(f, p, c0);
    584
    78
    585
    78
                // Scale by 2^n using integer bit manipulation
    586
    78
                // 2^n = reinterpret((n + 127) << 23) as float
    587
    78
                let bias = _mm256_set1_epi32(127);
    588
    78
                let n_biased = _mm256_add_epi32(n_i, bias);
    589
    78
                let exp_scale = _mm256_slli_epi32::<23>(n_biased);
    590
    78
                let exp_scale_f = _mm256_castsi256_ps(exp_scale);
    591
    78
    592
    78
                // exp(-g) = 2^n * p(f)
    593
    78
                let exp_neg_g = _mm256_mul_ps(p, exp_scale_f);
    594
    78
    595
    78
                // sigmoid(-(-g)) = 1 / (1 + exp(-g))
    596
    78
                // Use fast reciprocal approximation with Newton-Raphson refinement
    597
    78
                let denom = _mm256_add_ps(one, exp_neg_g);
    598
    78
                let rcp = _mm256_rcp_ps(denom); // ~12-bit precision
    599
    78
                                                // One Newton-Raphson iteration: x' = x * (2 - d*x)
    600
    78
                let sigmoid = _mm256_mul_ps(rcp, _mm256_fnmadd_ps(denom, rcp, two));
    601
    78
    602
    78
                // silu(g) = g * sigmoid(g)
    603
    78
                let silu_g = _mm256_mul_ps(g, sigmoid);
    604
    78
    605
    78
                // Result = silu(g) * u
    606
    78
                let result = _mm256_mul_ps(silu_g, u);
    607
    78
    608
    78
                // Store result
    609
    78
                _mm256_storeu_ps(gate.as_mut_ptr().add(i), result);
    610
    78
    611
    78
                i += 8;
    612
    78
            }
    613
    614
            // Handle remainder with scalar code
    615
    28
            while i < n {
    616
    13
                let g = gate[i];
    617
    13
                let silu_g = g / (1.0 + (-g).exp());
    618
    13
                gate[i] = silu_g * up[i];
    619
    13
                i += 1;
    620
    13
            }
    621
        }
    622
    15
    }
    623
    624
    /// SIMD-optimized in-place softmax
    625
    ///
    626
    /// Computes softmax(x) = exp(x - max) / sum(exp(x - max))
    627
    /// Uses AVX2/AVX-512 for vectorized exp and horizontal operations.
    628
    ///
    629
    /// # Arguments
    630
    /// * `x` - Slice to softmax in-place
    631
    #[inline]
    632
    2.80k
    pub fn softmax_simd(x: &mut [f32]) {
    633
    2.80k
        if x.is_empty() {
    634
    2
            return;
    635
    2.80k
        }
    636
    637
        #[cfg(target_arch = "x86_64")]
    638
        {
    639
    2.80k
            if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
    640
                // SAFETY: Memory safety ensured by bounds checking and alignment
    641
    2.80k
                unsafe {
    642
    2.80k
                    softmax_avx2(x);
    643
    2.80k
                }
    644
    2.80k
                return;
    645
    0
            }
    646
        }
    647
    648
        // Scalar fallback
    649
    0
        softmax_scalar(x);
    650
    2.80k
    }
    651
    652
    /// Scalar softmax
    653
    ///
    654
    /// Exposed as `pub(crate)` for direct testing on AVX2 machines.
    655
    #[inline]
    656
    0
    pub(crate) fn softmax_scalar(x: &mut [f32]) {
    657
        // Find max for numerical stability
    658
    0
        let max = x.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    659
    660
        // Compute exp(x - max) and sum
    661
    0
        let mut sum = 0.0f32;
    662
    0
        for v in x.iter_mut() {
    663
    0
            *v = (*v - max).exp();
    664
    0
            sum += *v;
    665
    0
        }
    666
    667
        // Normalize
    668
    0
        let inv_sum = 1.0 / sum;
    669
    0
        for v in x.iter_mut() {
    670
    0
            *v *= inv_sum;
    671
    0
        }
    672
    0
    }
    673
    674
    /// AVX2 SIMD softmax - only SIMD for max-find and normalization
    675
    /// (exp() uses libm which is faster than polynomial for short vectors)
    676
    #[cfg(target_arch = "x86_64")]
    677
    #[target_feature(enable = "avx2")]
    678
    #[inline]
    679
    #[allow(unsafe_op_in_unsafe_fn)]
    680
    2.80k
    unsafe fn softmax_avx2(x: &mut [f32]) {
    681
        use std::arch::x86_64::{
    682
            _mm256_loadu_ps, _mm256_max_ps, _mm256_mul_ps, _mm256_set1_ps, _mm256_storeu_ps,
    683
        };
    684
    685
    2.80k
        let n = x.len();
    686
    2.80k
        if n == 0 {
    687
    0
            return;
    688
    2.80k
        }
    689
    690
        // ============= Phase 1: Find max (SIMD) =============
    691
    2.80k
        let mut max_vec = _mm256_set1_ps(f32::NEG_INFINITY);
    692
    2.80k
        let mut i = 0;
    693
    694
    9.78k
        while i + 8 <= n {
    695
    6.97k
            let v = _mm256_loadu_ps(x.as_ptr().add(i));
    696
    6.97k
            max_vec = _mm256_max_ps(max_vec, v);
    697
    6.97k
            i += 8;
    698
    6.97k
        }
    699
    700
    2.80k
        let mut max_scalar = horizontal_max_avx2(max_vec);
    701
    9.95k
        for j in i2.80k..n2.80k {
    702
    9.95k
            max_scalar = max_scalar.max(x[j]);
    703
    9.95k
        }
    704
    705
        // ============= Phase 2: Compute exp(x - max) (scalar libm) =============
    706
    2.80k
        let mut sum_scalar = 0.0f32;
    707
    65.7k
        for j in 0..n2.80k {
    708
    65.7k
            let exp_v = (x[j] - max_scalar).exp();
    709
    65.7k
            x[j] = exp_v;
    710
    65.7k
            sum_scalar += exp_v;
    711
    65.7k
        }
    712
    713
        // ============= Phase 3: Normalize (SIMD) =============
    714
    2.80k
        let inv_sum = _mm256_set1_ps(1.0 / sum_scalar);
    715
    716
    2.80k
        i = 0;
    717
    9.78k
        while i + 8 <= n {
    718
    6.97k
            let v = _mm256_loadu_ps(x.as_ptr().add(i));
    719
    6.97k
            let normalized = _mm256_mul_ps(v, inv_sum);
    720
    6.97k
            _mm256_storeu_ps(x.as_mut_ptr().add(i), normalized);
    721
    6.97k
            i += 8;
    722
    6.97k
        }
    723
    724
    2.80k
        let inv_sum_scalar = 1.0 / sum_scalar;
    725
    9.95k
        for j in i2.80k..n2.80k {
    726
    9.95k
            x[j] *= inv_sum_scalar;
    727
    9.95k
        }
    728
    2.80k
    }
    729
    730
    /// Fast exp approximation using polynomial (AVX2)
    731
    #[cfg(target_arch = "x86_64")]
    732
    #[target_feature(enable = "avx2", enable = "fma")]
    733
    #[inline]
    734
    #[allow(unsafe_op_in_unsafe_fn)]
    735
    0
    unsafe fn fast_exp_avx2(
    736
    0
        x: std::arch::x86_64::__m256,
    737
    0
        ln2_inv: std::arch::x86_64::__m256,
    738
    0
        ln2: std::arch::x86_64::__m256,
    739
    0
        c0: std::arch::x86_64::__m256,
    740
    0
        c1: std::arch::x86_64::__m256,
    741
    0
        c2: std::arch::x86_64::__m256,
    742
    0
        c3: std::arch::x86_64::__m256,
    743
    0
        c4: std::arch::x86_64::__m256,
    744
    0
        c5: std::arch::x86_64::__m256,
    745
    0
        min_exp: std::arch::x86_64::__m256,
    746
    0
    ) -> std::arch::x86_64::__m256 {
    747
        use std::arch::x86_64::{
    748
            _mm256_add_epi32, _mm256_castsi256_ps, _mm256_cvtps_epi32, _mm256_floor_ps,
    749
            _mm256_fmadd_ps, _mm256_fnmadd_ps, _mm256_max_ps, _mm256_mul_ps, _mm256_set1_epi32,
    750
            _mm256_slli_epi32,
    751
        };
    752
    753
        {
    754
            // Clamp to avoid underflow
    755
    0
            let x_clamped = _mm256_max_ps(x, min_exp);
    756
    757
            // n = floor(x / ln2)
    758
    0
            let xln2 = _mm256_mul_ps(x_clamped, ln2_inv);
    759
    0
            let n_f = _mm256_floor_ps(xln2);
    760
    0
            let n_i = _mm256_cvtps_epi32(n_f);
    761
    762
            // f = x - n * ln2
    763
    0
            let f = _mm256_fnmadd_ps(n_f, ln2, x_clamped);
    764
    765
            // Polynomial: c0 + f*(c1 + f*(c2 + f*(c3 + f*(c4 + f*c5))))
    766
    0
            let p = _mm256_fmadd_ps(f, c5, c4);
    767
    0
            let p = _mm256_fmadd_ps(f, p, c3);
    768
    0
            let p = _mm256_fmadd_ps(f, p, c2);
    769
    0
            let p = _mm256_fmadd_ps(f, p, c1);
    770
    0
            let p = _mm256_fmadd_ps(f, p, c0);
    771
    772
            // 2^n via bit manipulation
    773
    0
            let bias = _mm256_set1_epi32(127);
    774
    0
            let n_biased = _mm256_add_epi32(n_i, bias);
    775
    0
            let exp_scale = _mm256_slli_epi32::<23>(n_biased);
    776
    0
            let exp_scale_f = _mm256_castsi256_ps(exp_scale);
    777
    778
    0
            _mm256_mul_ps(p, exp_scale_f)
    779
        }
    780
    0
    }
    781
    782
    /// Horizontal max of 8-wide AVX2 vector
    783
    #[cfg(target_arch = "x86_64")]
    784
    #[target_feature(enable = "avx2")]
    785
    #[inline]
    786
    #[allow(unsafe_op_in_unsafe_fn)]
    787
    2.80k
    unsafe fn horizontal_max_avx2(v: std::arch::x86_64::__m256) -> f32 {
    788
        use std::arch::x86_64::{
    789
            _mm256_extractf128_ps, _mm_cvtss_f32, _mm_max_ps, _mm_max_ss, _mm_movehl_ps, _mm_shuffle_ps,
    790
        };
    791
    792
        {
    793
            // Extract high and low 128-bit lanes
    794
    2.80k
            let hi = _mm256_extractf128_ps::<1>(v);
    795
    2.80k
            let lo = _mm256_extractf128_ps::<0>(v);
    796
    2.80k
            let max128 = _mm_max_ps(hi, lo);
    797
    798
            // Reduce 4 to 2
    799
    2.80k
            let max64 = _mm_max_ps(max128, _mm_movehl_ps(max128, max128));
    800
    801
            // Reduce 2 to 1
    802
    2.80k
            let max32 = _mm_max_ss(max64, _mm_shuffle_ps::<0x55>(max64, max64));
    803
    804
    2.80k
            _mm_cvtss_f32(max32)
    805
        }
    806
    2.80k
    }
    807
    808
    /// Horizontal sum of 8-wide AVX2 vector
    809
    #[cfg(target_arch = "x86_64")]
    810
    #[target_feature(enable = "avx2")]
    811
    #[inline]
    812
    #[allow(unsafe_op_in_unsafe_fn)]
    813
    0
    unsafe fn horizontal_sum_avx2(v: std::arch::x86_64::__m256) -> f32 {
    814
        use std::arch::x86_64::{
    815
            _mm256_extractf128_ps, _mm_add_ps, _mm_add_ss, _mm_cvtss_f32, _mm_movehl_ps, _mm_shuffle_ps,
    816
        };
    817
    818
        {
    819
            // Extract high and low 128-bit lanes
    820
    0
            let hi = _mm256_extractf128_ps::<1>(v);
    821
    0
            let lo = _mm256_extractf128_ps::<0>(v);
    822
    0
            let sum128 = _mm_add_ps(hi, lo);
    823
    824
            // Reduce 4 to 2
    825
    0
            let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
    826
    827
            // Reduce 2 to 1
    828
    0
            let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps::<0x55>(sum64, sum64));
    829
    830
    0
            _mm_cvtss_f32(sum32)
    831
        }
    832
    0
    }
    833
    834
    /// Quantize f32 activations to Q8_0 format for fast integer matmul
    835
    ///
    836
    /// Returns (scales, quantized_values) where each block of 32 values
    837
    /// has one f32 scale and 32 int8 quantized values.
    838
    #[inline]
    839
    806
    pub fn quantize_activations_q8_0(activations: &[f32]) -> (Vec<f32>, Vec<i8>) {
    840
    806
        let num_blocks = activations.len().div_ceil(32);
    841
    806
        let mut scales = Vec::with_capacity(num_blocks);
    842
    806
        let mut quants = Vec::with_capacity(num_blocks * 32);
    843
    844
    1.97k
        for block_idx in 0..num_blocks806 {
    845
    1.97k
            let start = block_idx * 32;
    846
    1.97k
            let end = (start + 32).min(activations.len());
    847
    848
            // Find max absolute value for symmetric quantization
    849
    1.97k
            let mut max_abs = 0.0f32;
    850
    63.0k
            for i in start1.97k..end1.97k {
    851
    63.0k
                let abs = activations[i].abs();
    852
    63.0k
                if abs > max_abs {
    853
    2.00k
                    max_abs = abs;
    854
    61.0k
                }
    855
            }
    856
    857
            // Compute scale (avoid division by zero)
    858
    1.97k
            let scale = if max_abs > 1e-10 {
    859
    936
                max_abs / 127.0
    860
            } else {
    861
    1.03k
                1.0 / 127.0
    862
            };
    863
    1.97k
            let inv_scale = 1.0 / scale;
    864
    1.97k
            scales.push(scale);
    865
    866
            // Quantize values
    867
    63.0k
            for i in start1.97k..end1.97k {
    868
    63.0k
                let q = (activations[i] * inv_scale).round();
    869
    63.0k
                quants.push(q.clamp(-128.0, 127.0) as i8);
    870
    63.0k
            }
    871
            // Pad to 32 if partial block
    872
    1.97k
            for _ in end..(start + 32) {
    873
    52
                quants.push(0i8);
    874
    52
            }
    875
        }
    876
    877
    806
        (scales, quants)
    878
    806
    }