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

    Line
    Count
    Source
    1
    //! Attention computation for OwnedQuantizedModel
    2
    //!
    3
    //! Contains apply_rope (rotary position embeddings), causal_attention,
    4
    //! and KV-cache based attention variants with GQA support.
    5
    6
    use crate::error::{RealizarError, Result};
    7
    use crate::gguf::{GGUFConfig, OwnedQuantizedModel};
    8
    9
    impl OwnedQuantizedModel {
    10
        /// Apply Rotary Position Embedding (RoPE) to Q or K vectors
    11
        ///
    12
        /// Supports two RoPE styles:
    13
        /// - NORM (type 0): Adjacent pairs rotation (LLaMA default)
    14
        /// - NEOX (type 2): Split halves rotation (GPT-NeoX)
    15
        ///
    16
        /// # Arguments
    17
        /// * `x` - Vector to rotate in-place [num_heads_in_x * head_dim]
    18
        /// * `position` - Position index for frequency calculation
    19
        /// * `num_heads_in_x` - Number of heads in x (num_heads for Q, num_kv_heads for K)
    20
        ///
    21
        /// # GQA Support
    22
        /// For GQA models, pass num_heads for Q vectors and num_kv_heads for K vectors.
    23
    1.43k
        pub(crate) fn apply_rope(&self, x: &mut [f32], position: usize, num_heads_in_x: usize) {
    24
    1.43k
            let head_dim = self.config.hidden_dim / self.config.num_heads;
    25
    1.43k
            let half_dim = head_dim / 2;
    26
    1.43k
            let theta = self.config.rope_theta;
    27
    1.43k
            let rope_type = self.config.rope_type;
    28
    29
            // Stack-based buffers (max 128 = 256 head_dim, covers all common models)
    30
            // Avoids heap allocation on every call
    31
    1.43k
            let mut cos_vals: [f32; 128] = [0.0; 128];
    32
    1.43k
            let mut sin_vals: [f32; 128] = [0.0; 128];
    33
    34
            // Pre-compute cos/sin for this position (reused across all heads)
    35
    1.43k
            let pos_f32 = position as f32;
    36
    1.43k
            let head_dim_f32 = head_dim as f32;
    37
    11.5k
            for i in 0..half_dim1.43k.min1.43k(128) {
    38
    11.5k
                let freq = 1.0 / theta.powf(2.0 * i as f32 / head_dim_f32);
    39
    11.5k
                let angle = pos_f32 * freq;
    40
    11.5k
                let (sin_v, cos_v) = angle.sin_cos();
    41
    11.5k
                cos_vals[i] = cos_v;
    42
    11.5k
                sin_vals[i] = sin_v;
    43
    11.5k
            }
    44
    45
            // Apply rotation to each head
    46
    5.79k
            for h in 0..num_heads_in_x1.43k {
    47
    5.79k
                let head_start = h * head_dim;
    48
    49
    5.79k
                if head_start + head_dim > x.len() {
    50
    0
                    continue;
    51
    5.79k
                }
    52
    53
    5.79k
                if rope_type == 2 {
    54
    0
                    // NEOX style: split halves (x[0..half], x[half..])
    55
    0
                    // Used by GPT-NeoX and some newer models
    56
    0
                    let (first_half, second_half) =
    57
    0
                        x[head_start..head_start + head_dim].split_at_mut(half_dim);
    58
    0
                    crate::quantize::apply_rope_rotation_simd(
    59
    0
                        first_half,
    60
    0
                        second_half,
    61
    0
                        &cos_vals[..half_dim],
    62
    0
                        &sin_vals[..half_dim],
    63
    0
                    );
    64
    0
                } else {
    65
                    // NORM style (type 0): adjacent pairs (x[0], x[1]), (x[2], x[3]), ...
    66
                    // This is the default for LLaMA-family models
    67
    5.79k
                    let head_slice = &mut x[head_start..head_start + head_dim];
    68
    52.5k
                    for i in 0..half_dim5.79k {
    69
    52.5k
                        let x0 = head_slice[2 * i];
    70
    52.5k
                        let x1 = head_slice[2 * i + 1];
    71
    52.5k
                        let cos_v = cos_vals[i];
    72
    52.5k
                        let sin_v = sin_vals[i];
    73
    52.5k
                        head_slice[2 * i] = x0 * cos_v - x1 * sin_v;
    74
    52.5k
                        head_slice[2 * i + 1] = x0 * sin_v + x1 * cos_v;
    75
    52.5k
                    }
    76
                }
    77
            }
    78
    1.43k
        }
    79
    80
        /// Compute scaled dot-product attention with causal mask (IMP-101b)
    81
        ///
    82
        /// Computes: softmax(QK^T / sqrt(d_k)) * V with causal masking
    83
        ///
    84
        /// # Arguments
    85
        /// * `q` - Query vectors [seq_len, q_dim] where q_dim = num_heads * head_dim
    86
        /// * `k` - Key vectors [seq_len, kv_dim] where kv_dim = num_kv_heads * head_dim
    87
        /// * `v` - Value vectors [seq_len, kv_dim] where kv_dim = num_kv_heads * head_dim
    88
        ///
    89
        /// # Returns
    90
        /// Attention output [seq_len, q_dim] where q_dim = num_heads * head_dim
    91
        ///
    92
        /// # GQA (Grouped Query Attention) Support
    93
        /// For models where num_kv_heads < num_heads (e.g., TinyLlama: 4 vs 32),
    94
        /// multiple Q heads share the same K/V head. The group size is num_heads / num_kv_heads.
    95
    7
        pub(crate) fn causal_attention(
    96
    7
            &self,
    97
    7
            q: &[f32],
    98
    7
            k: &[f32],
    99
    7
            v: &[f32],
    100
    7
            seq_len: usize,
    101
    7
        ) -> Vec<f32> {
    102
    7
            let num_heads = self.config.num_heads;
    103
    7
            let num_kv_heads = self.config.num_kv_heads;
    104
    7
            let head_dim = self.config.hidden_dim / num_heads;
    105
    7
            let scale = 1.0 / (head_dim as f32).sqrt();
    106
    107
            // GQA: multiple Q heads share each KV head
    108
            // group_size = num_heads / num_kv_heads (e.g., 32/4 = 8 for TinyLlama)
    109
    7
            let group_size = num_heads / num_kv_heads;
    110
    111
            // Q has num_heads heads, K/V have num_kv_heads heads
    112
    7
            let q_dim = num_heads * head_dim; // e.g., 32 * 64 = 2048
    113
    7
            let kv_dim = num_kv_heads * head_dim; // e.g., 4 * 64 = 256
    114
    115
    7
            let mut output = vec![0.0f32; seq_len * q_dim];
    116
    117
            // Process each Q head independently
    118
    20
            for head in 0..num_heads7 {
    119
                // Map Q head to corresponding KV head (GQA grouping)
    120
    20
                let kv_head = head / group_size;
    121
    122
    20
                let q_head_offset = head * head_dim;
    123
    20
                let kv_head_offset = kv_head * head_dim;
    124
    125
                // Process each query position
    126
    56
                for i in 0..seq_len20 {
    127
                    // Compute attention scores for this query against all keys up to position i (causal)
    128
    56
                    let mut scores = Vec::with_capacity(i + 1);
    129
    56
                    let q_start = i * q_dim + q_head_offset;
    130
    131
    135
                    for j in 0..=i56 {
    132
                        // Only attend to positions 0..=i (causal mask)
    133
    135
                        let k_start = j * kv_dim + kv_head_offset;
    134
    135
                        // Dot product Q[i] · K[j]
    136
    135
                        let mut score = 0.0f32;
    137
    1.54k
                        for d in 0..head_dim135 {
    138
    1.54k
                            score += q[q_start + d] * k[k_start + d];
    139
    1.54k
                        }
    140
    135
                        scores.push(score * scale);
    141
                    }
    142
    143
                    // Softmax (SIMD-optimized)
    144
    56
                    crate::quantize::softmax_simd(&mut scores);
    145
    146
                    // Weighted sum of values
    147
    56
                    let out_start = i * q_dim + q_head_offset;
    148
    135
                    for (j, &weight) in scores.iter()56.enumerate56() {
    149
    135
                        let v_start = j * kv_dim + kv_head_offset;
    150
    1.54k
                        for d in 0..head_dim135 {
    151
    1.54k
                            output[out_start + d] += weight * v[v_start + d];
    152
    1.54k
                        }
    153
                    }
    154
                }
    155
            }
    156
    157
    7
            output
    158
    7
        }
    159
    160
        /// Get model configuration
    161
    0
        pub fn config(&self) -> &GGUFConfig {
    162
    0
            &self.config
    163
    0
        }
    164
    165
        /// Check if CUDA is enabled
    166
        #[cfg(feature = "cuda")]
    167
        pub fn cuda_enabled(&self) -> bool {
    168
            self.cuda_executor.is_some()
    169
        }
    170
    171
        /// Check if CUDA acceleration is enabled (stub when cuda feature disabled)
    172
        #[cfg(not(feature = "cuda"))]
    173
    0
        pub fn cuda_enabled(&self) -> bool {
    174
    0
            false
    175
    0
        }
    176
    177
        // ============================================================================
    178
        // SIMD Helper Methods
    179
        // ============================================================================
    180
    181
        /// SIMD-optimized dot product for f32 slices
    182
        #[inline]
    183
    28.8k
        fn simd_dot_f32(a: &[f32], b: &[f32]) -> f32 {
    184
            #[cfg(target_arch = "x86_64")]
    185
            {
    186
    28.8k
                if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
    187
                    // SAFETY: We've verified AVX2+FMA support
    188
    28.8k
                    unsafe { Self::simd_dot_f32_avx2(a, b) }
    189
                } else {
    190
    0
                    Self::simd_dot_f32_scalar(a, b)
    191
                }
    192
            }
    193
            #[cfg(not(target_arch = "x86_64"))]
    194
            {
    195
                Self::simd_dot_f32_scalar(a, b)
    196
            }
    197
    28.8k
        }
    198
    199
        #[cfg(target_arch = "x86_64")]
    200
        #[target_feature(enable = "avx2", enable = "fma")]
    201
        #[inline]
    202
    28.8k
        unsafe fn simd_dot_f32_avx2(a: &[f32], b: &[f32]) -> f32 {
    203
            // SAFETY: Memory safety ensured by bounds checking and alignment
    204
            unsafe {
    205
                use std::arch::x86_64::{
    206
                    _mm256_castps256_ps128, _mm256_extractf128_ps, _mm256_fmadd_ps, _mm256_loadu_ps,
    207
                    _mm256_setzero_ps, _mm_add_ps, _mm_add_ss, _mm_cvtss_f32, _mm_movehdup_ps,
    208
                    _mm_movehl_ps,
    209
                };
    210
    211
    28.8k
                let len = a.len().min(b.len());
    212
    28.8k
                let mut acc = _mm256_setzero_ps();
    213
    28.8k
                let mut i = 0;
    214
    215
                // Process 16 floats at a time (2x unrolled for better ILP)
    216
    53.1k
                while i + 16 <= len {
    217
    24.2k
                    let va0 = _mm256_loadu_ps(a.as_ptr().add(i));
    218
    24.2k
                    let vb0 = _mm256_loadu_ps(b.as_ptr().add(i));
    219
    24.2k
                    let va1 = _mm256_loadu_ps(a.as_ptr().add(i + 8));
    220
    24.2k
                    let vb1 = _mm256_loadu_ps(b.as_ptr().add(i + 8));
    221
    24.2k
                    acc = _mm256_fmadd_ps(va0, vb0, acc);
    222
    24.2k
                    acc = _mm256_fmadd_ps(va1, vb1, acc);
    223
    24.2k
                    i += 16;
    224
    24.2k
                }
    225
                // Handle remaining 8-float chunk
    226
    28.8k
                if i + 8 <= len {
    227
    4.66k
                    let va = _mm256_loadu_ps(a.as_ptr().add(i));
    228
    4.66k
                    let vb = _mm256_loadu_ps(b.as_ptr().add(i));
    229
    4.66k
                    acc = _mm256_fmadd_ps(va, vb, acc);
    230
    4.66k
                    i += 8;
    231
    24.2k
                }
    232
    233
                // Horizontal sum
    234
    28.8k
                let hi = _mm256_extractf128_ps(acc, 1);
    235
    28.8k
                let lo = _mm256_castps256_ps128(acc);
    236
    28.8k
                let sum128 = _mm_add_ps(lo, hi);
    237
    28.8k
                let shuf = _mm_movehdup_ps(sum128);
    238
    28.8k
                let sums = _mm_add_ps(sum128, shuf);
    239
    28.8k
                let shuf2 = _mm_movehl_ps(sums, sums);
    240
    28.8k
                let result = _mm_add_ss(sums, shuf2);
    241
    28.8k
                let mut sum = _mm_cvtss_f32(result);
    242
    243
                // Handle remaining elements
    244
    28.9k
                while i < len {
    245
    56
                    sum += a[i] * b[i];
    246
    56
                    i += 1;
    247
    56
                }
    248
    249
    28.8k
                sum
    250
            }
    251
    28.8k
        }
    252
    253
        #[inline]
    254
    0
        fn simd_dot_f32_scalar(a: &[f32], b: &[f32]) -> f32 {
    255
    0
            a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
    256
    0
        }
    257
    258
        /// SIMD-optimized scaled accumulation: out[i] += weight * val[i]
    259
        #[inline]
    260
    28.8k
        fn simd_axpy_f32(out: &mut [f32], weight: f32, val: &[f32]) {
    261
            #[cfg(target_arch = "x86_64")]
    262
            {
    263
    28.8k
                if is_x86_feature_detected!("avx2") {
    264
                    // SAFETY: We've verified AVX2 support
    265
    28.8k
                    unsafe { Self::simd_axpy_f32_avx2(out, weight, val) }
    266
    0
                } else {
    267
    0
                    Self::simd_axpy_f32_scalar(out, weight, val);
    268
    0
                }
    269
            }
    270
            #[cfg(not(target_arch = "x86_64"))]
    271
            {
    272
                Self::simd_axpy_f32_scalar(out, weight, val);
    273
            }
    274
    28.8k
        }
    275
    276
        #[cfg(target_arch = "x86_64")]
    277
        #[target_feature(enable = "avx2", enable = "fma")]
    278
        #[inline]
    279
    28.8k
        unsafe fn simd_axpy_f32_avx2(out: &mut [f32], weight: f32, val: &[f32]) {
    280
            use std::arch::x86_64::{
    281
                _mm256_fmadd_ps, _mm256_loadu_ps, _mm256_set1_ps, _mm256_storeu_ps,
    282
            };
    283
    284
    28.8k
            let len = out.len().min(val.len());
    285
    28.8k
            let w = _mm256_set1_ps(weight);
    286
    28.8k
            let mut i = 0;
    287
    288
            // Process 8 floats at a time
    289
    81.9k
            while i + 8 <= len {
    290
    53.0k
                // SAFETY: bounds checked above, pointers valid
    291
    53.0k
                let v_out = unsafe { _mm256_loadu_ps(out.as_ptr().add(i)) };
    292
    53.0k
                // SAFETY: Memory safety ensured by bounds checking and alignment
    293
    53.0k
                let v_val = unsafe { _mm256_loadu_ps(val.as_ptr().add(i)) };
    294
    53.0k
                let result = _mm256_fmadd_ps(w, v_val, v_out);
    295
    53.0k
                // SAFETY: Memory safety ensured by bounds checking and alignment
    296
    53.0k
                unsafe { _mm256_storeu_ps(out.as_mut_ptr().add(i), result) };
    297
    53.0k
                i += 8;
    298
    53.0k
            }
    299
    300
            // Handle remaining elements
    301
    28.9k
            while i < len {
    302
    56
                out[i] += weight * val[i];
    303
    56
                i += 1;
    304
    56
            }
    305
    28.8k
        }
    306
    307
        #[inline]
    308
    0
        fn simd_axpy_f32_scalar(out: &mut [f32], weight: f32, val: &[f32]) {
    309
    0
            for (o, v) in out.iter_mut().zip(val.iter()) {
    310
    0
                *o += weight * *v;
    311
    0
            }
    312
    0
        }
    313
    314
        // ============================================================================
    315
        // KV Cache Attention Methods
    316
        // ============================================================================
    317
    318
        /// Attention with KV cache for autoregressive generation
    319
        ///
    320
        /// # Arguments
    321
        /// * `q` - Query vector for current position [hidden_dim]
    322
        /// * `k_cache` - Cached keys [cache_len, hidden_dim]
    323
        /// * `v_cache` - Cached values [cache_len, hidden_dim]
    324
        /// * `current_k` - Key for current position [hidden_dim]
    325
        /// * `current_v` - Value for current position [hidden_dim]
    326
        ///
    327
        /// # Returns
    328
        /// Attention output [hidden_dim]
    329
    4
        pub(crate) fn attention_with_cache(
    330
    4
            &self,
    331
    4
            q: &[f32],
    332
    4
            k_cache: &[f32],
    333
    4
            v_cache: &[f32],
    334
    4
            current_k: &[f32],
    335
    4
            current_v: &[f32],
    336
    4
        ) -> Vec<f32> {
    337
    4
            let hidden_dim = self.config.hidden_dim;
    338
    4
            let num_heads = self.config.num_heads;
    339
    4
            let head_dim = hidden_dim / num_heads;
    340
    4
            let scale = 1.0 / (head_dim as f32).sqrt();
    341
    342
            // Total sequence length = cached + 1 (current)
    343
    4
            let cache_len = k_cache.len() / hidden_dim;
    344
    4
            let total_len = cache_len + 1;
    345
    346
    4
            let mut output = vec![0.0f32; hidden_dim];
    347
    348
            // Process each head
    349
    13
            for head in 0..num_heads4 {
    350
    13
                let head_offset = head * head_dim;
    351
    13
                let q_head = &q[head_offset..head_offset + head_dim];
    352
    353
                // Compute attention scores against all positions (cached + current)
    354
    13
                let mut scores = Vec::with_capacity(total_len);
    355
    356
                // Scores against cached positions (SIMD-optimized)
    357
    257
                for pos in 0..cache_len13 {
    358
    257
                    let k_start = pos * hidden_dim + head_offset;
    359
    257
                    let cached_key = &k_cache[k_start..k_start + head_dim];
    360
    257
                    let score = Self::simd_dot_f32(q_head, cached_key) * scale;
    361
    257
                    scores.push(score);
    362
    257
                }
    363
    364
                // Score against current position (SIMD-optimized)
    365
    13
                let curr_key = &current_k[head_offset..head_offset + head_dim];
    366
    13
                let current_score = Self::simd_dot_f32(q_head, curr_key) * scale;
    367
    13
                scores.push(current_score);
    368
    369
                // Softmax (SIMD-optimized)
    370
    13
                crate::quantize::softmax_simd(&mut scores);
    371
    372
                // Weighted sum of values
    373
    13
                let out_head = &mut output[head_offset..head_offset + head_dim];
    374
    375
                // Sum over cached values (SIMD-optimized)
    376
    257
                for (pos, &weight) in scores.iter()13.enumerate13().take13(cache_len13) {
    377
    257
                    let v_start = pos * hidden_dim + head_offset;
    378
    257
                    let cached_val = &v_cache[v_start..v_start + head_dim];
    379
    257
                    Self::simd_axpy_f32(out_head, weight, cached_val);
    380
    257
                }
    381
    382
                // Add current value (SIMD-optimized)
    383
    13
                let curr_val = &current_v[head_offset..head_offset + head_dim];
    384
    13
                let current_weight = scores[cache_len];
    385
    13
                Self::simd_axpy_f32(out_head, current_weight, curr_val);
    386
            }
    387
    388
    4
            output
    389
    4
        }
    390
    391
        /// Compute attention with Grouped Query Attention (GQA) support (IMP-105)
    392
        ///
    393
        /// GQA uses fewer KV heads than Q heads, with multiple Q heads sharing each KV head.
    394
        /// This reduces memory bandwidth and KV cache size for large models.
    395
        ///
    396
        /// # Arguments
    397
        /// * `q` - Query vector for current position [hidden_dim] (num_heads Q heads)
    398
        /// * `k_cache` - Cached keys [cache_len, kv_dim] (num_kv_heads KV heads)
    399
        /// * `v_cache` - Cached values [cache_len, kv_dim] (num_kv_heads KV heads)
    400
        /// * `current_k` - Key for current position [kv_dim]
    401
        /// * `current_v` - Value for current position [kv_dim]
    402
        ///
    403
        /// # Returns
    404
        /// Attention output [hidden_dim]
    405
        ///
    406
        /// # GQA Mapping
    407
        /// Q head i uses KV head (i * num_kv_heads / num_heads)
    408
        /// Example: 8 Q heads, 2 KV heads → Q heads 0-3 use KV head 0, Q heads 4-7 use KV head 1
    409
    297
        pub fn attention_with_cache_gqa(
    410
    297
            &self,
    411
    297
            q: &[f32],
    412
    297
            k_cache: &[f32],
    413
    297
            v_cache: &[f32],
    414
    297
            current_k: &[f32],
    415
    297
            current_v: &[f32],
    416
    297
        ) -> Vec<f32> {
    417
    297
            let hidden_dim = self.config.hidden_dim;
    418
    297
            let num_heads = self.config.num_heads;
    419
    297
            let num_kv_heads = self.config.num_kv_heads;
    420
    297
            let head_dim = hidden_dim / num_heads;
    421
    297
            let kv_dim = num_kv_heads * head_dim;
    422
    297
            let scale = 1.0 / (head_dim as f32).sqrt();
    423
    424
            // Number of Q heads that share each KV head
    425
    297
            let q_per_kv = num_heads / num_kv_heads;
    426
    427
            // Total sequence length = cached + 1 (current)
    428
    297
            let cache_len = if kv_dim > 0 {
    429
    297
                k_cache.len() / kv_dim
    430
            } else {
    431
    0
                0
    432
            };
    433
    297
            let total_len = cache_len + 1;
    434
    435
    297
            let mut output = vec![0.0f32; hidden_dim];
    436
    437
            // Process each Q head
    438
    1.00k
            for q_head in 0..num_heads297 {
    439
    1.00k
                let q_head_offset = q_head * head_dim;
    440
    1.00k
                let q_head_data = &q[q_head_offset..q_head_offset + head_dim];
    441
    442
                // Map Q head to KV head (GQA mapping)
    443
    1.00k
                let kv_head = q_head / q_per_kv;
    444
    1.00k
                let kv_head_offset = kv_head * head_dim;
    445
    446
                // Compute attention scores against all positions (cached + current)
    447
    1.00k
                let mut scores = Vec::with_capacity(total_len);
    448
    449
                // Scores against cached positions (SIMD-optimized)
    450
    21.9k
                for pos in 0..cache_len1.00k {
    451
    21.9k
                    let k_start = pos * kv_dim + kv_head_offset;
    452
    21.9k
                    let cached_key = &k_cache[k_start..k_start + head_dim];
    453
    21.9k
                    let score = Self::simd_dot_f32(q_head_data, cached_key);
    454
    21.9k
                    scores.push(score * scale);
    455
    21.9k
                }
    456
    457
                // Score against current position (SIMD-optimized)
    458
    1.00k
                let curr_key = &current_k[kv_head_offset..kv_head_offset + head_dim];
    459
    1.00k
                let current_score = Self::simd_dot_f32(q_head_data, curr_key);
    460
    1.00k
                scores.push(current_score * scale);
    461
    462
                // Softmax (SIMD-optimized)
    463
    1.00k
                crate::quantize::softmax_simd(&mut scores);
    464
    465
                // Weighted sum of values
    466
    1.00k
                let out_head = &mut output[q_head_offset..q_head_offset + head_dim];
    467
    468
                // Sum over cached values (SIMD-optimized)
    469
    21.9k
                for (pos, &weight) in scores.iter()1.00k.enumerate1.00k().take1.00k(cache_len1.00k) {
    470
    21.9k
                    let v_start = pos * kv_dim + kv_head_offset;
    471
    21.9k
                    let cached_val = &v_cache[v_start..v_start + head_dim];
    472
    21.9k
                    Self::simd_axpy_f32(out_head, weight, cached_val);
    473
    21.9k
                }
    474
    475
                // Add current value (SIMD-optimized)
    476
    1.00k
                let curr_val = &current_v[kv_head_offset..kv_head_offset + head_dim];
    477
    1.00k
                let current_weight = scores[cache_len];
    478
    1.00k
                Self::simd_axpy_f32(out_head, current_weight, curr_val);
    479
            }
    480
    481
    297
            output
    482
    297
        }
    483
    484
        /// Attention with cache - writes to pre-allocated buffer (IMP-131)
    485
    223
        pub fn attention_with_cache_gqa_into(
    486
    223
            &self,
    487
    223
            q: &[f32],
    488
    223
            k_cache: &[f32],
    489
    223
            v_cache: &[f32],
    490
    223
            current_k: &[f32],
    491
    223
            current_v: &[f32],
    492
    223
            output: &mut [f32],
    493
    223
        ) {
    494
    223
            let hidden_dim = self.config.hidden_dim;
    495
    223
            let num_heads = self.config.num_heads;
    496
    223
            let num_kv_heads = self.config.num_kv_heads;
    497
    223
            let head_dim = hidden_dim / num_heads;
    498
    223
            let kv_dim = num_kv_heads * head_dim;
    499
    223
            let scale = 1.0 / (head_dim as f32).sqrt();
    500
    501
    223
            let q_per_kv = num_heads / num_kv_heads;
    502
    503
    223
            let cache_len = if kv_dim > 0 {
    504
    223
                k_cache.len() / kv_dim
    505
            } else {
    506
    0
                0
    507
            };
    508
    223
            let total_len = cache_len + 1;
    509
    510
            // Zero output buffer
    511
    223
            output[..hidden_dim].iter_mut().for_each(|x| *x = 0.0);
    512
    513
            // Stack-allocated scores buffer (max 8192 seq length)
    514
    223
            let mut scores_buf = [0.0f32; 8192];
    515
    223
            let scores = &mut scores_buf[..total_len];
    516
    517
    892
            for q_head in 0..num_heads223 {
    518
    892
                let q_head_offset = q_head * head_dim;
    519
    892
                let q_head_data = &q[q_head_offset..q_head_offset + head_dim];
    520
    521
    892
                let kv_head = q_head / q_per_kv;
    522
    892
                let kv_head_offset = kv_head * head_dim;
    523
    524
                // Compute attention scores
    525
    4.76k
                for pos in 0..cache_len892 {
    526
    4.76k
                    let k_start = pos * kv_dim + kv_head_offset;
    527
    4.76k
                    let cached_key = &k_cache[k_start..k_start + head_dim];
    528
    4.76k
                    scores[pos] = Self::simd_dot_f32(q_head_data, cached_key) * scale;
    529
    4.76k
                }
    530
    531
    892
                let curr_key = &current_k[kv_head_offset..kv_head_offset + head_dim];
    532
    892
                scores[cache_len] = Self::simd_dot_f32(q_head_data, curr_key) * scale;
    533
    534
                // Softmax
    535
    892
                crate::quantize::softmax_simd(scores);
    536
    537
                // Weighted sum of values
    538
    892
                let out_head = &mut output[q_head_offset..q_head_offset + head_dim];
    539
    540
    4.76k
                for (pos, &weight) in scores892.iter892().enumerate892().take892(cache_len892) {
    541
    4.76k
                    let v_start = pos * kv_dim + kv_head_offset;
    542
    4.76k
                    let cached_val = &v_cache[v_start..v_start + head_dim];
    543
    4.76k
                    Self::simd_axpy_f32(out_head, weight, cached_val);
    544
    4.76k
                }
    545
    546
    892
                let curr_val = &current_v[kv_head_offset..kv_head_offset + head_dim];
    547
    892
                Self::simd_axpy_f32(out_head, scores[cache_len], curr_val);
    548
            }
    549
    223
        }
    550
    551
        /// Adaptive attention with KV cache - auto-selects CPU or GPU backend (IMP-122)
    552
        ///
    553
        /// For short cache lengths (< 64), uses efficient CPU implementation.
    554
        /// For long cache lengths (>= 64), uses GPU-accelerated computation.
    555
        ///
    556
        /// # Arguments
    557
        /// * `q` - Query vector for current position [hidden_dim]
    558
        /// * `k_cache` - Cached keys [cache_len, hidden_dim]
    559
        /// * `v_cache` - Cached values [cache_len, hidden_dim]
    560
        /// * `current_k` - Key for current position [hidden_dim]
    561
        /// * `current_v` - Value for current position [hidden_dim]
    562
        ///
    563
        /// # Returns
    564
        /// Result containing attention output [hidden_dim]
    565
        ///
    566
        /// # Errors
    567
        /// Returns error if GPU operations fail (for GPU path)
    568
        #[cfg(feature = "gpu")]
    569
    65
        pub fn adaptive_attention_with_cache(
    570
    65
            &self,
    571
    65
            q: &[f32],
    572
    65
            k_cache: &[f32],
    573
    65
            v_cache: &[f32],
    574
    65
            current_k: &[f32],
    575
    65
            current_v: &[f32],
    576
    65
        ) -> Result<Vec<f32>> {
    577
    65
            let hidden_dim = self.config.hidden_dim;
    578
    579
            // Calculate cache length
    580
    65
            let cache_len = if hidden_dim > 0 {
    581
    65
                k_cache.len() / hidden_dim
    582
            } else {
    583
    0
                0
    584
            };
    585
    586
            // Threshold for GPU dispatch (matches IMP-119)
    587
            const GPU_CACHE_LEN_THRESHOLD: usize = 64;
    588
    589
    65
            if cache_len >= GPU_CACHE_LEN_THRESHOLD {
    590
                // GPU path for long sequences
    591
    63
                self.gpu_attention_with_cache(q, k_cache, v_cache, current_k, current_v)
    592
            } else {
    593
                // CPU path for short sequences - use existing implementation
    594
    2
                Ok(self.attention_with_cache(q, k_cache, v_cache, current_k, current_v))
    595
            }
    596
    65
        }
    597
    598
        /// CPU-only version of adaptive attention
    599
        #[cfg(not(feature = "gpu"))]
    600
        pub fn adaptive_attention_with_cache(
    601
            &self,
    602
            q: &[f32],
    603
            k_cache: &[f32],
    604
            v_cache: &[f32],
    605
            current_k: &[f32],
    606
            current_v: &[f32],
    607
        ) -> Result<Vec<f32>> {
    608
            Ok(self.attention_with_cache(q, k_cache, v_cache, current_k, current_v))
    609
        }
    610
    611
        /// GPU-accelerated attention with KV cache (IMP-122)
    612
        ///
    613
        /// Uses GPU for Q@K^T computation when cache is large enough.
    614
        #[cfg(feature = "gpu")]
    615
    63
        fn gpu_attention_with_cache(
    616
    63
            &self,
    617
    63
            q: &[f32],
    618
    63
            k_cache: &[f32],
    619
    63
            v_cache: &[f32],
    620
    63
            current_k: &[f32],
    621
    63
            current_v: &[f32],
    622
    63
        ) -> Result<Vec<f32>> {
    623
            use crate::gpu::HybridScheduler;
    624
    625
    63
            let hidden_dim = self.config.hidden_dim;
    626
    63
            let num_heads = self.config.num_heads;
    627
    63
            let head_dim = hidden_dim / num_heads;
    628
    63
            let scale = 1.0 / (head_dim as f32).sqrt();
    629
    630
            // Total sequence length = cached + 1 (current)
    631
    63
            let cache_len = k_cache.len() / hidden_dim;
    632
    63
            let total_len = cache_len + 1;
    633
    634
    63
            let mut output = vec![0.0f32; hidden_dim];
    635
    636
            // Create scheduler for GPU operations
    637
    63
            let mut scheduler = HybridScheduler::with_threshold(1000).map_err(|e| {0
    638
    0
                RealizarError::UnsupportedOperation {
    639
    0
                    operation: "gpu_attention_with_cache".to_string(),
    640
    0
                    reason: format!("Failed to create scheduler: {e}"),
    641
    0
                }
    642
    0
            })?;
    643
    644
            // Process each head
    645
    224
            for head in 0..num_heads63 {
    646
    224
                let head_offset = head * head_dim;
    647
    224
                let q_head = &q[head_offset..head_offset + head_dim];
    648
    649
                // Build full K matrix for this head: [total_len, head_dim]
    650
    224
                let mut k_full = Vec::with_capacity(total_len * head_dim);
    651
    18.2k
                for pos in 0..cache_len224 {
    652
    18.2k
                    let k_start = pos * hidden_dim + head_offset;
    653
    18.2k
                    k_full.extend_from_slice(&k_cache[k_start..k_start + head_dim]);
    654
    18.2k
                }
    655
    224
                k_full.extend_from_slice(&current_k[head_offset..head_offset + head_dim]);
    656
    657
                // Transpose K to [head_dim, total_len] for matmul
    658
    224
                let mut k_t = vec![0.0f32; head_dim * total_len];
    659
    18.4k
                for pos in 0..total_len224 {
    660
    277k
                    for d in 0..head_dim18.4k {
    661
    277k
                        k_t[d * total_len + pos] = k_full[pos * head_dim + d];
    662
    277k
                    }
    663
                }
    664
    665
                // GPU matmul: Q[1, head_dim] @ K_T[head_dim, total_len] -> [1, total_len]
    666
    224
                let scores_raw = scheduler
    667
    224
                    .matmul(q_head, &k_t, 1, head_dim, total_len)
    668
    224
                    .map_err(|e| RealizarError::UnsupportedOperation {
    669
    0
                        operation: "gpu_attention_with_cache".to_string(),
    670
    0
                        reason: format!("GPU matmul failed: {e}"),
    671
    0
                    })?;
    672
    673
                // Scale scores
    674
    18.4k
                let mut scores224: Vec<f32>224 = scores_raw.iter()224.map224(|&s| s * scale).collect224();
    675
    676
                // Softmax (SIMD-optimized)
    677
    224
                crate::quantize::softmax_simd(&mut scores);
    678
    679
                // Weighted sum of values
    680
    224
                let out_head = &mut output[head_offset..head_offset + head_dim];
    681
    682
                // Cached values
    683
    18.2k
                for (pos, &weight) in scores.iter()224.enumerate224().take224(cache_len224) {
    684
    18.2k
                    let v_start = pos * hidden_dim + head_offset;
    685
    18.2k
                    let cached_val = &v_cache[v_start..v_start + head_dim];
    686
    273k
                    for d in 0..head_dim18.2k {
    687
    273k
                        out_head[d] += weight * cached_val[d];
    688
    273k
                    }
    689
                }
    690
    691
                // Current value
    692
    224
                let curr_val = &current_v[head_offset..head_offset + head_dim];
    693
    224
                let current_weight = scores[cache_len];
    694
    3.32k
                for d in 0..head_dim224 {
    695
    3.32k
                    out_head[d] += current_weight * curr_val[d];
    696
    3.32k
                }
    697
            }
    698
    699
    63
            Ok(output)
    700
    63
        }
    701
    702
        /// FlashAttention: Tiled attention with O(N) memory (PARITY-026)
    703
        ///
    704
        /// Implements the FlashAttention algorithm from Dao et al. for memory-efficient attention.
    705
        /// Uses online softmax to process attention in tiles without materializing the N×N matrix.
    706
        ///
    707
        /// # Key Properties
    708
        /// - Memory: O(N) instead of O(N²)
    709
        /// - Numerically equivalent to standard attention
    710
        /// - 10-100x memory savings for long sequences
    711
        ///
    712
        /// # Arguments
    713
        /// * `q` - Query vector [hidden_dim]
    714
        /// * `k_cache` - Cached keys [cache_len, hidden_dim]
    715
        /// * `v_cache` - Cached values [cache_len, hidden_dim]
    716
        /// * `current_k` - Current key [hidden_dim]
    717
        /// * `current_v` - Current value [hidden_dim]
    718
        /// * `block_size` - Tile size for tiled computation (default: 64)
    719
        ///
    720
        /// # Returns
    721
        /// Attention output [hidden_dim]
    722
        #[cfg(feature = "gpu")]
    723
    0
        pub fn flash_attention_tiled(
    724
    0
            &self,
    725
    0
            q: &[f32],
    726
    0
            k_cache: &[f32],
    727
    0
            v_cache: &[f32],
    728
    0
            current_k: &[f32],
    729
    0
            current_v: &[f32],
    730
    0
            block_size: usize,
    731
    0
        ) -> Vec<f32> {
    732
    0
            let hidden_dim = self.config.hidden_dim;
    733
    0
            let num_heads = self.config.num_heads;
    734
    0
            let head_dim = hidden_dim / num_heads;
    735
    0
            let scale = 1.0 / (head_dim as f32).sqrt();
    736
    737
            // Total sequence length = cached + 1 (current)
    738
    0
            let cache_len = k_cache.len() / hidden_dim;
    739
    0
            let total_len = cache_len + 1;
    740
    741
    0
            let mut output = vec![0.0f32; hidden_dim];
    742
    743
            // Process each head with FlashAttention tiling
    744
    0
            for head in 0..num_heads {
    745
    0
                let head_offset = head * head_dim;
    746
    0
                let q_head = &q[head_offset..head_offset + head_dim];
    747
    748
                // Online softmax state for this head
    749
    0
                let mut m_i = f32::NEG_INFINITY; // Running max
    750
    0
                let mut l_i = 0.0f32; // Running sum of exp(score - max)
    751
    0
                let mut o_i = vec![0.0f32; head_dim]; // Accumulated output
    752
    753
                // Process KV cache in tiles
    754
    0
                let num_tiles = total_len.div_ceil(block_size);
    755
    756
    0
                for tile_idx in 0..num_tiles {
    757
    0
                    let tile_start = tile_idx * block_size;
    758
    0
                    let tile_end = (tile_start + block_size).min(total_len);
    759
    0
                    let tile_len = tile_end - tile_start;
    760
    761
                    // Compute scores for this tile
    762
    0
                    let mut tile_scores = Vec::with_capacity(tile_len);
    763
    0
                    let mut tile_values: Vec<&[f32]> = Vec::with_capacity(tile_len);
    764
    765
    0
                    for pos in tile_start..tile_end {
    766
    0
                        if pos < cache_len {
    767
                            // From cache
    768
    0
                            let k_start = pos * hidden_dim + head_offset;
    769
    0
                            let cached_key = &k_cache[k_start..k_start + head_dim];
    770
    771
                            // Compute Q·K score
    772
    0
                            let mut score = 0.0f32;
    773
    0
                            for d in 0..head_dim {
    774
    0
                                score += q_head[d] * cached_key[d];
    775
    0
                            }
    776
    0
                            tile_scores.push(score * scale);
    777
    778
    0
                            let v_start = pos * hidden_dim + head_offset;
    779
    0
                            tile_values.push(&v_cache[v_start..v_start + head_dim]);
    780
                        } else {
    781
                            // Current position
    782
    0
                            let curr_key = &current_k[head_offset..head_offset + head_dim];
    783
    784
    0
                            let mut score = 0.0f32;
    785
    0
                            for d in 0..head_dim {
    786
    0
                                score += q_head[d] * curr_key[d];
    787
    0
                            }
    788
    0
                            tile_scores.push(score * scale);
    789
    790
    0
                            tile_values.push(&current_v[head_offset..head_offset + head_dim]);
    791
                        }
    792
                    }
    793
    794
                    // Find max in this tile
    795
    0
                    let m_tile = tile_scores
    796
    0
                        .iter()
    797
    0
                        .cloned()
    798
    0
                        .fold(f32::NEG_INFINITY, f32::max);
    799
    800
                    // Update running max
    801
    0
                    let m_new = m_i.max(m_tile);
    802
    803
                    // Rescale factors for online softmax
    804
    0
                    let scale_old = (m_i - m_new).exp();
    805
    0
                    let scale_tile = (m_tile - m_new).exp();
    806
    807
                    // Compute local softmax sum for this tile
    808
    0
                    let l_tile: f32 = tile_scores.iter().map(|&s| (s - m_tile).exp()).sum();
    809
    810
                    // Update running sum
    811
    0
                    l_i = l_i * scale_old + l_tile * scale_tile;
    812
    813
                    // Update output: rescale old output and add new contribution
    814
    0
                    for o in &mut o_i {
    815
    0
                        *o *= scale_old;
    816
    0
                    }
    817
    818
                    // Add weighted values from this tile
    819
    0
                    for (j, &score) in tile_scores.iter().enumerate() {
    820
    0
                        let attn_weight = (score - m_tile).exp() * scale_tile;
    821
    0
                        let v = tile_values[j];
    822
    0
                        for d in 0..head_dim {
    823
    0
                            o_i[d] += attn_weight * v[d];
    824
    0
                        }
    825
                    }
    826
    827
    0
                    m_i = m_new;
    828
                }
    829
    830
                // Finalize: divide by sum
    831
    0
                if l_i > 0.0 {
    832
    0
                    for d in 0..head_dim {
    833
    0
                        output[head_offset + d] = o_i[d] / l_i;
    834
    0
                    }
    835
    0
                }
    836
            }
    837
    838
    0
            output
    839
    0
        }
    840
    841
        /// CPU fallback for flash_attention_tiled - uses standard attention
    842
        #[cfg(not(feature = "gpu"))]
    843
        #[allow(unused_variables)]
    844
        pub fn flash_attention_tiled(
    845
            &self,
    846
            q: &[f32],
    847
            k_cache: &[f32],
    848
            v_cache: &[f32],
    849
            current_k: &[f32],
    850
            current_v: &[f32],
    851
            block_size: usize,
    852
        ) -> Vec<f32> {
    853
            // FlashAttention is a GPU optimization; CPU path uses standard attention
    854
            self.attention_with_cache(q, k_cache, v_cache, current_k, current_v)
    855
        }
    856
    }