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

    Line
    Count
    Source
    1
    //! Fused Operations for Transformer Inference
    2
    //!
    3
    //! This module contains fused compute operations that combine multiple
    4
    //! operations into single passes for improved performance.
    5
    //!
    6
    //! # Operations
    7
    //!
    8
    //! - `FusedQKVOp`: Fused Query/Key/Value projection (3 matmuls → 1)
    9
    //! - `FusedGateUpOp`: Fused Gate+Up FFN projection with SiLU (SwiGLU)
    10
    //!
    11
    //! # Performance Impact
    12
    //!
    13
    //! Fusing operations provides:
    14
    //! - Reduced kernel launches (GPU)
    15
    //! - Better cache utilization (data loaded once)
    16
    //! - Eliminated intermediate memory traffic
    17
    18
    use super::{Backend, ComputeOp};
    19
    use crate::error::TruenoError;
    20
    21
    // ============================================================================
    22
    // Fused Q/K/V Projection (PMAT-PERF-009)
    23
    // ============================================================================
    24
    25
    /// Weights for fused QKV projection
    26
    #[derive(Debug, Clone)]
    27
    pub struct FusedQKVWeights {
    28
        /// Q projection weights [hidden_size, hidden_size]
    29
        pub q_weight: Vec<f32>,
    30
        /// K projection weights [hidden_size, kv_dim]
    31
        pub k_weight: Vec<f32>,
    32
        /// V projection weights [hidden_size, kv_dim]
    33
        pub v_weight: Vec<f32>,
    34
    }
    35
    36
    /// Fused Q/K/V projection operation for transformer attention.
    37
    ///
    38
    /// Computes Q, K, V projections in a single pass over the input:
    39
    /// - Q = x * W_q (hidden_size → hidden_size)
    40
    /// - K = x * W_k (hidden_size → kv_dim)
    41
    /// - V = x * W_v (hidden_size → kv_dim)
    42
    ///
    43
    /// # Performance Impact
    44
    ///
    45
    /// Fusing 3 separate matmuls into 1 operation provides:
    46
    /// - 3x reduction in kernel launches (GPU)
    47
    /// - Better cache utilization (input x loaded once)
    48
    /// - Expected speedup: 2-3x for decode phase
    49
    ///
    50
    /// # Five-Whys Root Cause (PMAT-PERF-009)
    51
    ///
    52
    /// ```text
    53
    /// Why 1: Why is decode throughput 131 tok/s vs 400 tok/s target?
    54
    /// → 280+ kernel launches per token (10+ per layer × 28 layers)
    55
    ///
    56
    /// Why 2: Why so many kernel launches?
    57
    /// → Q, K, V computed as 3 separate GEMV operations
    58
    ///
    59
    /// Why 3: Why separate operations?
    60
    /// → Original implementation didn't consider launch overhead
    61
    ///
    62
    /// Why 4: Why does launch overhead matter?
    63
    /// → GPU kernel launch: ~5-10µs, 280 launches = 1.4-2.8ms overhead/token
    64
    ///
    65
    /// Why 5: ROOT CAUSE
    66
    /// → Kernel launch overhead (2.8ms) exceeds compute time for small batch decode
    67
    /// → FIX: Fuse Q/K/V into single kernel, reducing launches by 2/3
    68
    /// ```
    69
    #[derive(Debug, Clone)]
    70
    pub struct FusedQKVOp {
    71
        /// Hidden dimension size
    72
        pub hidden_size: usize,
    73
        /// KV dimension (num_kv_heads * head_dim, may differ from hidden_size for GQA)
    74
        pub kv_dim: usize,
    75
        /// Number of attention heads
    76
        pub num_heads: usize,
    77
        /// Head dimension
    78
        pub head_dim: usize,
    79
    }
    80
    81
    impl FusedQKVOp {
    82
        /// Create a new fused QKV operation.
    83
        ///
    84
        /// # Arguments
    85
        /// * `hidden_size` - Hidden dimension (e.g., 3584 for Qwen 3B)
    86
        /// * `num_heads` - Number of attention heads
    87
        /// * `num_kv_heads` - Number of KV heads (may differ for GQA)
    88
    0
        pub fn new(hidden_size: usize, num_heads: usize, num_kv_heads: usize) -> Self {
    89
    0
            let head_dim = hidden_size / num_heads;
    90
    0
            let kv_dim = num_kv_heads * head_dim;
    91
    0
            Self {
    92
    0
                hidden_size,
    93
    0
                kv_dim,
    94
    0
                num_heads,
    95
    0
                head_dim,
    96
    0
            }
    97
    0
        }
    98
    }
    99
    100
    #[allow(clippy::needless_range_loop)] // Matrix indexing is clearer with explicit loops
    101
    impl ComputeOp for FusedQKVOp {
    102
        type Input = (Vec<f32>, FusedQKVWeights);
    103
        type Output = (Vec<f32>, Vec<f32>, Vec<f32>); // (Q, K, V)
    104
    105
    0
        fn name(&self) -> &'static str {
    106
    0
            "fused_qkv"
    107
    0
        }
    108
    109
    0
        fn execute(&self, input: Self::Input, _backend: Backend) -> Result<Self::Output, TruenoError> {
    110
    0
            let (x, weights) = input;
    111
    112
            // Validate input dimensions
    113
    0
            if x.len() != self.hidden_size {
    114
    0
                return Err(TruenoError::SizeMismatch {
    115
    0
                    expected: self.hidden_size,
    116
    0
                    actual: x.len(),
    117
    0
                });
    118
    0
            }
    119
    120
            // Q projection: x @ W_q^T -> [hidden_size]
    121
    0
            let mut q = vec![0.0f32; self.hidden_size];
    122
    0
            for i in 0..self.hidden_size {
    123
    0
                let mut sum = 0.0f32;
    124
    0
                for j in 0..self.hidden_size {
    125
    0
                    sum += x[j] * weights.q_weight[i * self.hidden_size + j];
    126
    0
                }
    127
    0
                q[i] = sum;
    128
            }
    129
    130
            // K projection: x @ W_k^T -> [kv_dim]
    131
    0
            let mut k = vec![0.0f32; self.kv_dim];
    132
    0
            for i in 0..self.kv_dim {
    133
    0
                let mut sum = 0.0f32;
    134
    0
                for j in 0..self.hidden_size {
    135
    0
                    sum += x[j] * weights.k_weight[i * self.hidden_size + j];
    136
    0
                }
    137
    0
                k[i] = sum;
    138
            }
    139
    140
            // V projection: x @ W_v^T -> [kv_dim]
    141
    0
            let mut v = vec![0.0f32; self.kv_dim];
    142
    0
            for i in 0..self.kv_dim {
    143
    0
                let mut sum = 0.0f32;
    144
    0
                for j in 0..self.hidden_size {
    145
    0
                    sum += x[j] * weights.v_weight[i * self.hidden_size + j];
    146
    0
                }
    147
    0
                v[i] = sum;
    148
            }
    149
    150
    0
            Ok((q, k, v))
    151
    0
        }
    152
    153
    0
        fn tokens(&self, _input: &Self::Input) -> usize {
    154
            // Output tokens = Q + K + V dimensions
    155
    0
            self.hidden_size + 2 * self.kv_dim
    156
    0
        }
    157
    }
    158
    159
    // ============================================================================
    160
    // Fused Gate+Up FFN Projection (PMAT-PERF-009)
    161
    // ============================================================================
    162
    163
    /// Weights for fused gate+up FFN projection
    164
    #[derive(Debug, Clone)]
    165
    pub struct FusedGateUpWeights {
    166
        /// Gate projection weights [hidden_size, intermediate_size]
    167
        pub gate_weight: Vec<f32>,
    168
        /// Up projection weights [hidden_size, intermediate_size]
    169
        pub up_weight: Vec<f32>,
    170
    }
    171
    172
    /// Fused Gate+Up FFN projection with SiLU activation.
    173
    ///
    174
    /// Computes gate and up projections in a single pass:
    175
    /// - gate = x * W_gate
    176
    /// - up = x * W_up
    177
    /// - output = SiLU(gate) * up (SwiGLU activation)
    178
    ///
    179
    /// # Performance Impact
    180
    ///
    181
    /// Fusing 2 separate matmuls + activation provides:
    182
    /// - 2x reduction in kernel launches (GPU)
    183
    /// - Fused SiLU avoids intermediate memory traffic
    184
    /// - Expected speedup: 1.5-2x for decode phase
    185
    ///
    186
    /// # Five-Whys Root Cause (PMAT-PERF-009)
    187
    ///
    188
    /// ```text
    189
    /// Why 1: Why is FFN phase slow?
    190
    /// → 3 kernel launches: gate_proj, up_proj, SiLU activation
    191
    ///
    192
    /// Why 2: Why separate kernels?
    193
    /// → Traditional implementation pattern from training frameworks
    194
    ///
    195
    /// Why 3: Why does this matter for inference?
    196
    /// → Inference is memory-bound; kernel launch overhead dominates
    197
    ///
    198
    /// Why 4: Why not fuse earlier?
    199
    /// → Requires custom kernel development
    200
    ///
    201
    /// Why 5: ROOT CAUSE
    202
    /// → SwiGLU requires gate*up pattern that naturally fuses
    203
    /// → FIX: Fuse gate+up+SiLU into single operation
    204
    /// ```
    205
    #[derive(Debug, Clone)]
    206
    pub struct FusedGateUpOp {
    207
        /// Hidden dimension size
    208
        pub hidden_size: usize,
    209
        /// Intermediate FFN dimension
    210
        pub intermediate_size: usize,
    211
    }
    212
    213
    impl FusedGateUpOp {
    214
        /// Create a new fused gate+up operation.
    215
        ///
    216
        /// # Arguments
    217
        /// * `hidden_size` - Hidden dimension (e.g., 3584 for Qwen 3B)
    218
        /// * `intermediate_size` - FFN intermediate dimension (e.g., 18944)
    219
    0
        pub fn new(hidden_size: usize, intermediate_size: usize) -> Self {
    220
    0
            Self {
    221
    0
                hidden_size,
    222
    0
                intermediate_size,
    223
    0
            }
    224
    0
        }
    225
    226
        /// SiLU activation: x * sigmoid(x)
    227
        #[inline]
    228
    0
        pub(crate) fn silu(x: f32) -> f32 {
    229
    0
            x / (1.0 + (-x).exp())
    230
    0
        }
    231
    }
    232
    233
    impl ComputeOp for FusedGateUpOp {
    234
        type Input = (Vec<f32>, FusedGateUpWeights);
    235
        type Output = Vec<f32>; // SwiGLU output [intermediate_size]
    236
    237
    0
        fn name(&self) -> &'static str {
    238
    0
            "fused_gate_up"
    239
    0
        }
    240
    241
    0
        fn execute(&self, input: Self::Input, _backend: Backend) -> Result<Self::Output, TruenoError> {
    242
    0
            let (x, weights) = input;
    243
    244
            // Validate input dimensions
    245
    0
            if x.len() != self.hidden_size {
    246
    0
                return Err(TruenoError::SizeMismatch {
    247
    0
                    expected: self.hidden_size,
    248
    0
                    actual: x.len(),
    249
    0
                });
    250
    0
            }
    251
    252
            // SIMD-optimized fused gate + up + SwiGLU
    253
            // Uses Vector dot product for ~4-8x speedup over scalar loops
    254
    0
            let mut output = vec![0.0f32; self.intermediate_size];
    255
    256
            // Select best SIMD backend (AVX2/AVX-512/NEON)
    257
    0
            let simd_backend = crate::Backend::select_best();
    258
    259
            // Create SIMD vector for input (reused for both gate and up projections)
    260
    0
            let x_vec = crate::Vector::from_slice_with_backend(&x, simd_backend);
    261
    262
    0
            for i in 0..self.intermediate_size {
    263
    0
                let row_start = i * self.hidden_size;
    264
    0
                let row_end = row_start + self.hidden_size;
    265
    0
    266
    0
                // Gate projection with SIMD dot product
    267
    0
                let gate_row = crate::Vector::from_slice_with_backend(
    268
    0
                    &weights.gate_weight[row_start..row_end],
    269
    0
                    simd_backend,
    270
    0
                );
    271
    0
                let gate_sum = x_vec.dot(&gate_row).unwrap_or(0.0);
    272
    0
    273
    0
                // Up projection with SIMD dot product
    274
    0
                let up_row = crate::Vector::from_slice_with_backend(
    275
    0
                    &weights.up_weight[row_start..row_end],
    276
    0
                    simd_backend,
    277
    0
                );
    278
    0
                let up_sum = x_vec.dot(&up_row).unwrap_or(0.0);
    279
    0
    280
    0
                // SwiGLU: SiLU(gate) * up
    281
    0
                output[i] = Self::silu(gate_sum) * up_sum;
    282
    0
            }
    283
    284
    0
            Ok(output)
    285
    0
        }
    286
    287
    0
        fn tokens(&self, _input: &Self::Input) -> usize {
    288
    0
            self.intermediate_size
    289
    0
        }
    290
    }
    291
    292
    #[cfg(test)]
    293
    mod tests {
    294
        use super::*;
    295
    296
        #[test]
    297
        fn test_fused_qkv_basic() {
    298
            // hidden=4, num_heads=2, kv_heads=1 → head_dim=2, kv_dim=2
    299
            let op = FusedQKVOp::new(4, 2, 1);
    300
    301
            let x = vec![1.0, 2.0, 3.0, 4.0];
    302
            let weights = FusedQKVWeights {
    303
                q_weight: vec![1.0; 16], // hidden_size x hidden_size = 4x4 = 16
    304
                k_weight: vec![1.0; 8],  // kv_dim x hidden_size = 2x4 = 8
    305
                v_weight: vec![1.0; 8],  // kv_dim x hidden_size = 2x4 = 8
    306
            };
    307
    308
            let (q, k, v) = op.execute((x, weights), Backend::Scalar).unwrap();
    309
    310
            assert_eq!(q.len(), 4);
    311
            assert_eq!(k.len(), 2);
    312
            assert_eq!(v.len(), 2);
    313
        }
    314
    315
        #[test]
    316
        fn test_fused_qkv_dimension_mismatch() {
    317
            let op = FusedQKVOp::new(4, 2, 2);
    318
            let x = vec![1.0, 2.0]; // Wrong size
    319
            let weights = FusedQKVWeights {
    320
                q_weight: vec![1.0; 16],
    321
                k_weight: vec![1.0; 8],
    322
                v_weight: vec![1.0; 8],
    323
            };
    324
    325
            let result = op.execute((x, weights), Backend::Scalar);
    326
            assert!(result.is_err());
    327
        }
    328
    329
        #[test]
    330
        fn test_fused_gate_up_basic() {
    331
            let op = FusedGateUpOp::new(4, 2);
    332
    333
            let x = vec![1.0, 2.0, 3.0, 4.0];
    334
            let weights = FusedGateUpWeights {
    335
                gate_weight: vec![1.0; 8], // 2x4
    336
                up_weight: vec![1.0; 8],   // 2x4
    337
            };
    338
    339
            let output = op.execute((x, weights), Backend::Scalar).unwrap();
    340
            assert_eq!(output.len(), 2);
    341
    342
            // Output should be SiLU(gate_sum) * up_sum
    343
            // gate_sum = up_sum = 1+2+3+4 = 10
    344
            // SiLU(10) ≈ 10 * sigmoid(10) ≈ 10 * 0.99995 ≈ 10
    345
            // output ≈ 10 * 10 = 100
    346
            assert!(output[0] > 90.0 && output[0] < 110.0);
    347
        }
    348
    349
        #[test]
    350
        fn test_fused_gate_up_dimension_mismatch() {
    351
            let op = FusedGateUpOp::new(4, 2);
    352
            let x = vec![1.0, 2.0]; // Wrong size
    353
            let weights = FusedGateUpWeights {
    354
                gate_weight: vec![1.0; 8],
    355
                up_weight: vec![1.0; 8],
    356
            };
    357
    358
            let result = op.execute((x, weights), Backend::Scalar);
    359
            assert!(result.is_err());
    360
        }
    361
    362
        #[test]
    363
        fn test_silu_values() {
    364
            // SiLU(0) = 0
    365
            assert!((FusedGateUpOp::silu(0.0) - 0.0).abs() < 1e-6);
    366
    367
            // SiLU(x) → x for large positive x
    368
            assert!((FusedGateUpOp::silu(10.0) - 10.0).abs() < 0.01);
    369
    370
            // SiLU(x) → 0 for large negative x
    371
            assert!(FusedGateUpOp::silu(-10.0).abs() < 0.01);
    372
        }
    373
    374
        #[test]
    375
        fn test_fused_qkv_tokens() {
    376
            // hidden=128, heads=8, kv_heads=4 → head_dim=16, kv_dim=64
    377
            let op = FusedQKVOp::new(128, 8, 4);
    378
            let weights = FusedQKVWeights {
    379
                q_weight: vec![],
    380
                k_weight: vec![],
    381
                v_weight: vec![],
    382
            };
    383
            // tokens = hidden + 2 * kv_dim = 128 + 2 * 64 = 256
    384
            assert_eq!(op.tokens(&(vec![], weights)), 256);
    385
        }
    386
    387
        #[test]
    388
        fn test_fused_gate_up_tokens() {
    389
            let op = FusedGateUpOp::new(128, 256);
    390
            let weights = FusedGateUpWeights {
    391
                gate_weight: vec![],
    392
                up_weight: vec![],
    393
            };
    394
            assert_eq!(op.tokens(&(vec![], weights)), 256);
    395
        }
    396
    }