inference_lab/kv_cache/
block.rs

1use crate::request::BlockId;
2
3/// Represents a single KV cache block
4#[derive(Debug, Clone)]
5pub struct Block {
6    /// Unique block ID
7    pub block_id: BlockId,
8
9    /// Reference count (number of requests using this block)
10    pub ref_count: u32,
11
12    /// Whether this block is free
13    pub is_free: bool,
14
15    /// For prefix caching: hash of the token sequence in this block
16    pub content_hash: Option<u64>,
17}
18
19impl Block {
20    /// Create a new free block
21    pub fn new(block_id: BlockId) -> Self {
22        Self {
23            block_id,
24            ref_count: 0,
25            is_free: true,
26            content_hash: None,
27        }
28    }
29
30    /// Allocate this block (increment ref count, mark as not free)
31    pub fn allocate(&mut self, content_hash: Option<u64>) -> Option<u64> {
32        self.ref_count += 1;
33        self.is_free = false;
34        std::mem::replace(&mut self.content_hash, content_hash)
35    }
36
37    /// Release this block (decrement ref count, mark as free if ref count reaches 0)
38    pub fn release(&mut self) {
39        if self.ref_count > 0 {
40            self.ref_count -= 1;
41        }
42
43        if self.ref_count == 0 {
44            // Leave the content hash in there. When we overwrite the block, we will update it.
45            self.is_free = true;
46        }
47    }
48}
49
50#[cfg(test)]
51mod tests {
52    use super::*;
53
54    #[test]
55    fn test_block_creation() {
56        let block = Block::new(0);
57        assert_eq!(block.block_id, 0);
58        assert_eq!(block.ref_count, 0);
59        assert!(block.is_free);
60        assert!(block.content_hash.is_none());
61    }
62
63    #[test]
64    fn test_block_allocate() {
65        let mut block = Block::new(0);
66
67        block.allocate(None);
68        assert_eq!(block.ref_count, 1);
69        assert!(!block.is_free);
70
71        block.allocate(None);
72        assert_eq!(block.ref_count, 2);
73        assert!(!block.is_free);
74    }
75
76    #[test]
77    fn test_block_release() {
78        let mut block = Block::new(0);
79        block.allocate(None);
80        block.allocate(None);
81
82        block.release();
83        assert_eq!(block.ref_count, 1);
84        assert!(!block.is_free);
85
86        block.release();
87        assert_eq!(block.ref_count, 0);
88        assert!(block.is_free);
89
90        // Releasing when already at 0 should be safe
91        block.release();
92        assert_eq!(block.ref_count, 0);
93        assert!(block.is_free);
94    }
95}
inference_lab/kv_cache/block.rs

inference_lab/kv_cache/
block.rs
