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