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:
2026-05-22 13:44:26 +08:00
parent 7d05ececa0
commit d8493bd70f
5 changed files with 348 additions and 100 deletions

View File

@@ -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();
}