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) {
32        self.ref_count += 1;
33        self.is_free = false;
34    }
35
36    /// Release this block (decrement ref count, mark as free if ref count reaches 0)
37    pub fn release(&mut self) {
38        if self.ref_count > 0 {
39            self.ref_count -= 1;
40        }
41
42        if self.ref_count == 0 {
43            self.is_free = true;
44            self.content_hash = None;
45        }
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52
53    #[test]
54    fn test_block_creation() {
55        let block = Block::new(0);
56        assert_eq!(block.block_id, 0);
57        assert_eq!(block.ref_count, 0);
58        assert!(block.is_free);
59        assert!(block.content_hash.is_none());
60    }
61
62    #[test]
63    fn test_block_allocate() {
64        let mut block = Block::new(0);
65
66        block.allocate();
67        assert_eq!(block.ref_count, 1);
68        assert!(!block.is_free);
69
70        block.allocate();
71        assert_eq!(block.ref_count, 2);
72        assert!(!block.is_free);
73    }
74
75    #[test]
76    fn test_block_release() {
77        let mut block = Block::new(0);
78        block.allocate();
79        block.allocate();
80
81        block.release();
82        assert_eq!(block.ref_count, 1);
83        assert!(!block.is_free);
84
85        block.release();
86        assert_eq!(block.ref_count, 0);
87        assert!(block.is_free);
88
89        // Releasing when already at 0 should be safe
90        block.release();
91        assert_eq!(block.ref_count, 0);
92        assert!(block.is_free);
93    }
94}
inference_lab/kv_cache/block.rs

inference_lab/kv_cache/
block.rs
