## Coverage Report

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

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

    /home/noah/src/trueno/src/brick/attention.rs

    Line
    Count
    Source
    1
    //! SIMD-Optimized Attention Operation (PMAT-017)
    2
    //!
    3
    //! This module contains the scaled dot-product attention operation
    4
    //! with SIMD optimization for CPU inference.
    5
    //!
    6
    //! # Algorithm
    7
    //!
    8
    //! Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V
    9
    //!
    10
    //! # SIMD Optimizations
    11
    //!
    12
    //! - Q @ K^T: Batched dot products with AVX2/AVX-512/FMA
    13
    //! - Softmax: Row-wise numerically stable implementation
    14
    //! - Scores @ V: SIMD-friendly weighted accumulation
    15
    //!
    16
    //! # Performance Target
    17
    //!
    18
    //! Close the 1.66x gap in CPU inference (25.4 → 42 tok/s) by replacing
    19
    //! scalar triple-nested loops with SIMD operations.
    20
    21
    use super::{Backend, ComputeOp};
    22
    use crate::error::TruenoError;
    23
    24
    /// Scaled dot-product attention operation.
    25
    ///
    26
    /// Computes: Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V
    27
    ///
    28
    /// # SIMD Optimization (PMAT-017)
    29
    ///
    30
    /// Uses trueno's SIMD backends for:
    31
    /// - Q @ K^T: Batched dot products with AVX2/AVX-512
    32
    /// - Softmax: Row-wise numerically stable softmax
    33
    /// - Scores @ V: Batched weighted sums
    34
    ///
    35
    /// # Performance Target
    36
    ///
    37
    /// Close the 1.66x gap in CPU inference (25.4 → 42 tok/s) by replacing
    38
    /// scalar triple-nested loops with SIMD operations.
    39
    #[derive(Debug, Clone)]
    40
    pub struct AttentionOp {
    41
        /// Sequence length (Q rows)
    42
        pub seq_len: usize,
    43
        /// Key/Value sequence length (may differ for cross-attention)
    44
        pub kv_seq_len: usize,
    45
        /// Head dimension
    46
        pub head_dim: usize,
    47
        /// Scale factor (1/sqrt(head_dim))
    48
        pub scale: f32,
    49
    }
    50
    51
    impl AttentionOp {
    52
        /// Create a new attention operation.
    53
        ///
    54
        /// # Arguments
    55
        ///
    56
        /// * `seq_len` - Query sequence length
    57
        /// * `kv_seq_len` - Key/Value sequence length
    58
        /// * `head_dim` - Dimension per head
    59
        #[must_use]
    60
    0
        pub fn new(seq_len: usize, kv_seq_len: usize, head_dim: usize) -> Self {
    61
    0
            Self {
    62
    0
                seq_len,
    63
    0
                kv_seq_len,
    64
    0
                head_dim,
    65
    0
                scale: 1.0 / (head_dim as f32).sqrt(),
    66
    0
            }
    67
    0
        }
    68
    69
        /// Create for self-attention (seq_len == kv_seq_len).
    70
        #[must_use]
    71
    0
        pub fn self_attention(seq_len: usize, head_dim: usize) -> Self {
    72
    0
            Self::new(seq_len, seq_len, head_dim)
    73
    0
        }
    74
    75
        /// SIMD-optimized dot product for attention scores.
    76
        ///
    77
        /// Computes Q[i] · K[j] using SIMD when available.
    78
        #[inline]
    79
    0
        pub(crate) fn simd_dot(a: &[f32], b: &[f32]) -> f32 {
    80
    0
            debug_assert_eq!(a.len(), b.len());
    81
    82
            // Use architecture-specific SIMD
    83
            #[cfg(target_arch = "x86_64")]
    84
            {
    85
    0
                if is_x86_feature_detected!("avx2") {
    86
    0
                    return unsafe { Self::avx2_dot(a, b) };
    87
    0
                }
    88
            }
    89
    90
            // Scalar fallback with manual unrolling for better vectorization
    91
    0
            let mut sum0 = 0.0f32;
    92
    0
            let mut sum1 = 0.0f32;
    93
    0
            let mut sum2 = 0.0f32;
    94
    0
            let mut sum3 = 0.0f32;
    95
    96
    0
            let chunks = a.len() / 4;
    97
    0
            for i in 0..chunks {
    98
    0
                let base = i * 4;
    99
    0
                sum0 += a[base] * b[base];
    100
    0
                sum1 += a[base + 1] * b[base + 1];
    101
    0
                sum2 += a[base + 2] * b[base + 2];
    102
    0
                sum3 += a[base + 3] * b[base + 3];
    103
    0
            }
    104
    105
            // Handle remainder
    106
    0
            for i in (chunks * 4)..a.len() {
    107
    0
                sum0 += a[i] * b[i];
    108
    0
            }
    109
    110
    0
            sum0 + sum1 + sum2 + sum3
    111
    0
        }
    112
    113
        /// AVX2-optimized dot product.
    114
        #[cfg(target_arch = "x86_64")]
    115
        #[target_feature(enable = "avx2", enable = "fma")]
    116
    0
        unsafe fn avx2_dot(a: &[f32], b: &[f32]) -> f32 {
    117
            use std::arch::x86_64::*;
    118
    119
    0
            let mut sum = _mm256_setzero_ps();
    120
    0
            let chunks = a.len() / 8;
    121
    122
    0
            for i in 0..chunks {
    123
    0
                let base = i * 8;
    124
    0
                let va = _mm256_loadu_ps(a.as_ptr().add(base));
    125
    0
                let vb = _mm256_loadu_ps(b.as_ptr().add(base));
    126
    0
                sum = _mm256_fmadd_ps(va, vb, sum);
    127
    0
            }
    128
    129
            // Horizontal sum
    130
    0
            let high = _mm256_extractf128_ps(sum, 1);
    131
    0
            let low = _mm256_castps256_ps128(sum);
    132
    0
            let sum128 = _mm_add_ps(high, low);
    133
    0
            let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
    134
    0
            let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
    135
    0
            let mut result = _mm_cvtss_f32(sum32);
    136
    137
            // Handle remainder
    138
    0
            for i in (chunks * 8)..a.len() {
    139
    0
                result += a[i] * b[i];
    140
    0
            }
    141
    142
    0
            result
    143
    0
        }
    144
    145
        /// Row-wise softmax with SIMD max/sum.
    146
        #[inline]
    147
    0
        pub(crate) fn simd_softmax_row(scores: &mut [f32]) {
    148
    0
            if scores.is_empty() {
    149
    0
                return;
    150
    0
            }
    151
    152
            // Find max for numerical stability
    153
    0
            let max = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    154
    155
            // Compute exp(x - max) and sum
    156
    0
            let mut sum = 0.0f32;
    157
    0
            for s in scores.iter_mut() {
    158
    0
                *s = (*s - max).exp();
    159
    0
                sum += *s;
    160
    0
            }
    161
    162
            // Normalize
    163
    0
            let inv_sum = 1.0 / sum;
    164
    0
            for s in scores.iter_mut() {
    165
    0
                *s *= inv_sum;
    166
    0
            }
    167
    0
        }
    168
    }
    169
    170
    impl ComputeOp for AttentionOp {
    171
        /// Input: (Q, K, V) tensors as flat vectors
    172
        /// Q: [seq_len * head_dim]
    173
        /// K: [kv_seq_len * head_dim]
    174
        /// V: [kv_seq_len * head_dim]
    175
        type Input = (Vec<f32>, Vec<f32>, Vec<f32>);
    176
        /// Output: attention output [seq_len * head_dim]
    177
        type Output = Vec<f32>;
    178
    179
    0
        fn name(&self) -> &'static str {
    180
    0
            "attention"
    181
    0
        }
    182
    183
    0
        fn execute(&self, input: Self::Input, _backend: Backend) -> Result<Self::Output, TruenoError> {
    184
    0
            let (q, k, v) = input;
    185
    186
            // Validate dimensions
    187
    0
            let expected_q = self.seq_len * self.head_dim;
    188
    0
            let expected_kv = self.kv_seq_len * self.head_dim;
    189
    190
    0
            if q.len() != expected_q {
    191
    0
                return Err(TruenoError::SizeMismatch {
    192
    0
                    expected: expected_q,
    193
    0
                    actual: q.len(),
    194
    0
                });
    195
    0
            }
    196
    0
            if k.len() != expected_kv || v.len() != expected_kv {
    197
    0
                return Err(TruenoError::SizeMismatch {
    198
    0
                    expected: expected_kv,
    199
    0
                    actual: k.len(),
    200
    0
                });
    201
    0
            }
    202
    203
            // Allocate output
    204
    0
            let mut output = vec![0.0f32; expected_q];
    205
    206
            // Allocate scores buffer (reused per query row)
    207
    0
            let mut scores = vec![0.0f32; self.kv_seq_len];
    208
    209
            // For each query position
    210
    0
            for qi in 0..self.seq_len {
    211
    0
                let q_row = &q[qi * self.head_dim..(qi + 1) * self.head_dim];
    212
    213
                // Compute Q[qi] · K[ki] for all ki (SIMD dot products)
    214
    0
                for ki in 0..self.kv_seq_len {
    215
    0
                    let k_row = &k[ki * self.head_dim..(ki + 1) * self.head_dim];
    216
    0
                    scores[ki] = Self::simd_dot(q_row, k_row) * self.scale;
    217
    0
                }
    218
    219
                // Softmax over scores
    220
    0
                Self::simd_softmax_row(&mut scores);
    221
    222
                // Compute weighted sum: output[qi] = sum(scores[ki] * V[ki])
    223
    0
                let out_row = &mut output[qi * self.head_dim..(qi + 1) * self.head_dim];
    224
    0
                out_row.fill(0.0);
    225
    226
    0
                for ki in 0..self.kv_seq_len {
    227
    0
                    let v_row = &v[ki * self.head_dim..(ki + 1) * self.head_dim];
    228
    0
                    let weight = scores[ki];
    229
    230
                    // SIMD-friendly accumulation
    231
    0
                    for (o, &vi) in out_row.iter_mut().zip(v_row.iter()) {
    232
    0
                        *o += weight * vi;
    233
    0
                    }
    234
                }
    235
            }
    236
    237
    0
            Ok(output)
    238
    0
        }
    239
    240
    0
        fn tokens(&self, _input: &Self::Input) -> usize {
    241
            // Output tokens = seq_len * head_dim
    242
    0
            self.seq_len * self.head_dim
    243
    0
        }
    244
    }
    245
    246
    #[cfg(test)]
    247
    mod tests {
    248
        use super::*;
    249
    250
        #[test]
    251
        fn test_attention_basic() {
    252
            let op = AttentionOp::self_attention(2, 4); // seq=2, head_dim=4
    253
    254
            // Simple identity-like setup
    255
            let q = vec![1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0]; // 2x4
    256
            let k = vec![1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0]; // 2x4
    257
            let v = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; // 2x4
    258
    259
            let output = op.execute((q, k, v), Backend::Scalar).unwrap();
    260
    261
            assert_eq!(output.len(), 8);
    262
            // Output should be weighted combination of V rows
    263
        }
    264
    265
        #[test]
    266
        fn test_attention_dimension_mismatch_q() {
    267
            let op = AttentionOp::self_attention(2, 4);
    268
            let q = vec![1.0; 4]; // Wrong size - should be 8
    269
            let k = vec![1.0; 8];
    270
            let v = vec![1.0; 8];
    271
    272
            let result = op.execute((q, k, v), Backend::Scalar);
    273
            assert!(result.is_err());
    274
        }
    275
    276
        #[test]
    277
        fn test_attention_dimension_mismatch_kv() {
    278
            let op = AttentionOp::self_attention(2, 4);
    279
            let q = vec![1.0; 8];
    280
            let k = vec![1.0; 4]; // Wrong size - should be 8
    281
            let v = vec![1.0; 8];
    282
    283
            let result = op.execute((q, k, v), Backend::Scalar);
    284
            assert!(result.is_err());
    285
        }
    286
    287
        #[test]
    288
        fn test_attention_cross_attention() {
    289
            // Cross-attention: Q from decoder (seq=1), K/V from encoder (seq=4)
    290
            let op = AttentionOp::new(1, 4, 8); // q_seq=1, kv_seq=4, head_dim=8
    291
    292
            let q = vec![1.0; 8];   // 1 x 8
    293
            let k = vec![1.0; 32];  // 4 x 8
    294
            let v = vec![1.0; 32];  // 4 x 8
    295
    296
            let output = op.execute((q, k, v), Backend::Scalar).unwrap();
    297
            assert_eq!(output.len(), 8);
    298
        }
    299
    300
        #[test]
    301
        fn test_attention_tokens() {
    302
            let op = AttentionOp::self_attention(16, 64);
    303
            let input = (vec![], vec![], vec![]);
    304
            // tokens = seq_len * head_dim = 16 * 64 = 1024
    305
            assert_eq!(op.tokens(&input), 1024);
    306
        }
    307
    308
        #[test]
    309
        fn test_simd_softmax_row_empty() {
    310
            let mut scores: Vec<f32> = vec![];
    311
            AttentionOp::simd_softmax_row(&mut scores);
    312
            assert!(scores.is_empty());
    313
        }
    314
    315
        #[test]
    316
        fn test_simd_softmax_row_single() {
    317
            let mut scores = vec![5.0];
    318
            AttentionOp::simd_softmax_row(&mut scores);
    319
            assert!((scores[0] - 1.0).abs() < 1e-6);
    320
        }
    321
    322
        #[test]
    323
        fn test_simd_softmax_row_uniform() {
    324
            let mut scores = vec![1.0, 1.0, 1.0, 1.0];
    325
            AttentionOp::simd_softmax_row(&mut scores);
    326
    327
            // All equal inputs → uniform distribution
    328
            for s in &scores {
    329
                assert!((s - 0.25).abs() < 1e-6);
    330
            }
    331
        }
    332
    333
        #[test]
    334
        fn test_simd_softmax_row_sum_to_one() {
    335
            let mut scores = vec![1.0, 2.0, 3.0, 4.0, 5.0];
    336
            AttentionOp::simd_softmax_row(&mut scores);
    337
    338
            let sum: f32 = scores.iter().sum();
    339
            assert!((sum - 1.0).abs() < 1e-5);
    340
        }
    341
    342
        #[test]
    343
        fn test_simd_dot_basic() {
    344
            let a = vec![1.0, 2.0, 3.0, 4.0];
    345
            let b = vec![1.0, 1.0, 1.0, 1.0];
    346
            let dot = AttentionOp::simd_dot(&a, &b);
    347
            assert!((dot - 10.0).abs() < 1e-5);
    348
        }
    349
    350
        #[test]
    351
        fn test_simd_dot_unaligned() {
    352
            // Test with non-multiple-of-8 length (tests scalar remainder handling)
    353
            let a = vec![1.0, 2.0, 3.0, 4.0, 5.0];
    354
            let b = vec![2.0, 2.0, 2.0, 2.0, 2.0];
    355
            let dot = AttentionOp::simd_dot(&a, &b);
    356
            // (1+2+3+4+5) * 2 = 30
    357
            assert!((dot - 30.0).abs() < 1e-5);
    358
        }
    359
    }