## 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/paged_kv/mod.rs

    Line
    Count
    Source
    1
    //! PagedAttention KV Cache Management
    2
    //!
    3
    //! Per spec §8.1: Efficient KV cache management based on vLLM's PagedAttention.
    4
    //! Reference: [4] Kwon et al. (2023) "Efficient Memory Management for LLM Serving"
    5
    //!
    6
    //! ## Key Features
    7
    //!
    8
    //! - **Physical Pages**: Fixed-size memory blocks for KV cache storage
    9
    //! - **Page Tables**: Logical to physical page mapping per sequence
    10
    //! - **Copy-on-Write**: Efficient prefix sharing between sequences
    11
    //! - **Dynamic Allocation**: Pages allocated on-demand during generation
    12
    //!
    13
    //! ## Memory Layout
    14
    //!
    15
    //! ```text
    16
    //! Physical Page (block_size tokens):
    17
    //! ┌─────────────────────────────────────────┐
    18
    //! │  K: [block_size, num_heads, head_dim]   │
    19
    //! │  V: [block_size, num_heads, head_dim]   │
    20
    //! └─────────────────────────────────────────┘
    21
    //! ```
    22
    23
    // Module-level clippy allows
    24
    #![allow(clippy::must_use_candidate)]
    25
    #![allow(clippy::return_self_not_must_use)]
    26
    #![allow(clippy::missing_errors_doc)]
    27
    28
    use serde::{Deserialize, Serialize};
    29
    use std::collections::{HashMap, VecDeque};
    30
    use std::sync::atomic::{AtomicU64, Ordering};
    31
    use thiserror::Error;
    32
    33
    /// Error type for PagedKvCache operations
    34
    #[derive(Debug, Error)]
    35
    pub enum PagedCacheError {
    36
        /// Out of memory - no free pages available
    37
        #[error("Out of memory: need {needed} pages, have {available}")]
    38
        OutOfMemory {
    39
            /// Number of pages needed
    40
            needed: usize,
    41
            /// Number of pages available
    42
            available: usize,
    43
        },
    44
    45
        /// Sequence not found in page table
    46
        #[error("Sequence not found: {0}")]
    47
        SequenceNotFound(u64),
    48
    49
        /// Invalid page access
    50
        #[error("Invalid page access: page {page_id} at offset {offset}")]
    51
        InvalidPageAccess {
    52
            /// Page ID accessed
    53
            page_id: u64,
    54
            /// Offset within page
    55
            offset: usize,
    56
        },
    57
    58
        /// Page table corruption
    59
        #[error("Page table corruption for sequence {seq_id}")]
    60
        PageTableCorruption {
    61
            /// Sequence ID
    62
            seq_id: u64,
    63
        },
    64
    }
    65
    66
    /// Unique sequence identifier
    67
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
    68
    pub struct SeqId(u64);
    69
    70
    impl SeqId {
    71
        /// Create a new unique sequence ID
    72
    91
        pub fn new() -> Self {
    73
            static COUNTER: AtomicU64 = AtomicU64::new(0);
    74
    91
            Self(COUNTER.fetch_add(1, Ordering::Relaxed))
    75
    91
        }
    76
    77
        /// Get the raw ID value
    78
    64
        pub fn value(&self) -> u64 {
    79
    64
            self.0
    80
    64
        }
    81
    }
    82
    83
    impl Default for SeqId {
    84
    2
        fn default() -> Self {
    85
    2
            Self::new()
    86
    2
        }
    87
    }
    88
    89
    /// Physical page identifier
    90
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
    91
    pub struct PageId(u64);
    92
    93
    impl PageId {
    94
        /// Create a new page ID
    95
    6.79k
        pub fn new(id: u64) -> Self {
    96
    6.79k
            Self(id)
    97
    6.79k
        }
    98
    99
        /// Get the raw ID value
    100
    196
        pub fn value(&self) -> u64 {
    101
    196
            self.0
    102
    196
        }
    103
    }
    104
    105
    /// KV cache data for a single page
    106
    #[derive(Debug, Clone)]
    107
    pub struct KvPage {
    108
        /// Page identifier
    109
        pub id: PageId,
    110
        /// Key cache: [block_size, num_heads, head_dim]
    111
        pub keys: Vec<f32>,
    112
        /// Value cache: [block_size, num_heads, head_dim]
    113
        pub values: Vec<f32>,
    114
        /// Number of tokens currently stored in this page
    115
        pub num_tokens: usize,
    116
        /// Reference count for copy-on-write
    117
        pub ref_count: usize,
    118
    }
    119
    120
    impl KvPage {
    121
        /// Create a new empty KV page
    122
    5.63k
        pub fn new(id: PageId, block_size: usize, num_heads: usize, head_dim: usize) -> Self {
    123
    5.63k
            let page_size = block_size * num_heads * head_dim;
    124
    5.63k
            Self {
    125
    5.63k
                id,
    126
    5.63k
                keys: vec![0.0; page_size],
    127
    5.63k
                values: vec![0.0; page_size],
    128
    5.63k
                num_tokens: 0,
    129
    5.63k
                ref_count: 1,
    130
    5.63k
            }
    131
    5.63k
        }
    132
    133
        /// Check if page is full
    134
    2
        pub fn is_full(&self, block_size: usize) -> bool {
    135
    2
            self.num_tokens >= block_size
    136
    2
        }
    137
    138
        /// Check if page is shared (copy-on-write)
    139
    6
        pub fn is_shared(&self) -> bool {
    140
    6
            self.ref_count > 1
    141
    6
        }
    142
    143
        /// Get remaining capacity
    144
    3
        pub fn remaining_capacity(&self, block_size: usize) -> usize {
    145
    3
            block_size.saturating_sub(self.num_tokens)
    146
    3
        }
    147
    }
    148
    149
    /// PagedAttention KV cache manager
    150
    /// Reference: [4] Kwon et al. (2023) "Efficient Memory Management for LLM Serving"
    151
    pub struct PagedKvCache {
    152
        /// Physical pages (fixed-size blocks)
    153
        physical_pages: Vec<KvPage>,
    154
        /// Logical to physical page mapping (per sequence)
    155
        page_tables: HashMap<SeqId, Vec<PageId>>,
    156
        /// Free page list
    157
        free_pages: VecDeque<PageId>,
    158
        /// Tokens per page (block size)
    159
        block_size: usize,
    160
        /// Number of attention heads
    161
        num_heads: usize,
    162
        /// Dimension per head
    163
        head_dim: usize,
    164
        /// Total pages allocated
    165
        total_pages: usize,
    166
        /// Statistics
    167
        stats: PagedCacheStats,
    168
    }
    169
    170
    /// Statistics for PagedKvCache
    171
    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    172
    pub struct PagedCacheStats {
    173
        /// Total sequences allocated
    174
        pub sequences_allocated: u64,
    175
        /// Total sequences freed
    176
        pub sequences_freed: u64,
    177
        /// Total pages allocated
    178
        pub pages_allocated: u64,
    179
        /// Total pages freed
    180
        pub pages_freed: u64,
    181
        /// Current active sequences
    182
        pub active_sequences: u64,
    183
        /// Current used pages
    184
        pub used_pages: u64,
    185
        /// Copy-on-write operations
    186
        pub cow_operations: u64,
    187
        /// Defragmentation operations performed
    188
        pub defrag_operations: u64,
    189
        /// Pages moved during defragmentation
    190
        pub pages_moved: u64,
    191
    }
    192
    193
    /// Fragmentation statistics for KV cache
    194
    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    195
    pub struct FragmentationStats {
    196
        /// Number of free page "holes" between used pages
    197
        pub holes: usize,
    198
        /// Total wasted capacity in partially-filled pages (tokens)
    199
        pub wasted_capacity: usize,
    200
        /// Fragmentation ratio (0.0 = no fragmentation, 1.0 = fully fragmented)
    201
        pub fragmentation_ratio: f32,
    202
        /// Largest contiguous free region (in pages)
    203
        pub largest_free_region: usize,
    204
        /// Average tokens per page (efficiency metric)
    205
        pub avg_tokens_per_page: f32,
    206
    }
    207
    208
    impl PagedKvCache {
    209
        /// Create a new PagedKvCache
    210
        ///
    211
        /// # Arguments
    212
        /// * `total_pages` - Total number of physical pages to allocate
    213
        /// * `block_size` - Tokens per page (typically 16 or 32)
    214
        /// * `num_heads` - Number of attention heads
    215
        /// * `head_dim` - Dimension per attention head
    216
    62
        pub fn new(total_pages: usize, block_size: usize, num_heads: usize, head_dim: usize) -> Self {
    217
    62
            let mut physical_pages = Vec::with_capacity(total_pages);
    218
    62
            let mut free_pages = VecDeque::with_capacity(total_pages);
    219
    220
            // Pre-allocate all pages
    221
    5.62k
            for i in 0..total_pages62 {
    222
    5.62k
                let page_id = PageId::new(i as u64);
    223
    5.62k
                physical_pages.push(KvPage::new(page_id, block_size, num_heads, head_dim));
    224
    5.62k
                free_pages.push_back(page_id);
    225
    5.62k
            }
    226
    227
    62
            Self {
    228
    62
                physical_pages,
    229
    62
                page_tables: HashMap::new(),
    230
    62
                free_pages,
    231
    62
                block_size,
    232
    62
                num_heads,
    233
    62
                head_dim,
    234
    62
                total_pages,
    235
    62
                stats: PagedCacheStats::default(),
    236
    62
            }
    237
    62
        }
    238
    239
        /// Allocate pages for a new sequence
    240
    61
        pub fn allocate_sequence(&mut self, num_tokens: usize) -> Result<SeqId, PagedCacheError> {
    241
    61
            let num_pages = self.tokens_to_pages(num_tokens);
    242
    243
    61
            if self.free_pages.len() < num_pages {
    244
    1
                return Err(PagedCacheError::OutOfMemory {
    245
    1
                    needed: num_pages,
    246
    1
                    available: self.free_pages.len(),
    247
    1
                });
    248
    60
            }
    249
    250
    60
            let seq_id = SeqId::new();
    251
    60
            let mut pages = Vec::with_capacity(num_pages);
    252
    253
    60
            for _ in 0..num_pages {
    254
    83
                if let Some(page_id) = self.free_pages.pop_front() {
    255
    83
                    // Reset the page
    256
    83
                    let page = &mut self.physical_pages[page_id.value() as usize];
    257
    83
                    page.num_tokens = 0;
    258
    83
                    page.ref_count = 1;
    259
    83
                    pages.push(page_id);
    260
    83
                }0
    261
            }
    262
    263
    60
            self.page_tables.insert(seq_id, pages);
    264
    60
            self.stats.sequences_allocated += 1;
    265
    60
            self.stats.pages_allocated += num_pages as u64;
    266
    60
            self.stats.active_sequences += 1;
    267
    60
            self.stats.used_pages += num_pages as u64;
    268
    269
    60
            Ok(seq_id)
    270
    61
        }
    271
    272
        /// Extend sequence by allocating more pages for generation
    273
    5
        pub fn extend(&mut self, seq_id: SeqId, num_tokens: usize) -> Result<(), PagedCacheError> {
    274
            // First, gather info without holding mutable borrow
    275
    4
            let (current_pages, current_tokens) = {
    276
    5
                let pages4 = self
    277
    5
                    .page_tables
    278
    5
                    .get(&seq_id)
    279
    5
                    .ok_or(PagedCacheError::SequenceNotFound(seq_id.value()))?1;
    280
    281
    4
                let mut total_tokens = 0;
    282
    8
                for page_id4 in pages {
    283
    4
                    let page = &self.physical_pages[page_id.value() as usize];
    284
    4
                    total_tokens += page.num_tokens;
    285
    4
                }
    286
    4
                (pages.len(), total_tokens)
    287
            };
    288
    289
    4
            let current_capacity = current_pages * self.block_size;
    290
    4
            let total_needed = current_tokens + num_tokens;
    291
    292
    4
            if total_needed <= current_capacity {
    293
                // No new pages needed
    294
    1
                return Ok(());
    295
    3
            }
    296
    297
    3
            let additional_pages = self.tokens_to_pages(total_needed) - current_pages;
    298
    299
    3
            if self.free_pages.len() < additional_pages {
    300
    1
                return Err(PagedCacheError::OutOfMemory {
    301
    1
                    needed: additional_pages,
    302
    1
                    available: self.free_pages.len(),
    303
    1
                });
    304
    2
            }
    305
    306
            // Collect new page IDs
    307
    2
            let mut new_pages = Vec::with_capacity(additional_pages);
    308
    2
            for _ in 0..additional_pages {
    309
    2
                if let Some(page_id) = self.free_pages.pop_front() {
    310
    2
                    let page = &mut self.physical_pages[page_id.value() as usize];
    311
    2
                    page.num_tokens = 0;
    312
    2
                    page.ref_count = 1;
    313
    2
                    new_pages.push(page_id);
    314
    2
                }0
    315
            }
    316
    317
            // Now update page table
    318
    2
            if let Some(pages) = self.page_tables.get_mut(&seq_id) {
    319
    2
                pages.extend(new_pages);
    320
    2
            }0
    321
    322
    2
            self.stats.pages_allocated += additional_pages as u64;
    323
    2
            self.stats.used_pages += additional_pages as u64;
    324
    325
    2
            Ok(())
    326
    5
        }
    327
    328
        /// Free sequence and return pages to pool
    329
    10
        pub fn free_sequence(&mut self, seq_id: SeqId) {
    330
    10
            if let Some(pages9) = self.page_tables.remove(&seq_id) {
    331
    20
                for page_id11 in pages {
    332
    11
                    let page = &mut self.physical_pages[page_id.value() as usize];
    333
    11
                    page.ref_count = page.ref_count.saturating_sub(1);
    334
    335
                    // Only return to free list if no references remain
    336
    11
                    if page.ref_count == 0 {
    337
    10
                        self.free_pages.push_back(page_id);
    338
    10
                        self.stats.pages_freed += 1;
    339
    10
                        self.stats.used_pages = self.stats.used_pages.saturating_sub(1);
    340
    10
                    }1
    341
                }
    342
    9
                self.stats.sequences_freed += 1;
    343
    9
                self.stats.active_sequences = self.stats.active_sequences.saturating_sub(1);
    344
    1
            }
    345
    10
        }
    346
    347
        /// Fork a sequence (copy-on-write for prefix sharing)
    348
    8
        pub fn fork_sequence(&mut self, parent_seq_id: SeqId) -> Result<SeqId, PagedCacheError> {
    349
    8
            let parent_pages7 = self
    350
    8
                .page_tables
    351
    8
                .get(&parent_seq_id)
    352
    8
                .ok_or(PagedCacheError::SequenceNotFound(parent_seq_id.value()))?1
    353
    7
                .clone();
    354
    355
            // Increment reference counts for shared pages
    356
    15
            for page_id8 in &parent_pages {
    357
    8
                self.physical_pages[page_id.value() as usize].ref_count += 1;
    358
    8
            }
    359
    360
    7
            let child_seq_id = SeqId::new();
    361
    7
            self.page_tables.insert(child_seq_id, parent_pages);
    362
    363
    7
            self.stats.sequences_allocated += 1;
    364
    7
            self.stats.active_sequences += 1;
    365
    7
            self.stats.cow_operations += 1;
    366
    367
    7
            Ok(child_seq_id)
    368
    8
        }
    369
    370
        /// Get the number of tokens stored for a sequence
    371
    6
        pub fn get_sequence_tokens(&self, seq_id: SeqId) -> Result<usize, PagedCacheError> {
    372
    6
            let pages5 = self
    373
    6
                .page_tables
    374
    6
                .get(&seq_id)
    375
    6
                .ok_or(PagedCacheError::SequenceNotFound(seq_id.value()))?1;
    376
    377
    5
            let mut total_tokens = 0;
    378
    13
            for page_id8 in pages {
    379
    8
                let page = &self.physical_pages[page_id.value() as usize];
    380
    8
                total_tokens += page.num_tokens;
    381
    8
            }
    382
    383
    5
            Ok(total_tokens)
    384
    6
        }
    385
    386
        /// Update token count for sequence (after writing KV data)
    387
    16
        pub fn update_tokens(
    388
    16
            &mut self,
    389
    16
            seq_id: SeqId,
    390
    16
            num_tokens: usize,
    391
    16
        ) -> Result<(), PagedCacheError> {
    392
    16
            let pages15 = self
    393
    16
                .page_tables
    394
    16
                .get(&seq_id)
    395
    16
                .ok_or(PagedCacheError::SequenceNotFound(seq_id.value()))?1;
    396
    397
    15
            let mut remaining = num_tokens;
    398
    21
            for page_id in pages {
    399
    21
                let page = &mut self.physical_pages[page_id.value() as usize];
    400
    21
                let tokens_in_page = remaining.min(self.block_size);
    401
    21
                page.num_tokens = tokens_in_page;
    402
    21
                remaining = remaining.saturating_sub(self.block_size);
    403
    21
                if remaining == 0 {
    404
    15
                    break;
    405
    6
                }
    406
            }
    407
    408
    15
            Ok(())
    409
    16
        }
    410
    411
        /// Get physical page for a logical position
    412
    6
        pub fn get_page(
    413
    6
            &self,
    414
    6
            seq_id: SeqId,
    415
    6
            token_position: usize,
    416
    6
        ) -> Result<&KvPage, PagedCacheError> {
    417
    6
            let pages5 = self
    418
    6
                .page_tables
    419
    6
                .get(&seq_id)
    420
    6
                .ok_or(PagedCacheError::SequenceNotFound(seq_id.value()))?1;
    421
    422
    5
            let page_index = token_position / self.block_size;
    423
    5
            let page_id4 = pages
    424
    5
                .get(page_index)
    425
    5
                .ok_or(PagedCacheError::InvalidPageAccess {
    426
    5
                    page_id: page_index as u64,
    427
    5
                    offset: token_position,
    428
    5
                })?1;
    429
    430
    4
            Ok(&self.physical_pages[page_id.value() as usize])
    431
    6
        }
    432
    433
        /// Get mutable physical page (handles copy-on-write)
    434
    6
        pub fn get_page_mut(
    435
    6
            &mut self,
    436
    6
            seq_id: SeqId,
    437
    6
            token_position: usize,
    438
    6
        ) -> Result<&mut KvPage, PagedCacheError> {
    439
    6
            let pages5 = self
    440
    6
                .page_tables
    441
    6
                .get(&seq_id)
    442
    6
                .ok_or(PagedCacheError::SequenceNotFound(seq_id.value()))?1;
    443
    444
    5
            let page_index = token_position / self.block_size;
    445
    5
            let page_id4 = *pages
    446
    5
                .get(page_index)
    447
    5
                .ok_or(PagedCacheError::InvalidPageAccess {
    448
    5
                    page_id: page_index as u64,
    449
    5
                    offset: token_position,
    450
    5
                })?1;
    451
    452
            // Handle copy-on-write if page is shared
    453
    4
            let page = &self.physical_pages[page_id.value() as usize];
    454
    4
            if page.is_shared() {
    455
                // Allocate a new page and copy data
    456
    3
                let new_page_id2 = self
    457
    3
                    .free_pages
    458
    3
                    .pop_front()
    459
    3
                    .ok_or(PagedCacheError::OutOfMemory {
    460
    3
                        needed: 1,
    461
    3
                        available: 0,
    462
    3
                    })?1;
    463
    464
                // Copy data to new page
    465
    2
                let old_page = &self.physical_pages[page_id.value() as usize];
    466
    2
                let keys = old_page.keys.clone();
    467
    2
                let values = old_page.values.clone();
    468
    2
                let num_tokens = old_page.num_tokens;
    469
    470
                // Update old page ref count
    471
    2
                self.physical_pages[page_id.value() as usize].ref_count -= 1;
    472
    473
                // Setup new page
    474
    2
                let new_page = &mut self.physical_pages[new_page_id.value() as usize];
    475
    2
                new_page.keys = keys;
    476
    2
                new_page.values = values;
    477
    2
                new_page.num_tokens = num_tokens;
    478
    2
                new_page.ref_count = 1;
    479
    480
                // Update page table
    481
    2
                let pages = self
    482
    2
                    .page_tables
    483
    2
                    .get_mut(&seq_id)
    484
    2
                    .ok_or(PagedCacheError::SequenceNotFound(seq_id.value()))?0;
    485
    2
                pages[page_index] = new_page_id;
    486
    487
    2
                self.stats.cow_operations += 1;
    488
    2
                self.stats.pages_allocated += 1;
    489
    2
                self.stats.used_pages += 1;
    490
    491
    2
                return Ok(&mut self.physical_pages[new_page_id.value() as usize]);
    492
    1
            }
    493
    494
    1
            Ok(&mut self.physical_pages[page_id.value() as usize])
    495
    6
        }
    496
    497
        /// Get cache statistics
    498
    15
        pub fn stats(&self) -> &PagedCacheStats {
    499
    15
            &self.stats
    500
    15
        }
    501
    502
        /// Get memory usage in bytes
    503
    2
        pub fn memory_usage(&self) -> usize {
    504
    2
            let page_size = self.block_size * self.num_heads * self.head_dim * 4 * 2; // f32 = 4 bytes, K+V = 2
    505
    2
            self.stats.used_pages as usize * page_size
    506
    2
        }
    507
    508
        /// Get total capacity in bytes
    509
    1
        pub fn total_capacity(&self) -> usize {
    510
    1
            let page_size = self.block_size * self.num_heads * self.head_dim * 4 * 2;
    511
    1
            self.total_pages * page_size
    512
    1
        }
    513
    514
        /// Get utilization percentage
    515
    3
        pub fn utilization(&self) -> f32 {
    516
    3
            if self.total_pages == 0 {
    517
    1
                return 0.0;
    518
    2
            }
    519
    2
            (self.stats.used_pages as f32 / self.total_pages as f32) * 100.0
    520
    3
        }
    521
    522
        /// Number of free pages available
    523
    11
        pub fn free_page_count(&self) -> usize {
    524
    11
            self.free_pages.len()
    525
    11
        }
    526
    527
        /// Number of pages needed for tokens
    528
    64
        fn tokens_to_pages(&self, num_tokens: usize) -> usize {
    529
    64
            num_tokens.div_ceil(self.block_size)
    530
    64
        }
    531
    532
        // ========================================================================
    533
        // Defragmentation (llama.cpp competitive feature)
    534
        // ========================================================================
    535
    536
        /// Calculate fragmentation statistics
    537
        ///
    538
        /// Per llama.cpp KV cache defrag: tracks holes, wasted capacity, and
    539
        /// fragmentation ratio to decide when defragmentation is beneficial.
    540
    10
        pub fn fragmentation_stats(&self) -> FragmentationStats {
    541
            // Build a usage map: true = used, false = free
    542
    10
            let mut usage_map = vec![false; self.total_pages];
    543
    10
            let mut total_tokens = 0usize;
    544
    10
            let mut pages_with_tokens = 0usize;
    545
    546
    10
            for pages8 in self.page_tables.values() {
    547
    21
                for page_id13 in pages {
    548
    13
                    let idx = page_id.value() as usize;
    549
    13
                    if idx < self.total_pages {
    550
    13
                        usage_map[idx] = true;
    551
    13
                        let page = &self.physical_pages[idx];
    552
    13
                        total_tokens += page.num_tokens;
    553
    13
                        if page.num_tokens > 0 {
    554
    7
                            pages_with_tokens += 1;
    555
    7
                        }6
    556
    0
                    }
    557
                }
    558
            }
    559
    560
            // Count holes (transitions from used to free in the middle of used regions)
    561
    10
            let mut holes = 0usize;
    562
    10
            let mut in_used_region = false;
    563
    10
            let mut current_free_run = 0usize;
    564
    10
            let mut largest_free_region = 0usize;
    565
    10
            let mut free_runs = Vec::new();
    566
    567
    920
            for &used910 in &usage_map {
    568
    910
                if used {
    569
    13
                    if in_used_region && current_free_run > 07 {
    570
    2
                        holes += 1;
    571
    2
                        free_runs.push(current_free_run);
    572
    11
                    }
    573
    13
                    in_used_region = true;
    574
    13
                    current_free_run = 0;
    575
    897
                } else {
    576
    897
                    current_free_run += 1;
    577
    897
                    largest_free_region = largest_free_region.max(current_free_run);
    578
    897
                }
    579
            }
    580
    581
            // Trailing free region
    582
    10
            if current_free_run > 0 {
    583
    10
                free_runs.push(current_free_run);
    584
    10
            }0
    585
    586
            // Calculate wasted capacity (unfilled slots in used pages)
    587
    10
            let used_pages = self.stats.used_pages as usize;
    588
    10
            let max_capacity = used_pages * self.block_size;
    589
    10
            let wasted_capacity = max_capacity.saturating_sub(total_tokens);
    590
    591
            // Fragmentation ratio: based on holes relative to used pages
    592
    10
            let fragmentation_ratio = if used_pages > 0 {
    593
    6
                (holes as f32) / (used_pages as f32).max(1.0)
    594
            } else {
    595
    4
                0.0
    596
            };
    597
    598
            // Average tokens per page
    599
    10
            let avg_tokens_per_page = if pages_with_tokens > 0 {
    600
    4
                total_tokens as f32 / pages_with_tokens as f32
    601
            } else {
    602
    6
                0.0
    603
            };
    604
    605
    10
            FragmentationStats {
    606
    10
                holes,
    607
    10
                wasted_capacity,
    608
    10
                fragmentation_ratio: fragmentation_ratio.min(1.0),
    609
    10
                largest_free_region,
    610
    10
                avg_tokens_per_page,
    611
    10
            }
    612
    10
        }
    613
    614
        /// Determine if defragmentation should be performed
    615
        ///
    616
        /// Heuristic based on:
    617
        /// - Fragmentation ratio > threshold (default 0.3)
    618
        /// - Wasted capacity > 25% of used capacity
    619
        /// - Free page count low but fragmented
    620
    2
        pub fn should_defragment(&self) -> bool {
    621
    2
            self.should_defragment_with_threshold(0.3)
    622
    2
        }
    623
    624
        /// Determine if defragmentation should be performed with custom threshold
    625
    4
        pub fn should_defragment_with_threshold(&self, threshold: f32) -> bool {
    626
    4
            let stats = self.fragmentation_stats();
    627
    628
            // High fragmentation ratio
    629
    4
            if stats.fragmentation_ratio > threshold {
    630
    0
                return true;
    631
    4
            }
    632
    633
            // Significant wasted capacity (>25% of block size)
    634
    4
            let used_pages = self.stats.used_pages as usize;
    635
    4
            if used_pages > 0 {
    636
    1
                let max_capacity = used_pages * self.block_size;
    637
    1
                let waste_ratio = stats.wasted_capacity as f32 / max_capacity as f32;
    638
    1
                if waste_ratio > 0.25 && stats.holes > 20 {
    639
    0
                    return true;
    640
    1
                }
    641
    3
            }
    642
    643
            // Low on free pages but have holes we can recover
    644
    4
            let free_ratio = self.free_pages.len() as f32 / self.total_pages as f32;
    645
    4
            if free_ratio < 0.1 && stats.holes > 00 {
    646
    0
                return true;
    647
    4
            }
    648
    649
    4
            false
    650
    4
        }
    651
    652
        /// Perform defragmentation - compact pages to reduce fragmentation
    653
        ///
    654
        /// This operation:
    655
        /// 1. Identifies fragmented sequences
    656
        /// 2. Moves pages to create contiguous allocations
    657
        /// 3. Updates page tables accordingly
    658
        /// 4. Returns number of pages moved
    659
        ///
    660
        /// Note: This is a relatively expensive operation and should be called
    661
        /// during low-activity periods or when `should_defragment()` returns true.
    662
    4
        pub fn defragment(&mut self) -> usize {
    663
    4
            let mut pages_moved = 0;
    664
    665
            // Collect sequences to defragment (those with non-contiguous pages)
    666
    4
            let seq_ids: Vec<SeqId> = self.page_tables.keys().copied().collect();
    667
    668
    6
            for seq_id2 in seq_ids {
    669
    2
                pages_moved += self.compact_sequence(seq_id);
    670
    2
            }
    671
    672
    4
            if pages_moved > 0 {
    673
    0
                self.stats.defrag_operations += 1;
    674
    0
                self.stats.pages_moved += pages_moved as u64;
    675
    4
            }
    676
    677
    4
            pages_moved
    678
    4
        }
    679
    680
        /// Compact a specific sequence's pages to be contiguous
    681
        ///
    682
        /// Returns number of pages moved.
    683
    6
        pub fn compact_sequence(&mut self, seq_id: SeqId) -> usize {
    684
    6
            let pages5 = match self.page_tables.get(&seq_id) {
    685
    5
                Some(p) => p.clone(),
    686
    1
                None => return 0,
    687
            };
    688
    689
    5
            if pages.is_empty() {
    690
    1
                return 0;
    691
    4
            }
    692
    693
            // Check if already contiguous
    694
    4
            let mut is_contiguous = true;
    695
    4
            for i3 in 1..pages.len() {
    696
    3
                let prev_id = pages[i - 1].value();
    697
    3
                let curr_id = pages[i].value();
    698
                // Check if pages are adjacent (within reasonable range for contiguity)
    699
    3
                if curr_id != prev_id + 1 {
    700
    0
                    is_contiguous = false;
    701
    0
                    break;
    702
    3
                }
    703
            }
    704
    705
    4
            if is_contiguous {
    706
    4
                return 0; // Already compact
    707
    0
            }
    708
    709
            // Find target region - look for contiguous free space or lowest-numbered free pages
    710
    0
            let mut pages_moved = 0;
    711
    712
            // Strategy: For each non-contiguous page, try to move it adjacent to previous
    713
    0
            let mut new_page_list = vec![pages[0]];
    714
    715
    0
            for i in 1..pages.len() {
    716
    0
                let prev_page_id = new_page_list[i - 1];
    717
    0
                let curr_page_id = pages[i];
    718
    719
                // Check if current page is already adjacent
    720
    0
                if curr_page_id.value() == prev_page_id.value() + 1 {
    721
    0
                    new_page_list.push(curr_page_id);
    722
    0
                    continue;
    723
    0
                }
    724
    725
                // Try to find a free page adjacent to previous
    726
    0
                let target_id = PageId::new(prev_page_id.value() + 1);
    727
    0
                let target_idx = target_id.value() as usize;
    728
    729
    0
                if target_idx < self.total_pages && self.is_page_free(target_id) {
    730
                    // Move data from curr_page to target_page
    731
    0
                    let curr_idx = curr_page_id.value() as usize;
    732
    733
                    // Copy data
    734
    0
                    let keys = self.physical_pages[curr_idx].keys.clone();
    735
    0
                    let values = self.physical_pages[curr_idx].values.clone();
    736
    0
                    let num_tokens = self.physical_pages[curr_idx].num_tokens;
    737
    0
                    let ref_count = self.physical_pages[curr_idx].ref_count;
    738
    739
                    // If current page is shared, we can't move it (COW semantics)
    740
    0
                    if ref_count > 1 {
    741
    0
                        new_page_list.push(curr_page_id);
    742
    0
                        continue;
    743
    0
                    }
    744
    745
                    // Remove target from free list
    746
    0
                    self.free_pages.retain(|&p| p != target_id);
    747
    748
                    // Setup target page
    749
    0
                    self.physical_pages[target_idx].keys = keys;
    750
    0
                    self.physical_pages[target_idx].values = values;
    751
    0
                    self.physical_pages[target_idx].num_tokens = num_tokens;
    752
    0
                    self.physical_pages[target_idx].ref_count = 1;
    753
    754
                    // Clear source page and return to free list
    755
    0
                    self.physical_pages[curr_idx].num_tokens = 0;
    756
    0
                    self.physical_pages[curr_idx].ref_count = 0;
    757
    0
                    self.free_pages.push_back(curr_page_id);
    758
    759
    0
                    new_page_list.push(target_id);
    760
    0
                    pages_moved += 1;
    761
    0
                } else {
    762
    0
                    // Can't move, keep current page
    763
    0
                    new_page_list.push(curr_page_id);
    764
    0
                }
    765
            }
    766
    767
            // Update page table
    768
    0
            if let Some(entry) = self.page_tables.get_mut(&seq_id) {
    769
    0
                *entry = new_page_list;
    770
    0
            }
    771
    772
    0
            pages_moved
    773
    6
        }
    774
    775
        /// Check if a page is free
    776
    0
        fn is_page_free(&self, page_id: PageId) -> bool {
    777
    0
            self.free_pages.contains(&page_id)
    778
    0
        }
    779
    780
        /// Get contiguity score for a sequence (1.0 = fully contiguous)
    781
    3
        pub fn sequence_contiguity(&self, seq_id: SeqId) -> Result<f32, PagedCacheError> {
    782
    3
            let pages2 = self
    783
    3
                .page_tables
    784
    3
                .get(&seq_id)
    785
    3
                .ok_or(PagedCacheError::SequenceNotFound(seq_id.value()))?1;
    786
    787
    2
            if pages.len() <= 1 {
    788
    1
                return Ok(1.0); // Single page is always contiguous
    789
    1
            }
    790
    791
    1
            let mut contiguous_pairs = 0;
    792
    1
            for i in 1..pages.len() {
    793
    1
                if pages[i].value() == pages[i - 1].value() + 1 {
    794
    1
                    contiguous_pairs += 1;
    795
    1
                }0
    796
            }
    797
    798
    1
            Ok(contiguous_pairs as f32 / (pages.len() - 1) as f32)
    799
    3
        }
    800
    }
    801
    802
    // ============================================================================
    803
    // PREFIX CACHING (per llama.cpp)
    804
    // ============================================================================
    805
    //
    806
    // Prefix caching allows reusing KV cache values for common prompt prefixes.
    807
    // When multiple requests share the same prefix tokens (e.g., system prompts),
    808
    // the KV cache for those tokens is computed once and reused.
    809
    //
    810
    // Benefits:
    811
    // - Reduces time-to-first-token for common prompts
    812
    // - Saves computation for repeated system instructions
    813
    // - Enables efficient multi-turn conversation handling
    814
    // ============================================================================
    815
    816
    /// Hash type for prefix cache lookup
    817
    pub type PrefixHash = u64;
    818
    819
    /// Compute hash for a token sequence (used for prefix lookup)
    820
    29
    pub fn compute_prefix_hash(tokens: &[u32]) -> PrefixHash {
    821
        // Simple FNV-1a hash for token sequences
    822
    29
        let mut hash: u64 = 0xcbf2_9ce4_8422_2325; // FNV offset basis
    823
    119
        for &token90 in tokens {
    824
    90
            hash ^= token as u64;
    825
    90
            hash = hash.wrapping_mul(0x0100_0000_01b3); // FNV prime
    826
    90
        }
    827
    29
        hash
    828
    29
    }
    829
    830
    /// Cached prefix entry
    831
    #[derive(Debug, Clone, Serialize, Deserialize)]
    832
    pub struct CachedPrefix {
    833
        /// Hash of the prefix tokens
    834
        pub hash: PrefixHash,
    835
        /// Number of tokens in prefix
    836
        pub num_tokens: usize,
    837
        /// Page IDs containing the cached KV values
    838
        pub page_ids: Vec<PageId>,
    839
        /// Reference count (number of sequences using this prefix)
    840
        pub ref_count: usize,
    841
        /// Last access timestamp (for LRU eviction)
    842
        pub last_access: u64,
    843
    }
    844
    845
    impl CachedPrefix {
    846
        /// Create new cached prefix
    847
    24
        pub fn new(hash: PrefixHash, num_tokens: usize, page_ids: Vec<PageId>) -> Self {
    848
    24
            Self {
    849
    24
                hash,
    850
    24
                num_tokens,
    851
    24
                page_ids,
    852
    24
                ref_count: 1,
    853
    24
                last_access: 0,
    854
    24
            }
    855
    24
        }
    856
    857
        /// Increment reference count
    858
    2
        pub fn add_ref(&mut self) {
    859
    2
            self.ref_count += 1;
    860
    2
        }
    861
    862
        /// Decrement reference count
    863
    6
        pub fn remove_ref(&mut self) -> bool {
    864
    6
            self.ref_count = self.ref_count.saturating_sub(1);
    865
    6
            self.ref_count == 0
    866
    6
        }
    867
    }
    868
    869
    /// Prefix cache for KV cache reuse
    870
    ///
    871
    /// Per llama.cpp's prompt cache: stores computed KV values for common
    872
    /// prompt prefixes, enabling fast cache hits for repeated system prompts.
    873
    pub struct PrefixCache {
    874
        /// Cached prefixes by hash
    875
        cache: HashMap<PrefixHash, CachedPrefix>,
    876
        /// Maximum number of cached prefixes
    877
        max_entries: usize,
    878
        /// Access counter for LRU
    879
        access_counter: u64,
    880
        /// Statistics
    881
        stats: PrefixCacheStats,
    882
    }
    883
    884
    /// Statistics for prefix cache
    885
    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    886
    pub struct PrefixCacheStats {
    887
        /// Cache hits
    888
        pub hits: u64,
    889
        /// Cache misses
    890
        pub misses: u64,
    891
        /// Total prefixes cached
    892
        pub prefixes_cached: u64,
    893
        /// Prefixes evicted
    894
        pub prefixes_evicted: u64,
    895
        /// Tokens saved (not recomputed)
    896
        pub tokens_saved: u64,
    897
    }
    898
    899
    impl PrefixCacheStats {
    900
        /// Hit rate (0.0 to 1.0)
    901
    2
        pub fn hit_rate(&self) -> f64 {
    902
    2
            let total = self.hits + self.misses;
    903
    2
            if total == 0 {
    904
    1
                0.0
    905
            } else {
    906
    1
                self.hits as f64 / total as f64
    907
            }
    908
    2
        }
    909
    }
    910
    911
    impl PrefixCache {
    912
        /// Create new prefix cache
    913
    18
        pub fn new(max_entries: usize) -> Self {
    914
    18
            Self {
    915
    18
                cache: HashMap::with_capacity(max_entries),
    916
    18
                max_entries,
    917
    18
                access_counter: 0,
    918
    18
                stats: PrefixCacheStats::default(),
    919
    18
            }
    920
    18
        }
    921
    922
        /// Look up cached prefix by hash
    923
    9
        pub fn lookup(&mut self, hash: PrefixHash) -> Option<&CachedPrefix> {
    924
    9
            if let Some(entry7) = self.cache.get_mut(&hash) {
    925
    7
                self.access_counter += 1;
    926
    7
                entry.last_access = self.access_counter;
    927
    7
                self.stats.hits += 1;
    928
                // Return immutable reference
    929
    7
                self.cache.get(&hash)
    930
            } else {
    931
    2
                self.stats.misses += 1;
    932
    2
                None
    933
            }
    934
    9
        }
    935
    936
        /// Look up cached prefix by tokens
    937
    1
        pub fn lookup_tokens(&mut self, tokens: &[u32]) -> Option<&CachedPrefix> {
    938
    1
            let hash = compute_prefix_hash(tokens);
    939
    1
            self.lookup(hash)
    940
    1
        }
    941
    942
        /// Check if prefix is cached (without updating stats)
    943
    21
        pub fn contains(&self, hash: PrefixHash) -> bool {
    944
    21
            self.cache.contains_key(&hash)
    945
    21
        }
    946
    947
        /// Insert cached prefix
    948
    21
        pub fn insert(&mut self, prefix: CachedPrefix) -> bool {
    949
    21
            let hash = prefix.hash;
    950
    951
            // Evict if at capacity
    952
    21
            if self.cache.len() >= self.max_entries && !2self.cache2.contains_key(&hash) {
    953
    2
                self.evict_lru();
    954
    19
            }
    955
    956
    21
            if self.cache.len() < self.max_entries {
    957
    20
                self.stats.prefixes_cached += 1;
    958
    20
                self.stats.tokens_saved += prefix.num_tokens as u64;
    959
    20
                self.cache.insert(hash, prefix);
    960
    20
                true
    961
            } else {
    962
    1
                false
    963
            }
    964
    21
        }
    965
    966
        /// Add reference to cached prefix
    967
    2
        pub fn add_ref(&mut self, hash: PrefixHash) -> bool {
    968
    2
            if let Some(entry1) = self.cache.get_mut(&hash) {
    969
    1
                entry.add_ref();
    970
    1
                self.access_counter += 1;
    971
    1
                entry.last_access = self.access_counter;
    972
    1
                true
    973
            } else {
    974
    1
                false
    975
            }
    976
    2
        }
    977
    978
        /// Remove reference from cached prefix
    979
        /// Returns true if prefix was removed (no more references)
    980
    4
        pub fn remove_ref(&mut self, hash: PrefixHash) -> bool {
    981
    4
            if let Some(entry3) = self.cache.get_mut(&hash) {
    982
    3
                if entry.remove_ref() {
    983
                    // No more references, remove from cache
    984
    2
                    self.cache.remove(&hash);
    985
    2
                    return true;
    986
    1
                }
    987
    1
            }
    988
    2
            false
    989
    4
        }
    990
    991
        /// Evict least recently used prefix
    992
    2
        fn evict_lru(&mut self) {
    993
    2
            if let Some((&hash1, _)) = self
    994
    2
                .cache
    995
    2
                .iter()
    996
    5
                .filter2(|(_, v)| v.ref_count == 0)
    997
    2
                .min_by_key(|(_, v)| v.last_access)
    998
    1
            {
    999
    1
                self.cache.remove(&hash);
    1000
    1
                self.stats.prefixes_evicted += 1;
    1001
    1
            }
    1002
    2
        }
    1003
    1004
        /// Get number of cached prefixes
    1005
    7
        pub fn len(&self) -> usize {
    1006
    7
            self.cache.len()
    1007
    7
        }
    1008
    1009
        /// Check if cache is empty
    1010
    4
        pub fn is_empty(&self) -> bool {
    1011
    4
            self.cache.is_empty()
    1012
    4
        }
    1013
    1014
        /// Get cache statistics
    1015
    4
        pub fn stats(&self) -> &PrefixCacheStats {
    1016
    4
            &self.stats
    1017
    4
        }
    1018
    1019
        /// Clear the cache
    1020
    1
        pub fn clear(&mut self) {
    1021
    1
            self.cache.clear();
    1022
    1
            self.access_counter = 0;
    1023
    1
        }
    1024
    1025
        /// Get cache utilization (0.0 to 1.0)
    1026
    5
        pub fn utilization(&self) -> f64 {
    1027
    5
            if self.max_entries == 0 {
    1028
    1
                0.0
    1029
            } else {
    1030
    4
                self.cache.len() as f64 / self.max_entries as f64
    1031
            }
    1032
    5
        }
    1033
    }
    1034
    1035
    impl Default for PrefixCache {
    1036
    1
        fn default() -> Self {
    1037
    1
            Self::new(100)
    1038
    1
        }
    1039
    }
    1040
    1041
    // ============================================================================
    1042
    // KV CACHE QUANTIZATION (per llama.cpp Q8/Q4 KV)
    1043
    // ============================================================================
    1044
    //
    1045
    // KV cache quantization reduces memory usage during inference:
    1046
    // - Q8_0: 8-bit quantization, ~2x memory reduction, minimal quality loss
    1047
    // - Q4_0: 4-bit quantization, ~4x memory reduction, some quality loss
    1048
    //
    1049
    // llama.cpp uses this for long-context inference where KV cache dominates memory.
    1050
    // ============================================================================
    1051
    1052
    /// KV cache quantization type
    1053
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
    1054
    pub enum KvQuantType {
    1055
        /// Full precision (32-bit float)
    1056
        #[default]
    1057
        FP32,
    1058
        /// 8-bit quantization (Q8_0 format)
    1059
        Q8,
    1060
        /// 4-bit quantization (Q4_0 format)
    1061
        Q4,
    1062
    }
    1063
    1064
    impl KvQuantType {
    1065
        /// Bytes per value for this quantization type
    1066
    6
        pub fn bytes_per_value(&self) -> f32 {
    1067
    6
            match self {
    1068
    2
                Self::FP32 => 4.0,
    1069
    2
                Self::Q8 => 1.0, // 8 bits = 1 byte
    1070
    2
                Self::Q4 => 0.5, // 4 bits = 0.5 bytes
    1071
            }
    1072
    6
        }
    1073
    1074
        /// Memory reduction factor compared to FP32
    1075
    3
        pub fn memory_reduction(&self) -> f32 {
    1076
    3
            4.0 / self.bytes_per_value()
    1077
    3
        }
    1078
    }
    1079
    1080
    /// Block size for KV quantization (matches GGML)
    1081
    pub const KV_QUANT_BLOCK_SIZE: usize = 32;
    1082
    1083
    /// Q8_0 quantized block for KV cache
    1084
    #[derive(Debug, Clone)]
    1085
    pub struct Q8KvBlock {
    1086
        /// Scale factor for the block
    1087
        pub scale: f32,
    1088
        /// Quantized values (int8, stored as i8)
    1089
        pub quants: [i8; KV_QUANT_BLOCK_SIZE],
    1090
    }
    1091
    1092
    impl Q8KvBlock {
    1093
        /// Create empty block
    1094
    1.61k
        pub fn new() -> Self {
    1095
    1.61k
            Self {
    1096
    1.61k
                scale: 0.0,
    1097
    1.61k
                quants: [0; KV_QUANT_BLOCK_SIZE],
    1098
    1.61k
            }
    1099
    1.61k
        }
    1100
    1101
        /// Quantize float values to Q8
    1102
    10
        pub fn quantize(values: &[f32; KV_QUANT_BLOCK_SIZE]) -> Self {
    1103
            // Find max absolute value for scale
    1104
    320
            let amax10 = values10.iter10().map10(|v| v.abs()).fold10(0.0f32, f32::max);
    1105
    1106
    10
            if amax < 1e-10 {
    1107
    1
                return Self::new();
    1108
    9
            }
    1109
    1110
    9
            let scale = amax / 127.0;
    1111
    9
            let inv_scale = 1.0 / scale;
    1112
    1113
    9
            let mut quants = [0i8; KV_QUANT_BLOCK_SIZE];
    1114
    288
            for (i, &v) in values9.iter9().enumerate9() {
    1115
    288
                let q = (v * inv_scale).round() as i32;
    1116
    288
                quants[i] = q.clamp(-127, 127) as i8;
    1117
    288
            }
    1118
    1119
    9
            Self { scale, quants }
    1120
    10
        }
    1121
    1122
        /// Dequantize to float values
    1123
    17
        pub fn dequantize(&self) -> [f32; KV_QUANT_BLOCK_SIZE] {
    1124
    17
            let mut result = [0.0f32; KV_QUANT_BLOCK_SIZE];
    1125
    544
            for (i, &q) in self.quants17.iter17().enumerate17() {
    1126
    544
                result[i] = q as f32 * self.scale;
    1127
    544
            }
    1128
    17
            result
    1129
    17
        }
    1130
    }
    1131
    1132
    impl Default for Q8KvBlock {
    1133
    1
        fn default() -> Self {
    1134
    1
            Self::new()
    1135
    1
        }
    1136
    }
    1137
    1138
    /// Q4_0 quantized block for KV cache
    1139
    #[derive(Debug, Clone)]
    1140
    pub struct Q4KvBlock {
    1141
        /// Scale factor for the block
    1142
        pub scale: f32,
    1143
        /// Quantized values (4-bit, packed 2 per byte)
    1144
        pub quants: [u8; KV_QUANT_BLOCK_SIZE / 2],
    1145
    }
    1146
    1147
    impl Q4KvBlock {
    1148
        /// Create empty block
    1149
    709
        pub fn new() -> Self {
    1150
    709
            Self {
    1151
    709
                scale: 0.0,
    1152
    709
                quants: [0; KV_QUANT_BLOCK_SIZE / 2],
    1153
    709
            }
    1154
    709
        }
    1155
    1156
        /// Quantize float values to Q4
    1157
    5
        pub fn quantize(values: &[f32; KV_QUANT_BLOCK_SIZE]) -> Self {
    1158
            // Find max absolute value for scale
    1159
    160
            let amax5 = values5.iter5().map5(|v| v.abs()).fold5(0.0f32, f32::max);
    1160
    1161
    5
            if amax < 1e-10 {
    1162
    1
                return Self::new();
    1163
    4
            }
    1164
    1165
            // Q4_0 uses signed 4-bit: -8 to 7
    1166
    4
            let scale = amax / 7.0;
    1167
    4
            let inv_scale = 1.0 / scale;
    1168
    1169
    4
            let mut quants = [0u8; KV_QUANT_BLOCK_SIZE / 2];
    1170
    64
            for i in 0..(KV_QUANT_BLOCK_SIZE / 2)4 {
    1171
    64
                let v0 = values[i * 2];
    1172
    64
                let v1 = values[i * 2 + 1];
    1173
    64
    1174
    64
                // Quantize to -8..7 range, then shift to 0..15 for unsigned storage
    1175
    64
                let q0 = ((v0 * inv_scale).round() as i32).clamp(-8, 7) + 8;
    1176
    64
                let q1 = ((v1 * inv_scale).round() as i32).clamp(-8, 7) + 8;
    1177
    64
    1178
    64
                // Pack two 4-bit values into one byte
    1179
    64
                quants[i] = ((q1 as u8) << 4) | (q0 as u8);
    1180
    64
            }
    1181
    1182
    4
            Self { scale, quants }
    1183
    5
        }
    1184
    1185
        /// Dequantize to float values
    1186
    7
        pub fn dequantize(&self) -> [f32; KV_QUANT_BLOCK_SIZE] {
    1187
    7
            let mut result = [0.0f32; KV_QUANT_BLOCK_SIZE];
    1188
    1189
    112
            for (i, &packed) in self.quants7.iter7().enumerate7() {
    1190
    112
                // Unpack two 4-bit values
    1191
    112
                let q0 = (packed & 0x0F) as i32 - 8;
    1192
    112
                let q1 = ((packed >> 4) & 0x0F) as i32 - 8;
    1193
    112
    1194
    112
                result[i * 2] = q0 as f32 * self.scale;
    1195
    112
                result[i * 2 + 1] = q1 as f32 * self.scale;
    1196
    112
            }
    1197
    1198
    7
            result
    1199
    7
        }
    1200
    }
    1201
    1202
    impl Default for Q4KvBlock {
    1203
    1
        fn default() -> Self {
    1204
    1
            Self::new()
    1205
    1
        }
    1206
    }
    1207
    1208
    /// Quantized KV cache data for a single page
    1209
    #[derive(Debug, Clone)]
    1210
    pub enum QuantizedKvData {
    1211
        /// Full precision storage
    1212
        FP32 {
    1213
            /// Key cache: [block_size, num_heads, head_dim]
    1214
            keys: Vec<f32>,
    1215
            /// Value cache: [block_size, num_heads, head_dim]
    1216
            values: Vec<f32>,
    1217
        },
    1218
        /// Q8 quantized storage
    1219
        Q8 {
    1220
            /// Quantized key blocks
    1221
            key_blocks: Vec<Q8KvBlock>,
    1222
            /// Quantized value blocks
    1223
            value_blocks: Vec<Q8KvBlock>,
    1224
        },
    1225
        /// Q4 quantized storage
    1226
        Q4 {
    1227
            /// Quantized key blocks
    1228
            key_blocks: Vec<Q4KvBlock>,
    1229
            /// Quantized value blocks
    1230
            value_blocks: Vec<Q4KvBlock>,
    1231
        },
    1232
    }
    1233
    1234
    impl QuantizedKvData {
    1235
        /// Create new quantized KV data with given precision
    1236
    1.16k
        pub fn new(
    1237
    1.16k
            quant_type: KvQuantType,
    1238
    1.16k
            block_size: usize,
    1239
    1.16k
            num_heads: usize,
    1240
    1.16k
            head_dim: usize,
    1241
    1.16k
        ) -> Self {
    1242
    1.16k
            let total_size = block_size * num_heads * head_dim;
    1243
    1.16k
            let num_quant_blocks = total_size.div_ceil(KV_QUANT_BLOCK_SIZE);
    1244
    1245
    1.16k
            match quant_type {
    1246
    6
                KvQuantType::FP32 => Self::FP32 {
    1247
    6
                    keys: vec![0.0; total_size],
    1248
    6
                    values: vec![0.0; total_size],
    1249
    6
                },
    1250
    808
                KvQuantType::Q8 => Self::Q8 {
    1251
    808
                    key_blocks: vec![Q8KvBlock::new(); num_quant_blocks],
    1252
    808
                    value_blocks: vec![Q8KvBlock::new(); num_quant_blocks],
    1253
    808
                },
    1254
    353
                KvQuantType::Q4 => Self::Q4 {
    1255
    353
                    key_blocks: vec![Q4KvBlock::new(); num_quant_blocks],
    1256
    353
                    value_blocks: vec![Q4KvBlock::new(); num_quant_blocks],
    1257
    353
                },
    1258
            }
    1259
    1.16k
        }
    1260
    1261
        /// Get quantization type
    1262
    6
        pub fn quant_type(&self) -> KvQuantType {
    1263
    6
            match self {
    1264
    1
                Self::FP32 { .. } => KvQuantType::FP32,
    1265
    4
                Self::Q8 { .. } => KvQuantType::Q8,
    1266
    1
                Self::Q4 { .. } => KvQuantType::Q4,
    1267
            }
    1268
    6
        }
    1269
    1270
        /// Memory usage in bytes
    1271
    11
        pub fn memory_bytes(&self) -> usize {
    1272
    11
            match self {
    1273
    3
                Self::FP32 { keys, values } => (keys.len() + values.len()) * 4,
    1274
                Self::Q8 {
    1275
    5
                    key_blocks,
    1276
    5
                    value_blocks,
    1277
                } => {
    1278
                    // Scale (4 bytes) + quants (32 bytes) = 36 bytes per block
    1279
    5
                    (key_blocks.len() + value_blocks.len()) * (4 + KV_QUANT_BLOCK_SIZE)
    1280
                },
    1281
                Self::Q4 {
    1282
    3
                    key_blocks,
    1283
    3
                    value_blocks,
    1284
                } => {
    1285
                    // Scale (4 bytes) + quants (16 bytes) = 20 bytes per block
    1286
    3
                    (key_blocks.len() + value_blocks.len()) * (4 + KV_QUANT_BLOCK_SIZE / 2)
    1287
                },
    1288
            }
    1289
    11
        }
    1290
    1291
        /// Write keys at given offset
    1292
    5
        pub fn write_keys(&mut self, offset: usize, data: &[f32]) {
    1293
    5
            match self {
    1294
    4
                Self::FP32 { keys, .. } => {
    1295
    4
                    let end = (offset + data.len()).min(keys.len());
    1296
    4
                    keys[offset..end].copy_from_slice(&data[..end - offset]);
    1297
    4
                },
    1298
    1
                Self::Q8 { key_blocks, .. } => {
    1299
    1
                    write_quantized_q8(key_blocks, offset, data);
    1300
    1
                },
    1301
    0
                Self::Q4 { key_blocks, .. } => {
    1302
    0
                    write_quantized_q4(key_blocks, offset, data);
    1303
    0
                },
    1304
            }
    1305
    5
        }
    1306
    1307
        /// Write values at given offset
    1308
    3
        pub fn write_values(&mut self, offset: usize, data: &[f32]) {
    1309
    3
            match self {
    1310
    1
                Self::FP32 { values, .. } => {
    1311
    1
                    let end = (offset + data.len()).min(values.len());
    1312
    1
                    values[offset..end].copy_from_slice(&data[..end - offset]);
    1313
    1
                },
    1314
    1
                Self::Q8 { value_blocks, .. } => {
    1315
    1
                    write_quantized_q8(value_blocks, offset, data);
    1316
    1
                },
    1317
    1
                Self::Q4 { value_blocks, .. } => {
    1318
    1
                    write_quantized_q4(value_blocks, offset, data);
    1319
    1
                },
    1320
            }
    1321
    3
        }
    1322
    1323
        /// Read keys at given offset
    1324
    5
        pub fn read_keys(&self, offset: usize, length: usize) -> Vec<f32> {
    1325
    5
            match self {
    1326
    4
                Self::FP32 { keys, .. } => {
    1327
    4
                    let end = (offset + length).min(keys.len());
    1328
    4
                    keys[offset..end].to_vec()
    1329
                },
    1330
    1
                Self::Q8 { key_blocks, .. } => read_quantized_q8(key_blocks, offset, length),
    1331
    0
                Self::Q4 { key_blocks, .. } => read_quantized_q4(key_blocks, offset, length),
    1332
            }
    1333
    5
        }
    1334
    1335
        /// Read values at given offset
    1336
    3
        pub fn read_values(&self, offset: usize, length: usize) -> Vec<f32> {
    1337
    3
            match self {
    1338
    1
                Self::FP32 { values, .. } => {
    1339
    1
                    let end = (offset + length).min(values.len());
    1340
    1
                    values[offset..end].to_vec()
    1341
                },
    1342
    1
                Self::Q8 { value_blocks, .. } => read_quantized_q8(value_blocks, offset, length),
    1343
    1
                Self::Q4 { value_blocks, .. } => read_quantized_q4(value_blocks, offset, length),
    1344
            }
    1345
    3
        }
    1346
    }
    1347
    1348
    // Helper: Write to Q8 blocks
    1349
    2
    fn write_quantized_q8(blocks: &mut [Q8KvBlock], offset: usize, data: &[f32]) {
    1350
    2
        let start_block = offset / KV_QUANT_BLOCK_SIZE;
    1351
    2
        let start_offset = offset % KV_QUANT_BLOCK_SIZE;
    1352
    1353
    2
        let mut data_idx = 0;
    1354
    2
        let mut block_idx = start_block;
    1355
    2
        let mut in_block_offset = start_offset;
    1356
    1357
    9
        while data_idx < data.len() && block_idx7 < blocks.len() {
    1358
            // Read existing block, modify, re-quantize
    1359
    7
            let mut values = blocks[block_idx].dequantize();
    1360
    1361
    199
            while in_block_offset < KV_QUANT_BLOCK_SIZE && data_idx193 < data.len() {
    1362
    192
                values[in_block_offset] = data[data_idx];
    1363
    192
                in_block_offset += 1;
    1364
    192
                data_idx += 1;
    1365
    192
            }
    1366
    1367
    7
            blocks[block_idx] = Q8KvBlock::quantize(&values);
    1368
    7
            block_idx += 1;
    1369
    7
            in_block_offset = 0;
    1370
        }
    1371
    2
    }
    1372
    1373
    // Helper: Write to Q4 blocks
    1374
    1
    fn write_quantized_q4(blocks: &mut [Q4KvBlock], offset: usize, data: &[f32]) {
    1375
    1
        let start_block = offset / KV_QUANT_BLOCK_SIZE;
    1376
    1
        let start_offset = offset % KV_QUANT_BLOCK_SIZE;
    1377
    1378
    1
        let mut data_idx = 0;
    1379
    1
        let mut block_idx = start_block;
    1380
    1
        let mut in_block_offset = start_offset;
    1381
    1382
    3
        while data_idx < data.len() && block_idx2 < blocks.len() {
    1383
    2
            let mut values = blocks[block_idx].dequantize();
    1384
    1385
    66
            while in_block_offset < KV_QUANT_BLOCK_SIZE && data_idx64 < data.len() {
    1386
    64
                values[in_block_offset] = data[data_idx];
    1387
    64
                in_block_offset += 1;
    1388
    64
                data_idx += 1;
    1389
    64
            }
    1390
    1391
    2
            blocks[block_idx] = Q4KvBlock::quantize(&values);
    1392
    2
            block_idx += 1;
    1393
    2
            in_block_offset = 0;
    1394
        }
    1395
    1
    }
    1396
    1397
    // Helper: Read from Q8 blocks
    1398
    2
    fn read_quantized_q8(blocks: &[Q8KvBlock], offset: usize, length: usize) -> Vec<f32> {
    1399
    2
        let mut result = Vec::with_capacity(length);
    1400
    2
        let start_block = offset / KV_QUANT_BLOCK_SIZE;
    1401
    2
        let start_offset = offset % KV_QUANT_BLOCK_SIZE;
    1402
    1403
    2
        let mut block_idx = start_block;
    1404
    2
        let mut in_block_offset = start_offset;
    1405
    2
        let mut remaining = length;
    1406
    1407
    9
        while remaining > 0 && block_idx7 < blocks.len() {
    1408
    7
            let values = blocks[block_idx].dequantize();
    1409
    1410
    199
            while in_block_offset < KV_QUANT_BLOCK_SIZE && remaining > 0193 {
    1411
    192
                result.push(values[in_block_offset]);
    1412
    192
                in_block_offset += 1;
    1413
    192
                remaining -= 1;
    1414
    192
            }
    1415
    1416
    7
            block_idx += 1;
    1417
    7
            in_block_offset = 0;
    1418
        }
    1419
    1420
    2
        result
    1421
    2
    }
    1422
    1423
    // Helper: Read from Q4 blocks
    1424
    1
    fn read_quantized_q4(blocks: &[Q4KvBlock], offset: usize, length: usize) -> Vec<f32> {
    1425
    1
        let mut result = Vec::with_capacity(length);
    1426
    1
        let start_block = offset / KV_QUANT_BLOCK_SIZE;
    1427
    1
        let start_offset = offset % KV_QUANT_BLOCK_SIZE;
    1428
    1429
    1
        let mut block_idx = start_block;
    1430
    1
        let mut in_block_offset = start_offset;
    1431
    1
        let mut remaining = length;
    1432
    1433
    3
        while remaining > 0 && block_idx2 < blocks.len() {
    1434
    2
            let values = blocks[block_idx].dequantize();
    1435
    1436
    66
            while in_block_offset < KV_QUANT_BLOCK_SIZE && remaining > 064 {
    1437
    64
                result.push(values[in_block_offset]);
    1438
    64
                in_block_offset += 1;
    1439
    64
                remaining -= 1;
    1440
    64
            }
    1441
    1442
    2
            block_idx += 1;
    1443
    2
            in_block_offset = 0;
    1444
        }
    1445
    1446
    1
        result
    1447
    1
    }
    1448
    1449
    /// Quantized KV page for memory-efficient cache
    1450
    #[derive(Debug, Clone)]
    1451
    pub struct QuantizedKvPage {
    1452
        /// Page identifier
    1453
        pub id: PageId,
    1454
        /// Quantized KV data
    1455
        pub data: QuantizedKvData,
    1456
        /// Number of tokens currently stored
    1457
        pub num_tokens: usize,
    1458
        /// Reference count for COW
    1459
        pub ref_count: usize,
    1460
        /// Block size (tokens per page)
    1461
        block_size: usize,
    1462
        /// Number of attention heads
    1463
        num_heads: usize,
    1464
        /// Head dimension
    1465
        head_dim: usize,
    1466
    }
    1467
    1468
    impl QuantizedKvPage {
    1469
        /// Create new quantized KV page
    1470
    1.15k
        pub fn new(
    1471
    1.15k
            id: PageId,
    1472
    1.15k
            quant_type: KvQuantType,
    1473
    1.15k
            block_size: usize,
    1474
    1.15k
            num_heads: usize,
    1475
    1.15k
            head_dim: usize,
    1476
    1.15k
        ) -> Self {
    1477
    1.15k
            Self {
    1478
    1.15k
                id,
    1479
    1.15k
                data: QuantizedKvData::new(quant_type, block_size, num_heads, head_dim),
    1480
    1.15k
                num_tokens: 0,
    1481
    1.15k
                ref_count: 0, // Pages start in free pool with ref_count 0
    1482
    1.15k
                block_size,
    1483
    1.15k
                num_heads,
    1484
    1.15k
                head_dim,
    1485
    1.15k
            }
    1486
    1.15k
        }
    1487
    1488
        /// Get quantization type
    1489
    3
        pub fn quant_type(&self) -> KvQuantType {
    1490
    3
            self.data.quant_type()
    1491
    3
        }
    1492
    1493
        /// Memory usage in bytes
    1494
    6
        pub fn memory_bytes(&self) -> usize {
    1495
    6
            self.data.memory_bytes()
    1496
    6
        }
    1497
    1498
        /// Check if page is full
    1499
    3
        pub fn is_full(&self) -> bool {
    1500
    3
            self.num_tokens >= self.block_size
    1501
    3
        }
    1502
    1503
        /// Check if page is shared (COW)
    1504
    1
        pub fn is_shared(&self) -> bool {
    1505
    1
            self.ref_count > 1
    1506
    1
        }
    1507
    1508
        /// Remaining capacity in tokens
    1509
    2
        pub fn remaining_capacity(&self) -> usize {
    1510
    2
            self.block_size.saturating_sub(self.num_tokens)
    1511
    2
        }
    1512
    1513
        /// Write keys for a token position
    1514
    3
        pub fn write_keys(&mut self, token_pos: usize, keys: &[f32]) {
    1515
    3
            let offset = token_pos * self.num_heads * self.head_dim;
    1516
    3
            self.data.write_keys(offset, keys);
    1517
    3
        }
    1518
    1519
        /// Write values for a token position
    1520
    1
        pub fn write_values(&mut self, token_pos: usize, values: &[f32]) {
    1521
    1
            let offset = token_pos * self.num_heads * self.head_dim;
    1522
    1
            self.data.write_values(offset, values);
    1523
    1
        }
    1524
    1525
        /// Read keys for a token position
    1526
    3
        pub fn read_keys(&self, token_pos: usize) -> Vec<f32> {
    1527
    3
            let offset = token_pos * self.num_heads * self.head_dim;
    1528
    3
            let length = self.num_heads * self.head_dim;
    1529
    3
            self.data.read_keys(offset, length)
    1530
    3
        }
    1531
    1532
        /// Read values for a token position
    1533
    1
        pub fn read_values(&self, token_pos: usize) -> Vec<f32> {
    1534
    1
            let offset = token_pos * self.num_heads * self.head_dim;
    1535
    1
            let length = self.num_heads * self.head_dim;
    1536
    1
            self.data.read_values(offset, length)
    1537
    1
        }
    1538
    }
    1539
    1540
    /// Quantized PagedKvCache with configurable precision
    1541
    pub struct QuantizedPagedKvCache {
    1542
        /// Physical pages with quantized storage
    1543
        physical_pages: Vec<QuantizedKvPage>,
    1544
        /// Page tables (same as regular PagedKvCache)
    1545
        page_tables: HashMap<SeqId, Vec<PageId>>,
    1546
        /// Free page list
    1547
        free_pages: VecDeque<PageId>,
    1548
        /// Quantization type
    1549
        quant_type: KvQuantType,
    1550
        /// Tokens per page
    1551
        block_size: usize,
    1552
        /// Number of attention heads
    1553
        num_heads: usize,
    1554
        /// Head dimension
    1555
        head_dim: usize,
    1556
        /// Total pages
    1557
        total_pages: usize,
    1558
        /// Statistics
    1559
        stats: PagedCacheStats,
    1560
    }
    1561
    1562
    impl QuantizedPagedKvCache {
    1563
        /// Create new quantized paged KV cache
    1564
    13
        pub fn new(
    1565
    13
            total_pages: usize,
    1566
    13
            block_size: usize,
    1567
    13
            num_heads: usize,
    1568
    13
            head_dim: usize,
    1569
    13
            quant_type: KvQuantType,
    1570
    13
        ) -> Self {
    1571
    13
            let mut physical_pages = Vec::with_capacity(total_pages);
    1572
    13
            let mut free_pages = VecDeque::with_capacity(total_pages);
    1573
    1574
    1.15k
            for i in 0..total_pages13 {
    1575
    1.15k
                let page_id = PageId::new(i as u64);
    1576
    1.15k
                physical_pages.push(QuantizedKvPage::new(
    1577
    1.15k
                    page_id, quant_type, block_size, num_heads, head_dim,
    1578
    1.15k
                ));
    1579
    1.15k
                free_pages.push_back(page_id);
    1580
    1.15k
            }
    1581
    1582
    13
            Self {
    1583
    13
                physical_pages,
    1584
    13
                page_tables: HashMap::new(),
    1585
    13
                free_pages,
    1586
    13
                quant_type,
    1587
    13
                block_size,
    1588
    13
                num_heads,
    1589
    13
                head_dim,
    1590
    13
                total_pages,
    1591
    13
                stats: PagedCacheStats::default(),
    1592
    13
            }
    1593
    13
        }
    1594
    1595
        /// Get quantization type
    1596
    1
        pub fn quant_type(&self) -> KvQuantType {
    1597
    1
            self.quant_type
    1598
    1
        }
    1599
    1600
        /// Allocate pages for a sequence
    1601
    10
        pub fn allocate_sequence(&mut self, num_tokens: usize) -> Result<SeqId, PagedCacheError> {
    1602
    10
            let num_pages = num_tokens.div_ceil(self.block_size);
    1603
    1604
    10
            if self.free_pages.len() < num_pages {
    1605
    1
                return Err(PagedCacheError::OutOfMemory {
    1606
    1
                    needed: num_pages,
    1607
    1
                    available: self.free_pages.len(),
    1608
    1
                });
    1609
    9
            }
    1610
    1611
    9
            let seq_id = SeqId::new();
    1612
    9
            let mut pages = Vec::with_capacity(num_pages);
    1613
    1614
    9
            for _ in 0..num_pages {
    1615
    11
                if let Some(page_id) = self.free_pages.pop_front() {
    1616
    11
                    let page = &mut self.physical_pages[page_id.value() as usize];
    1617
    11
                    page.num_tokens = 0;
    1618
    11
                    page.ref_count = 1;
    1619
    11
                    pages.push(page_id);
    1620
    11
                }0
    1621
            }
    1622
    1623
    9
            self.page_tables.insert(seq_id, pages);
    1624
    9
            self.stats.sequences_allocated += 1;
    1625
    9
            self.stats.pages_allocated += num_pages as u64;
    1626
    9
            self.stats.active_sequences += 1;
    1627
    9
            self.stats.used_pages += num_pages as u64;
    1628
    1629
    9
            Ok(seq_id)
    1630
    10
        }
    1631
    1632
        /// Free a sequence
    1633
    1
        pub fn free_sequence(&mut self, seq_id: SeqId) {
    1634
    1
            if let Some(pages) = self.page_tables.remove(&seq_id) {
    1635
    2
                for page_id1 in pages {
    1636
    1
                    let page = &mut self.physical_pages[page_id.value() as usize];
    1637
    1
                    page.ref_count = page.ref_count.saturating_sub(1);
    1638
    1639
    1
                    if page.ref_count == 0 {
    1640
    1
                        self.free_pages.push_back(page_id);
    1641
    1
                        self.stats.pages_freed += 1;
    1642
    1
                        self.stats.used_pages = self.stats.used_pages.saturating_sub(1);
    1643
    1
                    }0
    1644
                }
    1645
    1
                self.stats.sequences_freed += 1;
    1646
    1
                self.stats.active_sequences = self.stats.active_sequences.saturating_sub(1);
    1647
    0
            }
    1648
    1
        }
    1649
    1650
        /// Get page for a token position
    1651
    5
        pub fn get_page(
    1652
    5
            &self,
    1653
    5
            seq_id: SeqId,
    1654
    5
            token_position: usize,
    1655
    5
        ) -> Result<&QuantizedKvPage, PagedCacheError> {
    1656
    5
            let pages4 = self
    1657
    5
                .page_tables
    1658
    5
                .get(&seq_id)
    1659
    5
                .ok_or(PagedCacheError::SequenceNotFound(seq_id.value()))?1;
    1660
    1661
    4
            let page_index = token_position / self.block_size;
    1662
    4
            let page_id3 = pages
    1663
    4
                .get(page_index)
    1664
    4
                .ok_or(PagedCacheError::InvalidPageAccess {
    1665
    4
                    page_id: page_index as u64,
    1666
    4
                    offset: token_position,
    1667
    4
                })?1;
    1668
    1669
    3
            Ok(&self.physical_pages[page_id.value() as usize])
    1670
    5
        }
    1671
    1672
        /// Get mutable page for a token position
    1673
    2
        pub fn get_page_mut(
    1674
    2
            &mut self,
    1675
    2
            seq_id: SeqId,
    1676
    2
            token_position: usize,
    1677
    2
        ) -> Result<&mut QuantizedKvPage, PagedCacheError> {
    1678
    2
            let pages = self
    1679
    2
                .page_tables
    1680
    2
                .get(&seq_id)
    1681
    2
                .ok_or(PagedCacheError::SequenceNotFound(seq_id.value()))?0;
    1682
    1683
    2
            let page_index = token_position / self.block_size;
    1684
    2
            let page_id1 = *pages
    1685
    2
                .get(page_index)
    1686
    2
                .ok_or(PagedCacheError::InvalidPageAccess {
    1687
    2
                    page_id: page_index as u64,
    1688
    2
                    offset: token_position,
    1689
    2
                })?1;
    1690
    1691
    1
            Ok(&mut self.physical_pages[page_id.value() as usize])
    1692
    2
        }
    1693
    1694
        /// Get total pages capacity
    1695
    1
        pub fn total_pages(&self) -> usize {
    1696
    1
            self.total_pages
    1697
    1
        }
    1698
    1699
        /// Total memory usage in bytes
    1700
    2
        pub fn memory_usage(&self) -> usize {
    1701
    2
            self.physical_pages
    1702
    2
                .iter()
    1703
    200
                .filter2(|p| p.ref_count > 0)
    1704
    2
                .map(QuantizedKvPage::memory_bytes)
    1705
    2
                .sum()
    1706
    2
        }
    1707
    1708
        /// FP32 equivalent memory (for comparison)
    1709
    3
        pub fn fp32_equivalent_memory(&self) -> usize {
    1710
    3
            let page_size = self.block_size * self.num_heads * self.head_dim * 4 * 2;
    1711
    3
            self.stats.used_pages as usize * page_size
    1712
    3
        }
    1713
    1714
        /// Memory savings ratio (1.0 = no savings, 0.25 = 4x reduction)
    1715
    3
        pub fn memory_savings(&self) -> f32 {
    1716
    3
            let fp32_mem = self.fp32_equivalent_memory();
    1717
    3
            if fp32_mem == 0 {
    1718
    1
                return 1.0;
    1719
    2
            }
    1720
    2
            self.memory_usage() as f32 / fp32_mem as f32
    1721
    3
        }
    1722
    1723
        /// Get statistics
    1724
    3
        pub fn stats(&self) -> &PagedCacheStats {
    1725
    3
            &self.stats
    1726
    3
        }
    1727
    1728
        /// Free page count
    1729
    4
        pub fn free_page_count(&self) -> usize {
    1730
    4
            self.free_pages.len()
    1731
    4
        }
    1732
    }
    1733
    1734
    /// Find longest matching prefix in a sequence of tokens
    1735
    ///
    1736
    /// Returns (hash, num_matching_tokens) for the longest cached prefix
    1737
    3
    pub fn find_longest_prefix(cache: &mut PrefixCache, tokens: &[u32]) -> Option<(PrefixHash, usize)> {
    1738
    3
        let mut best_match = None;
    1739
    3
        let mut best_len = 0;
    1740
    1741
        // Try progressively longer prefixes (from 1 token up to full sequence)
    1742
    15
        for len in 1..=tokens3.len3() {
    1743
    15
            let prefix_hash = compute_prefix_hash(&tokens[..len]);
    1744
    15
            if cache.contains(prefix_hash) && len > best_len3 {
    1745
    3
                best_len = len;
    1746
    3
                best_match = Some((prefix_hash, len));
    1747
    12
            }
    1748
        }
    1749
    1750
        // Update stats if found
    1751
    3
        if let Some((hash2, _)) = best_match {
    1752
    2
            cache.lookup(hash); // Update access time
    1753
    2
        }1
    1754
    1755
    3
        best_match
    1756
    3
    }
    1757
    1758
    // ============================================================================
    1759
    // Tests
    1760
    // ============================================================================
    1761
    1762
    // Tests extracted to tests.rs (PMAT-802)
    1763
    #[cfg(test)]
    1764
    #[path = "tests.rs"]
    1765
    mod paged_kv_tests;