phase 12: implement real continuous batching scheduler
Rewrote engine.rs from scratch: - Scheduler loop: admit → prefill → decode → finish → check new requests - Multiple sequences run concurrently (max_batch_size configurable) - Each sequence has independent GpuKVCache - Non-blocking try_recv() for new requests during decode iterations - Dynamic join: new requests enter batch immediately, don't wait for others Verified with concurrent test (tools/test_concurrent.py): - 3 concurrent requests: wall_time=3.8s, concurrency_ratio=2.82x ✓ - 5 concurrent requests: wall_time=6.1s, concurrency_ratio=4.04x ✓ - All outputs are coherent and correct Design doc (docs/12-continuous-batching.md) fully rewritten with: - Detailed scheduler loop pseudocode - Data structures (Sequence, Scheduler) - Acceptance criteria with specific test cases - Clear separation from Phase 13 (HTTP layer) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -8,7 +8,7 @@ use engine::GenerateRequest;
|
||||
|
||||
pub struct AppState {
|
||||
pub model_name: String,
|
||||
pub engine_sender: Mutex<mpsc::SyncSender<GenerateRequest>>,
|
||||
pub engine_sender: Mutex<mpsc::Sender<GenerateRequest>>,
|
||||
pub engine_tokenizer: Mutex<xserv_tokenizer::Tokenizer>,
|
||||
}
|
||||
|
||||
@@ -16,7 +16,7 @@ pub struct AppState {
|
||||
async fn main() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
if args.len() < 2 {
|
||||
eprintln!("Usage: xserv-server <model-dir> [--port PORT]");
|
||||
eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N]");
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
@@ -26,21 +26,25 @@ async fn main() {
|
||||
.and_then(|i| args.get(i + 1))
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(8080);
|
||||
let max_batch: usize = args.iter()
|
||||
.position(|a| a == "--max-batch")
|
||||
.and_then(|i| args.get(i + 1))
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(4);
|
||||
|
||||
let model_name = model_dir.file_name()
|
||||
.map(|n| n.to_string_lossy().to_string())
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
|
||||
let tokenizer = xserv_tokenizer::Tokenizer::from_file(&model_dir.join("tokenizer.json"));
|
||||
let (tx, rx) = mpsc::sync_channel::<GenerateRequest>(1);
|
||||
|
||||
// Unbounded channel: allows multiple requests to queue up
|
||||
let (tx, rx) = mpsc::channel::<GenerateRequest>();
|
||||
|
||||
let model_dir_clone = model_dir.clone();
|
||||
std::thread::spawn(move || {
|
||||
let engine = engine::Engine::load(&model_dir_clone);
|
||||
eprintln!("[engine] Listening for requests...");
|
||||
while let Ok(req) = rx.recv() {
|
||||
engine.generate(req);
|
||||
}
|
||||
let engine = engine::Engine::load(&model_dir_clone, max_batch);
|
||||
engine.run(rx);
|
||||
});
|
||||
|
||||
let state = Arc::new(AppState {
|
||||
@@ -56,7 +60,7 @@ async fn main() {
|
||||
.layer(Extension(state));
|
||||
|
||||
let addr = format!("0.0.0.0:{port}");
|
||||
eprintln!("[server] Listening on {addr}");
|
||||
eprintln!("[server] Listening on {addr} (max_batch={max_batch})");
|
||||
let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user