Gahow Wang
986a289616
P0 fixes (blocking usability): - FIX-01: thread-local cuBLAS handle (was creating/destroying per matmul) - FIX-16: EOS token no longer leaks into API responses - FIX-17: max_seq_len configurable via --max-seq-len (default 2048, was hardcoded 256) - FIX-18: max_tokens clamped to available seq space, prompt overflow returns 400 P1 fixes (bugs & performance): - FIX-07: CachingAllocator wired into all hot paths (to_device, embedding, rope, concat) - FIX-08: CudaDeviceProp buffer increased to 32KB for CUDA 12.9 safety - FIX-09: tokenizer byte_fallback graceful degradation (was panic) - FIX-19: causal mask uses -INFINITY instead of -1e9 (BF16 supports inf) - FIX-20: LayerNorm rewritten to numerically stable two-pass algorithm - FIX-21: min block size guard (32 threads) for LayerNorm/RMSNorm launches P2 fixes (improvements): - FIX-22: Option<GpuKVCache> + take() eliminates dummy KV cache allocations - FIX-23: RoPE cache no longer artificially capped at 8192 positions Verified on dash5 (RTX 5090): 51 tok/s batch=1, 74 tok/s 2-concurrent, 1.7-3.3x HF transformers. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
74 lines
2.4 KiB
Rust
74 lines
2.4 KiB
Rust
mod api;
|
|
mod engine;
|
|
|
|
use axum::{routing::{get, post}, Extension, Router};
|
|
use std::path::PathBuf;
|
|
use std::sync::{mpsc, Arc, Mutex};
|
|
use engine::GenerateRequest;
|
|
|
|
pub struct AppState {
|
|
pub model_name: String,
|
|
pub engine_sender: Mutex<mpsc::Sender<GenerateRequest>>,
|
|
pub engine_tokenizer: Mutex<xserv_tokenizer::Tokenizer>,
|
|
pub max_seq_len: usize,
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
let args: Vec<String> = std::env::args().collect();
|
|
if args.len() < 2 {
|
|
eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N]");
|
|
std::process::exit(1);
|
|
}
|
|
|
|
let model_dir = PathBuf::from(&args[1]);
|
|
let port: u16 = args.iter()
|
|
.position(|a| a == "--port")
|
|
.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 max_seq_len: usize = args.iter()
|
|
.position(|a| a == "--max-seq-len")
|
|
.and_then(|i| args.get(i + 1))
|
|
.and_then(|s| s.parse().ok())
|
|
.unwrap_or(2048);
|
|
|
|
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"));
|
|
|
|
// 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, max_batch, max_seq_len);
|
|
engine.run(rx);
|
|
});
|
|
|
|
let state = Arc::new(AppState {
|
|
model_name,
|
|
engine_sender: Mutex::new(tx),
|
|
engine_tokenizer: Mutex::new(tokenizer),
|
|
max_seq_len,
|
|
});
|
|
|
|
let app = Router::new()
|
|
.route("/health", get(api::health))
|
|
.route("/v1/models", get(api::list_models))
|
|
.route("/v1/chat/completions", post(api::chat_completions))
|
|
.layer(Extension(state));
|
|
|
|
let addr = format!("0.0.0.0:{port}");
|
|
eprintln!("[server] Listening on {addr} (max_batch={max_batch}, max_seq_len={max_seq_len})");
|
|
let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
|
|
axum::serve(listener, app).await.unwrap();
|
|
}
|