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

    Line
    Count
    Source
    1
    //! Streaming KV Cache (PMAT-802)
    2
    //!
    3
    //! Memory-efficient key-value cache for transformer inference.
    4
    5
    // ============================================================================
    6
    // M6: Memory Efficiency - StreamingKVCache
    7
    // ============================================================================
    8
    9
    /// Streaming KV cache for memory-efficient inference
    10
    ///
    11
    /// Implements a bounded circular buffer for key-value cache that allows
    12
    /// efficient inference on long sequences without unbounded memory growth.
    13
    ///
    14
    /// ## Memory Bound
    15
    ///
    16
    /// Total memory = num_layers * max_positions * num_heads * head_dim * 2 (K+V) * sizeof(f32)
    17
    ///
    18
    /// For 7B model (32 layers, 2048 positions, 32 heads, 128 head_dim):
    19
    /// = 32 * 2048 * 32 * 128 * 2 * 4 = ~2GB
    20
    ///
    21
    /// ## Usage
    22
    ///
    23
    /// ```rust,ignore
    24
    /// let mut cache = StreamingKVCache::new(32, 2048, 32, 128);
    25
    /// cache.append(0, &key_vec, &value_vec);
    26
    /// let (keys, values) = cache.get_range(0, 0, 100);
    27
    /// ```
    28
    pub struct StreamingKVCache {
    29
        /// Number of transformer layers
    30
        num_layers: usize,
    31
        /// Maximum cached positions (context length)
    32
        max_positions: usize,
    33
        /// Number of attention heads
    34
        num_heads: usize,
    35
        /// Dimension per head
    36
        head_dim: usize,
    37
        /// Key cache per layer [num_layers][max_positions * num_heads * head_dim]
    38
        keys: Vec<Vec<f32>>,
    39
        /// Value cache per layer
    40
        values: Vec<Vec<f32>>,
    41
        /// Current write position (circular)
    42
        position: usize,
    43
        /// Number of valid positions cached
    44
        valid_positions: usize,
    45
    }
    46
    47
    impl StreamingKVCache {
    48
        /// Create a new streaming KV cache
    49
        ///
    50
        /// # Arguments
    51
        ///
    52
        /// * `num_layers` - Number of transformer layers
    53
        /// * `max_positions` - Maximum context length to cache
    54
        /// * `num_heads` - Number of attention heads
    55
        /// * `head_dim` - Dimension per attention head
    56
        #[must_use]
    57
    60
        pub fn new(num_layers: usize, max_positions: usize, num_heads: usize, head_dim: usize) -> Self {
    58
    60
            let kv_size = max_positions * num_heads * head_dim;
    59
    60
            Self {
    60
    60
                num_layers,
    61
    60
                max_positions,
    62
    60
                num_heads,
    63
    60
                head_dim,
    64
    60
                keys: vec![vec![0.0f32; kv_size]; num_layers],
    65
    60
                values: vec![vec![0.0f32; kv_size]; num_layers],
    66
    60
                position: 0,
    67
    60
                valid_positions: 0,
    68
    60
            }
    69
    60
        }
    70
    71
        /// Append key-value pair for a single position
    72
        ///
    73
        /// # Arguments
    74
        ///
    75
        /// * `layer` - Layer index (0-indexed)
    76
        /// * `key` - Key vector [num_heads * head_dim]
    77
        /// * `value` - Value vector [num_heads * head_dim]
    78
        ///
    79
        /// # Panics
    80
        ///
    81
        /// Panics if layer index is out of bounds or key/value dimensions are wrong.
    82
    461k
        pub fn append(&mut self, layer: usize, key: &[f32], value: &[f32]) {
    83
    461k
            let kv_dim = self.num_heads * self.head_dim;
    84
    461k
            assert!(layer < self.num_layers, "Layer index out of bounds"1);
    85
    461k
            assert_eq!(key.len(), kv_dim, "Key dimension mismatch"1);
    86
    461k
            assert_eq!(value.len(), kv_dim, "Value dimension mismatch"0);
    87
    88
    461k
            let offset = self.position * kv_dim;
    89
    461k
            self.keys[layer][offset..offset + kv_dim].copy_from_slice(key);
    90
    461k
            self.values[layer][offset..offset + kv_dim].copy_from_slice(value);
    91
    92
            // Update position only after last layer
    93
    461k
            if layer == self.num_layers - 1 {
    94
    115k
                self.position = (self.position + 1) % self.max_positions;
    95
    115k
                self.valid_positions = (self.valid_positions + 1).min(self.max_positions);
    96
    345k
            }
    97
    461k
        }
    98
    99
        /// Get keys and values for a range of positions
    100
        ///
    101
        /// # Arguments
    102
        ///
    103
        /// * `layer` - Layer index
    104
        /// * `start` - Start position (inclusive)
    105
        /// * `end` - End position (exclusive)
    106
        ///
    107
        /// # Returns
    108
        ///
    109
        /// Tuple of (keys, values) slices
    110
        #[must_use]
    111
    3.37k
        pub fn get_range(&self, layer: usize, start: usize, end: usize) -> (&[f32], &[f32]) {
    112
    3.37k
            let kv_dim = self.num_heads * self.head_dim;
    113
    3.37k
            let start_offset = start * kv_dim;
    114
    3.37k
            let end_offset = end * kv_dim;
    115
    116
    3.37k
            (
    117
    3.37k
                &self.keys[layer][start_offset..end_offset],
    118
    3.37k
                &self.values[layer][start_offset..end_offset],
    119
    3.37k
            )
    120
    3.37k
        }
    121
    122
        /// Get all valid cached keys and values for a layer
    123
        ///
    124
        /// # Arguments
    125
        ///
    126
        /// * `layer` - Layer index
    127
        ///
    128
        /// # Returns
    129
        ///
    130
        /// Tuple of (keys, values) for all valid positions
    131
        #[must_use]
    132
    3.36k
        pub fn get_valid(&self, layer: usize) -> (&[f32], &[f32]) {
    133
    3.36k
            self.get_range(layer, 0, self.valid_positions)
    134
    3.36k
        }
    135
    136
        /// Get current number of valid cached positions
    137
        #[must_use]
    138
    29
        pub fn len(&self) -> usize {
    139
    29
            self.valid_positions
    140
    29
        }
    141
    142
        /// Check if cache is empty
    143
        #[must_use]
    144
    5
        pub fn is_empty(&self) -> bool {
    145
    5
            self.valid_positions == 0
    146
    5
        }
    147
    148
        /// Get maximum positions (context length)
    149
        #[must_use]
    150
    4
        pub fn max_positions(&self) -> usize {
    151
    4
            self.max_positions
    152
    4
        }
    153
    154
        /// Reset the cache
    155
    2
        pub fn clear(&mut self) {
    156
    2
            self.position = 0;
    157
    2
            self.valid_positions = 0;
    158
            // Note: We don't zero the memory for performance
    159
    2
        }
    160
    161
        /// Calculate memory usage in bytes
    162
        #[must_use]
    163
    13
        pub fn memory_bytes(&self) -> usize {
    164
    13
            let kv_size = self.max_positions * self.num_heads * self.head_dim;
    165
            // Keys + Values, f32 = 4 bytes
    166
    13
            self.num_layers * kv_size * 2 * 4
    167
    13
        }
    168
    169
        /// Calculate memory usage in megabytes
    170
        #[must_use]
    171
    5
        pub fn memory_mb(&self) -> f64 {
    172
    5
            self.memory_bytes() as f64 / (1024.0 * 1024.0)
    173
    5
        }
    174
    }
    175
    176
    /// Streaming KV cache with FP16 storage for memory efficiency (M12)
    177
    ///
    178
    /// Uses half-precision (FP16) storage to halve memory usage compared to FP32,
    179
    /// enabling support for ultra-long contexts (65536+) on consumer GPUs.
    180
    ///
    181
    /// # Memory Efficiency
    182
    ///
    183
    /// For 65536 context with 7B model config:
    184
    /// - FP32: 32 layers × 65536 pos × 32 heads × 128 dim × 2 × 4 bytes = 68.72 GB
    185
    /// - FP16: 32 layers × 65536 pos × 32 heads × 128 dim × 2 × 2 bytes = 34.36 GB
    186
    ///
    187
    /// # Example
    188
    ///
    189
    /// ```
    190
    /// use realizar::gpu::StreamingKVCacheFp16;
    191
    ///
    192
    /// let mut cache = StreamingKVCacheFp16::new(32, 65536, 32, 128);
    193
    /// assert!(cache.memory_mb() < 36000.0); // < 36 GB
    194
    /// ```
    195
    pub struct StreamingKVCacheFp16 {
    196
        /// Number of transformer layers
    197
        num_layers: usize,
    198
        /// Maximum cached positions (context length)
    199
        max_positions: usize,
    200
        /// Number of attention heads
    201
        num_heads: usize,
    202
        /// Dimension per head
    203
        head_dim: usize,
    204
        /// Key cache per layer [num_layers][max_positions * num_heads * head_dim] stored as FP16 bits
    205
        keys: Vec<Vec<u16>>,
    206
        /// Value cache per layer stored as FP16 bits
    207
        values: Vec<Vec<u16>>,
    208
        /// Current write position (circular)
    209
        position: usize,
    210
        /// Number of valid positions cached
    211
        valid_positions: usize,
    212
    }
    213
    214
    impl StreamingKVCacheFp16 {
    215
        /// Create a new FP16 streaming KV cache
    216
        ///
    217
        /// # Arguments
    218
        ///
    219
        /// * `num_layers` - Number of transformer layers
    220
        /// * `max_positions` - Maximum context length to cache
    221
        /// * `num_heads` - Number of attention heads
    222
        /// * `head_dim` - Dimension per attention head
    223
        #[must_use]
    224
    9
        pub fn new(num_layers: usize, max_positions: usize, num_heads: usize, head_dim: usize) -> Self {
    225
    9
            let kv_size = max_positions * num_heads * head_dim;
    226
    9
            Self {
    227
    9
                num_layers,
    228
    9
                max_positions,
    229
    9
                num_heads,
    230
    9
                head_dim,
    231
    9
                keys: vec![vec![0u16; kv_size]; num_layers],
    232
    9
                values: vec![vec![0u16; kv_size]; num_layers],
    233
    9
                position: 0,
    234
    9
                valid_positions: 0,
    235
    9
            }
    236
    9
        }
    237
    238
        /// Convert f32 to FP16 bits
    239
        #[inline]
    240
    3.72k
        pub(crate) fn f32_to_f16(value: f32) -> u16 {
    241
    3.72k
            half::f16::from_f32(value).to_bits()
    242
    3.72k
        }
    243
    244
        /// Convert FP16 bits to f32
    245
        #[inline]
    246
    1.09k
        pub(crate) fn f16_to_f32(bits: u16) -> f32 {
    247
    1.09k
            half::f16::from_bits(bits).to_f32()
    248
    1.09k
        }
    249
    250
        /// Append key-value pair for a single position (FP32 input, stored as FP16)
    251
        ///
    252
        /// # Arguments
    253
        ///
    254
        /// * `layer` - Layer index (0-indexed)
    255
        /// * `key` - Key vector [num_heads * head_dim] as FP32
    256
        /// * `value` - Value vector [num_heads * head_dim] as FP32
    257
        ///
    258
        /// # Panics
    259
        ///
    260
        /// Panics if layer index is out of bounds or key/value dimensions are wrong.
    261
    16
        pub fn append(&mut self, layer: usize, key: &[f32], value: &[f32]) {
    262
    16
            let kv_dim = self.num_heads * self.head_dim;
    263
    16
            assert!(layer < self.num_layers, "Layer index out of bounds"0);
    264
    16
            assert_eq!(key.len(), kv_dim, "Key dimension mismatch"0);
    265
    16
            assert_eq!(value.len(), kv_dim, "Value dimension mismatch"0);
    266
    267
    16
            let offset = self.position * kv_dim;
    268
    269
            // Convert FP32 to FP16 and store
    270
    1.85k
            for (i, &k) in key16.iter16().enumerate16() {
    271
    1.85k
                self.keys[layer][offset + i] = Self::f32_to_f16(k);
    272
    1.85k
            }
    273
    1.85k
            for (i, &v) in value16.iter16().enumerate16() {
    274
    1.85k
                self.values[layer][offset + i] = Self::f32_to_f16(v);
    275
    1.85k
            }
    276
    277
            // Update position only after last layer
    278
    16
            if layer == self.num_layers - 1 {
    279
    8
                self.position = (self.position + 1) % self.max_positions;
    280
    8
                self.valid_positions = (self.valid_positions + 1).min(self.max_positions);
    281
    8
            }
    282
    16
        }
    283
    284
        /// Get keys and values for a range of positions (converted back to FP32)
    285
        ///
    286
        /// # Arguments
    287
        ///
    288
        /// * `layer` - Layer index
    289
        /// * `start` - Start position (inclusive)
    290
        /// * `end` - End position (exclusive)
    291
        ///
    292
        /// # Returns
    293
        ///
    294
        /// Tuple of (keys, values) as Vec<f32>
    295
        #[must_use]
    296
    3
        pub fn get_range_f32(&self, layer: usize, start: usize, end: usize) -> (Vec<f32>, Vec<f32>) {
    297
    3
            let kv_dim = self.num_heads * self.head_dim;
    298
    3
            let start_offset = start * kv_dim;
    299
    3
            let end_offset = end * kv_dim;
    300
    301
    3
            let keys: Vec<f32> = self.keys[layer][start_offset..end_offset]
    302
    3
                .iter()
    303
    544
                .map3(|&bits| Self::f16_to_f32(bits))
    304
    3
                .collect();
    305
    306
    3
            let values: Vec<f32> = self.values[layer][start_offset..end_offset]
    307
    3
                .iter()
    308
    544
                .map3(|&bits| Self::f16_to_f32(bits))
    309
    3
                .collect();
    310
    311
    3
            (keys, values)
    312
    3
        }
    313
    314
        /// Get raw FP16 keys and values for a range of positions
    315
        #[must_use]
    316
    1
        pub fn get_range_raw(&self, layer: usize, start: usize, end: usize) -> (&[u16], &[u16]) {
    317
    1
            let kv_dim = self.num_heads * self.head_dim;
    318
    1
            let start_offset = start * kv_dim;
    319
    1
            let end_offset = end * kv_dim;
    320
    321
    1
            (
    322
    1
                &self.keys[layer][start_offset..end_offset],
    323
    1
                &self.values[layer][start_offset..end_offset],
    324
    1
            )
    325
    1
        }
    326
    327
        /// Get all valid cached keys and values for a layer (as FP32)
    328
        #[must_use]
    329
    2
        pub fn get_valid_f32(&self, layer: usize) -> (Vec<f32>, Vec<f32>) {
    330
    2
            self.get_range_f32(layer, 0, self.valid_positions)
    331
    2
        }
    332
    333
        /// Get current number of valid cached positions
    334
        #[must_use]
    335
    7
        pub fn len(&self) -> usize {
    336
    7
            self.valid_positions
    337
    7
        }
    338
    339
        /// Check if cache is empty
    340
        #[must_use]
    341
    4
        pub fn is_empty(&self) -> bool {
    342
    4
            self.valid_positions == 0
    343
    4
        }
    344
    345
        /// Get maximum positions (context length)
    346
        #[must_use]
    347
    2
        pub fn max_positions(&self) -> usize {
    348
    2
            self.max_positions
    349
    2
        }
    350
    351
        /// Reset the cache
    352
    1
        pub fn clear(&mut self) {
    353
    1
            self.position = 0;
    354
    1
            self.valid_positions = 0;
    355
    1
        }
    356
    357
        /// Calculate memory usage in bytes (half of FP32 version)
    358
        #[must_use]
    359
    3
        pub fn memory_bytes(&self) -> usize {
    360
    3
            let kv_size = self.max_positions * self.num_heads * self.head_dim;
    361
            // Keys + Values, u16 (FP16) = 2 bytes
    362
    3
            self.num_layers * kv_size * 2 * 2
    363
    3
        }
    364
    365
        /// Calculate memory usage in megabytes
    366
        #[must_use]
    367
    1
        pub fn memory_mb(&self) -> f64 {
    368
    1
            self.memory_bytes() as f64 / (1024.0 * 1024.0)
    369
    1
        }
    370
    }