diff --git a/Cargo.toml b/Cargo.toml index 812c87b..986e4a0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ members = [ "crates/xserv-kernels", "crates/xserv-model", "crates/xserv-tokenizer", + "crates/xserv-server", ] [workspace.package] @@ -20,3 +21,6 @@ serde = { version = "1", features = ["derive"] } serde_json = "1" safetensors = "0.5" regex = "1" +tokio = { version = "1", features = ["full"] } +axum = "0.8" +uuid = { version = "1", features = ["v4"] } diff --git a/crates/xserv-server/Cargo.toml b/crates/xserv-server/Cargo.toml new file mode 100644 index 0000000..1c732e3 --- /dev/null +++ b/crates/xserv-server/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "xserv-server" +version.workspace = true +edition.workspace = true + +[[bin]] +name = "xserv-server" +path = "src/main.rs" + +[dependencies] +xserv-cuda = { path = "../xserv-cuda" } +xserv-tensor = { path = "../xserv-tensor" } +xserv-kernels = { path = "../xserv-kernels" } +xserv-model = { path = "../xserv-model" } +xserv-tokenizer = { path = "../xserv-tokenizer" } +half.workspace = true +serde.workspace = true +serde_json.workspace = true +tokio.workspace = true +axum.workspace = true +uuid.workspace = true diff --git a/crates/xserv-server/src/api.rs b/crates/xserv-server/src/api.rs new file mode 100644 index 0000000..e1c5dc1 --- /dev/null +++ b/crates/xserv-server/src/api.rs @@ -0,0 +1,115 @@ +use axum::Extension; +use axum::Json; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use uuid::Uuid; + +use crate::engine::{GenerateEvent, GenerateRequest}; +use crate::AppState; + +#[derive(Deserialize)] +pub struct ChatRequest { + #[serde(default)] + pub model: Option, + pub messages: Vec, + #[serde(default = "default_max_tokens")] + pub max_tokens: usize, +} + +#[derive(Deserialize)] +pub struct Message { + pub role: String, + pub content: String, +} + +fn default_max_tokens() -> usize { 256 } + +#[derive(Serialize)] +pub struct ModelsResponse { + object: &'static str, + data: Vec, +} + +#[derive(Serialize)] +pub struct ModelInfo { + id: String, + object: &'static str, + owned_by: &'static str, +} + +pub async fn health() -> &'static str { "ok" } + +pub async fn list_models(Extension(state): Extension>) -> Json { + Json(ModelsResponse { + object: "list", + data: vec![ModelInfo { + id: state.model_name.clone(), + object: "model", + owned_by: "xserv", + }], + }) +} + +pub async fn chat_completions( + Extension(state): Extension>, + Json(req): Json, +) -> Json { + let id = format!("chatcmpl-{}", Uuid::new_v4()); + let model_name = state.model_name.clone(); + let created = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + + // Prepare prompt tokens (MutexGuard scoped) + let prompt = build_prompt(&req.messages); + let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt); + + // Create channel and submit request (MutexGuard scoped) + let (tx, mut rx) = tokio::sync::mpsc::channel::(64); + let gen_req = GenerateRequest { + prompt_tokens, + max_tokens: req.max_tokens, + sender: tx, + }; + state.engine_sender.lock().unwrap().send(gen_req).unwrap(); + + // Now await — no MutexGuards held here + let mut content = String::new(); + let mut finish_reason = "length".to_string(); + while let Some(event) = rx.recv().await { + match event { + GenerateEvent::Token { text, .. } => content.push_str(&text), + GenerateEvent::Done { finish_reason: fr } => { finish_reason = fr; break; } + } + } + + Json(serde_json::json!({ + "id": id, + "object": "chat.completion", + "created": created, + "model": model_name, + "choices": [{ + "index": 0, + "message": { "role": "assistant", "content": content }, + "finish_reason": finish_reason, + }], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0 + } + })) +} + +fn build_prompt(messages: &[Message]) -> String { + let mut prompt = String::new(); + for msg in messages { + match msg.role.as_str() { + "system" => { prompt.push_str(&msg.content); prompt.push('\n'); } + "user" | "assistant" => { prompt.push_str(&msg.content); } + _ => {} + } + } + prompt +} diff --git a/crates/xserv-server/src/engine.rs b/crates/xserv-server/src/engine.rs new file mode 100644 index 0000000..79d6012 --- /dev/null +++ b/crates/xserv-server/src/engine.rs @@ -0,0 +1,76 @@ +use std::path::Path; +use xserv_model::{loader, GpuKVCache, ModelConfig, Qwen3}; +use xserv_model::qwen3::sample_greedy; +use xserv_tensor::{DType, Device}; +use xserv_tokenizer::Tokenizer; + +pub struct Engine { + model: Qwen3, + config: ModelConfig, + tokenizer: Tokenizer, +} + +pub struct GenerateRequest { + pub prompt_tokens: Vec, + pub max_tokens: usize, + pub sender: tokio::sync::mpsc::Sender, +} + +pub enum GenerateEvent { + Token { id: u32, text: String }, + Done { finish_reason: String }, +} + +impl Engine { + pub fn load(model_dir: &Path) -> Self { + xserv_cuda::device::set_device(0).unwrap(); + let config = ModelConfig::from_file(&model_dir.join("config.json")); + eprintln!("[engine] Loading weights..."); + let weights = loader::load_model_dir(model_dir, Device::Cuda(0)); + eprintln!("[engine] Loaded {} tensors", weights.len()); + let model = Qwen3::from_weights(config.clone(), weights); + let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json")); + eprintln!("[engine] Ready"); + Self { model, config, tokenizer } + } + + pub fn tokenizer(&self) -> &Tokenizer { + &self.tokenizer + } + + pub fn generate(&self, req: GenerateRequest) { + let max_seq = 256; + let mut cache = GpuKVCache::new(&self.config, max_seq, DType::BF16); + + let logits = self.model.forward_gpu_cache(&req.prompt_tokens, &mut cache); + let mut next = sample_greedy(&logits); + + for _ in 0..req.max_tokens { + let text = self.tokenizer.decode(&[next]); + if req.sender.blocking_send(GenerateEvent::Token { id: next, text }).is_err() { + return; + } + + if self.tokenizer.eos_token_id() == Some(next) { + let _ = req.sender.blocking_send(GenerateEvent::Done { + finish_reason: "stop".to_string(), + }); + return; + } + + if cache.seq_len() >= max_seq - 1 { + let _ = req.sender.blocking_send(GenerateEvent::Done { + finish_reason: "length".to_string(), + }); + return; + } + + let logits = self.model.forward_gpu_cache(&[next], &mut cache); + next = sample_greedy(&logits); + } + + let _ = req.sender.blocking_send(GenerateEvent::Done { + finish_reason: "length".to_string(), + }); + } +} diff --git a/crates/xserv-server/src/main.rs b/crates/xserv-server/src/main.rs new file mode 100644 index 0000000..9fcf688 --- /dev/null +++ b/crates/xserv-server/src/main.rs @@ -0,0 +1,62 @@ +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>, + pub engine_tokenizer: Mutex, +} + +#[tokio::main] +async fn main() { + let args: Vec = std::env::args().collect(); + if args.len() < 2 { + eprintln!("Usage: xserv-server [--port PORT]"); + 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 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::(1); + + 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 state = Arc::new(AppState { + model_name, + engine_sender: Mutex::new(tx), + engine_tokenizer: Mutex::new(tokenizer), + }); + + 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}"); + let listener = tokio::net::TcpListener::bind(&addr).await.unwrap(); + axum::serve(listener, app).await.unwrap(); +} diff --git a/docs/12-13-serving.md b/docs/12-13-serving.md new file mode 100644 index 0000000..3b63272 --- /dev/null +++ b/docs/12-13-serving.md @@ -0,0 +1,98 @@ +# Phase 12+13: Continuous Batching + HTTP API — Design Document (Milestone ③) + +## Goal + +实现 HTTP serving 层:接收请求、调度执行、streaming 返回结果。OpenAI 兼容 API。 + +由于当前是单请求引擎(无 multi-GPU、无并发),Phase 12 (continuous batching) 和 Phase 13 (HTTP API) 合并实现:先实现单请求 serving,scheduler 作为 placeholder 留待后续扩展。 + +## Architecture + +``` +Client (curl / OpenAI SDK) + │ + ▼ HTTP POST /v1/chat/completions +┌─────────────────────────────────────┐ +│ xserv-api (axum) │ +│ - Parse request │ +│ - Apply chat template │ +│ - Submit to engine via channel │ +│ - Stream SSE chunks from channel │ +└────────────┬────────────────────────┘ + │ InferenceRequest → mpsc channel + ▼ +┌─────────────────────────────────────┐ +│ xserv-engine (dedicated thread) │ +│ - Receive requests │ +│ - Run model forward (prefill+decode)│ +│ - Send tokens back via channel │ +└─────────────────────────────────────┘ +``` + +## Crates + +- `xserv-engine`: inference orchestration (model + cache + generate loop) +- `xserv-api`: HTTP server with axum + +Both in one binary: `xserv-server` + +## API Endpoints + +``` +POST /v1/chat/completions # main endpoint +GET /v1/models # list models +GET /health # health check +``` + +## Request/Response (OpenAI compatible) + +Request: +```json +{ + "model": "qwen3-8b", + "messages": [{"role": "user", "content": "Hello"}], + "stream": true, + "max_tokens": 256, + "temperature": 1.0 +} +``` + +SSE Response: +``` +data: {"id":"...","choices":[{"delta":{"content":"Hi"},"index":0}]} + +data: {"id":"...","choices":[{"delta":{},"finish_reason":"stop"}]} + +data: [DONE] +``` + +## Engine Design + +```rust +pub struct Engine { + model: Qwen3, + config: ModelConfig, + tokenizer: Tokenizer, +} + +impl Engine { + pub fn generate(&self, prompt_tokens: &[u32], params: &SamplingParams, + sender: mpsc::Sender) { ... } +} +``` + +Engine runs on a dedicated OS thread (avoids async/GPU conflicts). +API handlers communicate via `tokio::sync::mpsc` channels. + +## Sampling + +For this phase: greedy only (temperature=0 or 1 with argmax). +Top-k/top-p sampling added later. + +## Test Plan + +- [ ] curl streaming request → get SSE chunks +- [ ] Python OpenAI SDK client works +- [ ] /v1/models returns model info +- [ ] /health returns 200 +- [ ] Multiple sequential requests work