ferrum_engine/pipeline/
chunked_prefill.rs

1//! Chunked Prefill Implementation
2//!
3//! This module implements chunked prefill, which splits long prompts into
4//! smaller chunks for processing. Benefits include:
5//!
6//! - Better memory efficiency for long prompts
7//! - Ability to interleave prefill with decode for better GPU utilization
8//! - Reduced latency for first token when processing long contexts
9
10use ferrum_interfaces::{
11    model_executor::{PrefillInput, PrefillOutput},
12    KvCacheHandle, ModelExecutor, TensorRef,
13};
14use ferrum_models::CandleTensorWrapper;
15use ferrum_types::{FerrumError, Result, TokenId};
16use std::sync::Arc;
17use tracing::{debug, info};
18
19/// Configuration for chunked prefill
20#[derive(Debug, Clone)]
21pub struct ChunkedPrefillConfig {
22    /// Maximum tokens per chunk
23    pub chunk_size: usize,
24    /// Minimum tokens to trigger chunking (below this, process as single chunk)
25    pub min_sequence_for_chunking: usize,
26    /// Whether to overlap chunks for better context
27    pub enable_overlap: bool,
28    /// Number of tokens to overlap between chunks
29    pub overlap_size: usize,
30}
31
32impl Default for ChunkedPrefillConfig {
33    fn default() -> Self {
34        Self {
35            chunk_size: 512,
36            min_sequence_for_chunking: 128,
37            enable_overlap: false,
38            overlap_size: 16,
39        }
40    }
41}
42
43/// State for chunked prefill processing
44#[derive(Debug)]
45pub struct ChunkedPrefillState {
46    /// Original input tokens
47    pub tokens: Vec<TokenId>,
48    /// Current chunk index
49    pub current_chunk: usize,
50    /// Total number of chunks
51    pub total_chunks: usize,
52    /// Tokens processed so far
53    pub tokens_processed: usize,
54    /// KV cache handle
55    pub kv_cache: Option<Arc<dyn KvCacheHandle>>,
56    /// Configuration
57    pub config: ChunkedPrefillConfig,
58}
59
60impl ChunkedPrefillState {
61    /// Create new chunked prefill state
62    pub fn new(tokens: Vec<TokenId>, config: ChunkedPrefillConfig) -> Self {
63        let chunk_size = config.chunk_size;
64        let total_chunks = if tokens.len() <= config.min_sequence_for_chunking {
65            1
66        } else {
67            tokens.len().div_ceil(chunk_size)
68        };
69
70        Self {
71            tokens,
72            current_chunk: 0,
73            total_chunks,
74            tokens_processed: 0,
75            kv_cache: None,
76            config,
77        }
78    }
79
80    /// Check if chunking is needed
81    pub fn needs_chunking(&self) -> bool {
82        self.tokens.len() > self.config.min_sequence_for_chunking
83    }
84
85    /// Check if all chunks have been processed
86    pub fn is_complete(&self) -> bool {
87        self.tokens_processed >= self.tokens.len()
88    }
89
90    /// Get the next chunk of tokens
91    pub fn next_chunk(&self) -> Option<&[TokenId]> {
92        if self.is_complete() {
93            return None;
94        }
95
96        let start = self.tokens_processed;
97        let end = (start + self.config.chunk_size).min(self.tokens.len());
98
99        if start < self.tokens.len() {
100            Some(&self.tokens[start..end])
101        } else {
102            None
103        }
104    }
105
106    /// Mark current chunk as processed
107    pub fn advance(&mut self, tokens_processed: usize) {
108        self.tokens_processed += tokens_processed;
109        self.current_chunk += 1;
110    }
111
112    /// Get progress as a fraction
113    pub fn progress(&self) -> f32 {
114        if self.tokens.is_empty() {
115            1.0
116        } else {
117            self.tokens_processed as f32 / self.tokens.len() as f32
118        }
119    }
120}
121
122/// Executor for chunked prefill operations
123pub struct ChunkedPrefillExecutor {
124    /// Configuration
125    config: ChunkedPrefillConfig,
126    /// Underlying model executor
127    model_executor: Arc<dyn ModelExecutor + Send + Sync>,
128}
129
130impl ChunkedPrefillExecutor {
131    /// Create new chunked prefill executor
132    pub fn new(
133        model_executor: Arc<dyn ModelExecutor + Send + Sync>,
134        config: ChunkedPrefillConfig,
135    ) -> Self {
136        info!(
137            "Creating ChunkedPrefillExecutor: chunk_size={}, min_seq={}",
138            config.chunk_size, config.min_sequence_for_chunking
139        );
140        Self {
141            config,
142            model_executor,
143        }
144    }
145
146    /// Create with default config
147    pub fn with_defaults(model_executor: Arc<dyn ModelExecutor + Send + Sync>) -> Self {
148        Self::new(model_executor, ChunkedPrefillConfig::default())
149    }
150
151    /// Execute prefill with chunking if needed
152    pub async fn execute(&self, tokens: Vec<TokenId>) -> Result<PrefillOutput> {
153        let mut state = ChunkedPrefillState::new(tokens, self.config.clone());
154
155        if !state.needs_chunking() {
156            // Process as single chunk
157            debug!("Processing {} tokens as single chunk", state.tokens.len());
158            return self.process_single_chunk(&state.tokens).await;
159        }
160
161        debug!(
162            "Processing {} tokens in {} chunks",
163            state.tokens.len(),
164            state.total_chunks
165        );
166
167        let mut last_output: Option<PrefillOutput> = None;
168
169        while let Some(chunk) = state.next_chunk() {
170            let chunk_len = chunk.len();
171            debug!(
172                "Processing chunk {}/{}: {} tokens (progress: {:.1}%)",
173                state.current_chunk + 1,
174                state.total_chunks,
175                chunk_len,
176                state.progress() * 100.0
177            );
178
179            let output = if state.kv_cache.is_some() {
180                // Continue from existing KV cache
181                self.process_continuation_chunk(chunk, state.kv_cache.as_ref().unwrap().clone())
182                    .await?
183            } else {
184                // First chunk
185                self.process_single_chunk(chunk).await?
186            };
187
188            state.kv_cache = Some(output.kv_cache.clone());
189            state.advance(chunk_len);
190            last_output = Some(output);
191        }
192
193        last_output.ok_or_else(|| FerrumError::internal("No output from chunked prefill"))
194    }
195
196    /// Process a single chunk as prefill
197    async fn process_single_chunk(&self, tokens: &[TokenId]) -> Result<PrefillOutput> {
198        // Convert tokens to tensor
199        let token_u32s: Vec<u32> = tokens.iter().map(|t| t.get()).collect();
200        let tensor = candle_core::Tensor::new(&token_u32s[..], &candle_core::Device::Cpu)
201            .map_err(|e| FerrumError::model(format!("Tensor error: {}", e)))?
202            .unsqueeze(0)
203            .map_err(|e| FerrumError::model(format!("Unsqueeze error: {}", e)))?;
204
205        let tensor_ref: TensorRef = Arc::new(CandleTensorWrapper::new(tensor));
206        let input = PrefillInput::new(tensor_ref);
207
208        self.model_executor.prefill(&input).await
209    }
210
211    /// Process a continuation chunk using existing KV cache
212    async fn process_continuation_chunk(
213        &self,
214        tokens: &[TokenId],
215        _kv_cache: Arc<dyn KvCacheHandle>,
216    ) -> Result<PrefillOutput> {
217        // For continuation, we use decode-like processing but with multiple tokens
218        // This is an optimization - some models support batched decode which is
219        // essentially the same as prefill with existing KV cache
220
221        // Convert tokens to tensor
222        let token_u32s: Vec<u32> = tokens.iter().map(|t| t.get()).collect();
223        let tensor = candle_core::Tensor::new(&token_u32s[..], &candle_core::Device::Cpu)
224            .map_err(|e| FerrumError::model(format!("Tensor error: {}", e)))?
225            .unsqueeze(0)
226            .map_err(|e| FerrumError::model(format!("Unsqueeze error: {}", e)))?;
227
228        let tensor_ref: TensorRef = Arc::new(CandleTensorWrapper::new(tensor));
229
230        // Use prefill with the input tensor and pass position info
231        // In a full implementation, we'd also pass the existing KV cache
232        let input = PrefillInput::new(tensor_ref);
233        let output = self.model_executor.prefill(&input).await?;
234
235        // Return with the updated KV cache
236        Ok(PrefillOutput::new(output.logits, output.kv_cache))
237    }
238
239    /// Get configuration
240    pub fn config(&self) -> &ChunkedPrefillConfig {
241        &self.config
242    }
243}
244
245impl std::fmt::Debug for ChunkedPrefillExecutor {
246    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
247        f.debug_struct("ChunkedPrefillExecutor")
248            .field("chunk_size", &self.config.chunk_size)
249            .field(
250                "min_sequence_for_chunking",
251                &self.config.min_sequence_for_chunking,
252            )
253            .finish()
254    }
255}
256
257// ============================================================================
258// Tests
259// ============================================================================
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264
265    #[test]
266    fn test_chunked_state_creation() {
267        let tokens: Vec<TokenId> = (0..100).map(|i| TokenId::new(i as u32)).collect();
268        let config = ChunkedPrefillConfig {
269            chunk_size: 32,
270            min_sequence_for_chunking: 50,
271            ..Default::default()
272        };
273
274        let state = ChunkedPrefillState::new(tokens, config);
275
276        assert!(state.needs_chunking());
277        assert_eq!(state.total_chunks, 4); // 100 / 32 = 3.125, rounded up to 4
278        assert!(!state.is_complete());
279    }
280
281    #[test]
282    fn test_chunked_state_no_chunking() {
283        let tokens: Vec<TokenId> = (0..30).map(|i| TokenId::new(i as u32)).collect();
284        let config = ChunkedPrefillConfig {
285            chunk_size: 32,
286            min_sequence_for_chunking: 50,
287            ..Default::default()
288        };
289
290        let state = ChunkedPrefillState::new(tokens, config);
291
292        assert!(!state.needs_chunking());
293        assert_eq!(state.total_chunks, 1);
294    }
295
296    #[test]
297    fn test_chunked_state_iteration() {
298        let tokens: Vec<TokenId> = (0..100).map(|i| TokenId::new(i as u32)).collect();
299        let config = ChunkedPrefillConfig {
300            chunk_size: 32,
301            min_sequence_for_chunking: 50,
302            ..Default::default()
303        };
304
305        let mut state = ChunkedPrefillState::new(tokens, config);
306
307        // First chunk: 0-31
308        let chunk1 = state.next_chunk().unwrap();
309        assert_eq!(chunk1.len(), 32);
310        assert_eq!(chunk1[0].get(), 0);
311        state.advance(32);
312
313        // Second chunk: 32-63
314        let chunk2 = state.next_chunk().unwrap();
315        assert_eq!(chunk2.len(), 32);
316        assert_eq!(chunk2[0].get(), 32);
317        state.advance(32);
318
319        // Third chunk: 64-95
320        let chunk3 = state.next_chunk().unwrap();
321        assert_eq!(chunk3.len(), 32);
322        state.advance(32);
323
324        // Fourth chunk: 96-99 (4 tokens)
325        let chunk4 = state.next_chunk().unwrap();
326        assert_eq!(chunk4.len(), 4);
327        state.advance(4);
328
329        // No more chunks
330        assert!(state.is_complete());
331        assert!(state.next_chunk().is_none());
332    }
333
334    #[test]
335    fn test_progress() {
336        let tokens: Vec<TokenId> = (0..100).map(|i| TokenId::new(i as u32)).collect();
337        let config = ChunkedPrefillConfig::default();
338        let mut state = ChunkedPrefillState::new(tokens, config);
339
340        assert_eq!(state.progress(), 0.0);
341
342        state.advance(50);
343        assert!((state.progress() - 0.5).abs() < 0.01);
344
345        state.advance(50);
346        assert!((state.progress() - 1.0).abs() < 0.01);
347    }
348}
ferrum_engine/pipeline/chunked_prefill.rs

ferrum_engine/pipeline/
chunked_prefill.rs
