## 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_dequant.rs

    Line
    Count
    Source
    1
    //! Parallel dequantization functions (PMAT-802)
    2
    //!
    3
    //! Implements parallel dequantization using rayon for multi-core acceleration:
    4
    //! - `dequantize_q4_k_parallel` - Parallel Q4_K dequantization
    5
    //! - `dequantize_q4_k_simd` - SIMD-accelerated Q4_K dequantization
    6
    //! - `dequantize_q8_0_parallel` - Parallel Q8_0 dequantization
    7
    //! - `dequantize_q8_0_simd` - SIMD-accelerated Q8_0 dequantization
    8
    9
    use crate::error::{RealizarError, Result};
    10
    use super::dequant::{read_f16, f16_to_f32};
    11
    use super::simd::extract_scale_min;
    12
    use super::types::QK_K;
    13
    14
    // ============================================================================
    15
    16
    /// Parallel Q4_K dequantization using rayon
    17
    ///
    18
    /// Processes super-blocks in parallel for faster model loading.
    19
    /// Each super-block (256 values) is dequantized independently.
    20
    ///
    21
    /// # Arguments
    22
    ///
    23
    /// * `data` - Raw Q4_K quantized data (144 bytes per super-block)
    24
    ///
    25
    /// # Returns
    26
    ///
    27
    /// Dequantized f32 values
    28
    ///
    29
    /// # Errors
    30
    ///
    31
    /// Returns error if data length is not a multiple of super-block size
    32
    ///
    33
    /// # Performance
    34
    ///
    35
    /// On a 16-core system, achieves ~10x speedup over serial dequantization
    36
    /// for large models (1B+ parameters).
    37
    7
    pub fn dequantize_q4_k_parallel(data: &[u8]) -> Result<Vec<f32>> {
    38
        use rayon::prelude::*;
    39
    40
        const SUPER_BLOCK_BYTES: usize = 144;
    41
    42
    7
        if !data.len().is_multiple_of(SUPER_BLOCK_BYTES) {
    43
    2
            return Err(RealizarError::InvalidShape {
    44
    2
                reason: format!(
    45
    2
                    "Q4_K data length {} is not a multiple of super-block size {}",
    46
    2
                    data.len(),
    47
    2
                    SUPER_BLOCK_BYTES
    48
    2
                ),
    49
    2
            });
    50
    5
        }
    51
    52
    5
        let num_super_blocks = data.len() / SUPER_BLOCK_BYTES;
    53
    54
        // Process super-blocks in parallel
    55
    5
        let result: Vec<f32> = (0..num_super_blocks)
    56
    5
            .into_par_iter()
    57
    8
            .flat_map5(|sb_idx| {
    58
    8
                let sb_start = sb_idx * SUPER_BLOCK_BYTES;
    59
    8
                let sb_data = &data[sb_start..sb_start + SUPER_BLOCK_BYTES];
    60
    8
                dequantize_q4_k_superblock(sb_data)
    61
    8
            })
    62
    5
            .collect();
    63
    64
    5
        Ok(result)
    65
    7
    }
    66
    67
    /// Dequantize a single Q4_K super-block (256 values)
    68
    ///
    69
    /// Internal helper for parallel processing.
    70
    ///
    71
    /// Exposed as `pub(crate)` for direct testing.
    72
    #[inline]
    73
    8
    pub(crate) fn dequantize_q4_k_superblock(sb_data: &[u8]) -> Vec<f32> {
    74
    8
        let mut result = vec![0.0f32; QK_K];
    75
    76
        // Read d (f16 -> f32)
    77
    8
        let d = read_f16(&sb_data[0..2]);
    78
    79
        // Read dmin (f16 -> f32)
    80
    8
        let dmin = read_f16(&sb_data[2..4]);
    81
    82
        // Read scales (12 bytes)
    83
    8
        let mut scales = [0u8; 12];
    84
    8
        scales.copy_from_slice(&sb_data[4..16]);
    85
    86
        // Read qs (128 bytes)
    87
    8
        let qs = &sb_data[16..144];
    88
    89
        // Dequantize following candle's layout:
    90
        // For each 64-value chunk, output 32 low nibbles then 32 high nibbles
    91
    8
        let mut ys_index = 0;
    92
    93
    32
        for j in (0..QK_K)8.step_by8(64) {
    94
    32
            let q = &qs[j / 2..j / 2 + 32];
    95
    96
            // Get scales for the two 32-value halves
    97
    32
            let is = j / 32;
    98
    32
            let (sc1, m1) = extract_scale_min(&scales, is);
    99
    32
            let d1 = d * sc1;
    100
    32
            let dm1 = dmin * m1;
    101
    102
    32
            let (sc2, m2) = extract_scale_min(&scales, is + 1);
    103
    32
            let d2 = d * sc2;
    104
    32
            let dm2 = dmin * m2;
    105
    106
            // First pass: 32 low nibbles
    107
    1.05k
            for &byte1.02k in q {
    108
    1.02k
                result[ys_index] = d1 * (byte & 0xF) as f32 - dm1;
    109
    1.02k
                ys_index += 1;
    110
    1.02k
            }
    111
    112
            // Second pass: 32 high nibbles
    113
    1.05k
            for &byte1.02k in q {
    114
    1.02k
                result[ys_index] = d2 * (byte >> 4) as f32 - dm2;
    115
    1.02k
                ys_index += 1;
    116
    1.02k
            }
    117
        }
    118
    119
    8
        result
    120
    8
    }
    121
    122
    /// SIMD-accelerated Q4_K dequantization with runtime feature detection
    123
    ///
    124
    /// Uses AVX2 when available, falls back to parallel scalar otherwise.
    125
    ///
    126
    /// # Arguments
    127
    ///
    128
    /// * `data` - Raw Q4_K quantized data (144 bytes per super-block)
    129
    ///
    130
    /// # Returns
    131
    ///
    132
    /// Dequantized f32 values
    133
    ///
    134
    /// # Errors
    135
    ///
    136
    /// Returns error if data length is not a multiple of super-block size
    137
    ///
    138
    /// # Performance
    139
    ///
    140
    /// AVX2: ~4x speedup over scalar per super-block
    141
    /// Combined with rayon: ~40x speedup on 16-core system
    142
    6.48k
    pub fn dequantize_q4_k_simd(data: &[u8]) -> Result<Vec<f32>> {
    143
        #[cfg(target_arch = "x86_64")]
    144
        {
    145
    6.48k
            if is_x86_feature_detected!("avx2") {
    146
                // SAFETY: AVX2 verified at runtime
    147
    6.48k
                return unsafe { dequantize_q4_k_avx2_parallel(data) };
    148
    0
            }
    149
        }
    150
    151
        // Fallback to parallel scalar
    152
    0
        dequantize_q4_k_parallel(data)
    153
    6.48k
    }
    154
    155
    /// AVX2-accelerated parallel Q4_K dequantization
    156
    ///
    157
    /// # Safety
    158
    ///
    159
    /// Caller must ensure AVX2 is available (use runtime feature detection)
    160
    #[cfg(target_arch = "x86_64")]
    161
    #[target_feature(enable = "avx2")]
    162
    6.48k
    unsafe fn dequantize_q4_k_avx2_parallel(data: &[u8]) -> Result<Vec<f32>> {
    163
        use rayon::prelude::*;
    164
    165
        const SUPER_BLOCK_BYTES: usize = 144;
    166
        // Process 64 super-blocks per parallel task to reduce scheduling overhead
    167
        const CHUNK_SIZE: usize = 64;
    168
        const CHUNK_BYTES: usize = SUPER_BLOCK_BYTES * CHUNK_SIZE;
    169
    170
    6.48k
        if !data.len().is_multiple_of(SUPER_BLOCK_BYTES) {
    171
    2
            return Err(RealizarError::InvalidShape {
    172
    2
                reason: format!(
    173
    2
                    "Q4_K data length {} is not a multiple of super-block size {}",
    174
    2
                    data.len(),
    175
    2
                    SUPER_BLOCK_BYTES
    176
    2
                ),
    177
    2
            });
    178
    6.48k
        }
    179
    180
    6.48k
        let num_super_blocks = data.len() / SUPER_BLOCK_BYTES;
    181
    182
        // For small data, skip parallelism overhead
    183
    6.48k
        if num_super_blocks < CHUNK_SIZE * 2 {
    184
    6.47k
            let mut result = Vec::with_capacity(num_super_blocks * QK_K);
    185
    6.54k
            for sb_idx in 0..num_super_blocks6.47k {
    186
    6.54k
                let sb_start = sb_idx * SUPER_BLOCK_BYTES;
    187
    6.54k
                let sb_data = &data[sb_start..sb_start + SUPER_BLOCK_BYTES];
    188
    6.54k
                // SAFETY: AVX2 availability verified by caller
    189
    6.54k
                result.extend(unsafe { dequantize_q4_k_superblock_avx2(sb_data) });
    190
    6.54k
            }
    191
    6.47k
            return Ok(result);
    192
    5
        }
    193
    194
        // Process chunks of super-blocks in parallel
    195
    5
        let result: Vec<f32> = data
    196
    5
            .par_chunks(CHUNK_BYTES)
    197
    136
            .flat_map5(|chunk| {
    198
    136
                let mut chunk_result = Vec::with_capacity(chunk.len() / SUPER_BLOCK_BYTES * QK_K);
    199
    8.70k
                for sb_data in chunk136.chunks_exact136(SUPER_BLOCK_BYTES) {
    200
    8.70k
                    // SAFETY: AVX2 availability verified by caller
    201
    8.70k
                    chunk_result.extend(unsafe { dequantize_q4_k_superblock_avx2(sb_data) });
    202
    8.70k
                }
    203
    136
                chunk_result
    204
    136
            })
    205
    5
            .collect();
    206
    207
    5
        Ok(result)
    208
    6.48k
    }
    209
    210
    /// AVX2 SIMD dequantization for a single Q4_K super-block
    211
    ///
    212
    /// Uses 256-bit SIMD to process 8 values simultaneously.
    213
    #[cfg(target_arch = "x86_64")]
    214
    #[target_feature(enable = "avx2", enable = "fma")]
    215
    15.2k
    unsafe fn dequantize_q4_k_superblock_avx2(sb_data: &[u8]) -> Vec<f32> {
    216
        #[allow(clippy::wildcard_imports)]
    217
        use std::arch::x86_64::*;
    218
    219
    15.2k
        let mut result = vec![0.0f32; QK_K];
    220
    221
        // Read d and dmin
    222
    15.2k
        let d = read_f16(&sb_data[0..2]);
    223
    15.2k
        let dmin = read_f16(&sb_data[2..4]);
    224
    225
        // SAFETY: AVX2 availability verified by caller's target_feature
    226
        unsafe {
    227
            // Read scales
    228
    15.2k
            let mut scales = [0u8; 12];
    229
    15.2k
            scales.copy_from_slice(&sb_data[4..16]);
    230
    231
    15.2k
            let qs = &sb_data[16..144];
    232
    233
            // Dequantize following candle's layout:
    234
            // For each 64-value chunk, output 32 low nibbles then 32 high nibbles
    235
    15.2k
            let mut ys_index = 0;
    236
    237
    60.9k
            for j in (0..QK_K)15.2k.step_by15.2k(64) {
    238
    60.9k
                let q = &qs[j / 2..j / 2 + 32];
    239
    240
                // Get scales for the two 32-value halves
    241
    60.9k
                let is = j / 32;
    242
    60.9k
                let (sc1, m1) = extract_scale_min(&scales, is);
    243
    60.9k
                let d1 = d * sc1;
    244
    60.9k
                let dm1 = dmin * m1;
    245
    60.9k
                let d1_vec = _mm256_set1_ps(d1);
    246
    60.9k
                let dm1_vec = _mm256_set1_ps(dm1);
    247
    248
    60.9k
                let (sc2, m2) = extract_scale_min(&scales, is + 1);
    249
    60.9k
                let d2 = d * sc2;
    250
    60.9k
                let dm2 = dmin * m2;
    251
    60.9k
                let d2_vec = _mm256_set1_ps(d2);
    252
    60.9k
                let dm2_vec = _mm256_set1_ps(dm2);
    253
    254
                // First pass: 32 low nibbles in 4 iterations of 8
    255
    304k
                for chunk243k in 0..4 {
    256
    243k
                    let byte_start = chunk * 8;
    257
    243k
    258
    243k
                    // Extract 8 low nibbles from 8 bytes
    259
    243k
                    let q0 = (q[byte_start] & 0x0F) as i32;
    260
    243k
                    let q1 = (q[byte_start + 1] & 0x0F) as i32;
    261
    243k
                    let q2 = (q[byte_start + 2] & 0x0F) as i32;
    262
    243k
                    let q3 = (q[byte_start + 3] & 0x0F) as i32;
    263
    243k
                    let q4 = (q[byte_start + 4] & 0x0F) as i32;
    264
    243k
                    let q5 = (q[byte_start + 5] & 0x0F) as i32;
    265
    243k
                    let q6 = (q[byte_start + 6] & 0x0F) as i32;
    266
    243k
                    let q7 = (q[byte_start + 7] & 0x0F) as i32;
    267
    243k
    268
    243k
                    let q_vec = _mm256_setr_epi32(q0, q1, q2, q3, q4, q5, q6, q7);
    269
    243k
                    let q_f32 = _mm256_cvtepi32_ps(q_vec);
    270
    243k
                    let dequant = _mm256_fmsub_ps(d1_vec, q_f32, dm1_vec);
    271
    243k
    272
    243k
                    _mm256_storeu_ps(result.as_mut_ptr().add(ys_index), dequant);
    273
    243k
                    ys_index += 8;
    274
    243k
                }
    275
    276
                // Second pass: 32 high nibbles in 4 iterations of 8
    277
    304k
                for chunk243k in 0..4 {
    278
    243k
                    let byte_start = chunk * 8;
    279
    243k
    280
    243k
                    // Extract 8 high nibbles from 8 bytes
    281
    243k
                    let q0 = (q[byte_start] >> 4) as i32;
    282
    243k
                    let q1 = (q[byte_start + 1] >> 4) as i32;
    283
    243k
                    let q2 = (q[byte_start + 2] >> 4) as i32;
    284
    243k
                    let q3 = (q[byte_start + 3] >> 4) as i32;
    285
    243k
                    let q4 = (q[byte_start + 4] >> 4) as i32;
    286
    243k
                    let q5 = (q[byte_start + 5] >> 4) as i32;
    287
    243k
                    let q6 = (q[byte_start + 6] >> 4) as i32;
    288
    243k
                    let q7 = (q[byte_start + 7] >> 4) as i32;
    289
    243k
    290
    243k
                    let q_vec = _mm256_setr_epi32(q0, q1, q2, q3, q4, q5, q6, q7);
    291
    243k
                    let q_f32 = _mm256_cvtepi32_ps(q_vec);
    292
    243k
                    let dequant = _mm256_fmsub_ps(d2_vec, q_f32, dm2_vec);
    293
    243k
    294
    243k
                    _mm256_storeu_ps(result.as_mut_ptr().add(ys_index), dequant);
    295
    243k
                    ys_index += 8;
    296
    243k
                }
    297
            }
    298
        }
    299
    300
    15.2k
        result
    301
    15.2k
    }
    302
    303
    /// Parallel Q8_0 dequantization using rayon
    304
    ///
    305
    /// Q8_0 is simpler than Q4_K (no scale packing), making SIMD even more effective.
    306
    ///
    307
    /// # Arguments
    308
    ///
    309
    /// * `data` - Raw Q8_0 quantized data (36 bytes per block: 4 scale + 32 quants)
    310
    ///
    311
    /// # Returns
    312
    ///
    313
    /// Dequantized f32 values
    314
    ///
    315
    /// # Errors
    316
    ///
    317
    /// Returns error if data length is not a multiple of block size
    318
    7
    pub fn dequantize_q8_0_parallel(data: &[u8]) -> Result<Vec<f32>> {
    319
        use rayon::prelude::*;
    320
    321
        const BLOCK_BYTES: usize = 34; // 2 (f16 scale) + 32 (i8 quants)
    322
    323
    7
        if !data.len().is_multiple_of(BLOCK_BYTES) {
    324
    2
            return Err(RealizarError::InvalidShape {
    325
    2
                reason: format!(
    326
    2
                    "Q8_0 data length {} is not a multiple of block size {}",
    327
    2
                    data.len(),
    328
    2
                    BLOCK_BYTES
    329
    2
                ),
    330
    2
            });
    331
    5
        }
    332
    333
    5
        let num_blocks = data.len() / BLOCK_BYTES;
    334
    335
        // Process blocks in parallel
    336
    5
        let result: Vec<f32> = (0..num_blocks)
    337
    5
            .into_par_iter()
    338
    1.00k
            .flat_map5(|block_idx| {
    339
    1.00k
                let block_start = block_idx * BLOCK_BYTES;
    340
    1.00k
                let block_data = &data[block_start..block_start + BLOCK_BYTES];
    341
    1.00k
                dequantize_q8_0_block(block_data)
    342
    1.00k
            })
    343
    5
            .collect();
    344
    345
    5
        Ok(result)
    346
    7
    }
    347
    348
    /// Dequantize a single Q8_0 block (32 values)
    349
    ///
    350
    /// Exposed as `pub(crate)` for direct testing.
    351
    #[inline]
    352
    1.00k
    pub(crate) fn dequantize_q8_0_block(block_data: &[u8]) -> Vec<f32> {
    353
    1.00k
        let mut result = Vec::with_capacity(32);
    354
    355
        // Read scale (f16 -> f32)
    356
    1.00k
        let scale_bits = u16::from_le_bytes([block_data[0], block_data[1]]);
    357
    1.00k
        let scale = f16_to_f32(scale_bits);
    358
    359
        // Dequantize 32 int8 values
    360
    32.1k
        for &byte in &block_data[2..34]1.00k {
    361
    32.1k
            let value = i8::from_le_bytes([byte]);
    362
    32.1k
            result.push(scale * f32::from(value));
    363
    32.1k
        }
    364
    365
    1.00k
        result
    366
    1.00k
    }
    367
    368
    /// SIMD-accelerated Q8_0 dequantization
    369
    ///
    370
    /// Uses AVX2 when available for 8x parallel i8→f32 conversion.
    371
    ///
    372
    /// # Errors
    373
    ///
    374
    /// Returns error if data length is not a multiple of block size (36 bytes)
    375
    7
    pub fn dequantize_q8_0_simd(data: &[u8]) -> Result<Vec<f32>> {
    376
        #[cfg(target_arch = "x86_64")]
    377
        {
    378
    7
            if is_x86_feature_detected!("avx2") {
    379
                // SAFETY: AVX2 verified at runtime
    380
    7
                return unsafe { dequantize_q8_0_avx2_parallel(data) };
    381
    0
            }
    382
        }
    383
    384
        // Fallback to parallel scalar
    385
    0
        dequantize_q8_0_parallel(data)
    386
    7
    }
    387
    388
    /// AVX2-accelerated parallel Q8_0 dequantization
    389
    #[cfg(target_arch = "x86_64")]
    390
    #[target_feature(enable = "avx2")]
    391
    7
    unsafe fn dequantize_q8_0_avx2_parallel(data: &[u8]) -> Result<Vec<f32>> {
    392
        use rayon::prelude::*;
    393
    394
        const BLOCK_BYTES: usize = 34; // 2 (f16 scale) + 32 (i8 quants)
    395
    396
    7
        if !data.len().is_multiple_of(BLOCK_BYTES) {
    397
    2
            return Err(RealizarError::InvalidShape {
    398
    2
                reason: format!(
    399
    2
                    "Q8_0 data length {} is not a multiple of block size {}",
    400
    2
                    data.len(),
    401
    2
                    BLOCK_BYTES
    402
    2
                ),
    403
    2
            });
    404
    5
        }
    405
    406
    5
        let num_blocks = data.len() / BLOCK_BYTES;
    407
    408
    5
        let result: Vec<f32> = (0..num_blocks)
    409
    5
            .into_par_iter()
    410
    5
            .flat_map(|block_idx| {
    411
    5
                let block_start = block_idx * BLOCK_BYTES;
    412
    5
                let block_data = &data[block_start..block_start + BLOCK_BYTES];
    413
                // SAFETY: AVX2 availability verified by caller
    414
    5
                unsafe { dequantize_q8_0_block_avx2(block_data) }
    415
    5
            })
    416
    5
            .collect();
    417
    418
    5
        Ok(result)
    419
    7
    }
    420
    421
    /// AVX2 SIMD dequantization for a single Q8_0 block
    422
    #[cfg(target_arch = "x86_64")]
    423
    #[target_feature(enable = "avx2")]
    424
    5
    unsafe fn dequantize_q8_0_block_avx2(block_data: &[u8]) -> Vec<f32> {
    425
        #[allow(clippy::wildcard_imports)]
    426
        use std::arch::x86_64::*;
    427
    428
    5
        let mut result = vec![0.0f32; 32];
    429
    430
        // Read scale (f16 -> f32)
    431
    5
        let scale_bits = u16::from_le_bytes([block_data[0], block_data[1]]);
    432
    5
        let scale = f16_to_f32(scale_bits);
    433
    434
        // SAFETY: AVX2 availability verified by caller's target_feature
    435
        unsafe {
    436
    5
            let scale_vec = _mm256_set1_ps(scale);
    437
    438
            // Process 32 i8 values in 4 iterations of 8
    439
    25
            for chunk20 in 0..4 {
    440
    20
                let byte_start = 2 + chunk * 8; // Start at offset 2 (after f16 scale)
    441
    20
    442
    20
                // Load 8 i8 values and sign-extend to i32
    443
    20
                let q0 = block_data[byte_start] as i8 as i32;
    444
    20
                let q1 = block_data[byte_start + 1] as i8 as i32;
    445
    20
                let q2 = block_data[byte_start + 2] as i8 as i32;
    446
    20
                let q3 = block_data[byte_start + 3] as i8 as i32;
    447
    20
                let q4 = block_data[byte_start + 4] as i8 as i32;
    448
    20
                let q5 = block_data[byte_start + 5] as i8 as i32;
    449
    20
                let q6 = block_data[byte_start + 6] as i8 as i32;
    450
    20
                let q7 = block_data[byte_start + 7] as i8 as i32;
    451
    20
    452
    20
                let q_vec = _mm256_setr_epi32(q0, q1, q2, q3, q4, q5, q6, q7);
    453
    20
                let q_f32 = _mm256_cvtepi32_ps(q_vec);
    454
    20
    455
    20
                // Multiply by scale
    456
    20
                let dequant = _mm256_mul_ps(scale_vec, q_f32);
    457
    20
    458
    20
                // Store 8 results
    459
    20
                _mm256_storeu_ps(result.as_mut_ptr().add(chunk * 8), dequant);
    460
    20
            }
    461
        }
    462
    463
    5
        result
    464
    5
    }
    465
    466
    // DequantStats, SimdBackend, detect_simd_backend moved to types.rs (PMAT-802)
    467
    468
    /// SIMD-optimized RoPE rotation for a single head
    469
    ///
    470
    /// Applies rotary position embedding rotation to a single attention head:
    471
    /// x1[i] = x1[i] * cos[i] - x2[i] * sin[i]
    472
    /// x2[i] = x1[i] * sin[i] + x2[i] * cos[i]
    473
    ///
    474
    /// # Arguments
    475
    /// * `x1` - First half of head (will be modified in-place)
    476
    /// * `x2` - Second half of head (will be modified in-place)
    477
    /// * `cos_vals` - Precomputed cosine values
    478
    /// * `sin_vals` - Precomputed sine values
    479
    #[inline]
    480
    9
    pub fn apply_rope_rotation_simd(
    481
    9
        x1: &mut [f32],
    482
    9
        x2: &mut [f32],
    483
    9
        cos_vals: &[f32],
    484
    9
        sin_vals: &[f32],
    485
    9
    ) {
    486
    9
        debug_assert_eq!(x1.len(), x2.len());
    487
    9
        debug_assert_eq!(x1.len(), cos_vals.len());
    488
    9
        debug_assert_eq!(x1.len(), sin_vals.len());
    489
    490
        #[cfg(target_arch = "x86_64")]
    491
        {
    492
    9
            if is_x86_feature_detected!("avx512f") {
    493
                // SAFETY: Memory safety ensured by bounds checking and alignment
    494
    9
                unsafe {
    495
    9
                    apply_rope_rotation_avx512(x1, x2, cos_vals, sin_vals);
    496
    9
                }
    497
    9
                return;
    498
    0
            }
    499
    0
            if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
    500
                // SAFETY: Memory safety ensured by bounds checking and alignment
    501
    0
                unsafe {
    502
    0
                    apply_rope_rotation_avx2(x1, x2, cos_vals, sin_vals);
    503
    0
                }
    504
    0
                return;
    505
    0
            }
    506
        }
    507
    508
        // Scalar fallback
    509
    0
        apply_rope_rotation_scalar(x1, x2, cos_vals, sin_vals);
    510
    9
    }
    511
    512
    /// Scalar fallback for RoPE rotation
    513
    ///
    514
    /// Exposed as `pub(crate)` for direct testing on AVX2 machines.
    515
    #[inline]
    516
    0
    pub(crate) fn apply_rope_rotation_scalar(
    517
    0
        x1: &mut [f32],
    518
    0
        x2: &mut [f32],
    519
    0
        cos_vals: &[f32],
    520
    0
        sin_vals: &[f32],
    521
    0
    ) {
    522
    0
        for i in 0..x1.len() {
    523
    0
            let v1 = x1[i];
    524
    0
            let v2 = x2[i];
    525
    0
            let cos_v = cos_vals[i];
    526
    0
            let sin_v = sin_vals[i];
    527
    0
            x1[i] = v1 * cos_v - v2 * sin_v;
    528
    0
            x2[i] = v1 * sin_v + v2 * cos_v;
    529
    0
        }
    530
    0
    }
    531
    532
    #[cfg(target_arch = "x86_64")]
    533
    #[target_feature(enable = "avx2", enable = "fma")]
    534
    #[inline]
    535
    #[allow(unsafe_op_in_unsafe_fn)]
    536
    0
    unsafe fn apply_rope_rotation_avx2(
    537
    0
        x1: &mut [f32],
    538
    0
        x2: &mut [f32],
    539
    0
        cos_vals: &[f32],
    540
    0
        sin_vals: &[f32],
    541
    0
    ) {
    542
        use std::arch::x86_64::{
    543
            _mm256_fmadd_ps, _mm256_fnmadd_ps, _mm256_loadu_ps, _mm256_mul_ps, _mm256_storeu_ps,
    544
        };
    545
    546
    0
        let n = x1.len();
    547
    0
        let mut i = 0;
    548
    549
        // Process 8 elements at a time
    550
    0
        while i + 8 <= n {
    551
    0
            let v1 = _mm256_loadu_ps(x1.as_ptr().add(i));
    552
    0
            let v2 = _mm256_loadu_ps(x2.as_ptr().add(i));
    553
    0
            let cos_v = _mm256_loadu_ps(cos_vals.as_ptr().add(i));
    554
    0
            let sin_v = _mm256_loadu_ps(sin_vals.as_ptr().add(i));
    555
    0
    556
    0
            // r1 = v1 * cos - v2 * sin
    557
    0
            let v1_cos = _mm256_mul_ps(v1, cos_v);
    558
    0
            let r1 = _mm256_fnmadd_ps(v2, sin_v, v1_cos);
    559
    0
    560
    0
            // r2 = v1 * sin + v2 * cos
    561
    0
            let v1_sin = _mm256_mul_ps(v1, sin_v);
    562
    0
            let r2 = _mm256_fmadd_ps(v2, cos_v, v1_sin);
    563
    0
    564
    0
            _mm256_storeu_ps(x1.as_mut_ptr().add(i), r1);
    565
    0
            _mm256_storeu_ps(x2.as_mut_ptr().add(i), r2);
    566
    0
    567
    0
            i += 8;
    568
    0
        }
    569
    570
        // Handle remainder
    571
    0
        while i < n {
    572
    0
            let v1 = x1[i];
    573
    0
            let v2 = x2[i];
    574
    0
            let cos_v = cos_vals[i];
    575
    0
            let sin_v = sin_vals[i];
    576
    0
            x1[i] = v1 * cos_v - v2 * sin_v;
    577
    0
            x2[i] = v1 * sin_v + v2 * cos_v;
    578
    0
            i += 1;
    579
    0
        }
    580
    0
    }
    581
    582
    #[cfg(target_arch = "x86_64")]
    583
    #[target_feature(enable = "avx512f")]
    584
    #[inline]
    585
    #[allow(unsafe_op_in_unsafe_fn)]
    586
    9
    unsafe fn apply_rope_rotation_avx512(
    587
    9
        x1: &mut [f32],
    588
    9
        x2: &mut [f32],
    589
    9
        cos_vals: &[f32],
    590
    9
        sin_vals: &[f32],
    591
    9
    ) {
    592
        use std::arch::x86_64::{
    593
            _mm512_fmadd_ps, _mm512_fnmadd_ps, _mm512_loadu_ps, _mm512_mul_ps, _mm512_storeu_ps,
    594
        };
    595
    596
    9
        let n = x1.len();
    597
    9
        let mut i = 0;
    598
    599
        // Process 16 elements at a time with AVX-512
    600
    16
        while i + 16 <= n {
    601
    7
            let v1 = _mm512_loadu_ps(x1.as_ptr().add(i));
    602
    7
            let v2 = _mm512_loadu_ps(x2.as_ptr().add(i));
    603
    7
            let cos_v = _mm512_loadu_ps(cos_vals.as_ptr().add(i));
    604
    7
            let sin_v = _mm512_loadu_ps(sin_vals.as_ptr().add(i));
    605
    7
    606
    7
            // r1 = v1 * cos - v2 * sin
    607
    7
            let v1_cos = _mm512_mul_ps(v1, cos_v);
    608
    7
            let r1 = _mm512_fnmadd_ps(v2, sin_v, v1_cos);
    609
    7
    610
    7
            // r2 = v1 * sin + v2 * cos
    611
    7
            let v1_sin = _mm512_mul_ps(v1, sin_v);
    612
    7
            let r2 = _mm512_fmadd_ps(v2, cos_v, v1_sin);
    613
    7
    614
    7
            _mm512_storeu_ps(x1.as_mut_ptr().add(i), r1);
    615
    7
            _mm512_storeu_ps(x2.as_mut_ptr().add(i), r2);
    616
    7
    617
    7
            i += 16;
    618
    7
        }
    619
    620
        // Handle remainder with AVX2 or scalar
    621
    20
        while i < n {
    622
    11
            let v1 = x1[i];
    623
    11
            let v2 = x2[i];
    624
    11
            let cos_v = cos_vals[i];
    625
    11
            let sin_v = sin_vals[i];
    626
    11
            x1[i] = v1 * cos_v - v2 * sin_v;
    627
    11
            x2[i] = v1 * sin_v + v2 * cos_v;
    628
    11
            i += 1;
    629
    11
        }
    630
    9
    }