gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: 12 bug fixes from comprehensive review — 51 tok/s verified on RTX 5090

Browse Source 
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>
This commit is contained in:
Gahow Wang
2026-05-23 14:13:43 +08:00
parent a67e724119
commit 986a289616
14 changed files with 285 additions and 292 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
crates/xserv-server/src/api.rs
View File

@@ -1,5 +1,6 @@
use axum::Extension;
use axum::Json;
use axum::http::StatusCode;
use axum::response::sse::{Event, KeepAlive, Sse};
use axum::response::{IntoResponse, Response};
use serde::{Deserialize, Serialize};
@@ -73,13 +74,13 @@ pub async fn chat_completions(
Json(req): Json<ChatRequest>,
) -> Response {
if req.stream == Some(true) {
chat_stream(state, req).into_response()
chat_stream(state, req)
} else {
chat_non_stream(state, req).await.into_response()
chat_non_stream(state, req).await
}
}
async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Json<serde_json::Value> {
async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
let id = format!("chatcmpl-{}", Uuid::new_v4());
let model_name = state.model_name.clone();
let created = unix_timestamp();
@@ -88,10 +89,21 @@ async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Json<serde_j
let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
let prompt_token_count = prompt_tokens.len();
let max_seq_len = state.max_seq_len;
if prompt_token_count >= max_seq_len {
return (StatusCode::BAD_REQUEST, Json(serde_json::json!({
"error": {
"message": format!("prompt is {} tokens, exceeds max_seq_len {}", prompt_token_count, max_seq_len),
"type": "invalid_request_error"
}
}))).into_response();
}
let max_tokens = req.max_tokens.min(max_seq_len - prompt_token_count);
let (tx, mut rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
let gen_req = GenerateRequest {
prompt_tokens,
max_tokens: req.max_tokens,
max_tokens,
sampling: sampling_params(&req),
sender: tx,
};
@@ -133,13 +145,13 @@ async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Json<serde_j
"completion_tokens": completion_token_count,
"total_tokens": prompt_token_count + completion_token_count
}
}))
})).into_response()
}
fn chat_stream(
state: Arc<AppState>,
req: ChatRequest,
) -> Sse<impl tokio_stream::Stream<Item = Result<Event, Infallible>>> {
) -> Response {
let id = format!("chatcmpl-{}", Uuid::new_v4());
let model_name = state.model_name.clone();
let created = unix_timestamp();
@@ -147,10 +159,21 @@ fn chat_stream(
let prompt = build_prompt(&req.messages);
let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
let max_seq_len = state.max_seq_len;
if prompt_tokens.len() >= max_seq_len {
return (StatusCode::BAD_REQUEST, Json(serde_json::json!({
"error": {
"message": format!("prompt is {} tokens, exceeds max_seq_len {}", prompt_tokens.len(), max_seq_len),
"type": "invalid_request_error"
}
}))).into_response();
}
let max_tokens = req.max_tokens.min(max_seq_len - prompt_tokens.len());
let (engine_tx, engine_rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
let gen_req = GenerateRequest {
prompt_tokens,
max_tokens: req.max_tokens,
max_tokens,
sampling: sampling_params(&req),
sender: engine_tx,
};
@@ -202,7 +225,7 @@ fn chat_stream(
}
});
Sse::new(ReceiverStream::new(sse_rx)).keep_alive(KeepAlive::default())
Sse::new(ReceiverStream::new(sse_rx)).keep_alive(KeepAlive::default()).into_response()
}
fn make_chunk(

29
crates/xserv-server/src/engine.rs
View File

@@ -34,7 +34,7 @@ struct Sequence {
generated_tokens: Vec<u32>,
max_tokens: usize,
sampling: SamplingParams,
kv_cache: GpuKVCache,
kv_cache: Option<GpuKVCache>,
sender: tokio::sync::mpsc::Sender<GenerateEvent>,
prefilled: bool,
eos_token_id: Option<u32>,
@@ -42,7 +42,7 @@ struct Sequence {
}
impl Engine {
pub fn load(model_dir: &Path, max_batch_size: usize) -> Self {
pub fn load(model_dir: &Path, max_batch_size: usize, max_seq_len: usize) -> Self {
xserv_cuda::device::set_device(0).unwrap();
let config = ModelConfig::from_file(&model_dir.join("config.json"));
eprintln!("[engine] Loading weights...");
@@ -50,13 +50,14 @@ impl Engine {
eprintln!("[engine] Loaded {} tensors", weights.len());
let model = Qwen3::from_weights(config.clone(), weights);
let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
let max_seq_len = 256;
eprintln!("[engine] Ready (max_batch_size={max_batch_size}, max_seq_len={max_seq_len})");
Self { model, config, tokenizer, max_batch_size, max_seq_len }
}
pub fn tokenizer(&self) -> &Tokenizer { &self.tokenizer }
pub fn max_seq_len(&self) -> usize { self.max_seq_len }
/// Main scheduler loop. Receives requests from channel, manages concurrent sequences.
pub fn run(&self, rx: mpsc::Receiver<GenerateRequest>) {
let mut waiting: VecDeque<Sequence> = VecDeque::new();
@@ -95,7 +96,7 @@ impl Engine {
let mut newly_prefilled = Vec::new();
for seq in running.iter_mut() {
if !seq.prefilled {
let logits = self.model.forward_gpu_cache(&seq.prompt_tokens, &mut seq.kv_cache);
let logits = self.model.forward_gpu_cache(&seq.prompt_tokens, seq.kv_cache.as_mut().unwrap());
let next = sample(&logits, &seq.sampling);
seq.generated_tokens.push(next);
seq.prefilled = true;
@@ -122,7 +123,7 @@ impl Engine {
// Single sequence: use per-seq path (no batching overhead)
let i = decode_indices[0];
let last = *running[i].generated_tokens.last().unwrap();
let logits = self.model.forward_gpu_cache(&[last], &mut running[i].kv_cache);
let logits = self.model.forward_gpu_cache(&[last], running[i].kv_cache.as_mut().unwrap());
let next = sample(&logits, &running[i].sampling);
running[i].generated_tokens.push(next);
self.emit_token(&running[i], next);
@@ -132,19 +133,12 @@ impl Engine {
.map(|&i| *running[i].generated_tokens.last().unwrap())
.collect();
let positions: Vec<usize> = decode_indices.iter()
.map(|&i| running[i].kv_cache.seq_len())
.map(|&i| running[i].kv_cache.as_ref().unwrap().seq_len())
.collect();
// Take caches out of sequences temporarily to satisfy borrow checker.
// The dummy caches (max_seq_len=1) are never used and dropped immediately
// after the real caches are restored. Minor alloc overhead (~microseconds).
// Take caches out of sequences via Option::take (no dummy allocation).
let mut caches: Vec<GpuKVCache> = decode_indices.iter()
.map(|&i| {
std::mem::replace(
&mut running[i].kv_cache,
GpuKVCache::new(&self.config, 1, DType::BF16, 0),
)
})
.map(|&i| running[i].kv_cache.take().unwrap())
.collect();
let mut cache_refs: Vec<&mut GpuKVCache> = caches.iter_mut().collect();
@@ -153,7 +147,7 @@ impl Engine {
// Put caches back: pop from end while iterating in reverse
drop(cache_refs);
for &i in decode_indices.iter().rev() {
running[i].kv_cache = caches.pop().unwrap();
running[i].kv_cache = Some(caches.pop().unwrap());
}
// Sample per-sequence from batched logits [B, vocab_size]
@@ -203,7 +197,7 @@ impl Engine {
generated_tokens: Vec::new(),
max_tokens: req.max_tokens,
sampling: req.sampling,
kv_cache,
kv_cache: Some(kv_cache),
sender: req.sender,
prefilled: false,
eos_token_id: self.tokenizer.eos_token_id(),
@@ -215,7 +209,6 @@ impl Engine {
let text = self.tokenizer.decode(&[token_id]);
if self.tokenizer.eos_token_id() == Some(token_id) {
let _ = seq.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
let _ = seq.sender.blocking_send(GenerateEvent::Done {
finish_reason: "stop".to_string(),
});

13
crates/xserv-server/src/main.rs
View File

@@ -10,13 +10,14 @@ 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]");
eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N]");
std::process::exit(1);
}
@@ -31,6 +32,11 @@ async fn main() {
.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())
@@ -43,7 +49,7 @@ async fn main() {
let model_dir_clone = model_dir.clone();
std::thread::spawn(move || {
let engine = engine::Engine::load(&model_dir_clone, max_batch);
let engine = engine::Engine::load(&model_dir_clone, max_batch, max_seq_len);
engine.run(rx);
});
@@ -51,6 +57,7 @@ async fn main() {
model_name,
engine_sender: Mutex::new(tx),
engine_tokenizer: Mutex::new(tokenizer),
max_seq_len,
});
let app = Router::new()
@@ -60,7 +67,7 @@ async fn main() {
.layer(Extension(state));
let addr = format!("0.0.0.0:{port}");
eprintln!("[server] Listening on {addr} (max_batch={max_batch})");
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();
}