推理引擎中的Continuous Batching实现:从请求级到迭代级调度的演进与性能提升

📅 2026/7/14 10:34:08
推理引擎中的Continuous Batching实现:从请求级到迭代级调度的演进与性能提升
推理引擎中的Continuous Batching实现从请求级到迭代级调度的演进与性能提升一、当5个并发请求只使用了60%的GPU静态批处理的效率黑洞推理服务的默认调度方式是静态批处理等待N个请求到达后组成一个batch一起执行前向传播。问题在于当第N个请求迟迟不来时前N-1个请求的GPU计算资源处于闲置。更严重的是批次内不同请求的生成长短不一——短请求生成10个token完成后仍需等待长请求生成100个token才能释放batch槽位。实测数据显示5并发、平均生成长度50 token的场景下GPU利用率仅62%。37%的时间花在等待慢请求完成和填充padding token。Continuous Batching通过迭代级调度和动态批次重组将GPU利用率提升至85%以上。二、请求级与迭代级调度的架构对比sequenceDiagram participant R1 as 请求1(50 tokens) participant R2 as 请求2(10 tokens) participant R3 as 请求3(80 tokens) participant S as 调度器 participant G as GPU Note over R1,G: 静态批处理请求级 R1-S: 到达 R2-S: 到达 S-S: 等待窗口填满 R3-S: 到达 S-G: Batch([R1,R2,R3]) G-G: 迭代1-10 Note over R2: 完成(10 tokens) G-G: 迭代11-50(padding R2) Note over R1: 完成(50 tokens) G-G: 迭代51-80(padding R1,R2) Note over R3: 完成(80 tokens) Note over R1,G: Continuous Batching迭代级 R1-S: 到达 S-G: Batch([R1]) R2-S: 到达 S-G: Batch([R1,R2]) G-G: 迭代1-10 Note over R2: 完成 R3-S: 到达 S-G: Batch([R1,R3]) G-G: 迭代11-50 Note over R1: 完成 S-G: Batch([R3]) G-G: 迭代51-80 Note over R3: 完成静态批处理的效率公式GPU利用率 ∑实际计算时间 / (最大生成长度 × 迭代时间) 140 / (80 × 1.0) 140/240 58%Continuous Batching的优化效果GPU利用率 (104030) / (80 × 1.0) 80/80 100%理想情况三、Continuous Batching的Rust实现use std::collections::{BTreeMap, VecDeque}; use std::time::Instant; /// 单个推理请求的状态 struct InferenceRequest { id: u64, // 已生成的token数 generated_tokens: usize, // 最大生成token数 max_tokens: usize, // KV Cache引用 kv_cache: ArcKVCacheBlock, // 到达时间用于延迟SLA计算 arrival_time: Instant, // 优先级高优先级请求不等待batch填满 priority: RequestPriority, } #[derive(Debug, Clone, PartialEq, Eq)] enum RequestPriority { High, // 在线对话 Normal, // 普通请求 Low, // 批量离线推理 } /// KV Cache块迭代级别的缓存单元 /// 每个请求的KV Cache由多个块组成 struct KVCacheBlock { block_index: usize, // 是否正在被使用 in_use: AtomicBool, // 所属请求ID owner_id: AtomicU64, } /// 调度状态决定是否在本次迭代中包含请求 enum ScheduleDecision { /// 包含请求被加入当前batch Include, /// 跳过请求等待如等待KV Cache分配 Skip, /// 驱逐请求驱逐超时或显存不足 Evict, } /// Continuous Batching调度器 struct ContinuousBatchingScheduler { // 等待队列新到达的请求 waiting_queue: VecDequeInferenceRequest, // 运行中请求按优先级到达时间排列 running_requests: BTreeMapRequestPriority, VecInferenceRequest, // KV Cache块管理器 kv_cache_manager: KVCacheManager, // 最大batch大小由显存容量限制 max_batch_size: usize, // 最大token数所有请求的token总和上限 max_tokens_per_batch: usize, } impl ContinuousBatchingScheduler { /// 每个迭代调用构建当前batch /// 核心逻辑遍历运行中请求按优先级填充batch fn schedule_iteration(mut self) - VecInferenceRequest { let mut batch Vec::with_capacity(self.max_batch_size); let mut total_tokens 0usize; let now Instant::now(); // 第一阶段从运行中请求选择 // 按优先级顺序High → Normal → Low for priority in [RequestPriority::High, RequestPriority::Normal, RequestPriority::Low] { if let Some(requests) self.running_requests.get(priority) { for request in requests { // 检查batch容量 if batch.len() self.max_batch_size { break; } // 检查token预算 let new_tokens request.generated_tokens 1; if total_tokens new_tokens self.max_tokens_per_batch { break; } // 检查KV Cache是否可用 if self.kv_cache_manager.ensure_block(request.id) { batch.push(request); total_tokens new_tokens; } } } } // 第二阶段从等待队列补充 // 仅在batch未满时新请求才有机会加入 while batch.len() self.max_batch_size { if let Some(request) self.waiting_queue.pop_front() { // 新请求需要prefill阶段 let prefill_tokens 128; // 假设prompt长度 if total_tokens prefill_tokens self.max_tokens_per_batch { if self.kv_cache_manager.allocate_blocks(request.id, prefill_tokens) { batch.push(request); total_tokens prefill_tokens; // 移入运行中队列 self.add_to_running(request); } } else { // token预算不足放回等待队列 self.waiting_queue.push_front(request); break; } } else { break; } } batch } /// 迭代完成后更新请求状态 fn post_iteration(mut self, batch_results: VecBatchResult) { let mut completed Vec::new(); for result in batch_results { if result.is_finished { completed.push(result.request_id); // 释放KV Cache块 self.kv_cache_manager.release_blocks(result.request_id); } else { // 更新token计数 self.update_token_count(result.request_id); } } // 移除完成的请求 self.remove_completed(completed); } /// 驱逐策略显存不足时驱逐低优先级请求 fn evict_if_needed(mut self, target_free_blocks: usize) - Vecu64 { let free_blocks self.kv_cache_manager.free_block_count(); if free_blocks target_free_blocks { return vec![]; } let mut evicted Vec::new(); let needed target_free_blocks - free_blocks; let mut reclaimed 0usize; // 驱逐顺序Low → Normal → 等待队列中最老的 for priority in [RequestPriority::Low, RequestPriority::Normal] { if reclaimed needed { break; } if let Some(requests) self.running_requests.get_mut(priority) { // 驱逐最老的请求 requests.sort_by_key(|r| r.arrival_time); while let Some(request) requests.first() { if reclaimed needed { break; } let blocks self.kv_cache_manager.release_blocks(request.id); reclaimed blocks; evicted.push(request.id); requests.remove(0); // 记录驱逐事件用于监控 metrics::counter!(evicted_requests_total, priority format!({:?}, priority) ).increment(1); } } } evicted } fn add_to_running(mut self, request: InferenceRequest) { let priority request.priority.clone(); self.running_requests .entry(priority) .or_insert_with(Vec::new) .push(request); } fn remove_completed(mut self, ids: [u64]) { for requests in self.running_requests.values_mut() { requests.retain(|r| !ids.contains(r.id)); } } fn update_token_count(mut self, _request_id: u64) { // 更新请求的generated_tokens } } /// KV Cache块管理器 struct KVCacheManager { // 总块数由显存大小 / 单块大小决定 total_blocks: usize, // 空闲块队列 free_blocks: VecDequeusize, // 已分配映射request_id → allocated_blocks allocations: HashMapu64, Vecusize, } impl KVCacheManager { /// 为新请求分配KV Cache块 fn allocate_blocks(mut self, request_id: u64, token_count: usize) - bool { let blocks_needed (token_count BLOCK_SIZE - 1) / BLOCK_SIZE; if self.free_blocks.len() blocks_needed { return false; } let mut allocated Vec::with_capacity(blocks_needed); for _ in 0..blocks_needed { if let Some(block) self.free_blocks.pop_front() { allocated.push(block); } } self.allocations.insert(request_id, allocated); true } /// 确保请求有足够块用于下一次迭代 /// 每个迭代需要至少1个新块 fn ensure_block(mut self, request_id: u64) - bool { if let Some(blocks) self.allocations.get(request_id) { // 已有分配每个迭代消费一个块 return true; } // 首次分配至少需要1个块 self.allocate_blocks(request_id, 1) } fn release_blocks(mut self, request_id: u64) - usize { if let Some(blocks) self.allocations.remove(request_id) { let count blocks.len(); self.free_blocks.extend(blocks); count } else { 0 } } fn free_block_count(self) - usize { self.free_blocks.len() } } const BLOCK_SIZE: usize 256; // 每个块覆盖256个token struct BatchResult { request_id: u64, is_finished: bool, token_count: usize, } use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; use std::collections::HashMap;核心设计迭代级调度每个token生成步骤后重新组合batch优先级驱逐Low→Normal→等待队列——保护高优先级请求KV Cache块管理固定大小块支持分配和回收token预算控制防止单batch的token总数超限Prefill与Decode的解耦调度——Chunked Prefill的深层优化。Continuous Batching的传统实现将Prefill和Decode放在同一迭代中混合调度但这忽略了二者在计算特征上的根本差异Prefill是compute-bound处理整段prompt需要大量矩阵乘法Decode是memory-bound每次只生成一个token瓶颈在KV Cache读写。混合调度的后果是一个长prompt4096 tokens的Prefill会占用GPU数百毫秒阻塞所有Decode请求的响应。Chunked Prefill将长prompt的Prefill拆分为多个chunk如512 tokens/chunk让Prefill和Decode交替执行——每个chunk的Prefill耗时控制在10ms以内使Decode请求的最大等待延迟从最长Prefill时间降为单个chunk时间。工程实现上Chunk拆分不只是均匀切分还需要在Transformer的Layer边界对齐——因为Attention计算的中间结果QKV投影不能跨chunk分割必须在Layer之间拆分才能保证计算正确性。更深一层的优化是Prefill和Decode使用不同的CUDA StreamPrefill使用高优先级Stream抢占计算资源Decode使用低优先级Stream在空闲时填充实现计算和内存访问的流水线化。这一优化在vLLM中被称为Prefix Caching Chunked Prefill实测可将长prompt场景的TTFTTime To First Token降低40-60%。另外值得注意的是调度抖动Scheduling JitterLinux内核的调度器可能在任意时刻抢占GPU推理进程的CPU核心导致一次迭代的延迟在0.5ms基线基础上增加0-2ms——对于高频迭代每10ms一次这种抖动会放大为20%的延迟波动。解决方案是通过taskset将推理进程绑定到隔离的CPU核心并使用SCHED_FIFO实时调度策略消除内核抢占。四、Continuous Batching的边界条件优势GPU利用率从60%提升至85%无padding浪费减少短请求的尾部延迟——完成即释放槽位支持新请求随时加入——无需等待batch窗口填满代价调度开销每迭代需重组batch约0.1-0.5ms/迭代KV Cache碎片化块级别分配导致内碎片实现复杂度跃升调度器需维护请求状态机不适用场景单请求推理服务无需batch调度请求生成长度高度一致静态batch的效率接近continuous极度延迟敏感1ms级调度开销不可接受五、总结Continuous Batching通过迭代级调度消除padding浪费GPU利用率提升约20-30%KV Cache块管理是调度的核心基础设施——分配/回收延迟直接决定batch重组效率优先级驱逐保护SLA关键请求Low→Normal→Waiting的驱逐顺序最小化业务影响max_tokens_per_batch的软限制兼顾GPU利用率和显存安全调度开销控制在0.5ms/迭代以下是工程可行性的前提