mirror of
https://github.com/ruvnet/RuView
synced 2026-08-11 20:41:44 +00:00
ruv
501 lines
14 KiB
Rust
501 lines
14 KiB
Rust
//! Batch management for continuous batching
|
|
//!
|
|
//! This module provides structures for organizing requests into
|
|
//! efficient batches that can be processed together by the model.
|
|
|
|
use super::request::{RequestId, RunningRequest};
|
|
use std::collections::HashMap;
|
|
|
|
/// A request that has been prepared for batch processing
|
|
#[derive(Debug, Clone)]
|
|
pub struct BatchedRequest {
|
|
/// Request identifier
|
|
pub request_id: RequestId,
|
|
/// Token IDs to process in this batch iteration
|
|
pub token_ids: Vec<u32>,
|
|
/// Position offset for this request's tokens
|
|
pub position_offset: usize,
|
|
/// KV cache slot assignment
|
|
pub kv_cache_slot: usize,
|
|
/// Block table for paged attention
|
|
pub block_table: Vec<usize>,
|
|
/// Whether this is a prefill (true) or decode (false) request
|
|
pub is_prefill: bool,
|
|
/// Sequence length including new tokens
|
|
pub seq_len: usize,
|
|
/// Context length (tokens already in cache)
|
|
pub context_len: usize,
|
|
}
|
|
|
|
impl BatchedRequest {
|
|
/// Create a prefill batch request
|
|
pub fn prefill(
|
|
request_id: RequestId,
|
|
token_ids: Vec<u32>,
|
|
kv_cache_slot: usize,
|
|
block_table: Vec<usize>,
|
|
) -> Self {
|
|
let seq_len = token_ids.len();
|
|
Self {
|
|
request_id,
|
|
token_ids,
|
|
position_offset: 0,
|
|
kv_cache_slot,
|
|
block_table,
|
|
is_prefill: true,
|
|
seq_len,
|
|
context_len: 0,
|
|
}
|
|
}
|
|
|
|
/// Create a decode batch request
|
|
pub fn decode(
|
|
request_id: RequestId,
|
|
token_id: u32,
|
|
position_offset: usize,
|
|
kv_cache_slot: usize,
|
|
block_table: Vec<usize>,
|
|
context_len: usize,
|
|
) -> Self {
|
|
Self {
|
|
request_id,
|
|
token_ids: vec![token_id],
|
|
position_offset,
|
|
kv_cache_slot,
|
|
block_table,
|
|
is_prefill: false,
|
|
seq_len: context_len + 1,
|
|
context_len,
|
|
}
|
|
}
|
|
|
|
/// Get the number of tokens in this batch request
|
|
pub fn num_tokens(&self) -> usize {
|
|
self.token_ids.len()
|
|
}
|
|
}
|
|
|
|
/// A scheduled batch ready for model execution
|
|
#[derive(Debug)]
|
|
pub struct ScheduledBatch {
|
|
/// Batched requests
|
|
pub requests: Vec<BatchedRequest>,
|
|
/// Total tokens in this batch
|
|
pub total_tokens: usize,
|
|
/// Whether this batch contains any prefill operations
|
|
pub has_prefill: bool,
|
|
/// Whether this batch contains any decode operations
|
|
pub has_decode: bool,
|
|
/// Maximum sequence length in the batch
|
|
pub max_seq_len: usize,
|
|
/// Batch ID for tracking
|
|
pub batch_id: u64,
|
|
}
|
|
|
|
impl ScheduledBatch {
|
|
/// Create an empty batch
|
|
pub fn new(batch_id: u64) -> Self {
|
|
Self {
|
|
requests: Vec::new(),
|
|
total_tokens: 0,
|
|
has_prefill: false,
|
|
has_decode: false,
|
|
max_seq_len: 0,
|
|
batch_id,
|
|
}
|
|
}
|
|
|
|
/// Add a batched request
|
|
pub fn add(&mut self, request: BatchedRequest) {
|
|
self.total_tokens += request.num_tokens();
|
|
self.has_prefill |= request.is_prefill;
|
|
self.has_decode |= !request.is_prefill;
|
|
self.max_seq_len = self.max_seq_len.max(request.seq_len);
|
|
self.requests.push(request);
|
|
}
|
|
|
|
/// Check if the batch is empty
|
|
pub fn is_empty(&self) -> bool {
|
|
self.requests.is_empty()
|
|
}
|
|
|
|
/// Get the number of requests in the batch
|
|
pub fn len(&self) -> usize {
|
|
self.requests.len()
|
|
}
|
|
|
|
/// Get request IDs in the batch
|
|
pub fn request_ids(&self) -> Vec<RequestId> {
|
|
self.requests.iter().map(|r| r.request_id).collect()
|
|
}
|
|
|
|
/// Merge prefill and decode requests into a single batch
|
|
///
|
|
/// This is key for continuous batching efficiency - we can process
|
|
/// both prefill and decode requests in a single forward pass.
|
|
pub fn merge_prefill_decode(
|
|
prefill: Vec<BatchedRequest>,
|
|
decode: Vec<BatchedRequest>,
|
|
batch_id: u64,
|
|
) -> Self {
|
|
let mut batch = Self::new(batch_id);
|
|
|
|
// Add all prefill requests first
|
|
for req in prefill {
|
|
batch.add(req);
|
|
}
|
|
|
|
// Then add decode requests
|
|
for req in decode {
|
|
batch.add(req);
|
|
}
|
|
|
|
batch
|
|
}
|
|
|
|
/// Get the batch as separated prefill and decode requests
|
|
pub fn split_by_type(&self) -> (Vec<&BatchedRequest>, Vec<&BatchedRequest>) {
|
|
let prefill: Vec<_> = self.requests.iter().filter(|r| r.is_prefill).collect();
|
|
let decode: Vec<_> = self.requests.iter().filter(|r| !r.is_prefill).collect();
|
|
(prefill, decode)
|
|
}
|
|
|
|
/// Collect all input token IDs (padded for batch processing)
|
|
pub fn collect_input_ids(&self) -> Vec<Vec<u32>> {
|
|
self.requests.iter().map(|r| r.token_ids.clone()).collect()
|
|
}
|
|
|
|
/// Collect position offsets
|
|
pub fn collect_positions(&self) -> Vec<usize> {
|
|
self.requests.iter().map(|r| r.position_offset).collect()
|
|
}
|
|
|
|
/// Collect KV cache slots
|
|
pub fn collect_kv_slots(&self) -> Vec<usize> {
|
|
self.requests.iter().map(|r| r.kv_cache_slot).collect()
|
|
}
|
|
|
|
/// Calculate batch statistics
|
|
pub fn stats(&self) -> BatchStats {
|
|
let prefill_count = self.requests.iter().filter(|r| r.is_prefill).count();
|
|
let decode_count = self.requests.len() - prefill_count;
|
|
|
|
let prefill_tokens: usize = self
|
|
.requests
|
|
.iter()
|
|
.filter(|r| r.is_prefill)
|
|
.map(|r| r.num_tokens())
|
|
.sum();
|
|
|
|
BatchStats {
|
|
batch_id: self.batch_id,
|
|
total_requests: self.requests.len(),
|
|
prefill_requests: prefill_count,
|
|
decode_requests: decode_count,
|
|
total_tokens: self.total_tokens,
|
|
prefill_tokens,
|
|
decode_tokens: self.total_tokens - prefill_tokens,
|
|
max_seq_len: self.max_seq_len,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Statistics for a scheduled batch
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct BatchStats {
|
|
/// Batch identifier
|
|
pub batch_id: u64,
|
|
/// Total number of requests
|
|
pub total_requests: usize,
|
|
/// Number of prefill requests
|
|
pub prefill_requests: usize,
|
|
/// Number of decode requests
|
|
pub decode_requests: usize,
|
|
/// Total tokens in batch
|
|
pub total_tokens: usize,
|
|
/// Tokens from prefill operations
|
|
pub prefill_tokens: usize,
|
|
/// Tokens from decode operations
|
|
pub decode_tokens: usize,
|
|
/// Maximum sequence length
|
|
pub max_seq_len: usize,
|
|
}
|
|
|
|
/// Prefill task for iteration scheduling
|
|
#[derive(Debug, Clone)]
|
|
pub struct PrefillTask {
|
|
/// Request ID
|
|
pub request_id: RequestId,
|
|
/// Tokens to prefill
|
|
pub tokens: Vec<u32>,
|
|
/// Starting position
|
|
pub start_position: usize,
|
|
/// KV cache slot
|
|
pub kv_cache_slot: usize,
|
|
/// Block table
|
|
pub block_table: Vec<usize>,
|
|
}
|
|
|
|
/// Decode task for iteration scheduling
|
|
#[derive(Debug, Clone)]
|
|
pub struct DecodeTask {
|
|
/// Request ID
|
|
pub request_id: RequestId,
|
|
/// Token to decode from
|
|
pub input_token: u32,
|
|
/// Position offset
|
|
pub position: usize,
|
|
/// KV cache slot
|
|
pub kv_cache_slot: usize,
|
|
/// Block table
|
|
pub block_table: Vec<usize>,
|
|
/// Context length
|
|
pub context_len: usize,
|
|
}
|
|
|
|
/// Plan for a single iteration of the serving loop
|
|
#[derive(Debug)]
|
|
pub struct IterationPlan {
|
|
/// Prefill tasks to execute
|
|
pub prefill_tasks: Vec<PrefillTask>,
|
|
/// Decode tasks to execute
|
|
pub decode_tasks: Vec<DecodeTask>,
|
|
/// Requests that were evicted due to preemption
|
|
pub evicted_requests: Vec<RequestId>,
|
|
/// Requests that should be swapped out
|
|
pub swap_out_requests: Vec<RequestId>,
|
|
/// Requests that should be swapped in
|
|
pub swap_in_requests: Vec<RequestId>,
|
|
}
|
|
|
|
impl IterationPlan {
|
|
/// Create an empty iteration plan
|
|
pub fn empty() -> Self {
|
|
Self {
|
|
prefill_tasks: Vec::new(),
|
|
decode_tasks: Vec::new(),
|
|
evicted_requests: Vec::new(),
|
|
swap_out_requests: Vec::new(),
|
|
swap_in_requests: Vec::new(),
|
|
}
|
|
}
|
|
|
|
/// Check if there's work to do
|
|
pub fn has_work(&self) -> bool {
|
|
!self.prefill_tasks.is_empty() || !self.decode_tasks.is_empty()
|
|
}
|
|
|
|
/// Total number of requests to process
|
|
pub fn total_requests(&self) -> usize {
|
|
self.prefill_tasks.len() + self.decode_tasks.len()
|
|
}
|
|
|
|
/// Total tokens to process
|
|
pub fn total_tokens(&self) -> usize {
|
|
let prefill_tokens: usize = self.prefill_tasks.iter().map(|t| t.tokens.len()).sum();
|
|
let decode_tokens = self.decode_tasks.len(); // Each decode is 1 token
|
|
prefill_tokens + decode_tokens
|
|
}
|
|
|
|
/// Convert to a scheduled batch
|
|
pub fn to_scheduled_batch(&self, batch_id: u64) -> ScheduledBatch {
|
|
let prefill: Vec<BatchedRequest> = self
|
|
.prefill_tasks
|
|
.iter()
|
|
.map(|t| {
|
|
BatchedRequest::prefill(
|
|
t.request_id,
|
|
t.tokens.clone(),
|
|
t.kv_cache_slot,
|
|
t.block_table.clone(),
|
|
)
|
|
})
|
|
.collect();
|
|
|
|
let decode: Vec<BatchedRequest> = self
|
|
.decode_tasks
|
|
.iter()
|
|
.map(|t| {
|
|
BatchedRequest::decode(
|
|
t.request_id,
|
|
t.input_token,
|
|
t.position,
|
|
t.kv_cache_slot,
|
|
t.block_table.clone(),
|
|
t.context_len,
|
|
)
|
|
})
|
|
.collect();
|
|
|
|
ScheduledBatch::merge_prefill_decode(prefill, decode, batch_id)
|
|
}
|
|
}
|
|
|
|
/// Token budget for iteration scheduling
|
|
#[derive(Debug, Clone)]
|
|
pub struct TokenBudget {
|
|
/// Maximum tokens for prefill operations
|
|
pub max_prefill_tokens: usize,
|
|
/// Maximum tokens for decode operations (usually = max_batch_size)
|
|
pub max_decode_tokens: usize,
|
|
/// Maximum total tokens per iteration
|
|
pub max_total_tokens: usize,
|
|
/// Current prefill tokens allocated
|
|
pub prefill_tokens: usize,
|
|
/// Current decode tokens allocated
|
|
pub decode_tokens: usize,
|
|
}
|
|
|
|
impl TokenBudget {
|
|
/// Create a new token budget
|
|
pub fn new(max_prefill: usize, max_decode: usize, max_total: usize) -> Self {
|
|
Self {
|
|
max_prefill_tokens: max_prefill,
|
|
max_decode_tokens: max_decode,
|
|
max_total_tokens: max_total,
|
|
prefill_tokens: 0,
|
|
decode_tokens: 0,
|
|
}
|
|
}
|
|
|
|
/// Reset the budget for a new iteration
|
|
pub fn reset(&mut self) {
|
|
self.prefill_tokens = 0;
|
|
self.decode_tokens = 0;
|
|
}
|
|
|
|
/// Total tokens currently allocated
|
|
pub fn total_tokens(&self) -> usize {
|
|
self.prefill_tokens + self.decode_tokens
|
|
}
|
|
|
|
/// Remaining capacity for prefill tokens
|
|
pub fn remaining_prefill(&self) -> usize {
|
|
let from_prefill_limit = self.max_prefill_tokens.saturating_sub(self.prefill_tokens);
|
|
let from_total_limit = self.max_total_tokens.saturating_sub(self.total_tokens());
|
|
from_prefill_limit.min(from_total_limit)
|
|
}
|
|
|
|
/// Remaining capacity for decode tokens
|
|
pub fn remaining_decode(&self) -> usize {
|
|
let from_decode_limit = self.max_decode_tokens.saturating_sub(self.decode_tokens);
|
|
let from_total_limit = self.max_total_tokens.saturating_sub(self.total_tokens());
|
|
from_decode_limit.min(from_total_limit)
|
|
}
|
|
|
|
/// Try to allocate prefill tokens
|
|
pub fn try_allocate_prefill(&mut self, tokens: usize) -> bool {
|
|
if tokens <= self.remaining_prefill() {
|
|
self.prefill_tokens += tokens;
|
|
true
|
|
} else {
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Try to allocate a decode token
|
|
pub fn try_allocate_decode(&mut self) -> bool {
|
|
if self.remaining_decode() > 0 {
|
|
self.decode_tokens += 1;
|
|
true
|
|
} else {
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Check if budget is exhausted
|
|
pub fn is_exhausted(&self) -> bool {
|
|
self.total_tokens() >= self.max_total_tokens
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_batched_request() {
|
|
let prefill = BatchedRequest::prefill(RequestId::new(), vec![1, 2, 3, 4], 0, vec![0, 1]);
|
|
assert!(prefill.is_prefill);
|
|
assert_eq!(prefill.num_tokens(), 4);
|
|
assert_eq!(prefill.seq_len, 4);
|
|
|
|
let decode = BatchedRequest::decode(RequestId::new(), 5, 10, 1, vec![0, 1, 2], 10);
|
|
assert!(!decode.is_prefill);
|
|
assert_eq!(decode.num_tokens(), 1);
|
|
assert_eq!(decode.context_len, 10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_scheduled_batch() {
|
|
let mut batch = ScheduledBatch::new(1);
|
|
|
|
batch.add(BatchedRequest::prefill(
|
|
RequestId::new(),
|
|
vec![1, 2, 3],
|
|
0,
|
|
vec![],
|
|
));
|
|
batch.add(BatchedRequest::decode(RequestId::new(), 4, 5, 1, vec![], 5));
|
|
|
|
assert_eq!(batch.len(), 2);
|
|
assert!(batch.has_prefill);
|
|
assert!(batch.has_decode);
|
|
assert_eq!(batch.total_tokens, 4); // 3 prefill + 1 decode
|
|
|
|
let (prefill, decode) = batch.split_by_type();
|
|
assert_eq!(prefill.len(), 1);
|
|
assert_eq!(decode.len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_token_budget() {
|
|
let mut budget = TokenBudget::new(100, 32, 128);
|
|
|
|
assert!(budget.try_allocate_prefill(50));
|
|
assert_eq!(budget.prefill_tokens, 50);
|
|
assert_eq!(budget.remaining_prefill(), 50);
|
|
|
|
assert!(budget.try_allocate_decode());
|
|
assert_eq!(budget.decode_tokens, 1);
|
|
|
|
// Should fail - exceeds prefill limit
|
|
assert!(!budget.try_allocate_prefill(60));
|
|
|
|
budget.reset();
|
|
assert_eq!(budget.total_tokens(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_iteration_plan() {
|
|
let plan = IterationPlan {
|
|
prefill_tasks: vec![PrefillTask {
|
|
request_id: RequestId::new(),
|
|
tokens: vec![1, 2, 3, 4, 5],
|
|
start_position: 0,
|
|
kv_cache_slot: 0,
|
|
block_table: vec![],
|
|
}],
|
|
decode_tasks: vec![DecodeTask {
|
|
request_id: RequestId::new(),
|
|
input_token: 6,
|
|
position: 10,
|
|
kv_cache_slot: 1,
|
|
block_table: vec![],
|
|
context_len: 10,
|
|
}],
|
|
evicted_requests: vec![],
|
|
swap_out_requests: vec![],
|
|
swap_in_requests: vec![],
|
|
};
|
|
|
|
assert!(plan.has_work());
|
|
assert_eq!(plan.total_requests(), 2);
|
|
assert_eq!(plan.total_tokens(), 6); // 5 prefill + 1 decode
|
|
|
|
let batch = plan.to_scheduled_batch(42);
|
|
assert_eq!(batch.batch_id, 42);
|
|
assert_eq!(batch.len(), 2);
|
|
}
|
|
}
|