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

    Line
    Count
    Source
    1
    //! Key-Value cache for efficient autoregressive generation
    2
    //!
    3
    //! Stores past key and value tensors to avoid recomputation during generation.
    4
    //! Enables O(1) per-token computation instead of O(n) for sequence of length n.
    5
    //!
    6
    //! ## Cache Types
    7
    //!
    8
    //! - [`KVCache`] - Basic KV cache with row-major storage
    9
    //! - [`OptimizedKVCache`] - Optimized cache with transposed V for better memory access
    10
    //!
    11
    //! ## Attention Functions
    12
    //!
    13
    //! - [`attention_with_cache`] - Scaled dot-product attention using cached K/V
    14
    //! - [`attention_with_transposed_v`] - Attention with transposed V storage
    15
    16
    use super::{simd_dot, simd_softmax};
    17
    18
    /// Key-Value cache for autoregressive generation
    19
    ///
    20
    /// Stores past key and value tensors to avoid recomputation during generation.
    21
    /// Enables O(1) per-token computation instead of O(n) for sequence of length n.
    22
    ///
    23
    /// # Example
    24
    ///
    25
    /// ```
    26
    /// use realizar::inference::KVCache;
    27
    ///
    28
    /// let mut cache = KVCache::new(12, 768, 2048);  // 12 layers, 768 hidden, 2048 max seq
    29
    /// assert!(cache.is_empty());
    30
    ///
    31
    /// // Store KV for position 0
    32
    /// let k = vec![0.1; 768];
    33
    /// let v = vec![0.2; 768];
    34
    /// cache.store(0, &k, &v);
    35
    /// cache.advance();
    36
    ///
    37
    /// assert_eq!(cache.len(), 1);
    38
    /// ```
    39
    #[derive(Clone)]
    40
    pub struct KVCache {
    41
        /// Key cache: [num_layers, max_seq_len, kv_dim]
    42
        k_cache: Vec<Vec<f32>>,
    43
        /// Value cache: [num_layers, max_seq_len, kv_dim]
    44
        v_cache: Vec<Vec<f32>>,
    45
        /// Current sequence length
    46
        seq_len: usize,
    47
        /// Hidden dimension per head
    48
        hidden_dim: usize,
    49
    }
    50
    51
    impl KVCache {
    52
        /// Create a new KV cache
    53
        ///
    54
        /// # Arguments
    55
        ///
    56
        /// * `num_layers` - Number of transformer layers
    57
        /// * `hidden_dim` - Hidden dimension (total across all heads)
    58
        /// * `max_seq_len` - Maximum sequence length to cache
    59
        #[must_use]
    60
    14
        pub fn new(num_layers: usize, hidden_dim: usize, max_seq_len: usize) -> Self {
    61
    14
            let k_cache = vec![vec![0.0; max_seq_len * hidden_dim]; num_layers];
    62
    14
            let v_cache = vec![vec![0.0; max_seq_len * hidden_dim]; num_layers];
    63
    14
            Self {
    64
    14
                k_cache,
    65
    14
                v_cache,
    66
    14
                seq_len: 0,
    67
    14
                hidden_dim,
    68
    14
            }
    69
    14
        }
    70
    71
        /// Store a new KV pair for a layer
    72
        ///
    73
        /// Stores the key and value vectors at the current sequence position.
    74
        /// Does nothing if the cache is full.
    75
        ///
    76
        /// # Arguments
    77
        ///
    78
        /// * `layer` - Layer index (0 to num_layers-1)
    79
        /// * `k` - Key vector of length `hidden_dim`
    80
        /// * `v` - Value vector of length `hidden_dim`
    81
    151
        pub fn store(&mut self, layer: usize, k: &[f32], v: &[f32]) {
    82
    151
            let start = self.seq_len * self.hidden_dim;
    83
    151
            let end = start + self.hidden_dim;
    84
    85
    151
            if end <= self.k_cache[layer].len() {
    86
    150
                self.k_cache[layer][start..end].copy_from_slice(k);
    87
    150
                self.v_cache[layer][start..end].copy_from_slice(v);
    88
    150
            }1
    89
    151
        }
    90
    91
        /// Advance the sequence position
    92
        ///
    93
        /// Call after storing all layers for the current position.
    94
    50
        pub fn advance(&mut self) {
    95
    50
            self.seq_len += 1;
    96
    50
        }
    97
    98
        /// Get cached keys for a layer
    99
        ///
    100
        /// Returns all keys from position 0 to current seq_len.
    101
        #[must_use]
    102
    14
        pub fn get_k(&self, layer: usize) -> &[f32] {
    103
    14
            &self.k_cache[layer][..self.seq_len * self.hidden_dim]
    104
    14
        }
    105
    106
        /// Get cached values for a layer
    107
        ///
    108
        /// Returns all values from position 0 to current seq_len.
    109
        #[must_use]
    110
    12
        pub fn get_v(&self, layer: usize) -> &[f32] {
    111
    12
            &self.v_cache[layer][..self.seq_len * self.hidden_dim]
    112
    12
        }
    113
    114
        /// Get current sequence length
    115
        #[must_use]
    116
    10
        pub fn len(&self) -> usize {
    117
    10
            self.seq_len
    118
    10
        }
    119
    120
        /// Check if cache is empty
    121
        #[must_use]
    122
    6
        pub fn is_empty(&self) -> bool {
    123
    6
            self.seq_len == 0
    124
    6
        }
    125
    126
        /// Reset the cache for a new sequence
    127
        ///
    128
        /// Clears the sequence position but keeps allocated memory.
    129
    3
        pub fn reset(&mut self) {
    130
    3
            self.seq_len = 0;
    131
    3
        }
    132
    }
    133
    134
    /// Compute attention with KV cache
    135
    ///
    136
    /// Computes scaled dot-product attention: softmax(QK^T / sqrt(d)) V
    137
    /// Uses cached K and V from previous positions.
    138
    ///
    139
    /// # Arguments
    140
    ///
    141
    /// * `q` - Query vector [hidden_dim]
    142
    /// * `k_cache` - Cached keys [seq_len × hidden_dim]
    143
    /// * `v_cache` - Cached values [seq_len × hidden_dim]
    144
    /// * `current_k` - Key for current position [hidden_dim]
    145
    /// * `current_v` - Value for current position [hidden_dim]
    146
    /// * `num_heads` - Number of attention heads
    147
    ///
    148
    /// # Returns
    149
    ///
    150
    /// Attention output [hidden_dim]
    151
    ///
    152
    /// # Example
    153
    ///
    154
    /// ```
    155
    /// use realizar::inference::attention_with_cache;
    156
    ///
    157
    /// let hidden_dim = 64;
    158
    /// let num_heads = 2;
    159
    ///
    160
    /// let q = vec![0.1; hidden_dim];
    161
    /// let k_cache: Vec<f32> = vec![];  // No cached positions
    162
    /// let v_cache: Vec<f32> = vec![];
    163
    /// let current_k = vec![0.1; hidden_dim];
    164
    /// let current_v = vec![0.2; hidden_dim];
    165
    ///
    166
    /// let output = attention_with_cache(&q, &k_cache, &v_cache, &current_k, &current_v, num_heads);
    167
    /// assert_eq!(output.len(), hidden_dim);
    168
    /// ```
    169
    #[must_use]
    170
    22
    pub fn attention_with_cache(
    171
    22
        q: &[f32],
    172
    22
        k_cache: &[f32],
    173
    22
        v_cache: &[f32],
    174
    22
        current_k: &[f32],
    175
    22
        current_v: &[f32],
    176
    22
        num_heads: usize,
    177
    22
    ) -> Vec<f32> {
    178
    22
        let hidden_dim = q.len();
    179
    22
        let head_dim = hidden_dim / num_heads;
    180
    22
        let cache_len = if hidden_dim > 0 {
    181
    22
            k_cache.len() / hidden_dim
    182
        } else {
    183
    0
            0
    184
        };
    185
    22
        let total_len = cache_len + 1;
    186
    22
        let scale = 1.0 / (head_dim as f32).sqrt();
    187
    188
    22
        let mut output = vec![0.0; hidden_dim];
    189
    190
        // Process each head independently
    191
    87
        for h in 0..num_heads22 {
    192
    87
            let head_offset = h * head_dim;
    193
    87
            let q_head = &q[head_offset..head_offset + head_dim];
    194
    195
            // Compute attention scores for all positions
    196
    87
            let mut scores = Vec::with_capacity(total_len);
    197
    198
            // Scores against cached K
    199
    128
            for pos in 0..cache_len87 {
    200
    128
                let k_start = pos * hidden_dim + head_offset;
    201
    128
                let k_head = &k_cache[k_start..k_start + head_dim];
    202
    128
                let score = simd_dot(q_head, k_head) * scale;
    203
    128
                scores.push(score);
    204
    128
            }
    205
    206
            // Score against current K
    207
    87
            let current_k_head = &current_k[head_offset..head_offset + head_dim];
    208
    87
            scores.push(simd_dot(q_head, current_k_head) * scale);
    209
    210
            // Softmax
    211
    87
            simd_softmax(&mut scores);
    212
    213
            // Weighted sum of V
    214
    87
            let out_head = &mut output[head_offset..head_offset + head_dim];
    215
    216
    128
            for (pos, &weight) in scores.iter()87.enumerate87().take87(cache_len87) {
    217
    128
                let v_start = pos * hidden_dim + head_offset;
    218
    128
                let v_head = &v_cache[v_start..v_start + head_dim];
    219
    7.69k
                for (i, &v) in v_head128.iter128().enumerate128() {
    220
    7.69k
                    out_head[i] += weight * v;
    221
    7.69k
                }
    222
            }
    223
    224
            // Add contribution from current V
    225
    87
            let current_v_head = &current_v[head_offset..head_offset + head_dim];
    226
    87
            let current_weight = scores[cache_len];
    227
    3.92k
            for (i, &v) in current_v_head87.iter87().enumerate87() {
    228
    3.92k
                out_head[i] += current_weight * v;
    229
    3.92k
            }
    230
        }
    231
    232
    22
        output
    233
    22
    }
    234
    235
    /// Optimized KV cache with contiguous storage
    236
    ///
    237
    /// Uses transposed V storage for better memory access during attention.
    238
    /// The value cache is stored as [hidden_dim × max_seq_len] instead of
    239
    /// [max_seq_len × hidden_dim], enabling better cache locality when
    240
    /// computing attention.
    241
    #[derive(Clone)]
    242
    pub struct OptimizedKVCache {
    243
        /// Key cache: [num_layers][seq_len × hidden_dim]
    244
        k_cache: Vec<Vec<f32>>,
    245
        /// Value cache (transposed): [num_layers][hidden_dim × seq_len]
    246
        v_cache: Vec<Vec<f32>>,
    247
        /// Current sequence length
    248
        seq_len: usize,
    249
        /// Hidden dimension
    250
        hidden_dim: usize,
    251
        /// Maximum sequence length
    252
        max_seq_len: usize,
    253
    }
    254
    255
    impl OptimizedKVCache {
    256
        /// Create a new optimized KV cache
    257
        #[must_use]
    258
    6
        pub fn new(num_layers: usize, hidden_dim: usize, max_seq_len: usize) -> Self {
    259
    6
            let k_cache = vec![vec![0.0; max_seq_len * hidden_dim]; num_layers];
    260
    6
            let v_cache = vec![vec![0.0; hidden_dim * max_seq_len]; num_layers];
    261
    6
            Self {
    262
    6
                k_cache,
    263
    6
                v_cache,
    264
    6
                seq_len: 0,
    265
    6
                hidden_dim,
    266
    6
                max_seq_len,
    267
    6
            }
    268
    6
        }
    269
    270
        /// Store a new KV pair with transposed V storage
    271
        ///
    272
        /// Does nothing if the cache is at maximum capacity.
    273
    10
        pub fn store(&mut self, layer: usize, k: &[f32], v: &[f32]) {
    274
    10
            if self.seq_len >= self.max_seq_len {
    275
    1
                return;
    276
    9
            }
    277
    278
            // Store K in normal format
    279
    9
            let k_start = self.seq_len * self.hidden_dim;
    280
    9
            let k_end = k_start + self.hidden_dim;
    281
    9
            self.k_cache[layer][k_start..k_end].copy_from_slice(k);
    282
    283
            // Store V transposed: v[i] goes to v_cache[i * max_seq_len + seq_len]
    284
    26
            for (i, &val) in v9.iter9().enumerate9() {
    285
    26
                self.v_cache[layer][i * self.max_seq_len + self.seq_len] = val;
    286
    26
            }
    287
    10
        }
    288
    289
        /// Advance the sequence position
    290
    10
        pub fn advance(&mut self) {
    291
    10
            if self.seq_len < self.max_seq_len {
    292
    9
                self.seq_len += 1;
    293
    9
            }1
    294
    10
        }
    295
    296
        /// Get cached keys for a layer
    297
        #[must_use]
    298
    2
        pub fn get_k(&self, layer: usize) -> &[f32] {
    299
    2
            &self.k_cache[layer][..self.seq_len * self.hidden_dim]
    300
    2
        }
    301
    302
        /// Get cached values (transposed) for a layer
    303
        #[must_use]
    304
    3
        pub fn get_v_transposed(&self, layer: usize) -> &[f32] {
    305
    3
            &self.v_cache[layer]
    306
    3
        }
    307
    308
        /// Get current sequence length
    309
        #[must_use]
    310
    4
        pub fn len(&self) -> usize {
    311
    4
            self.seq_len
    312
    4
        }
    313
    314
        /// Check if cache is empty
    315
        #[must_use]
    316
    2
        pub fn is_empty(&self) -> bool {
    317
    2
            self.seq_len == 0
    318
    2
        }
    319
    320
        /// Reset the cache
    321
    1
        pub fn reset(&mut self) {
    322
    1
            self.seq_len = 0;
    323
    1
        }
    324
    325
        /// Get maximum sequence length
    326
        #[must_use]
    327
    1
        pub fn max_len(&self) -> usize {
    328
    1
            self.max_seq_len
    329
    1
        }
    330
    }
    331
    332
    /// Attention with transposed V cache for better memory access
    333
    ///
    334
    /// Uses transposed V storage for improved cache locality during
    335
    /// the weighted sum computation.
    336
    #[must_use]
    337
    3
    pub fn attention_with_transposed_v(
    338
    3
        q: &[f32],
    339
    3
        k_cache: &[f32],
    340
    3
        v_cache_transposed: &[f32],
    341
    3
        current_k: &[f32],
    342
    3
        current_v: &[f32],
    343
    3
        num_heads: usize,
    344
    3
        max_seq_len: usize,
    345
    3
    ) -> Vec<f32> {
    346
    3
        let hidden_dim = q.len();
    347
    3
        let head_dim = hidden_dim / num_heads;
    348
    3
        let cache_len = if hidden_dim > 0 {
    349
    3
            k_cache.len() / hidden_dim
    350
        } else {
    351
    0
            0
    352
        };
    353
    3
        let total_len = cache_len + 1;
    354
    3
        let scale = 1.0 / (head_dim as f32).sqrt();
    355
    356
    3
        let mut output = vec![0.0; hidden_dim];
    357
    358
    6
        for h in 0..num_heads3 {
    359
    6
            let head_offset = h * head_dim;
    360
    6
            let q_head = &q[head_offset..head_offset + head_dim];
    361
    362
            // Compute attention scores
    363
    6
            let mut scores = Vec::with_capacity(total_len);
    364
    365
    6
            for pos in 0..cache_len {
    366
    6
                let k_start = pos * hidden_dim + head_offset;
    367
    6
                let k_head = &k_cache[k_start..k_start + head_dim];
    368
    6
                scores.push(simd_dot(q_head, k_head) * scale);
    369
    6
            }
    370
    371
    6
            let current_k_head = &current_k[head_offset..head_offset + head_dim];
    372
    6
            scores.push(simd_dot(q_head, current_k_head) * scale);
    373
    374
    6
            simd_softmax(&mut scores);
    375
    376
            // Weighted sum with transposed V (better cache locality)
    377
    6
            let out_head = &mut output[head_offset..head_offset + head_dim];
    378
    379
    12
            for i in 0..head_dim6 {
    380
    12
                let v_idx = (head_offset + i) * max_seq_len;
    381
    12
                let mut sum = 0.0;
    382
    12
                for (pos, &weight) in scores.iter().enumerate().take(cache_len) {
    383
    12
                    sum += weight * v_cache_transposed[v_idx + pos];
    384
    12
                }
    385
    12
                sum += scores[cache_len] * current_v[head_offset + i];
    386
    12
                out_head[i] = sum;
    387
            }
    388
        }
    389
    390
    3
        output
    391
    3
    }
    392
    393
    // ============================================================================
    394
    // EXTREME TDD: Comprehensive Tests
    395
    // ============================================================================
    396
    397
    #[cfg(test)]
    398
    mod tests {
    399
        use super::*;
    400
    401
        // ------------------------------------------------------------------------
    402
        // KVCache Tests
    403
        // ------------------------------------------------------------------------
    404
    405
        #[test]
    406
    1
        fn test_kv_cache_new() {
    407
    1
            let cache = KVCache::new(12, 768, 2048);
    408
    1
            assert!(cache.is_empty());
    409
    1
            assert_eq!(cache.len(), 0);
    410
    1
        }
    411
    412
        #[test]
    413
    1
        fn test_kv_cache_store_and_retrieve() {
    414
    1
            let mut cache = KVCache::new(2, 4, 10);
    415
    416
    1
            let k = vec![1.0, 2.0, 3.0, 4.0];
    417
    1
            let v = vec![5.0, 6.0, 7.0, 8.0];
    418
    419
    1
            cache.store(0, &k, &v);
    420
    1
            cache.advance();
    421
    422
    1
            assert_eq!(cache.len(), 1);
    423
    1
            assert_eq!(cache.get_k(0), &[1.0, 2.0, 3.0, 4.0]);
    424
    1
            assert_eq!(cache.get_v(0), &[5.0, 6.0, 7.0, 8.0]);
    425
    1
        }
    426
    427
        #[test]
    428
    1
        fn test_kv_cache_multiple_positions() {
    429
    1
            let mut cache = KVCache::new(1, 2, 10);
    430
    431
            // Store position 0
    432
    1
            cache.store(0, &[1.0, 2.0], &[3.0, 4.0]);
    433
    1
            cache.advance();
    434
    435
            // Store position 1
    436
    1
            cache.store(0, &[5.0, 6.0], &[7.0, 8.0]);
    437
    1
            cache.advance();
    438
    439
    1
            assert_eq!(cache.len(), 2);
    440
    1
            assert_eq!(cache.get_k(0), &[1.0, 2.0, 5.0, 6.0]);
    441
    1
            assert_eq!(cache.get_v(0), &[3.0, 4.0, 7.0, 8.0]);
    442
    1
        }
    443
    444
        #[test]
    445
    1
        fn test_kv_cache_multiple_layers() {
    446
    1
            let mut cache = KVCache::new(2, 2, 10);
    447
    448
    1
            cache.store(0, &[1.0, 2.0], &[3.0, 4.0]);
    449
    1
            cache.store(1, &[5.0, 6.0], &[7.0, 8.0]);
    450
    1
            cache.advance();
    451
    452
    1
            assert_eq!(cache.get_k(0), &[1.0, 2.0]);
    453
    1
            assert_eq!(cache.get_k(1), &[5.0, 6.0]);
    454
    1
            assert_eq!(cache.get_v(0), &[3.0, 4.0]);
    455
    1
            assert_eq!(cache.get_v(1), &[7.0, 8.0]);
    456
    1
        }
    457
    458
        #[test]
    459
    1
        fn test_kv_cache_reset() {
    460
    1
            let mut cache = KVCache::new(1, 4, 10);
    461
    462
    1
            cache.store(0, &[1.0; 4], &[2.0; 4]);
    463
    1
            cache.advance();
    464
    1
            assert_eq!(cache.len(), 1);
    465
    466
    1
            cache.reset();
    467
    1
            assert!(cache.is_empty());
    468
    1
            assert_eq!(cache.len(), 0);
    469
    1
        }
    470
    471
        #[test]
    472
    1
        fn test_kv_cache_is_empty() {
    473
    1
            let mut cache = KVCache::new(1, 4, 10);
    474
    1
            assert!(cache.is_empty());
    475
    476
    1
            cache.store(0, &[1.0; 4], &[1.0; 4]);
    477
    1
            cache.advance();
    478
    1
            assert!(!cache.is_empty());
    479
    1
        }
    480
    481
        #[test]
    482
    1
        fn test_kv_cache_clone() {
    483
    1
            let mut cache = KVCache::new(1, 2, 10);
    484
    1
            cache.store(0, &[1.0, 2.0], &[3.0, 4.0]);
    485
    1
            cache.advance();
    486
    487
    1
            let cloned = cache.clone();
    488
    1
            assert_eq!(cloned.len(), 1);
    489
    1
            assert_eq!(cloned.get_k(0), &[1.0, 2.0]);
    490
    1
        }
    491
    492
        // ------------------------------------------------------------------------
    493
        // OptimizedKVCache Tests
    494
        // ------------------------------------------------------------------------
    495
    496
        #[test]
    497
    1
        fn test_optimized_cache_new() {
    498
    1
            let cache = OptimizedKVCache::new(12, 768, 2048);
    499
    1
            assert!(cache.is_empty());
    500
    1
            assert_eq!(cache.len(), 0);
    501
    1
            assert_eq!(cache.max_len(), 2048);
    502
    1
        }
    503
    504
        #[test]
    505
    1
        fn test_optimized_cache_store_and_retrieve() {
    506
    1
            let mut cache = OptimizedKVCache::new(1, 4, 10);
    507
    508
    1
            let k = vec![1.0, 2.0, 3.0, 4.0];
    509
    1
            let v = vec![5.0, 6.0, 7.0, 8.0];
    510
    511
    1
            cache.store(0, &k, &v);
    512
    1
            cache.advance();
    513
    514
    1
            assert_eq!(cache.len(), 1);
    515
    1
            assert_eq!(cache.get_k(0), &[1.0, 2.0, 3.0, 4.0]);
    516
    517
            // V is transposed: v[i] at position 0 is at index i * max_seq_len
    518
    1
            let v_transposed = cache.get_v_transposed(0);
    519
    1
            assert_eq!(v_transposed[0], 5.0); // v[0] at pos 0
    520
    1
            assert_eq!(v_transposed[10], 6.0); // v[1] at pos 0 (stride = max_seq_len = 10)
    521
    1
            assert_eq!(v_transposed[20], 7.0); // v[2] at pos 0
    522
    1
            assert_eq!(v_transposed[30], 8.0); // v[3] at pos 0
    523
    1
        }
    524
    525
        #[test]
    526
    1
        fn test_optimized_cache_transposed_v_layout() {
    527
    1
            let mut cache = OptimizedKVCache::new(1, 2, 5);
    528
    529
            // Store position 0
    530
    1
            cache.store(0, &[1.0, 2.0], &[10.0, 20.0]);
    531
    1
            cache.advance();
    532
    533
            // Store position 1
    534
    1
            cache.store(0, &[3.0, 4.0], &[30.0, 40.0]);
    535
    1
            cache.advance();
    536
    537
    1
            let v_transposed = cache.get_v_transposed(0);
    538
            // v[0] positions: indices 0, 1, 2, ... (stride 1)
    539
            // v[1] positions: indices 5, 6, 7, ... (stride 1, offset = hidden_dim * max_seq_len / hidden_dim = max_seq_len)
    540
    1
            assert_eq!(v_transposed[0], 10.0); // v[0] at pos 0
    541
    1
            assert_eq!(v_transposed[1], 30.0); // v[0] at pos 1
    542
    1
            assert_eq!(v_transposed[5], 20.0); // v[1] at pos 0
    543
    1
            assert_eq!(v_transposed[6], 40.0); // v[1] at pos 1
    544
    1
        }
    545
    546
        #[test]
    547
    1
        fn test_optimized_cache_max_capacity() {
    548
    1
            let mut cache = OptimizedKVCache::new(1, 2, 3);
    549
    550
            // Fill to capacity
    551
    4
            for i3 in 0..3 {
    552
    3
                cache.store(0, &[i as f32; 2], &[i as f32; 2]);
    553
    3
                cache.advance();
    554
    3
            }
    555
    1
            assert_eq!(cache.len(), 3);
    556
    557
            // Should not advance beyond max
    558
    1
            cache.store(0, &[99.0; 2], &[99.0; 2]);
    559
    1
            cache.advance();
    560
    1
            assert_eq!(cache.len(), 3); // Still at max
    561
    1
        }
    562
    563
        #[test]
    564
    1
        fn test_optimized_cache_reset() {
    565
    1
            let mut cache = OptimizedKVCache::new(1, 4, 10);
    566
    567
    1
            cache.store(0, &[1.0; 4], &[2.0; 4]);
    568
    1
            cache.advance();
    569
    1
            cache.reset();
    570
    571
    1
            assert!(cache.is_empty());
    572
    1
        }
    573
    574
        // ------------------------------------------------------------------------
    575
        // attention_with_cache Tests
    576
        // ------------------------------------------------------------------------
    577
    578
        #[test]
    579
    1
        fn test_attention_with_cache_no_history() {
    580
    1
            let hidden_dim = 4;
    581
    1
            let num_heads = 2;
    582
    583
    1
            let q = vec![1.0; hidden_dim];
    584
    1
            let k_cache: Vec<f32> = vec![];
    585
    1
            let v_cache: Vec<f32> = vec![];
    586
    1
            let current_k = vec![1.0; hidden_dim];
    587
    1
            let current_v = vec![2.0; hidden_dim];
    588
    589
    1
            let output =
    590
    1
                attention_with_cache(&q, &k_cache, &v_cache, &current_k, &current_v, num_heads);
    591
    592
    1
            assert_eq!(output.len(), hidden_dim);
    593
            // With no history and uniform attention, output should equal current_v
    594
    5
            for &v4 in &output {
    595
    4
                assert!((v - 2.0).abs() < 1e-5);
    596
            }
    597
    1
        }
    598
    599
        #[test]
    600
    1
        fn test_attention_with_cache_one_cached() {
    601
    1
            let hidden_dim = 4;
    602
    1
            let num_heads = 2;
    603
    604
    1
            let q = vec![1.0; hidden_dim];
    605
            // One cached position
    606
    1
            let k_cache = vec![1.0; hidden_dim];
    607
    1
            let v_cache = vec![1.0; hidden_dim];
    608
    1
            let current_k = vec![1.0; hidden_dim];
    609
    1
            let current_v = vec![3.0; hidden_dim];
    610
    611
    1
            let output =
    612
    1
                attention_with_cache(&q, &k_cache, &v_cache, &current_k, &current_v, num_heads);
    613
    614
    1
            assert_eq!(output.len(), hidden_dim);
    615
            // With uniform K, attention is 0.5 to each position
    616
            // output = 0.5 * 1.0 + 0.5 * 3.0 = 2.0
    617
    5
            for &v4 in &output {
    618
    4
                assert!((v - 2.0).abs() < 1e-5);
    619
            }
    620
    1
        }
    621
    622
        #[test]
    623
    1
        fn test_attention_with_cache_multi_head() {
    624
    1
            let hidden_dim = 8;
    625
    1
            let num_heads = 4;
    626
    1
            let head_dim = hidden_dim / num_heads;
    627
    628
            // Each head has different Q
    629
    1
            let mut q = vec![0.0; hidden_dim];
    630
    4
            for h in 0..num_heads1 {
    631
    8
                for i in 0..head_dim4 {
    632
    8
                    q[h * head_dim + i] = (h + 1) as f32;
    633
    8
                }
    634
            }
    635
    636
    1
            let k_cache: Vec<f32> = vec![];
    637
    1
            let v_cache: Vec<f32> = vec![];
    638
    1
            let current_k = vec![1.0; hidden_dim];
    639
    1
            let current_v = vec![1.0; hidden_dim];
    640
    641
    1
            let output =
    642
    1
                attention_with_cache(&q, &k_cache, &v_cache, &current_k, &current_v, num_heads);
    643
    644
    1
            assert_eq!(output.len(), hidden_dim);
    645
            // All outputs should be 1.0 (current_v with softmax weight 1.0)
    646
    9
            for &v8 in &output {
    647
    8
                assert!((v - 1.0).abs() < 1e-5);
    648
            }
    649
    1
        }
    650
    651
        #[test]
    652
    1
        fn test_attention_preserves_dimension() {
    653
    5
            for hidden_dim4 in [64, 128, 256, 512] {
    654
    20
                for num_heads16 in [1, 2, 4, 8] {
    655
    16
                    if hidden_dim % num_heads != 0 {
    656
    0
                        continue;
    657
    16
                    }
    658
    659
    16
                    let q = vec![0.1; hidden_dim];
    660
    16
                    let k_cache = vec![0.1; hidden_dim * 2];
    661
    16
                    let v_cache = vec![0.2; hidden_dim * 2];
    662
    16
                    let current_k = vec![0.1; hidden_dim];
    663
    16
                    let current_v = vec![0.3; hidden_dim];
    664
    665
    16
                    let output =
    666
    16
                        attention_with_cache(&q, &k_cache, &v_cache, &current_k, &current_v, num_heads);
    667
    668
    16
                    assert_eq!(output.len(), hidden_dim);
    669
                }
    670
            }
    671
    1
        }
    672
    673
        // ------------------------------------------------------------------------
    674
        // attention_with_transposed_v Tests
    675
        // ------------------------------------------------------------------------
    676
    677
        #[test]
    678
    1
        fn test_transposed_attention_no_history() {
    679
    1
            let hidden_dim = 4;
    680
    1
            let num_heads = 2;
    681
    1
            let max_seq_len = 10;
    682
    683
    1
            let q = vec![1.0; hidden_dim];
    684
    1
            let k_cache: Vec<f32> = vec![];
    685
    1
            let v_cache_transposed = vec![0.0; hidden_dim * max_seq_len];
    686
    1
            let current_k = vec![1.0; hidden_dim];
    687
    1
            let current_v = vec![2.0; hidden_dim];
    688
    689
    1
            let output = attention_with_transposed_v(
    690
    1
                &q,
    691
    1
                &k_cache,
    692
    1
                &v_cache_transposed,
    693
    1
                &current_k,
    694
    1
                &current_v,
    695
    1
                num_heads,
    696
    1
                max_seq_len,
    697
            );
    698
    699
    1
            assert_eq!(output.len(), hidden_dim);
    700
            // Output should equal current_v when no history
    701
    5
            for &v4 in &output {
    702
    4
                assert!((v - 2.0).abs() < 1e-5);
    703
            }
    704
    1
        }
    705
    706
        #[test]
    707
    1
        fn test_transposed_attention_one_cached() {
    708
    1
            let hidden_dim = 4;
    709
    1
            let num_heads = 2;
    710
    1
            let max_seq_len = 10;
    711
    712
    1
            let q = vec![1.0; hidden_dim];
    713
    1
            let k_cache = vec![1.0; hidden_dim]; // 1 cached position
    714
    715
            // Transposed V: v[i] at pos j is at index i * max_seq_len + j
    716
    1
            let mut v_cache_transposed = vec![0.0; hidden_dim * max_seq_len];
    717
    4
            for i in 0..hidden_dim1 {
    718
    4
                v_cache_transposed[i * max_seq_len] = 1.0; // v[i] = 1.0 at pos 0
    719
    4
            }
    720
    721
    1
            let current_k = vec![1.0; hidden_dim];
    722
    1
            let current_v = vec![3.0; hidden_dim];
    723
    724
    1
            let output = attention_with_transposed_v(
    725
    1
                &q,
    726
    1
                &k_cache,
    727
    1
                &v_cache_transposed,
    728
    1
                &current_k,
    729
    1
                &current_v,
    730
    1
                num_heads,
    731
    1
                max_seq_len,
    732
            );
    733
    734
            // Uniform attention: 0.5 * 1.0 + 0.5 * 3.0 = 2.0
    735
    5
            for &v4 in &output {
    736
    4
                assert!((v - 2.0).abs() < 1e-5);
    737
            }
    738
    1
        }
    739
    740
        // ------------------------------------------------------------------------
    741
        // Integration Tests
    742
        // ------------------------------------------------------------------------
    743
    744
        #[test]
    745
    1
        fn test_cache_and_attention_integration() {
    746
    1
            let num_layers = 2;
    747
    1
            let hidden_dim = 4;
    748
    1
            let max_seq_len = 10;
    749
    1
            let num_heads = 2;
    750
    751
    1
            let mut cache = KVCache::new(num_layers, hidden_dim, max_seq_len);
    752
    753
            // Simulate 3 tokens
    754
    4
            for pos3 in 0..3 {
    755
    3
                let k = vec![pos as f32; hidden_dim];
    756
    3
                let v = vec![(pos * 2) as f32; hidden_dim];
    757
    758
    6
                for layer in 0..num_layers3 {
    759
    6
                    cache.store(layer, &k, &v);
    760
    6
                }
    761
    3
                cache.advance();
    762
            }
    763
    764
            // Now compute attention for token 4
    765
    1
            let q = vec![1.0; hidden_dim];
    766
    1
            let current_k = vec![3.0; hidden_dim];
    767
    1
            let current_v = vec![6.0; hidden_dim];
    768
    769
    1
            let k_cached = cache.get_k(0);
    770
    1
            let v_cached = cache.get_v(0);
    771
    772
    1
            let output =
    773
    1
                attention_with_cache(&q, k_cached, v_cached, &current_k, &current_v, num_heads);
    774
    775
    1
            assert_eq!(output.len(), hidden_dim);
    776
            // Output should be some weighted combination
    777
    4
            assert!1(output.iter()1.all1(|&x| x.is_finite()));
    778
    1
        }
    779
    780
        #[test]
    781
    1
        fn test_optimized_cache_and_attention_integration() {
    782
    1
            let num_layers = 1;
    783
    1
            let hidden_dim = 4;
    784
    1
            let max_seq_len = 10;
    785
    1
            let num_heads = 2;
    786
    787
    1
            let mut cache = OptimizedKVCache::new(num_layers, hidden_dim, max_seq_len);
    788
    789
            // Store 2 positions
    790
    1
            cache.store(0, &[1.0; 4], &[2.0; 4]);
    791
    1
            cache.advance();
    792
    1
            cache.store(0, &[1.0; 4], &[4.0; 4]);
    793
    1
            cache.advance();
    794
    795
    1
            let q = vec![1.0; hidden_dim];
    796
    1
            let current_k = vec![1.0; hidden_dim];
    797
    1
            let current_v = vec![6.0; hidden_dim];
    798
    799
    1
            let output = attention_with_transposed_v(
    800
    1
                &q,
    801
    1
                cache.get_k(0),
    802
    1
                cache.get_v_transposed(0),
    803
    1
                &current_k,
    804
    1
                &current_v,
    805
    1
                num_heads,
    806
    1
                max_seq_len,
    807
            );
    808
    809
    1
            assert_eq!(output.len(), hidden_dim);
    810
            // Uniform attention: (2 + 4 + 6) / 3 = 4.0
    811
    5
            for &v4 in &output {
    812
    4
                assert!((v - 4.0).abs() < 1e-5);
    813
            }
    814
    1
        }
    815
    816
        // ------------------------------------------------------------------------
    817
        // Edge Cases
    818
        // ------------------------------------------------------------------------
    819
    820
        #[test]
    821
    1
        fn test_single_head_attention() {
    822
    1
            let hidden_dim = 4;
    823
    1
            let num_heads = 1;
    824
    825
    1
            let q = vec![1.0; hidden_dim];
    826
    1
            let k_cache: Vec<f32> = vec![];
    827
    1
            let v_cache: Vec<f32> = vec![];
    828
    1
            let current_k = vec![1.0; hidden_dim];
    829
    1
            let current_v = vec![5.0; hidden_dim];
    830
    831
    1
            let output =
    832
    1
                attention_with_cache(&q, &k_cache, &v_cache, &current_k, &current_v, num_heads);
    833
    834
    5
            for &v4 in &output {
    835
    4
                assert!((v - 5.0).abs() < 1e-5);
    836
            }
    837
    1
        }
    838
    839
        #[test]
    840
    1
        fn test_many_heads_attention() {
    841
    1
            let hidden_dim = 64;
    842
    1
            let num_heads = 16;
    843
    844
    1
            let q = vec![0.1; hidden_dim];
    845
    1
            let k_cache: Vec<f32> = vec![];
    846
    1
            let v_cache: Vec<f32> = vec![];
    847
    1
            let current_k = vec![0.1; hidden_dim];
    848
    1
            let current_v = vec![1.0; hidden_dim];
    849
    850
    1
            let output =
    851
    1
                attention_with_cache(&q, &k_cache, &v_cache, &current_k, &current_v, num_heads);
    852
    853
    1
            assert_eq!(output.len(), hidden_dim);
    854
    65
            for &v64 in &output {
    855
    64
                assert!((v - 1.0).abs() < 1e-5);
    856
            }
    857
    1
        }
    858
    }