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
}