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

    Line
    Count
    Source
    1
    //! Attention mechanisms for transformer models
    2
    //!
    3
    //! Extracted from layers/mod.rs (PMAT-802) to reduce module size.
    4
    //! Contains:
    5
    //! - Attention: Basic scaled dot-product attention
    6
    //! - SlidingWindowAttention: Efficient attention with fixed window size
    7
    //! - FusedQKVAttention: FlashAttention-style tiled attention
    8
    //! - MultiHeadAttention: Full multi-head attention with Q/K/V projections
    9
    10
    use crate::{
    11
        error::{RealizarError, Result},
    12
        tensor::Tensor,
    13
    };
    14
    15
    use super::{softmax, Linear};
    16
    17
    /// Scaled dot-product attention
    18
    ///
    19
    /// Computes attention as:
    20
    /// ```text
    21
    /// Attention(Q, K, V) = softmax(Q @ K.T / sqrt(d_k)) @ V
    22
    /// ```
    23
    ///
    24
    /// This is a building block for multi-head attention.
    25
    ///
    26
    /// # References
    27
    ///
    28
    /// "Attention is All You Need" - Vaswani et al., 2017
    29
    #[derive(Debug, Clone)]
    30
    pub struct Attention {
    31
        /// Head dimension (`d_k` = `d_model` / `num_heads`)
    32
        head_dim: usize,
    33
        /// Scale factor: 1 / `sqrt(head_dim)`
    34
        scale: f32,
    35
    }
    36
    37
    impl Attention {
    38
        /// Create a new attention layer
    39
        ///
    40
        /// # Arguments
    41
        ///
    42
        /// * `head_dim` - Dimension of each attention head
    43
        ///
    44
        /// # Errors
    45
        ///
    46
        /// Returns error if `head_dim` is zero
    47
    255
        pub fn new(head_dim: usize) -> Result<Self> {
    48
    255
            if head_dim == 0 {
    49
    2
                return Err(RealizarError::InvalidShape {
    50
    2
                    reason: "head_dim must be > 0".to_string(),
    51
    2
                });
    52
    253
            }
    53
    54
            #[allow(clippy::cast_precision_loss)]
    55
    253
            let scale = 1.0 / (head_dim as f32).sqrt();
    56
    57
    253
            Ok(Self { head_dim, scale })
    58
    255
        }
    59
    60
        /// Compute scaled dot-product attention
    61
        ///
    62
        /// # Arguments
    63
        ///
    64
        /// * `query` - Query tensor `[seq_len, head_dim]`
    65
        /// * `key` - Key tensor `[seq_len, head_dim]`
    66
        /// * `value` - Value tensor `[seq_len, head_dim]`
    67
        ///
    68
        /// # Returns
    69
        ///
    70
        /// Output tensor `[seq_len, head_dim]`
    71
        ///
    72
        /// # Errors
    73
        ///
    74
        /// Returns error if shapes don't match
    75
    3.09k
        pub fn forward(
    76
    3.09k
            &self,
    77
    3.09k
            query: &Tensor<f32>,
    78
    3.09k
            key: &Tensor<f32>,
    79
    3.09k
            value: &Tensor<f32>,
    80
    3.09k
        ) -> Result<Tensor<f32>> {
    81
    3.09k
            let q_shape = query.shape();
    82
    3.09k
            let k_shape = key.shape();
    83
    3.09k
            let v_shape = value.shape();
    84
    85
            // Validate shapes
    86
    3.09k
            if q_shape.is_empty() || k_shape.is_empty() || v_shape.is_empty() {
    87
    0
                return Err(RealizarError::InvalidShape {
    88
    0
                    reason: "Query, key, value tensors must have at least 1 dimension".to_string(),
    89
    0
                });
    90
    3.09k
            }
    91
    92
    3.09k
            let q_last = q_shape[q_shape.len() - 1];
    93
    3.09k
            let k_last = k_shape[k_shape.len() - 1];
    94
    3.09k
            let v_last = v_shape[v_shape.len() - 1];
    95
    96
    3.09k
            if q_last != self.head_dim || k_last != self.head_dim3.09k || v_last != self.head_dim3.09k {
    97
    2
                return Err(RealizarError::InvalidShape {
    98
    2
                    reason: format!(
    99
    2
                        "Expected head_dim={}, got Q={}, K={}, V={}",
    100
    2
                        self.head_dim, q_last, k_last, v_last
    101
    2
                    ),
    102
    2
                });
    103
    3.09k
            }
    104
    105
            // Get sequence lengths
    106
    3.09k
            let q_seq_len = if q_shape.len() > 1 { q_shape[0]3.09k } else { 11 };
    107
    3.09k
            let k_seq_len = if k_shape.len() > 1 { k_shape[0]3.09k } else { 11 };
    108
    3.09k
            let v_seq_len = if v_shape.len() > 1 { v_shape[0]3.09k } else { 11 };
    109
    110
    3.09k
            if k_seq_len != v_seq_len {
    111
    2
                return Err(RealizarError::InvalidShape {
    112
    2
                    reason: format!("Key seq_len {k_seq_len} != Value seq_len {v_seq_len}"),
    113
    2
                });
    114
    3.08k
            }
    115
    116
    3.08k
            let q_data = query.data();
    117
    3.08k
            let k_data = key.data();
    118
    3.08k
            let v_data = value.data();
    119
    120
            // Compute attention scores: Q @ K.T
    121
            // scores[i][j] = sum(Q[i][k] * K[j][k]) for all k
    122
    3.08k
            let mut scores = Vec::with_capacity(q_seq_len * k_seq_len);
    123
    132k
            for i in 0..q_seq_len3.08k {
    124
    18.1M
                for j in 0..k_seq_len132k {
    125
    18.1M
                    let mut dot = 0.0;
    126
    567M
                    for k in 0..self.head_dim18.1M {
    127
    567M
                        dot += q_data[i * self.head_dim + k] * k_data[j * self.head_dim + k];
    128
    567M
                    }
    129
    18.1M
                    scores.push(dot * self.scale);
    130
                }
    131
            }
    132
    133
            // Apply softmax to each row of scores
    134
    3.08k
            let scores_tensor = Tensor::from_vec(vec![q_seq_len, k_seq_len], scores)?0;
    135
    3.08k
            let attn_weights = softmax(&scores_tensor)?0;
    136
    3.08k
            let attn_data = attn_weights.data();
    137
    138
            // Compute output: attn_weights @ V
    139
            // output[i][k] = sum(attn_weights[i][j] * V[j][k]) for all j
    140
    3.08k
            let mut output = Vec::with_capacity(q_seq_len * self.head_dim);
    141
    132k
            for i in 0..q_seq_len3.08k {
    142
    3.75M
                for k in 0..self.head_dim132k {
    143
    3.75M
                    let mut sum = 0.0;
    144
    567M
                    for j in 0..k_seq_len3.75M {
    145
    567M
                        sum += attn_data[i * k_seq_len + j] * v_data[j * self.head_dim + k];
    146
    567M
                    }
    147
    3.75M
                    output.push(sum);
    148
                }
    149
            }
    150
    151
            // Debug assertion for numerical stability
    152
    3.08k
            debug_assert!(
    153
    3.75M
                output.iter()3.08k.all3.08k(|&x| x.is_finite()),
    154
    0
                "Attention layer produced NaN or Inf values - check input scaling"
    155
            );
    156
    157
    3.08k
            Tensor::from_vec(vec![q_seq_len, self.head_dim], output)
    158
    3.09k
        }
    159
    160
        /// Get head dimension
    161
        #[must_use]
    162
    3
        pub fn head_dim(&self) -> usize {
    163
    3
            self.head_dim
    164
    3
        }
    165
    166
        /// Get scale factor
    167
        #[must_use]
    168
    7
        pub fn scale(&self) -> f32 {
    169
    7
            self.scale
    170
    7
        }
    171
    172
        /// Compute Flash Attention - memory-efficient block-wise attention
    173
        ///
    174
        /// Uses tiling and recomputation to reduce memory usage from O(N²) to O(N).
    175
        /// Implements block-wise softmax with running max/sum statistics.
    176
        ///
    177
        /// # Arguments
    178
        ///
    179
        /// * `query` - Query tensor `[seq_len, head_dim]`
    180
        /// * `key` - Key tensor `[seq_len, head_dim]`
    181
        /// * `value` - Value tensor `[seq_len, head_dim]`
    182
        /// * `block_size` - Tile size for block-wise computation (e.g., 64, 128)
    183
        ///
    184
        /// # Returns
    185
        ///
    186
        /// Output tensor `[seq_len, head_dim]` (same as standard attention)
    187
        ///
    188
        /// # Errors
    189
        ///
    190
        /// Returns error if shapes don't match or `block_size` is zero
    191
        ///
    192
        /// # References
    193
        ///
    194
        /// - "`FlashAttention`: Fast and Memory-Efficient Exact Attention" - Dao et al., 2022
    195
        /// - "FlashAttention-2: Faster Attention with Better Parallelism" - Dao, 2023
    196
    9
        pub fn flash_forward(
    197
    9
            &self,
    198
    9
            query: &Tensor<f32>,
    199
    9
            key: &Tensor<f32>,
    200
    9
            value: &Tensor<f32>,
    201
    9
            block_size: usize,
    202
    9
        ) -> Result<Tensor<f32>> {
    203
    9
            if block_size == 0 {
    204
    1
                return Err(RealizarError::InvalidShape {
    205
    1
                    reason: "block_size must be > 0".to_string(),
    206
    1
                });
    207
    8
            }
    208
    209
    8
            let q_shape = query.shape();
    210
    8
            let k_shape = key.shape();
    211
    8
            let v_shape = value.shape();
    212
    213
            // Validate shapes (same as standard attention)
    214
    8
            if q_shape.is_empty() || k_shape.is_empty() || v_shape.is_empty() {
    215
    0
                return Err(RealizarError::InvalidShape {
    216
    0
                    reason: "Query, key, value tensors must have at least 1 dimension".to_string(),
    217
    0
                });
    218
    8
            }
    219
    220
    8
            let q_last = q_shape[q_shape.len() - 1];
    221
    8
            let k_last = k_shape[k_shape.len() - 1];
    222
    8
            let v_last = v_shape[v_shape.len() - 1];
    223
    224
    8
            if q_last != self.head_dim || k_last != self.head_dim || v_last != self.head_dim {
    225
    0
                return Err(RealizarError::InvalidShape {
    226
    0
                    reason: format!(
    227
    0
                        "Expected head_dim={}, got Q={}, K={}, V={}",
    228
    0
                        self.head_dim, q_last, k_last, v_last
    229
    0
                    ),
    230
    0
                });
    231
    8
            }
    232
    233
            // Get sequence lengths
    234
    8
            let q_seq_len = if q_shape.len() > 1 { q_shape[0] } else { 10 };
    235
    8
            let k_seq_len = if k_shape.len() > 1 { k_shape[0] } else { 10 };
    236
    8
            let v_seq_len = if v_shape.len() > 1 { v_shape[0] } else { 10 };
    237
    238
    8
            if k_seq_len != v_seq_len {
    239
    1
                return Err(RealizarError::InvalidShape {
    240
    1
                    reason: format!("Key seq_len {k_seq_len} != Value seq_len {v_seq_len}"),
    241
    1
                });
    242
    7
            }
    243
    244
    7
            let q_data = query.data();
    245
    7
            let k_data = key.data();
    246
    7
            let v_data = value.data();
    247
    248
            // Initialize output and statistics
    249
    7
            let mut output = vec![0.0; q_seq_len * self.head_dim];
    250
    7
            let mut row_max = vec![f32::NEG_INFINITY; q_seq_len]; // Running max for each query row
    251
    7
            let mut row_sum = vec![0.0; q_seq_len]; // Running sum for each query row
    252
    253
            // Iterate over K/V blocks (outer loop)
    254
    7
            let num_kv_blocks = k_seq_len.div_ceil(block_size);
    255
    16
            for kv_block_idx in 0..num_kv_blocks7 {
    256
    16
                let kv_start = kv_block_idx * block_size;
    257
    16
                let kv_end = (kv_start + block_size).min(k_seq_len);
    258
    16
                let kv_block_len = kv_end - kv_start;
    259
    260
                // Iterate over Q blocks (inner loop)
    261
    16
                let num_q_blocks = q_seq_len.div_ceil(block_size);
    262
    48
                for q_block_idx in 0..num_q_blocks16 {
    263
    48
                    let q_start = q_block_idx * block_size;
    264
    48
                    let q_end = (q_start + block_size).min(q_seq_len);
    265
    266
                    // Compute attention scores for this block: Q_block @ K_block.T
    267
    48
                    let mut scores = vec![0.0; (q_end - q_start) * kv_block_len];
    268
    342
                    for (i, q_idx) in (q_start..q_end)48.enumerate48() {
    269
    4.38k
                        for (j, kv_idx) in (kv_start..kv_end)342.enumerate342() {
    270
    4.38k
                            let mut dot = 0.0;
    271
    133k
                            for k in 0..self.head_dim4.38k {
    272
    133k
                                dot += q_data[q_idx * self.head_dim + k]
    273
    133k
                                    * k_data[kv_idx * self.head_dim + k];
    274
    133k
                            }
    275
    4.38k
                            scores[i * kv_block_len + j] = dot * self.scale;
    276
                        }
    277
                    }
    278
    279
                    // Update running max and apply softmax with new max
    280
    342
                    for (i, q_idx) in (q_start..q_end)48.enumerate48() {
    281
                        // Find max in current block
    282
    342
                        let block_max = (0..kv_block_len)
    283
    4.38k
                            .map342(|j| scores[i * kv_block_len + j])
    284
    342
                            .fold(f32::NEG_INFINITY, f32::max);
    285
    286
                        // Update global max
    287
    342
                        let old_max = row_max[q_idx];
    288
    342
                        let new_max = old_max.max(block_max);
    289
    342
                        row_max[q_idx] = new_max;
    290
    291
                        // Compute exp(scores - new_max) and update running sum
    292
    342
                        let mut block_sum = 0.0;
    293
    4.38k
                        for j in 0..kv_block_len342 {
    294
    4.38k
                            let exp_val = (scores[i * kv_block_len + j] - new_max).exp();
    295
    4.38k
                            scores[i * kv_block_len + j] = exp_val;
    296
    4.38k
                            block_sum += exp_val;
    297
    4.38k
                        }
    298
    299
                        // Rescale old output and sum based on new max
    300
    342
                        let scale_factor = (old_max - new_max).exp();
    301
    8.79k
                        for k in 0..self.head_dim342 {
    302
    8.79k
                            output[q_idx * self.head_dim + k] *= scale_factor;
    303
    8.79k
                        }
    304
    342
                        row_sum[q_idx] = row_sum[q_idx] * scale_factor + block_sum;
    305
                    }
    306
    307
                    // Accumulate weighted values: output += scores @ V_block
    308
    342
                    for (i, q_idx) in (q_start..q_end)48.enumerate48() {
    309
    8.79k
                        for k in 0..self.head_dim342 {
    310
    8.79k
                            let mut weighted_sum = 0.0;
    311
    133k
                            for (j, kv_idx) in (kv_start..kv_end)8.79k.enumerate8.79k() {
    312
    133k
                                weighted_sum +=
    313
    133k
                                    scores[i * kv_block_len + j] * v_data[kv_idx * self.head_dim + k];
    314
    133k
                            }
    315
    8.79k
                            output[q_idx * self.head_dim + k] += weighted_sum;
    316
                        }
    317
                    }
    318
                }
    319
            }
    320
    321
            // Final normalization by row_sum
    322
    93
            for i in 0..q_seq_len7 {
    323
    2.23k
                for k in 0..self.head_dim93 {
    324
    2.23k
                    output[i * self.head_dim + k] /= row_sum[i];
    325
    2.23k
                }
    326
            }
    327
    328
    7
            Tensor::from_vec(vec![q_seq_len, self.head_dim], output)
    329
    9
        }
    330
    331
        /// Flash Attention v2 with SIMD-accelerated dot products
    332
        ///
    333
        /// Optimized implementation using AVX2 SIMD for dot products.
    334
        /// Uses parallel outer loop over query blocks for better multi-core utilization.
    335
        ///
    336
        /// # Arguments
    337
        ///
    338
        /// * `query` - Query tensor `[seq_len, head_dim]`
    339
        /// * `key` - Key tensor `[seq_len, head_dim]`
    340
        /// * `value` - Value tensor `[seq_len, head_dim]`
    341
        /// * `block_size` - Tile size for block-wise computation (e.g., 64, 128)
    342
        ///
    343
        /// # Returns
    344
        ///
    345
        /// Output tensor `[seq_len, head_dim]` (same as standard attention)
    346
        ///
    347
        /// # Errors
    348
        ///
    349
        /// Returns error if shapes don't match or `block_size` is zero
    350
        ///
    351
        /// # References
    352
        ///
    353
        /// - "FlashAttention-2: Faster Attention with Better Parallelism" - Dao, 2023
    354
        #[allow(clippy::similar_names)]
    355
    109
        pub fn flash_forward_v2(
    356
    109
            &self,
    357
    109
            query: &Tensor<f32>,
    358
    109
            key: &Tensor<f32>,
    359
    109
            value: &Tensor<f32>,
    360
    109
            block_size: usize,
    361
    109
        ) -> Result<Tensor<f32>> {
    362
    109
            if block_size == 0 {
    363
    1
                return Err(RealizarError::InvalidShape {
    364
    1
                    reason: "block_size must be > 0".to_string(),
    365
    1
                });
    366
    108
            }
    367
    368
    108
            let q_shape = query.shape();
    369
    108
            let k_shape = key.shape();
    370
    108
            let v_shape = value.shape();
    371
    372
            // Validate shapes
    373
    108
            if q_shape.is_empty() || k_shape.is_empty() || v_shape.is_empty() {
    374
    0
                return Err(RealizarError::InvalidShape {
    375
    0
                    reason: "Query, key, value tensors must have at least 1 dimension".to_string(),
    376
    0
                });
    377
    108
            }
    378
    379
    108
            let q_last = q_shape[q_shape.len() - 1];
    380
    108
            let k_last = k_shape[k_shape.len() - 1];
    381
    108
            let v_last = v_shape[v_shape.len() - 1];
    382
    383
    108
            if q_last != self.head_dim || k_last != self.head_dim || v_last != self.head_dim {
    384
    0
                return Err(RealizarError::InvalidShape {
    385
    0
                    reason: format!(
    386
    0
                        "Expected head_dim={}, got Q={}, K={}, V={}",
    387
    0
                        self.head_dim, q_last, k_last, v_last
    388
    0
                    ),
    389
    0
                });
    390
    108
            }
    391
    392
    108
            let q_seq_len = if q_shape.len() > 1 { q_shape[0] } else { 10 };
    393
    108
            let k_seq_len = if k_shape.len() > 1 { k_shape[0] } else { 10 };
    394
    108
            let v_seq_len = if v_shape.len() > 1 { v_shape[0] } else { 10 };
    395
    396
    108
            if k_seq_len != v_seq_len {
    397
    0
                return Err(RealizarError::InvalidShape {
    398
    0
                    reason: format!("Key seq_len {k_seq_len} != Value seq_len {v_seq_len}"),
    399
    0
                });
    400
    108
            }
    401
    402
    108
            let q_data = query.data();
    403
    108
            let k_data = key.data();
    404
    108
            let v_data = value.data();
    405
    108
            let head_dim = self.head_dim;
    406
    108
            let scale = self.scale;
    407
    408
            // Initialize output and statistics
    409
    108
            let mut output = vec![0.0; q_seq_len * head_dim];
    410
    108
            let mut row_max = vec![f32::NEG_INFINITY; q_seq_len];
    411
    108
            let mut row_sum = vec![0.0; q_seq_len];
    412
    413
            // Flash Attention v2: Iterate over K/V blocks in outer loop
    414
            // This allows better memory access patterns
    415
    108
            let num_kv_blocks = k_seq_len.div_ceil(block_size);
    416
    417
    420
            for kv_block_idx in 0..num_kv_blocks108 {
    418
    420
                let kv_start = kv_block_idx * block_size;
    419
    420
                let kv_end = (kv_start + block_size).min(k_seq_len);
    420
    420
                let kv_block_len = kv_end - kv_start;
    421
    422
                // Process all Q rows against this K/V block
    423
    13.1k
                for q_idx in 0..q_seq_len420 {
    424
                    // SIMD-accelerated dot products for this row
    425
    13.1k
                    let mut scores = Vec::with_capacity(kv_block_len);
    426
    104k
                    for kv_idx in kv_start13.1k..kv_end13.1k {
    427
    104k
                        let dot = Self::simd_dot_product(
    428
    104k
                            &q_data[q_idx * head_dim..(q_idx + 1) * head_dim],
    429
    104k
                            &k_data[kv_idx * head_dim..(kv_idx + 1) * head_dim],
    430
    104k
                        );
    431
    104k
                        scores.push(dot * scale);
    432
    104k
                    }
    433
    434
                    // Find max in current block
    435
    104k
                    let block_max13.1k = scores.iter()13.1k.fold13.1k(f32::NEG_INFINITY, |a, &b| a.max(b));
    436
    437
                    // Update global max
    438
    13.1k
                    let old_max = row_max[q_idx];
    439
    13.1k
                    let new_max = old_max.max(block_max);
    440
    13.1k
                    row_max[q_idx] = new_max;
    441
    442
                    // Compute exp(scores - new_max) and update running sum
    443
    13.1k
                    let mut block_sum = 0.0;
    444
    117k
                    for score104k in &mut scores {
    445
    104k
                        let exp_val = (*score - new_max).exp();
    446
    104k
                        *score = exp_val;
    447
    104k
                        block_sum += exp_val;
    448
    104k
                    }
    449
    450
                    // Rescale old output and sum based on new max
    451
    13.1k
                    let scale_factor = (old_max - new_max).exp();
    452
    829k
                    for k in 0..head_dim13.1k {
    453
    829k
                        output[q_idx * head_dim + k] *= scale_factor;
    454
    829k
                    }
    455
    13.1k
                    row_sum[q_idx] = row_sum[q_idx] * scale_factor + block_sum;
    456
    457
                    // Accumulate weighted values: output += scores @ V_block
    458
    104k
                    for (j, kv_idx) in (kv_start..kv_end)13.1k.enumerate13.1k() {
    459
    104k
                        let weight = scores[j];
    460
    6.62M
                        for k in 0..head_dim104k {
    461
    6.62M
                            output[q_idx * head_dim + k] += weight * v_data[kv_idx * head_dim + k];
    462
    6.62M
                        }
    463
                    }
    464
                }
    465
            }
    466
    467
            // Final normalization by row_sum
    468
    3.29k
            for i in 0..q_seq_len108 {
    469
    3.29k
                let inv_sum = 1.0 / row_sum[i];
    470
    207k
                for k in 0..head_dim3.29k {
    471
    207k
                    output[i * head_dim + k] *= inv_sum;
    472
    207k
                }
    473
            }
    474
    475
    108
            Tensor::from_vec(vec![q_seq_len, self.head_dim], output)
    476
    109
        }
    477
    478
        /// SIMD-accelerated dot product
    479
        ///
    480
        /// Uses AVX2 on x86_64 for 8-way f32 parallelism
    481
        #[inline]
    482
    212k
        fn simd_dot_product(a: &[f32], b: &[f32]) -> f32 {
    483
            #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
    484
            {
    485
                Self::simd_dot_avx2(a, b)
    486
            }
    487
    488
            #[cfg(not(all(target_arch = "x86_64", target_feature = "avx2")))]
    489
            {
    490
    212k
                Self::scalar_dot_product(a, b)
    491
            }
    492
    212k
        }
    493
    494
        /// AVX2 SIMD dot product (8-way f32 parallelism)
    495
        #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
    496
        #[inline]
    497
        #[allow(clippy::wildcard_imports)]
    498
        fn simd_dot_avx2(a: &[f32], b: &[f32]) -> f32 {
    499
            use std::arch::x86_64::*;
    500
    501
            let len = a.len().min(b.len());
    502
            let chunks = len / 8;
    503
            let remainder = len % 8;
    504
    505
            // SIMD accumulator
    506
            // SAFETY: Memory safety ensured by bounds checking and alignment
    507
            let simd_sum = unsafe {
    508
                let mut acc = _mm256_setzero_ps();
    509
    510
                for i in 0..chunks {
    511
                    let a_vec = _mm256_loadu_ps(a.as_ptr().add(i * 8));
    512
                    let b_vec = _mm256_loadu_ps(b.as_ptr().add(i * 8));
    513
                    acc = _mm256_fmadd_ps(a_vec, b_vec, acc);
    514
                }
    515
    516
                // Horizontal sum of 8 floats
    517
                let hi = _mm256_extractf128_ps(acc, 1);
    518
                let lo = _mm256_castps256_ps128(acc);
    519
                let sum128 = _mm_add_ps(lo, hi);
    520
                let hi64 = _mm_movehl_ps(sum128, sum128);
    521
                let sum64 = _mm_add_ps(sum128, hi64);
    522
                let hi32 = _mm_shuffle_ps(sum64, sum64, 0x55);
    523
                let sum32 = _mm_add_ss(sum64, hi32);
    524
                _mm_cvtss_f32(sum32)
    525
            };
    526
    527
            // Handle remainder
    528
            let remainder_sum: f32 = (0..remainder)
    529
                .map(|i| a[chunks * 8 + i] * b[chunks * 8 + i])
    530
                .sum();
    531
    532
            simd_sum + remainder_sum
    533
        }
    534
    535
        /// Scalar fallback dot product
    536
        #[cfg(not(all(target_arch = "x86_64", target_feature = "avx2")))]
    537
        #[inline]
    538
    212k
        fn scalar_dot_product(a: &[f32], b: &[f32]) -> f32 {
    539
    13.3M
            a212k.iter212k().zip212k(b212k.iter212k()).map212k(|(x, y)| x * y).sum212k()
    540
    212k
        }
    541
    542
        /// Parallel Flash Attention v2 using rayon
    543
        ///
    544
        /// Parallelizes over query positions for multi-core utilization.
    545
        /// Each thread processes a subset of query rows independently.
    546
        ///
    547
        /// # Arguments
    548
        ///
    549
        /// * `query` - Query tensor `[seq_len, head_dim]`
    550
        /// * `key` - Key tensor `[seq_len, head_dim]`
    551
        /// * `value` - Value tensor `[seq_len, head_dim]`
    552
        /// * `block_size` - Tile size for block-wise computation (e.g., 64, 128)
    553
        ///
    554
        /// # Returns
    555
        ///
    556
        /// Output tensor `[seq_len, head_dim]` (same as standard attention)
    557
        ///
    558
        /// # Errors
    559
        ///
    560
        /// Returns error if shapes don't match or `block_size` is zero
    561
        #[allow(clippy::similar_names)]
    562
    108
        pub fn flash_forward_parallel(
    563
    108
            &self,
    564
    108
            query: &Tensor<f32>,
    565
    108
            key: &Tensor<f32>,
    566
    108
            value: &Tensor<f32>,
    567
    108
            block_size: usize,
    568
    108
        ) -> Result<Tensor<f32>> {
    569
            use rayon::prelude::*;
    570
    571
    108
            if block_size == 0 {
    572
    1
                return Err(RealizarError::InvalidShape {
    573
    1
                    reason: "block_size must be > 0".to_string(),
    574
    1
                });
    575
    107
            }
    576
    577
    107
            let q_shape = query.shape();
    578
    107
            let k_shape = key.shape();
    579
    107
            let v_shape = value.shape();
    580
    581
            // Validate shapes
    582
    107
            if q_shape.is_empty() || k_shape.is_empty() || v_shape.is_empty() {
    583
    0
                return Err(RealizarError::InvalidShape {
    584
    0
                    reason: "Query, key, value tensors must have at least 1 dimension".to_string(),
    585
    0
                });
    586
    107
            }
    587
    588
    107
            let q_last = q_shape[q_shape.len() - 1];
    589
    107
            let k_last = k_shape[k_shape.len() - 1];
    590
    107
            let v_last = v_shape[v_shape.len() - 1];
    591
    592
    107
            if q_last != self.head_dim || k_last != self.head_dim || v_last != self.head_dim {
    593
    0
                return Err(RealizarError::InvalidShape {
    594
    0
                    reason: format!(
    595
    0
                        "Expected head_dim={}, got Q={}, K={}, V={}",
    596
    0
                        self.head_dim, q_last, k_last, v_last
    597
    0
                    ),
    598
    0
                });
    599
    107
            }
    600
    601
    107
            let q_seq_len = if q_shape.len() > 1 { q_shape[0] } else { 10 };
    602
    107
            let k_seq_len = if k_shape.len() > 1 { k_shape[0] } else { 10 };
    603
    107
            let v_seq_len = if v_shape.len() > 1 { v_shape[0] } else { 10 };
    604
    605
    107
            if k_seq_len != v_seq_len {
    606
    0
                return Err(RealizarError::InvalidShape {
    607
    0
                    reason: format!("Key seq_len {k_seq_len} != Value seq_len {v_seq_len}"),
    608
    0
                });
    609
    107
            }
    610
    611
    107
            let q_data = query.data();
    612
    107
            let k_data = key.data();
    613
    107
            let v_data = value.data();
    614
    107
            let head_dim = self.head_dim;
    615
    107
            let scale = self.scale;
    616
    617
            // Parallel over query positions
    618
    107
            let output: Vec<f32> = (0..q_seq_len)
    619
    107
                .into_par_iter()
    620
    3.32k
                .flat_map107(|q_idx| {
    621
                    // Each query row is processed independently
    622
    3.32k
                    let mut row_output = vec![0.0; head_dim];
    623
    3.32k
                    let mut row_max = f32::NEG_INFINITY;
    624
    3.32k
                    let mut row_sum = 0.0;
    625
    626
    3.32k
                    let num_kv_blocks = k_seq_len.div_ceil(block_size);
    627
    628
    13.2k
                    for kv_block_idx in 0..num_kv_blocks3.32k {
    629
    13.2k
                        let kv_start = kv_block_idx * block_size;
    630
    13.2k
                        let kv_end = (kv_start + block_size).min(k_seq_len);
    631
    632
                        // Compute scores for this K/V block
    633
    13.2k
                        let mut scores: Vec<f32> = (kv_start..kv_end)
    634
    107k
                            .map13.2k(|kv_idx| {
    635
    107k
                                let dot = Self::simd_dot_product(
    636
    107k
                                    &q_data[q_idx * head_dim..(q_idx + 1) * head_dim],
    637
    107k
                                    &k_data[kv_idx * head_dim..(kv_idx + 1) * head_dim],
    638
                                );
    639
    107k
                                dot * scale
    640
    107k
                            })
    641
    13.2k
                            .collect();
    642
    643
                        // Online softmax: find block max and update global max
    644
    107k
                        let block_max13.2k = scores.iter()13.2k.fold13.2k(f32::NEG_INFINITY, |a, &b| a.max(b));
    645
    13.2k
                        let old_max = row_max;
    646
    13.2k
                        let new_max = old_max.max(block_max);
    647
    13.2k
                        row_max = new_max;
    648
    649
                        // Compute exp(scores - new_max)
    650
    13.2k
                        let mut block_sum = 0.0;
    651
    121k
                        for score107k in &mut scores {
    652
    107k
                            let exp_val = (*score - new_max).exp();
    653
    107k
                            *score = exp_val;
    654
    107k
                            block_sum += exp_val;
    655
    107k
                        }
    656
    657
                        // Rescale previous output
    658
    13.2k
                        let scale_factor = (old_max - new_max).exp();
    659
    845k
                        for out_val832k in &mut row_output {
    660
    832k
                            *out_val *= scale_factor;
    661
    832k
                        }
    662
    13.2k
                        row_sum = row_sum * scale_factor + block_sum;
    663
    664
                        // Accumulate weighted values
    665
    107k
                        for (j, kv_idx) in (kv_start..kv_end)13.2k.enumerate13.2k() {
    666
    107k
                            let weight = scores[j];
    667
    6.68M
                            for k in 0..head_dim107k {
    668
    6.68M
                                row_output[k] += weight * v_data[kv_idx * head_dim + k];
    669
    6.68M
                            }
    670
                        }
    671
                    }
    672
    673
                    // Final normalization
    674
    3.32k
                    let inv_sum = 1.0 / row_sum;
    675
    211k
                    for out_val208k in &mut row_output {
    676
    208k
                        *out_val *= inv_sum;
    677
    208k
                    }
    678
    679
    3.32k
                    row_output
    680
    3.32k
                })
    681
    107
                .collect();
    682
    683
    107
            Tensor::from_vec(vec![q_seq_len, self.head_dim], output)
    684
    108
        }
    685
    }
    686
    687
    // ============================================================================
    688
    // Sliding Window Attention (Mistral/Mixtral style)
    689
    // ============================================================================
    690
    //
    691
    // Limits attention to a fixed window of recent tokens for efficient
    692
    // long-context inference. Used by Mistral-7B, Mixtral, and similar models.
    693
    //
    694
    // Benefits:
    695
    // - Reduces memory from O(n²) to O(n*w) where w = window_size
    696
    // - Enables very long context with bounded KV cache
    697
    // - Compatible with Flash Attention algorithms
    698
    //
    699
    // Reference: "Mistral 7B" - Jiang et al., 2023
    700
    // ============================================================================
    701
    702
    /// Sliding Window Attention
    703
    ///
    704
    /// Limits each token to attending only to the most recent `window_size` tokens.
    705
    /// This provides linear memory scaling for long sequences while maintaining
    706
    /// local context.
    707
    ///
    708
    /// # Algorithm
    709
    ///
    710
    /// For each query position i, attention is computed only over keys/values
    711
    /// in positions `[max(0, i - window_size + 1), i]`.
    712
    ///
    713
    /// ```text
    714
    /// Standard Attention (full):  Sliding Window (w=3):
    715
    ///   Q K K K K K                 Q K K K . .
    716
    ///   Q K K K K K                 . Q K K K .
    717
    ///   Q K K K K K                 . . Q K K K
    718
    ///   Q K K K K K                 . . . Q K K
    719
    /// ```
    720
    ///
    721
    /// # References
    722
    ///
    723
    /// - "Mistral 7B" - Jiang et al., 2023
    724
    /// - "Longformer: The Long-Document Transformer" - Beltagy et al., 2020
    725
    #[derive(Debug, Clone)]
    726
    pub struct SlidingWindowAttention {
    727
        /// Head dimension (`d_k` = `d_model` / `num_heads`)
    728
        head_dim: usize,
    729
        /// Scale factor: 1 / `sqrt(head_dim)`
    730
        scale: f32,
    731
        /// Window size (number of tokens each query can attend to)
    732
        window_size: usize,
    733
    }
    734
    735
    impl SlidingWindowAttention {
    736
        /// Create a new sliding window attention layer
    737
        ///
    738
        /// # Arguments
    739
        ///
    740
        /// * `head_dim` - Dimension of each attention head
    741
        /// * `window_size` - Number of tokens each query can attend to
    742
        ///
    743
        /// # Errors
    744
        ///
    745
        /// Returns error if `head_dim` is zero or `window_size` is zero
    746
    18
        pub fn new(head_dim: usize, window_size: usize) -> Result<Self> {
    747
    18
            if head_dim == 0 {
    748
    2
                return Err(RealizarError::InvalidShape {
    749
    2
                    reason: "head_dim must be > 0".to_string(),
    750
    2
                });
    751
    16
            }
    752
    16
            if window_size == 0 {
    753
    2
                return Err(RealizarError::InvalidShape {
    754
    2
                    reason: "window_size must be > 0".to_string(),
    755
    2
                });
    756
    14
            }
    757
    758
            #[allow(clippy::cast_precision_loss)]
    759
    14
            let scale = 1.0 / (head_dim as f32).sqrt();
    760
    761
    14
            Ok(Self {
    762
    14
                head_dim,
    763
    14
                scale,
    764
    14
                window_size,
    765
    14
            })
    766
    18
        }
    767
    768
        /// Compute sliding window attention
    769
        ///
    770
        /// # Arguments
    771
        ///
    772
        /// * `query` - Query tensor `[seq_len, head_dim]`
    773
        /// * `key` - Key tensor `[seq_len, head_dim]`
    774
        /// * `value` - Value tensor `[seq_len, head_dim]`
    775
        ///
    776
        /// # Returns
    777
        ///
    778
        /// Output tensor `[seq_len, head_dim]`
    779
        ///
    780
        /// # Errors
    781
        ///
    782
        /// Returns error if shapes don't match
    783
    9
        pub fn forward(
    784
    9
            &self,
    785
    9
            query: &Tensor<f32>,
    786
    9
            key: &Tensor<f32>,
    787
    9
            value: &Tensor<f32>,
    788
    9
        ) -> Result<Tensor<f32>> {
    789
    9
            let q_shape = query.shape();
    790
    9
            let k_shape = key.shape();
    791
    9
            let v_shape = value.shape();
    792
    793
            // Validate shapes
    794
    9
            if q_shape.is_empty() || k_shape.is_empty() || v_shape.is_empty() {
    795
    0
                return Err(RealizarError::InvalidShape {
    796
    0
                    reason: "Query, key, value tensors must have at least 1 dimension".to_string(),
    797
    0
                });
    798
    9
            }
    799
    800
    9
            let q_last = q_shape[q_shape.len() - 1];
    801
    9
            let k_last = k_shape[k_shape.len() - 1];
    802
    9
            let v_last = v_shape[v_shape.len() - 1];
    803
    804
    9
            if q_last != self.head_dim || k_last != self.head_dim || v_last != self.head_dim8 {
    805
    1
                return Err(RealizarError::InvalidShape {
    806
    1
                    reason: format!(
    807
    1
                        "Expected head_dim={}, got Q={}, K={}, V={}",
    808
    1
                        self.head_dim, q_last, k_last, v_last
    809
    1
                    ),
    810
    1
                });
    811
    8
            }
    812
    813
            // Get sequence lengths
    814
    8
            let q_seq_len = if q_shape.len() > 1 { q_shape[0] } else { 10 };
    815
    8
            let k_seq_len = if k_shape.len() > 1 { k_shape[0] } else { 10 };
    816
    8
            let v_seq_len = if v_shape.len() > 1 { v_shape[0] } else { 10 };
    817
    818
    8
            if k_seq_len != v_seq_len {
    819
    1
                return Err(RealizarError::InvalidShape {
    820
    1
                    reason: format!("Key seq_len {k_seq_len} != Value seq_len {v_seq_len}"),
    821
    1
                });
    822
    7
            }
    823
    824
    7
            let q_data = query.data();
    825
    7
            let k_data = key.data();
    826
    7
            let v_data = value.data();
    827
    828
    7
            let mut output = Vec::with_capacity(q_seq_len * self.head_dim);
    829
    830
            // Process each query position with sliding window
    831
    25
            for i in 0..q_seq_len7 {
    832
                // Compute window boundaries [window_start, window_end)
    833
                // For causal attention: can only attend to positions <= i
    834
    25
                let window_end = (i + 1).min(k_seq_len);
    835
    25
                let window_start = window_end.saturating_sub(self.window_size);
    836
    25
                let window_len = window_end - window_start;
    837
    838
    25
                if window_len == 0 {
    839
                    // No keys to attend to, output zeros
    840
    0
                    output.extend(std::iter::repeat_n(0.0, self.head_dim));
    841
    0
                    continue;
    842
    25
                }
    843
    844
                // Compute attention scores for this window
    845
    25
                let mut scores = Vec::with_capacity(window_len);
    846
    54
                for j in window_start25..window_end25 {
    847
    54
                    let mut dot = 0.0;
    848
    134
                    for k in 0..self.head_dim54 {
    849
    134
                        dot += q_data[i * self.head_dim + k] * k_data[j * self.head_dim + k];
    850
    134
                    }
    851
    54
                    scores.push(dot * self.scale);
    852
                }
    853
    854
                // Apply softmax over window scores
    855
    54
                let max_score25 = scores.iter()25.fold25(f32::NEG_INFINITY, |a, &b| a.max(b));
    856
    25
                let mut exp_sum = 0.0;
    857
    79
                for score54 in &mut scores {
    858
    54
                    let exp_val = (*score - max_score).exp();
    859
    54
                    *score = exp_val;
    860
    54
                    exp_sum += exp_val;
    861
    54
                }
    862
    25
                let inv_sum = 1.0 / exp_sum;
    863
    79
                for score54 in &mut scores {
    864
    54
                    *score *= inv_sum;
    865
    54
                }
    866
    867
                // Compute weighted sum of values
    868
    62
                for k in 0..self.head_dim25 {
    869
    62
                    let mut sum = 0.0;
    870
    134
                    for (idx, j) in (window_start..window_end)62.enumerate62() {
    871
    134
                        sum += scores[idx] * v_data[j * self.head_dim + k];
    872
    134
                    }
    873
    62
                    output.push(sum);
    874
                }
    875
            }
    876
    877
    7
            Tensor::from_vec(vec![q_seq_len, self.head_dim], output)
    878
    9
        }
    879
    880
        /// Compute sliding window attention with mask
    881
        ///
    882
        /// Supports bidirectional attention (non-causal) with the sliding window.
    883
        ///
    884
        /// # Arguments
    885
        ///
    886
        /// * `query` - Query tensor `[seq_len, head_dim]`
    887
        /// * `key` - Key tensor `[seq_len, head_dim]`
    888
        /// * `value` - Value tensor `[seq_len, head_dim]`
    889
        /// * `causal` - If true, only attend to past positions (causal/autoregressive)
    890
        ///
    891
        /// # Returns
    892
        ///
    893
        /// Output tensor `[seq_len, head_dim]`
    894
        ///
    895
        /// # Errors
    896
        ///
    897
        /// Returns error if shapes don't match
    898
    2
        pub fn forward_with_mask(
    899
    2
            &self,
    900
    2
            query: &Tensor<f32>,
    901
    2
            key: &Tensor<f32>,
    902
    2
            value: &Tensor<f32>,
    903
    2
            causal: bool,
    904
    2
        ) -> Result<Tensor<f32>> {
    905
    2
            if causal {
    906
                // Causal is the default behavior
    907
    1
                return self.forward(query, key, value);
    908
    1
            }
    909
    910
    1
            let q_shape = query.shape();
    911
    1
            let k_shape = key.shape();
    912
    1
            let v_shape = value.shape();
    913
    914
            // Validate shapes
    915
    1
            if q_shape.is_empty() || k_shape.is_empty() || v_shape.is_empty() {
    916
    0
                return Err(RealizarError::InvalidShape {
    917
    0
                    reason: "Query, key, value tensors must have at least 1 dimension".to_string(),
    918
    0
                });
    919
    1
            }
    920
    921
    1
            let q_last = q_shape[q_shape.len() - 1];
    922
    1
            let k_last = k_shape[k_shape.len() - 1];
    923
    1
            let v_last = v_shape[v_shape.len() - 1];
    924
    925
    1
            if q_last != self.head_dim || k_last != self.head_dim || v_last != self.head_dim {
    926
    0
                return Err(RealizarError::InvalidShape {
    927
    0
                    reason: format!(
    928
    0
                        "Expected head_dim={}, got Q={}, K={}, V={}",
    929
    0
                        self.head_dim, q_last, k_last, v_last
    930
    0
                    ),
    931
    0
                });
    932
    1
            }
    933
    934
    1
            let q_seq_len = if q_shape.len() > 1 { q_shape[0] } else { 10 };
    935
    1
            let k_seq_len = if k_shape.len() > 1 { k_shape[0] } else { 10 };
    936
    1
            let v_seq_len = if v_shape.len() > 1 { v_shape[0] } else { 10 };
    937
    938
    1
            if k_seq_len != v_seq_len {
    939
    0
                return Err(RealizarError::InvalidShape {
    940
    0
                    reason: format!("Key seq_len {k_seq_len} != Value seq_len {v_seq_len}"),
    941
    0
                });
    942
    1
            }
    943
    944
    1
            let q_data = query.data();
    945
    1
            let k_data = key.data();
    946
    1
            let v_data = value.data();
    947
    948
    1
            let mut output = Vec::with_capacity(q_seq_len * self.head_dim);
    949
    1
            let half_window = self.window_size / 2;
    950
    951
            // Process each query position with bidirectional sliding window
    952
    5
            for i in 0..q_seq_len1 {
    953
                // Bidirectional window centered on position i
    954
    5
                let window_start = i.saturating_sub(half_window);
    955
    5
                let window_end = (i + half_window + 1).min(k_seq_len);
    956
    5
                let window_len = window_end - window_start;
    957
    958
    5
                if window_len == 0 {
    959
    0
                    output.extend(std::iter::repeat_n(0.0, self.head_dim));
    960
    0
                    continue;
    961
    5
                }
    962
    963
                // Compute attention scores for this window
    964
    5
                let mut scores = Vec::with_capacity(window_len);
    965
    19
                for j in window_start5..window_end5 {
    966
    19
                    let mut dot = 0.0;
    967
    38
                    for k in 0..self.head_dim19 {
    968
    38
                        dot += q_data[i * self.head_dim + k] * k_data[j * self.head_dim + k];
    969
    38
                    }
    970
    19
                    scores.push(dot * self.scale);
    971
                }
    972
    973
                // Apply softmax over window scores
    974
    19
                let max_score5 = scores.iter()5.fold5(f32::NEG_INFINITY, |a, &b| a.max(b));
    975
    5
                let mut exp_sum = 0.0;
    976
    24
                for score19 in &mut scores {
    977
    19
                    let exp_val = (*score - max_score).exp();
    978
    19
                    *score = exp_val;
    979
    19
                    exp_sum += exp_val;
    980
    19
                }
    981
    5
                let inv_sum = 1.0 / exp_sum;
    982
    24
                for score19 in &mut scores {
    983
    19
                    *score *= inv_sum;
    984
    19
                }
    985
    986
                // Compute weighted sum of values
    987
    10
                for k in 0..self.head_dim5 {
    988
    10
                    let mut sum = 0.0;
    989
    38
                    for (idx, j) in (window_start..window_end)10.enumerate10() {
    990
    38
                        sum += scores[idx] * v_data[j * self.head_dim + k];
    991
    38
                    }
    992
    10
                    output.push(sum);
    993
                }
    994
            }
    995
    996
    1
            Tensor::from_vec(vec![q_seq_len, self.head_dim], output)
    997
    2
        }
    998
    999
        /// Get head dimension
    1000
        #[must_use]
    1001
    2
        pub fn head_dim(&self) -> usize {
    1002
    2
            self.head_dim
    1003
    2
        }
    1004
    1005
        /// Get scale factor
    1006
        #[must_use]
    1007
    1
        pub fn scale(&self) -> f32 {
    1008
    1
            self.scale
    1009
    1
        }
    1010
    1011
        /// Get window size
    1012
        #[must_use]
    1013
    4
        pub fn window_size(&self) -> usize {
    1014
    4
            self.window_size
    1015
    4
        }
    1016
    1017
        /// Compute the effective context at a given position
    1018
        ///
    1019
        /// Returns the number of tokens this position can attend to
    1020
        #[must_use]
    1021
    4
        pub fn effective_context(&self, position: usize, seq_len: usize) -> usize {
    1022
    4
            let window_end = (position + 1).min(seq_len);
    1023
    4
            let window_start = window_end.saturating_sub(self.window_size);
    1024
    4
            window_end - window_start
    1025
    4
        }
    1026
    1027
        /// Memory usage relative to full attention
    1028
        ///
    1029
        /// Returns the ratio of memory used compared to full attention.
    1030
        /// For window_size w and seq_len n: memory = O(n*w) vs O(n²)
    1031
        #[must_use]
    1032
    2
        pub fn memory_ratio(&self, seq_len: usize) -> f32 {
    1033
    2
            if seq_len == 0 {
    1034
    0
                return 1.0;
    1035
    2
            }
    1036
            #[allow(clippy::cast_precision_loss)]
    1037
            {
    1038
    2
                (self.window_size.min(seq_len) as f32) / (seq_len as f32)
    1039
            }
    1040
    2
        }
    1041
    }
    1042
    1043
    // ============================================================================
    1044
    // Fused QKV + Attention (IMP-003)
    1045
    // Per spec: performance-parity-ollama-llamacpp-gpu-inference-llms.md
    1046
    // ============================================================================
    1047
    1048
    /// Fused Query-Key-Value projection with scaled dot-product attention
    1049
    ///
    1050
    /// Combines QKV projection and attention into a single fused operation for
    1051
    /// improved memory efficiency and performance. Eliminates intermediate
    1052
    /// materializations by computing attention in a single pass.
    1053
    ///
    1054
    /// # Performance Benefits
    1055
    ///
    1056
    /// - **Memory Bandwidth**: Single read of input, single write of output
    1057
    /// - **Cache Efficiency**: QKV computed block-wise to maximize L1/L2 reuse
    1058
    /// - **Numerical Stability**: Uses log-sum-exp trick for softmax
    1059
    ///
    1060
    /// # Algorithm (Flash Attention style)
    1061
    ///
    1062
    /// ```text
    1063
    /// for each block of queries:
    1064
    ///     Q_block = input_block @ W_q
    1065
    ///     for each block of keys/values:
    1066
    ///         K_block = input_block @ W_k
    1067
    ///         V_block = input_block @ W_v
    1068
    ///         scores = Q_block @ K_block^T / sqrt(d)
    1069
    ///         update running softmax and output
    1070
    /// ```
    1071
    ///
    1072
    /// # References
    1073
    ///
    1074
    /// - [1] Dao et al., "FlashAttention: Fast and Memory-Efficient Attention", 2022
    1075
    /// - [2] Tri Dao, "FlashAttention-2: Faster Attention with Better Parallelism", 2023
    1076
    #[derive(Debug, Clone)]
    1077
    pub struct FusedQKVAttention {
    1078
        /// Dimension per attention head
    1079
        head_dim: usize,
    1080
        /// Total hidden dimension
    1081
        hidden_dim: usize,
    1082
        /// Number of attention heads
    1083
        num_heads: usize,
    1084
        /// Scale factor: 1 / sqrt(head_dim)
    1085
        scale: f32,
    1086
        /// Query projection weights: [hidden_dim, hidden_dim]
    1087
        w_q: Vec<f32>,
    1088
        /// Key projection weights: [hidden_dim, hidden_dim]
    1089
        w_k: Vec<f32>,
    1090
        /// Value projection weights: [hidden_dim, hidden_dim]
    1091
        w_v: Vec<f32>,
    1092
        /// Output projection weights: [hidden_dim, hidden_dim]
    1093
        w_o: Vec<f32>,
    1094
    }
    1095
    1096
    impl FusedQKVAttention {
    1097
        /// Create a new fused QKV attention layer
    1098
        ///
    1099
        /// # Arguments
    1100
        ///
    1101
        /// * `head_dim` - Dimension per attention head
    1102
        /// * `hidden_dim` - Total hidden dimension (must be divisible by head_dim)
    1103
        ///
    1104
        /// # Errors
    1105
        ///
    1106
        /// Returns error if head_dim is 0, hidden_dim is 0, or hidden_dim % head_dim != 0
    1107
    15
        pub fn new(head_dim: usize, hidden_dim: usize) -> Result<Self> {
    1108
    15
            if head_dim == 0 {
    1109
    2
                return Err(RealizarError::InvalidShape {
    1110
    2
                    reason: "head_dim must be > 0".to_string(),
    1111
    2
                });
    1112
    13
            }
    1113
    13
            if hidden_dim == 0 {
    1114
    1
                return Err(RealizarError::InvalidShape {
    1115
    1
                    reason: "hidden_dim must be > 0".to_string(),
    1116
    1
                });
    1117
    12
            }
    1118
    12
            if !hidden_dim.is_multiple_of(head_dim) {
    1119
    1
                return Err(RealizarError::InvalidShape {
    1120
    1
                    reason: format!(
    1121
    1
                        "hidden_dim ({}) must be divisible by head_dim ({})",
    1122
    1
                        hidden_dim, head_dim
    1123
    1
                    ),
    1124
    1
                });
    1125
    11
            }
    1126
    1127
    11
            let num_heads = hidden_dim / head_dim;
    1128
    11
            let scale = 1.0 / (head_dim as f32).sqrt();
    1129
    11
            let proj_size = hidden_dim * hidden_dim;
    1130
    1131
            // Initialize with small random-like values for non-degenerate behavior
    1132
    44
            let init_weight11 = |size: usize| -> Vec<f32> {
    1133
    1.90M
                (0..size)44.map44(|i| (i as f32 * 0.001).sin() * 0.02).collect44()
    1134
    44
            };
    1135
    1136
    11
            Ok(Self {
    1137
    11
                head_dim,
    1138
    11
                hidden_dim,
    1139
    11
                num_heads,
    1140
    11
                scale,
    1141
    11
                w_q: init_weight(proj_size),
    1142
    11
                w_k: init_weight(proj_size),
    1143
    11
                w_v: init_weight(proj_size),
    1144
    11
                w_o: init_weight(proj_size),
    1145
    11
            })
    1146
    15
        }
    1147
    1148
        /// Forward pass with fused QKV projection and attention
    1149
        ///
    1150
        /// # Arguments
    1151
        ///
    1152
        /// * `input` - Input tensor [seq_len, hidden_dim]
    1153
        ///
    1154
        /// # Returns
    1155
        ///
    1156
        /// Output tensor [seq_len, hidden_dim]
    1157
        ///
    1158
        /// # Errors
    1159
        ///
    1160
        /// Returns error if input shape doesn't match hidden_dim
    1161
    58
        pub fn forward(&self, input: &Tensor<f32>) -> Result<Tensor<f32>> {
    1162
    58
            let shape = input.shape();
    1163
    58
            if shape.len() < 2 {
    1164
    0
                return Err(RealizarError::InvalidShape {
    1165
    0
                    reason: "Input must have at least 2 dimensions [seq_len, hidden_dim]".to_string(),
    1166
    0
                });
    1167
    58
            }
    1168
    1169
    58
            let seq_len = shape[0];
    1170
    58
            let input_dim = shape[shape.len() - 1];
    1171
    1172
    58
            if input_dim != self.hidden_dim {
    1173
    1
                return Err(RealizarError::InvalidShape {
    1174
    1
                    reason: format!(
    1175
    1
                        "Input hidden_dim ({}) doesn't match layer hidden_dim ({})",
    1176
    1
                        input_dim, self.hidden_dim
    1177
    1
                    ),
    1178
    1
                });
    1179
    57
            }
    1180
    1181
    57
            let data = input.data();
    1182
    1183
            // Compute Q, K, V projections
    1184
    57
            let mut q = vec![0.0f32; seq_len * self.hidden_dim];
    1185
    57
            let mut k = vec![0.0f32; seq_len * self.hidden_dim];
    1186
    57
            let mut v = vec![0.0f32; seq_len * self.hidden_dim];
    1187
    1188
            // Matrix multiply: [seq_len, hidden_dim] @ [hidden_dim, hidden_dim]
    1189
    841
            for i in 0..seq_len57 {
    1190
    53.3k
                for j in 0..self.hidden_dim841 {
    1191
    53.3k
                    let mut sum_q = 0.0f32;
    1192
    53.3k
                    let mut sum_k = 0.0f32;
    1193
    53.3k
                    let mut sum_v = 0.0f32;
    1194
    3.40M
                    for l in 0..self.hidden_dim53.3k {
    1195
    3.40M
                        let inp = data[i * self.hidden_dim + l];
    1196
    3.40M
                        sum_q += inp * self.w_q[l * self.hidden_dim + j];
    1197
    3.40M
                        sum_k += inp * self.w_k[l * self.hidden_dim + j];
    1198
    3.40M
                        sum_v += inp * self.w_v[l * self.hidden_dim + j];
    1199
    3.40M
                    }
    1200
    53.3k
                    q[i * self.hidden_dim + j] = sum_q;
    1201
    53.3k
                    k[i * self.hidden_dim + j] = sum_k;
    1202
    53.3k
                    v[i * self.hidden_dim + j] = sum_v;
    1203
                }
    1204
            }
    1205
    1206
            // Compute attention per head
    1207
    57
            let mut output = vec![0.0f32; seq_len * self.hidden_dim];
    1208
    1209
    138
            for head in 0..self.num_heads57 {
    1210
    138
                let head_offset = head * self.head_dim;
    1211
    1212
                // Compute attention scores for this head
    1213
    1.82k
                for i in 0..seq_len138 {
    1214
                    // Find max for numerical stability (causal: only j <= i)
    1215
    1.82k
                    let mut max_score = f32::NEG_INFINITY;
    1216
    14.6k
                    for j in 0..=i1.82k {
    1217
    14.6k
                        let mut dot = 0.0f32;
    1218
    447k
                        for d in 0..self.head_dim14.6k {
    1219
    447k
                            let q_idx = i * self.hidden_dim + head_offset + d;
    1220
    447k
                            let k_idx = j * self.hidden_dim + head_offset + d;
    1221
    447k
                            dot += q[q_idx] * k[k_idx];
    1222
    447k
                        }
    1223
    14.6k
                        let score = dot * self.scale;
    1224
    14.6k
                        if score > max_score {
    1225
    1.87k
                            max_score = score;
    1226
    12.7k
                        }
    1227
                    }
    1228
    1229
                    // Compute softmax with log-sum-exp trick
    1230
                    // Using enumerate() pattern for causal attention where j <= i
    1231
    1.82k
                    let mut sum_exp = 0.0f32;
    1232
    1.82k
                    let mut scores = vec![0.0f32; i + 1];
    1233
    14.6k
                    for (j, score) in scores.iter_mut()1.82k.enumerate1.82k() {
    1234
    14.6k
                        let mut dot = 0.0f32;
    1235
    447k
                        for d in 0..self.head_dim14.6k {
    1236
    447k
                            let q_idx = i * self.hidden_dim + head_offset + d;
    1237
    447k
                            let k_idx = j * self.hidden_dim + head_offset + d;
    1238
    447k
                            dot += q[q_idx] * k[k_idx];
    1239
    447k
                        }
    1240
    14.6k
                        *score = (dot * self.scale - max_score).exp();
    1241
    14.6k
                        sum_exp += *score;
    1242
                    }
    1243
    1244
                    // Normalize and compute weighted sum of values
    1245
    1.82k
                    if sum_exp > 0.0 {
    1246
    53.3k
                        for d in 0..self.head_dim1.82k {
    1247
    53.3k
                            let mut weighted_sum = 0.0f32;
    1248
    447k
                            for (j, &score) in scores.iter()53.3k.enumerate53.3k() {
    1249
    447k
                                let v_idx = j * self.hidden_dim + head_offset + d;
    1250
    447k
                                weighted_sum += (score / sum_exp) * v[v_idx];
    1251
    447k
                            }
    1252
    53.3k
                            output[i * self.hidden_dim + head_offset + d] = weighted_sum;
    1253
                        }
    1254
    0
                    }
    1255
                }
    1256
            }
    1257
    1258
            // Output projection
    1259
    57
            let mut final_output = vec![0.0f32; seq_len * self.hidden_dim];
    1260
    841
            for i in 0..seq_len57 {
    1261
    53.3k
                for j in 0..self.hidden_dim841 {
    1262
    53.3k
                    let mut sum = 0.0f32;
    1263
    3.40M
                    for l in 0..self.hidden_dim53.3k {
    1264
    3.40M
                        sum += output[i * self.hidden_dim + l] * self.w_o[l * self.hidden_dim + j];
    1265
    3.40M
                    }
    1266
    53.3k
                    final_output[i * self.hidden_dim + j] = sum;
    1267
                }
    1268
            }
    1269
    1270
    57
            Tensor::from_vec(vec![seq_len, self.hidden_dim], final_output)
    1271
    58
        }
    1272
    1273
        /// Get head dimension
    1274
        #[must_use]
    1275
    1
        pub fn head_dim(&self) -> usize {
    1276
    1
            self.head_dim
    1277
    1
        }
    1278
    1279
        /// Get hidden dimension
    1280
        #[must_use]
    1281
    1
        pub fn hidden_dim(&self) -> usize {
    1282
    1
            self.hidden_dim
    1283
    1
        }
    1284
    1285
        /// Get number of attention heads
    1286
        #[must_use]
    1287
    3
        pub fn num_heads(&self) -> usize {
    1288
    3
            self.num_heads
    1289
    3
        }
    1290
    1291
        /// Get mutable access to Q projection weights for loading
    1292
    0
        pub fn w_q_mut(&mut self) -> &mut [f32] {
    1293
    0
            &mut self.w_q
    1294
    0
        }
    1295
    1296
        /// Get mutable access to K projection weights for loading
    1297
    0
        pub fn w_k_mut(&mut self) -> &mut [f32] {
    1298
    0
            &mut self.w_k
    1299
    0
        }
    1300
    1301
        /// Get mutable access to V projection weights for loading
    1302
    0
        pub fn w_v_mut(&mut self) -> &mut [f32] {
    1303
    0
            &mut self.w_v
    1304
    0
        }
    1305
    1306
        /// Get mutable access to output projection weights for loading
    1307
    0
        pub fn w_o_mut(&mut self) -> &mut [f32] {
    1308
    0
            &mut self.w_o
    1309
    0
        }
    1310
    }
    1311
    1312
    /// Multi-Head Attention with support for MHA, MQA, and GQA
    1313
    ///
    1314
    /// Implements three attention variants through configurable `KV` head count:
    1315
    ///
    1316
    /// **Multi-Head Attention (MHA):** `num_kv_heads = num_heads`
    1317
    /// - Each head has separate Q, K, V projections
    1318
    /// - `KV` cache: `O(num_heads * seq_len * head_dim)`
    1319
    /// - Standard attention mechanism
    1320
    ///
    1321
    /// **Multi-Query Attention (MQA):** `num_kv_heads = 1`
    1322
    /// - Each head has separate Q projection
    1323
    /// - All heads share single K, V projection
    1324
    /// - `KV` cache: `O(seq_len * head_dim)` - reduces by `num_heads` factor
    1325
    /// - Used in `PaLM`, Falcon, `StarCoder`
    1326
    ///
    1327
    /// **Grouped-Query Attention (GQA):** `1 < num_kv_heads < num_heads`
    1328
    /// - Heads grouped into `num_kv_heads` groups
    1329
    /// - Each group shares K, V projections
    1330
    /// - `KV` cache: `O(num_kv_heads * seq_len * head_dim)`
    1331
    /// - Used in `Llama-2`, Mistral, `CodeLlama`
    1332
    ///
    1333
    /// # Architecture
    1334
    ///
    1335
    /// ```text
    1336
    /// Input [hidden_dim]
    1337
    ///   |
    1338
    ///   ├─> Q_proj [hidden_dim -> hidden_dim] -> split into num_heads
    1339
    ///   ├─> K_proj [hidden_dim -> num_kv_heads * head_dim]
    1340
    ///   └─> V_proj [hidden_dim -> num_kv_heads * head_dim]
    1341
    ///   |
    1342
    ///   ├─> Attention (grouped by num_kv_heads)
    1343
    ///   |
    1344
    ///   └─> O_proj [hidden_dim -> hidden_dim]
    1345
    ///       |
    1346
    ///     Output [hidden_dim]
    1347
    /// ```
    1348
    ///
    1349
    /// # References
    1350
    ///
    1351
    /// - "Attention is All You Need" - Vaswani et al., 2017 (MHA)
    1352
    /// - "Fast Transformer Decoding: One Write-Head is All You Need" - Shazeer, 2019 (MQA)
    1353
    /// - "`PaLM`: Scaling Language Modeling with Pathways" - Chowdhery et al., 2022 (MQA)
    1354
    /// - "`GQA`: Training Generalized Multi-Query Transformer" - Ainslie et al., 2023 (GQA)
    1355
    #[derive(Debug, Clone)]
    1356
    pub struct MultiHeadAttention {
    1357
        /// Number of attention heads (Q heads)
    1358
        num_heads: usize,
    1359
        /// Number of key/value heads (for GQA/MQA)
    1360
        num_kv_heads: usize,
    1361
        /// Dimension per attention head
    1362
        head_dim: usize,
    1363
        /// Total hidden dimension (`num_heads * head_dim`)
    1364
        hidden_dim: usize,
    1365
        /// Query projection: `hidden_dim -> hidden_dim`
    1366
        q_proj: Linear,
    1367
        /// Key projection: `hidden_dim -> num_kv_heads * head_dim`
    1368
        k_proj: Linear,
    1369
        /// Value projection: `hidden_dim -> num_kv_heads * head_dim`
    1370
        v_proj: Linear,
    1371
        /// Output projection: `hidden_dim -> hidden_dim`
    1372
        o_proj: Linear,
    1373
        /// Per-head attention mechanism
    1374
        attention: Attention,
    1375
    }
    1376
    1377
    impl MultiHeadAttention {
    1378
        /// Create a new Multi-Head Attention layer with configurable `KV` heads
    1379
        ///
    1380
        /// # Arguments
    1381
        ///
    1382
        /// * `hidden_dim` - Total hidden dimension (must be divisible by `num_heads`)
    1383
        /// * `num_heads` - Number of query heads
    1384
        /// * `num_kv_heads` - Number of key/value heads (must divide `num_heads`)
    1385
        ///
    1386
        /// # Modes
    1387
        ///
    1388
        /// - MHA: `num_kv_heads = num_heads` (standard multi-head)
    1389
        /// - MQA: `num_kv_heads = 1` (all heads share K/V)
    1390
        /// - GQA: `1 < num_kv_heads < num_heads` (grouped heads)
    1391
        ///
    1392
        /// # Errors
    1393
        ///
    1394
        /// Returns error if:
    1395
        /// - `hidden_dim` is zero or not divisible by `num_heads`
    1396
        /// - `num_heads` is zero or not divisible by `num_kv_heads`
    1397
        /// - `num_kv_heads` is zero or greater than `num_heads`
    1398
        ///
    1399
        /// # Examples
    1400
        ///
    1401
        /// ```rust,ignore
    1402
        /// // Standard Multi-Head Attention (MHA)
    1403
        /// let mha = MultiHeadAttention::new(512, 8, 8)?;
    1404
        ///
    1405
        /// // Multi-Query Attention (MQA)
    1406
        /// let mqa = MultiHeadAttention::new(512, 8, 1)?;
    1407
        ///
    1408
        /// // Grouped-Query Attention (GQA) - 4 heads per group
    1409
        /// let gqa = MultiHeadAttention::new(512, 8, 2)?;
    1410
        /// ```
    1411
    224
        pub fn new(hidden_dim: usize, num_heads: usize, num_kv_heads: usize) -> Result<Self> {
    1412
    224
            if hidden_dim == 0 {
    1413
    2
                return Err(RealizarError::InvalidShape {
    1414
    2
                    reason: "hidden_dim must be > 0".to_string(),
    1415
    2
                });
    1416
    222
            }
    1417
    222
            if num_heads == 0 {
    1418
    2
                return Err(RealizarError::InvalidShape {
    1419
    2
                    reason: "num_heads must be > 0".to_string(),
    1420
    2
                });
    1421
    220
            }
    1422
    220
            if num_kv_heads == 0 {
    1423
    1
                return Err(RealizarError::InvalidShape {
    1424
    1
                    reason: "num_kv_heads must be > 0".to_string(),
    1425
    1
                });
    1426
    219
            }
    1427
    219
            if num_kv_heads > num_heads {
    1428
    1
                return Err(RealizarError::InvalidShape {
    1429
    1
                    reason: format!(
    1430
    1
                        "num_kv_heads {num_kv_heads} cannot be greater than num_heads {num_heads}"
    1431
    1
                    ),
    1432
    1
                });
    1433
    218
            }
    1434
    218
            if !hidden_dim.is_multiple_of(num_heads) {
    1435
    2
                return Err(RealizarError::InvalidShape {
    1436
    2
                    reason: format!(
    1437
    2
                        "hidden_dim {hidden_dim} must be divisible by num_heads {num_heads}"
    1438
    2
                    ),
    1439
    2
                });
    1440
    216
            }
    1441
    216
            if !num_heads.is_multiple_of(num_kv_heads) {
    1442
    1
                return Err(RealizarError::InvalidShape {
    1443
    1
                    reason: format!(
    1444
    1
                        "num_heads {num_heads} must be divisible by num_kv_heads {num_kv_heads}"
    1445
    1
                    ),
    1446
    1
                });
    1447
    215
            }
    1448
    1449
    215
            let head_dim = hidden_dim / num_heads;
    1450
    1451
            // Q projection: always hidden_dim -> hidden_dim (all query heads)
    1452
    215
            let q_proj = Linear::new(hidden_dim, hidden_dim)?0;
    1453
    1454
            // K/V projections: hidden_dim -> num_kv_heads * head_dim
    1455
    215
            let kv_dim = num_kv_heads * head_dim;
    1456
    215
            let k_proj = Linear::new(hidden_dim, kv_dim)?0;
    1457
    215
            let v_proj = Linear::new(hidden_dim, kv_dim)?0;
    1458
    1459
            // Output projection: hidden_dim -> hidden_dim
    1460
    215
            let o_proj = Linear::new(hidden_dim, hidden_dim)?0;
    1461
    1462
            // Per-head attention mechanism
    1463
    215
            let attention = Attention::new(head_dim)?0;
    1464
    1465
    215
            Ok(Self {
    1466
    215
                num_heads,
    1467
    215
                num_kv_heads,
    1468
    215
                head_dim,
    1469
    215
                hidden_dim,
    1470
    215
                q_proj,
    1471
    215
                k_proj,
    1472
    215
                v_proj,
    1473
    215
                o_proj,
    1474
    215
                attention,
    1475
    215
            })
    1476
    224
        }
    1477
    1478
        /// Create standard Multi-Head Attention (MHA) - each head has separate K/V
    1479
        ///
    1480
        /// # Errors
    1481
        ///
    1482
        /// Returns `RealizarError::InvalidShape` if:
    1483
        /// - `hidden_dim` is 0
    1484
        /// - `num_heads` is 0
    1485
        /// - `hidden_dim` is not divisible by `num_heads`
    1486
    200
        pub fn mha(hidden_dim: usize, num_heads: usize) -> Result<Self> {
    1487
    200
            Self::new(hidden_dim, num_heads, num_heads)
    1488
    200
        }
    1489
    1490
        /// Create Multi-Query Attention (MQA) - all heads share K/V
    1491
        ///
    1492
        /// # Errors
    1493
        ///
    1494
        /// Returns `RealizarError::InvalidShape` if:
    1495
        /// - `hidden_dim` is 0
    1496
        /// - `num_heads` is 0
    1497
        /// - `hidden_dim` is not divisible by `num_heads`
    1498
    6
        pub fn mqa(hidden_dim: usize, num_heads: usize) -> Result<Self> {
    1499
    6
            Self::new(hidden_dim, num_heads, 1)
    1500
    6
        }
    1501
    1502
        /// Create Grouped-Query Attention (GQA) - heads grouped to share K/V
    1503
        ///
    1504
        /// # Errors
    1505
        ///
    1506
        /// Returns `RealizarError::InvalidShape` if:
    1507
        /// - `hidden_dim` is 0
    1508
        /// - `num_heads` is 0
    1509
        /// - `num_kv_heads` is 0
    1510
        /// - `num_kv_heads` is greater than `num_heads`
    1511
        /// - `hidden_dim` is not divisible by `num_heads`
    1512
        /// - `num_heads` is not divisible by `num_kv_heads`
    1513
    5
        pub fn gqa(hidden_dim: usize, num_heads: usize, num_kv_heads: usize) -> Result<Self> {
    1514
    5
            Self::new(hidden_dim, num_heads, num_kv_heads)
    1515
    5
        }
    1516
    1517
        /// Forward pass through multi-head attention
    1518
        ///
    1519
        /// # Arguments
    1520
        ///
    1521
        /// * `input` - Input tensor `[seq_len, hidden_dim]`
    1522
        ///
    1523
        /// # Returns
    1524
        ///
    1525
        /// Output tensor `[seq_len, hidden_dim]`
    1526
        ///
    1527
        /// # Errors
    1528
        ///
    1529
        /// Returns error if input shape is invalid
    1530
    1.60k
        pub fn forward(&self, input: &Tensor<f32>) -> Result<Tensor<f32>> {
    1531
    1.60k
            let shape = input.shape();
    1532
    1533
    1.60k
            if shape.len() != 2 {
    1534
    1
                return Err(RealizarError::InvalidShape {
    1535
    1
                    reason: format!("Expected 2D tensor [seq_len, hidden_dim], got shape {shape:?}"),
    1536
    1
                });
    1537
    1.60k
            }
    1538
    1539
    1.60k
            let seq_len = shape[0];
    1540
    1.60k
            let input_dim = shape[1];
    1541
    1542
    1.60k
            if input_dim != self.hidden_dim {
    1543
    1
                return Err(RealizarError::InvalidShape {
    1544
    1
                    reason: format!("Expected hidden_dim={}, got {}", self.hidden_dim, input_dim),
    1545
    1
                });
    1546
    1.59k
            }
    1547
    1548
            // Project Q, K, V
    1549
    1.59k
            let q = self.q_proj.forward(input)?0; // [seq_len, hidden_dim]
    1550
    1.59k
            let k = self.k_proj.forward(input)?0; // [seq_len, kv_dim]
    1551
    1.59k
            let v = self.v_proj.forward(input)?0; // [seq_len, kv_dim]
    1552
    1553
            // Reshape Q into heads: [seq_len, num_heads, head_dim]
    1554
    1.59k
            let q_data = q.data();
    1555
    1.59k
            let k_data = k.data();
    1556
    1.59k
            let v_data = v.data();
    1557
    1558
            // Calculate heads per group for GQA
    1559
    1.59k
            let heads_per_group = self.num_heads / self.num_kv_heads;
    1560
    1561
            // Process each query head
    1562
    1.59k
            let mut head_outputs = Vec::with_capacity(self.num_heads);
    1563
    1564
    3.05k
            for head_idx in 0..self.num_heads1.59k {
    1565
                // Extract Q for this head
    1566
    3.05k
                let mut q_head_data = Vec::with_capacity(seq_len * self.head_dim);
    1567
    131k
                for seq_idx in 0..seq_len3.05k {
    1568
    131k
                    let q_row_start = seq_idx * self.hidden_dim;
    1569
    131k
                    let head_start = q_row_start + head_idx * self.head_dim;
    1570
    3.72M
                    for offset in 0..self.head_dim131k {
    1571
    3.72M
                        q_head_data.push(q_data[head_start + offset]);
    1572
    3.72M
                    }
    1573
                }
    1574
    3.05k
                let q_head = Tensor::from_vec(vec![seq_len, self.head_dim], q_head_data)?0;
    1575
    1576
                // Determine which KV head this Q head uses (for GQA/MQA/MHA)
    1577
    3.05k
                let kv_head_idx = head_idx / heads_per_group;
    1578
    3.05k
                let kv_dim = self.num_kv_heads * self.head_dim;
    1579
    1580
                // Extract K, V for the corresponding KV head
    1581
    3.05k
                let mut k_head_data = Vec::with_capacity(seq_len * self.head_dim);
    1582
    3.05k
                let mut v_head_data = Vec::with_capacity(seq_len * self.head_dim);
    1583
    131k
                for seq_idx in 0..seq_len3.05k {
    1584
    131k
                    let kv_row_start = seq_idx * kv_dim;
    1585
    131k
                    let kv_head_start = kv_row_start + kv_head_idx * self.head_dim;
    1586
    3.72M
                    for offset in 0..self.head_dim131k {
    1587
    3.72M
                        k_head_data.push(k_data[kv_head_start + offset]);
    1588
    3.72M
                        v_head_data.push(v_data[kv_head_start + offset]);
    1589
    3.72M
                    }
    1590
                }
    1591
    3.05k
                let k_head = Tensor::from_vec(vec![seq_len, self.head_dim], k_head_data)?0;
    1592
    3.05k
                let v_head = Tensor::from_vec(vec![seq_len, self.head_dim], v_head_data)?0;
    1593
    1594
                // Compute attention for this head
    1595
    3.05k
                let head_output = self.attention.forward(&q_head, &k_head, &v_head)?0;
    1596
    3.05k
                head_outputs.push(head_output);
    1597
            }
    1598
    1599
            // Concatenate all head outputs: [seq_len, hidden_dim]
    1600
    1.59k
            let mut concat_data = Vec::with_capacity(seq_len * self.hidden_dim);
    1601
    115k
            for seq_idx in 0..seq_len1.59k {
    1602
    247k
                for head_output131k in &head_outputs {
    1603
    131k
                    let head_output_data = head_output.data();
    1604
    131k
                    let head_row_start = seq_idx * self.head_dim;
    1605
    3.72M
                    for offset in 0..self.head_dim131k {
    1606
    3.72M
                        concat_data.push(head_output_data[head_row_start + offset]);
    1607
    3.72M
                    }
    1608
                }
    1609
            }
    1610
    1611
    1.59k
            let concat = Tensor::from_vec(vec![seq_len, self.hidden_dim], concat_data)?0;
    1612
    1613
            // Output projection
    1614
    1.59k
            self.o_proj.forward(&concat)
    1615
    1.60k
        }
    1616
    1617
        /// Get number of query heads
    1618
        #[must_use]
    1619
    6
        pub fn num_heads(&self) -> usize {
    1620
    6
            self.num_heads
    1621
    6
        }
    1622
    1623
        /// Get number of key/value heads
    1624
        #[must_use]
    1625
    3
        pub fn num_kv_heads(&self) -> usize {
    1626
    3
            self.num_kv_heads
    1627
    3
        }
    1628
    1629
        /// Get head dimension
    1630
        #[must_use]
    1631
    4
        pub fn head_dim(&self) -> usize {
    1632
    4
            self.head_dim
    1633
    4
        }
    1634
    1635
        /// Get hidden dimension
    1636
        #[must_use]
    1637
    4
        pub fn hidden_dim(&self) -> usize {
    1638
    4
            self.hidden_dim
    1639
    4
        }
    1640
    1641
        /// Check if using Multi-Query Attention (MQA)
    1642
        #[must_use]
    1643
    3
        pub fn is_mqa(&self) -> bool {
    1644
    3
            self.num_kv_heads == 1
    1645
    3
        }
    1646
    1647
        /// Check if using Grouped-Query Attention (GQA)
    1648
        #[must_use]
    1649
    3
        pub fn is_gqa(&self) -> bool {
    1650
    3
            self.num_kv_heads > 1 && self.num_kv_heads < self.num_heads2
    1651
    3
        }
    1652
    1653
        /// Check if using standard Multi-Head Attention (MHA)
    1654
        #[must_use]
    1655
    3
        pub fn is_mha(&self) -> bool {
    1656
    3
            self.num_kv_heads == self.num_heads
    1657
    3
        }
    1658
    }