## 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/scheduler/chunked_prefill.rs

    Line
    Count
    Source
    1
    //! Chunked Prefill Scheduler Component
    2
    //!
    3
    //! Implements chunked prefill per Sarathi-Serve/vLLM for reduced TTFT variance.
    4
    //! Extracted from scheduler/mod.rs (PMAT-802).
    5
    6
    #![allow(dead_code)]
    7
    #![allow(clippy::too_many_arguments)]
    8
    9
    use serde::{Deserialize, Serialize};
    10
    use std::collections::{HashMap, VecDeque};
    11
    12
    13
    // ============================================================================
    14
    // CHUNKED PREFILL (per Sarathi-Serve / vLLM)
    15
    // ============================================================================
    16
    //
    17
    // Chunked prefill breaks long prompt processing into smaller chunks, allowing
    18
    // decode steps for other requests to interleave. This reduces TTFT variance
    19
    // and prevents head-of-line blocking from long prompts.
    20
    //
    21
    // Reference: Agrawal et al. (2024) "Sarathi-Serve: Large Language Model Inference
    22
    // with Chunked Prefill and Decode Disaggregation"
    23
    24
    /// Configuration for chunked prefill
    25
    ///
    26
    /// Controls how long prompts are split into chunks for interleaved processing.
    27
    #[derive(Debug, Clone, Serialize, Deserialize)]
    28
    pub struct ChunkedPrefillConfig {
    29
        /// Enable chunked prefill (default: true)
    30
        pub enabled: bool,
    31
        /// Maximum tokens per prefill chunk (default: 512)
    32
        pub chunk_size: usize,
    33
        /// Minimum prompt length to trigger chunking (default: 256)
    34
        pub min_prompt_length: usize,
    35
        /// Allow decode interleaving between chunks (default: true)
    36
        pub allow_decode_interleave: bool,
    37
        /// Priority boost for partially-prefilled sequences (default: true)
    38
        pub boost_partial_prefill: bool,
    39
        /// Maximum chunks before forcing completion (default: 16)
    40
        pub max_chunks: usize,
    41
    }
    42
    43
    impl Default for ChunkedPrefillConfig {
    44
    22
        fn default() -> Self {
    45
    22
            Self {
    46
    22
                enabled: true,
    47
    22
                chunk_size: 512,
    48
    22
                min_prompt_length: 256,
    49
    22
                allow_decode_interleave: true,
    50
    22
                boost_partial_prefill: true,
    51
    22
                max_chunks: 16,
    52
    22
            }
    53
    22
        }
    54
    }
    55
    56
    impl ChunkedPrefillConfig {
    57
        /// Create config with specified chunk size
    58
    2
        pub fn with_chunk_size(mut self, size: usize) -> Self {
    59
    2
            self.chunk_size = size;
    60
    2
            self
    61
    2
        }
    62
    63
        /// Disable chunked prefill
    64
    2
        pub fn disabled() -> Self {
    65
    2
            Self {
    66
    2
                enabled: false,
    67
    2
                ..Default::default()
    68
    2
            }
    69
    2
        }
    70
    71
        /// Create config for low-latency scenarios (smaller chunks)
    72
    1
        pub fn low_latency() -> Self {
    73
    1
            Self {
    74
    1
                enabled: true,
    75
    1
                chunk_size: 128,
    76
    1
                min_prompt_length: 64,
    77
    1
                allow_decode_interleave: true,
    78
    1
                boost_partial_prefill: true,
    79
    1
                max_chunks: 32,
    80
    1
            }
    81
    1
        }
    82
    83
        /// Create config for high-throughput scenarios (larger chunks)
    84
    1
        pub fn high_throughput() -> Self {
    85
    1
            Self {
    86
    1
                enabled: true,
    87
    1
                chunk_size: 1024,
    88
    1
                min_prompt_length: 512,
    89
    1
                allow_decode_interleave: false,
    90
    1
                boost_partial_prefill: false,
    91
    1
                max_chunks: 8,
    92
    1
            }
    93
    1
        }
    94
    }
    95
    96
    /// State of chunked prefill for a sequence
    97
    #[derive(Debug, Clone, Serialize, Deserialize)]
    98
    pub struct ChunkedPrefillState {
    99
        /// Sequence ID
    100
        pub seq_id: u64,
    101
        /// Total prompt tokens
    102
        pub total_tokens: usize,
    103
        /// Tokens processed so far
    104
        pub processed_tokens: usize,
    105
        /// Current chunk index
    106
        pub current_chunk: usize,
    107
        /// Total number of chunks
    108
        pub total_chunks: usize,
    109
        /// Start time of prefill (for latency tracking)
    110
        pub start_time_ms: u64,
    111
        /// Time spent on each chunk (milliseconds)
    112
        pub chunk_latencies: Vec<u64>,
    113
    }
    114
    115
    impl ChunkedPrefillState {
    116
        /// Create new chunked prefill state
    117
    21
        pub fn new(seq_id: u64, total_tokens: usize, chunk_size: usize) -> Self {
    118
    21
            let total_chunks = total_tokens.div_ceil(chunk_size);
    119
    21
            Self {
    120
    21
                seq_id,
    121
    21
                total_tokens,
    122
    21
                processed_tokens: 0,
    123
    21
                current_chunk: 0,
    124
    21
                total_chunks,
    125
    21
                start_time_ms: 0,
    126
    21
                chunk_latencies: Vec::with_capacity(total_chunks),
    127
    21
            }
    128
    21
        }
    129
    130
        /// Get next chunk of tokens to process
    131
    6
        pub fn next_chunk(&self, chunk_size: usize) -> std::ops::Range<usize> {
    132
    6
            let start = self.processed_tokens;
    133
    6
            let end = (start + chunk_size).min(self.total_tokens);
    134
    6
            start..end
    135
    6
        }
    136
    137
        /// Advance to next chunk
    138
    13
        pub fn advance(&mut self, tokens_processed: usize, latency_ms: u64) {
    139
    13
            self.processed_tokens += tokens_processed;
    140
    13
            self.current_chunk += 1;
    141
    13
            self.chunk_latencies.push(latency_ms);
    142
    13
        }
    143
    144
        /// Check if prefill is complete
    145
    17
        pub fn is_complete(&self) -> bool {
    146
    17
            self.processed_tokens >= self.total_tokens
    147
    17
        }
    148
    149
        /// Get progress as percentage
    150
    4
        pub fn progress(&self) -> f64 {
    151
    4
            if self.total_tokens == 0 {
    152
    1
                100.0
    153
            } else {
    154
    3
                (self.processed_tokens as f64 / self.total_tokens as f64) * 100.0
    155
            }
    156
    4
        }
    157
    158
        /// Remaining tokens to prefill
    159
    2
        pub fn remaining_tokens(&self) -> usize {
    160
    2
            self.total_tokens.saturating_sub(self.processed_tokens)
    161
    2
        }
    162
    163
        /// Average chunk latency
    164
    3
        pub fn avg_chunk_latency_ms(&self) -> f64 {
    165
    3
            if self.chunk_latencies.is_empty() {
    166
    1
                0.0
    167
            } else {
    168
    2
                self.chunk_latencies.iter().sum::<u64>() as f64 / self.chunk_latencies.len() as f64
    169
            }
    170
    3
        }
    171
    }
    172
    173
    /// Statistics for chunked prefill scheduler
    174
    #[derive(Debug, Clone, Default, Serialize, Deserialize)]
    175
    pub struct ChunkedPrefillStats {
    176
        /// Total sequences that used chunked prefill
    177
        pub chunked_sequences: u64,
    178
        /// Total sequences that bypassed chunking (short prompts)
    179
        pub bypassed_sequences: u64,
    180
        /// Total chunks processed
    181
        pub chunks_processed: u64,
    182
        /// Total decode interleaves (decode steps between prefill chunks)
    183
        pub decode_interleaves: u64,
    184
        /// Sum of all chunk latencies (for averaging)
    185
        pub total_chunk_latency_ms: u64,
    186
        /// Maximum single chunk latency
    187
        pub max_chunk_latency_ms: u64,
    188
        /// Total tokens saved from prefix cache hits during chunked prefill
    189
        pub prefix_cache_hits: u64,
    190
    }
    191
    192
    impl ChunkedPrefillStats {
    193
        /// Average chunk latency
    194
    2
        pub fn avg_chunk_latency_ms(&self) -> f64 {
    195
    2
            if self.chunks_processed == 0 {
    196
    1
                0.0
    197
            } else {
    198
    1
                self.total_chunk_latency_ms as f64 / self.chunks_processed as f64
    199
            }
    200
    2
        }
    201
    202
        /// Chunking rate (what % of sequences used chunking)
    203
    2
        pub fn chunking_rate(&self) -> f64 {
    204
    2
            let total = self.chunked_sequences + self.bypassed_sequences;
    205
    2
            if total == 0 {
    206
    1
                0.0
    207
            } else {
    208
    1
                self.chunked_sequences as f64 / total as f64
    209
            }
    210
    2
        }
    211
    }
    212
    213
    /// Chunked prefill scheduler
    214
    ///
    215
    /// Manages the chunked processing of long prompts, interleaving with decode
    216
    /// operations for improved TTFT across all requests.
    217
    pub struct ChunkedPrefillScheduler {
    218
        /// Configuration
    219
        config: ChunkedPrefillConfig,
    220
        /// Active chunked prefill states (seq_id -> state)
    221
        active_prefills: HashMap<u64, ChunkedPrefillState>,
    222
        /// Queue of sequences waiting for prefill
    223
        prefill_queue: VecDeque<u64>,
    224
        /// Statistics
    225
        stats: ChunkedPrefillStats,
    226
        /// Next sequence ID
    227
        next_seq_id: u64,
    228
    }
    229
    230
    impl ChunkedPrefillScheduler {
    231
        /// Create new chunked prefill scheduler
    232
    19
        pub fn new(config: ChunkedPrefillConfig) -> Self {
    233
    19
            Self {
    234
    19
                config,
    235
    19
                active_prefills: HashMap::new(),
    236
    19
                prefill_queue: VecDeque::new(),
    237
    19
                stats: ChunkedPrefillStats::default(),
    238
    19
                next_seq_id: 0,
    239
    19
            }
    240
    19
        }
    241
    242
        /// Submit a new sequence for prefill
    243
        ///
    244
        /// Returns the sequence ID and whether it will use chunked prefill.
    245
    15
        pub fn submit(&mut self, prompt_tokens: usize) -> (u64, bool) {
    246
    15
            let seq_id = self.next_seq_id;
    247
    15
            self.next_seq_id += 1;
    248
    249
    15
            let use_chunking = self.config.enabled && prompt_tokens >= self.config.min_prompt_length14;
    250
    251
    15
            if use_chunking {
    252
    13
                let state = ChunkedPrefillState::new(seq_id, prompt_tokens, self.config.chunk_size);
    253
    13
                self.active_prefills.insert(seq_id, state);
    254
    13
                self.prefill_queue.push_back(seq_id);
    255
    13
                self.stats.chunked_sequences += 1;
    256
    13
            } else {
    257
    2
                self.stats.bypassed_sequences += 1;
    258
    2
            }
    259
    260
    15
            (seq_id, use_chunking)
    261
    15
        }
    262
    263
        /// Get the next chunk to process
    264
        ///
    265
        /// Returns (seq_id, token_range) for the next prefill chunk, or None if
    266
        /// all prefills are complete or queue is empty.
    267
    6
        pub fn next_chunk(&mut self) -> Option<(u64, std::ops::Range<usize>)> {
    268
            // Find next sequence with remaining prefill work
    269
    6
            while let Some(&seq_id4) = self.prefill_queue.front() {
    270
    4
                if let Some(state) = self.active_prefills.get(&seq_id) {
    271
    4
                    if !state.is_complete() {
    272
    4
                        let range = state.next_chunk(self.config.chunk_size);
    273
    4
                        return Some((seq_id, range));
    274
    0
                    }
    275
    0
                }
    276
                // Remove completed or missing sequences from queue
    277
    0
                self.prefill_queue.pop_front();
    278
            }
    279
    2
            None
    280
    6
        }
    281
    282
        /// Record completion of a prefill chunk
    283
    6
        pub fn complete_chunk(&mut self, seq_id: u64, tokens_processed: usize, latency_ms: u64) {
    284
    6
            if let Some(state5) = self.active_prefills.get_mut(&seq_id) {
    285
    5
                state.advance(tokens_processed, latency_ms);
    286
    5
                self.stats.chunks_processed += 1;
    287
    5
                self.stats.total_chunk_latency_ms += latency_ms;
    288
    5
                self.stats.max_chunk_latency_ms = self.stats.max_chunk_latency_ms.max(latency_ms);
    289
    290
                // If complete, clean up
    291
    5
                if state.is_complete() {
    292
                    // Move to back of queue or remove
    293
    2
                    if let Some(pos) = self.prefill_queue.iter().position(|&id| id == seq_id) {
    294
    2
                        self.prefill_queue.remove(pos);
    295
    2
                    }0
    296
    3
                } else if self.config.boost_partial_prefill {
    297
    2
                    // Keep at front for priority
    298
    2
                } else {
    299
                    // Move to back of queue for round-robin
    300
    1
                    if let Some(pos) = self.prefill_queue.iter().position(|&id| id == seq_id) {
    301
    1
                        self.prefill_queue.remove(pos);
    302
    1
                        self.prefill_queue.push_back(seq_id);
    303
    1
                    }0
    304
                }
    305
    1
            }
    306
    6
        }
    307
    308
        /// Record a decode interleave (decode step between prefill chunks)
    309
    1
        pub fn record_decode_interleave(&mut self) {
    310
    1
            self.stats.decode_interleaves += 1;
    311
    1
        }
    312
    313
        /// Check if we should allow decode interleaving
    314
    3
        pub fn should_interleave_decode(&self) -> bool {
    315
    3
            self.config.allow_decode_interleave && !self.prefill_queue.is_empty()2
    316
    3
        }
    317
    318
        /// Get state for a sequence
    319
    2
        pub fn get_state(&self, seq_id: u64) -> Option<&ChunkedPrefillState> {
    320
    2
            self.active_prefills.get(&seq_id)
    321
    2
        }
    322
    323
        /// Check if sequence has pending prefill work
    324
    4
        pub fn has_pending_prefill(&self, seq_id: u64) -> bool {
    325
    4
            self.active_prefills
    326
    4
                .get(&seq_id)
    327
    4
                .is_some_and(|s| !s3.is_complete3())
    328
    4
        }
    329
    330
        /// Remove completed sequence
    331
    2
        pub fn remove(&mut self, seq_id: u64) -> Option<ChunkedPrefillState> {
    332
    2
            if let Some(pos1) = self.prefill_queue.iter().position(|&id| id1 == seq_id1) {
    333
    1
                self.prefill_queue.remove(pos);
    334
    1
            }
    335
    2
            self.active_prefills.remove(&seq_id)
    336
    2
        }
    337
    338
        /// Number of sequences with pending prefill
    339
    2
        pub fn pending_count(&self) -> usize {
    340
    2
            self.active_prefills
    341
    2
                .values()
    342
    2
                .filter(|s| !s0.is_complete0())
    343
    2
                .count()
    344
    2
        }
    345
    346
        /// Total prefill queue length
    347
    5
        pub fn queue_len(&self) -> usize {
    348
    5
            self.prefill_queue.len()
    349
    5
        }
    350
    351
        /// Get statistics
    352
    12
        pub fn stats(&self) -> &ChunkedPrefillStats {
    353
    12
            &self.stats
    354
    12
        }
    355
    356
        /// Get configuration
    357
    1
        pub fn config(&self) -> &ChunkedPrefillConfig {
    358
    1
            &self.config
    359
    1
        }
    360
    361
        /// Clear all state
    362
    1
        pub fn clear(&mut self) {
    363
    1
            self.active_prefills.clear();
    364
    1
            self.prefill_queue.clear();
    365
    1
        }
    366
    367
        /// Record prefix cache hit during prefill
    368
    2
        pub fn record_prefix_cache_hit(&mut self, tokens_saved: usize) {
    369
    2
            self.stats.prefix_cache_hits += tokens_saved as u64;
    370
    2
        }
    371
    }
    372
    373
    impl Default for ChunkedPrefillScheduler {
    374
    1
        fn default() -> Self {
    375
    1
            Self::new(ChunkedPrefillConfig::default())
    376
    1
        }
    377
    }
    378