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

3
crates/xserv-cuda/src/device.rs
View File

@@ -46,7 +46,8 @@ pub fn current_device() -> Result<u32> {
pub fn device_info(device: u32) -> Result<DeviceInfo> {
// Heap-allocate oversized buffer for cudaDeviceProp (layout varies by CUDA version).
let mut prop_buf = vec![0u8; 16384];
// CUDA 12.x struct is ~5-6 KB; use 32 KB to guard against future growth.
let mut prop_buf = vec![0u8; 32768];
error::check(unsafe {
ffi::cudaGetDeviceProperties(prop_buf.as_mut_ptr(), device as i32)
})?;

2
crates/xserv-kernels/src/embedding.rs
View File

@@ -26,7 +26,7 @@ pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor {
num_tokens * std::mem::size_of::<u32>(),
)
};
let mut ids_gpu = GpuBuffer::alloc(ids_bytes.len()).expect("alloc token_ids");
let mut ids_gpu = xserv_cuda::allocator::cached_alloc(ids_bytes.len()).expect("alloc token_ids");
ids_gpu.copy_from_host(ids_bytes).unwrap();
let out = Tensor::empty(&[num_tokens, hidden_size], table.dtype(), table.device());

36
crates/xserv-kernels/src/gemm.rs
View File

@@ -1,3 +1,4 @@
use std::cell::RefCell;
use std::ffi::c_void;
use xserv_cuda::error::{self, Result};
use xserv_tensor::{DType, Device, Tensor};
@@ -82,6 +83,23 @@ impl Drop for CublasContext {
}
}
thread_local! {
static CUBLAS_CTX: RefCell<CublasContext> = RefCell::new(
CublasContext::new().expect("failed to create thread-local cuBLAS handle")
);
}
/// Borrow the thread-local cuBLAS handle for the duration of a closure.
fn with_cublas<F, R>(f: F) -> R
where
F: FnOnce(CublasHandle) -> R,
{
CUBLAS_CTX.with(|cell| {
let ctx = cell.borrow();
f(ctx.handle)
})
}
/// Matrix multiplication: C = A @ B
/// A: [M, K], B: [K, N], C: [M, N]
/// All tensors must be contiguous and on the same GPU.
@@ -143,7 +161,6 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
// cuBLAS uses column-major, but we have row-major tensors.
// Trick: compute C^T = B^T @ A^T, which gives us C in row-major.
// cuBLAS sees our row-major data as column-major transposed.
let ctx = CublasContext::new().unwrap();
let alpha = 1.0f32;
let beta = 0.0f32;
@@ -153,12 +170,12 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
_ => panic!("unsupported dtype for cuBLAS GEMM"),
};
unsafe {
cublasSetStream_v2(ctx.handle, null_stream);
with_cublas(|handle| unsafe {
cublasSetStream_v2(handle, null_stream);
// Row-major trick: swap A/B and transpose flags
// C(row-major) = A @ B <=> C^T(col-major) = B^T @ A^T
error::check(cublasGemmEx(
ctx.handle,
handle,
CUBLAS_OP_N, CUBLAS_OP_N,
n as i32, m as i32, k as i32,
&alpha as *const f32 as *const c_void,
@@ -169,7 +186,7 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
CUBLAS_COMPUTE_32F,
-1, // default algo
)).expect("cuBLAS GEMM failed");
}
});
}
}
}
@@ -221,12 +238,11 @@ pub fn batched_matmul(a: &Tensor, b: &Tensor) -> Tensor {
let stride_b = (k * n) as i64;
let stride_c = (m * n) as i64;
let ctx = CublasContext::new().unwrap();
unsafe {
cublasSetStream_v2(ctx.handle, std::ptr::null_mut());
with_cublas(|handle| unsafe {
cublasSetStream_v2(handle, std::ptr::null_mut());
// Row-major trick: C = A @ B ⟺ C^T = B^T @ A^T (col-major)
error::check(cublasGemmStridedBatchedEx(
ctx.handle,
handle,
CUBLAS_OP_N, CUBLAS_OP_N,
n as i32, m as i32, k as i32,
&alpha as *const f32 as *const c_void,
@@ -238,6 +254,6 @@ pub fn batched_matmul(a: &Tensor, b: &Tensor) -> Tensor {
CUBLAS_COMPUTE_32F,
-1,
)).expect("cuBLAS batched GEMM failed");
}
});
c
}

2
crates/xserv-kernels/src/rope.rs
View File

@@ -58,7 +58,7 @@ pub fn rope_inplace(x: &Tensor, cache: &RopeCache, positions: &[u32]) {
num_tokens * std::mem::size_of::<u32>(),
)
};
let mut pos_gpu = GpuBuffer::alloc(pos_bytes.len()).expect("alloc positions");
let mut pos_gpu = xserv_cuda::allocator::cached_alloc(pos_bytes.len()).expect("alloc positions");
pos_gpu.copy_from_host(pos_bytes).unwrap();
unsafe {

4
crates/xserv-model/src/qwen3.rs
View File

@@ -42,7 +42,7 @@ impl Qwen3 {
let lm_head_raw = take(&mut w, "lm_head.weight");
let rope_cache = RopeCache::new(
config.max_seq_len().min(8192), // limit for memory
config.max_seq_len(),
config.head_dim(),
config.rope_theta.unwrap_or(1_000_000.0) as f32,
);
@@ -453,7 +453,7 @@ fn concat_rows(rows: &[Tensor]) -> Tensor {
// Allocate output [B, cols] and copy each row into it
let total_bytes = batch * row_bytes;
let mut out_buf = xserv_cuda::GpuBuffer::alloc(total_bytes).expect("alloc concat_rows");
let mut out_buf = xserv_cuda::allocator::cached_alloc(total_bytes).expect("alloc concat_rows");
for (b, row) in rows.iter().enumerate() {
assert_eq!(row.shape(), &[1, cols]);

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

6
crates/xserv-tensor/src/storage.rs
View File

@@ -73,7 +73,7 @@ impl Storage {
match (current, target) {
(Device::Cpu, Device::Cuda(dev)) => {
let cpu_data = self.as_cpu_bytes();
let mut buf = GpuBuffer::alloc(cpu_data.len())?;
let mut buf = xserv_cuda::allocator::cached_alloc(cpu_data.len())?;
buf.copy_from_host(cpu_data)?;
Ok(Storage::cuda(buf, dev))
}
@@ -85,7 +85,7 @@ impl Storage {
}
(Device::Cuda(_), Device::Cuda(dev)) => {
let src = self.gpu_buffer();
let mut dst = GpuBuffer::alloc(src.len())?;
let mut dst = xserv_cuda::allocator::cached_alloc(src.len())?;
dst.copy_from_device(src)?;
Ok(Storage::cuda(dst, dev))
}
@@ -98,7 +98,7 @@ impl Storage {
match self.0.as_ref() {
StorageInner::Cpu { data } => Ok(Storage::cpu(data.clone())),
StorageInner::Cuda { buffer, device } => {
let mut dst = GpuBuffer::alloc(buffer.len())?;
let mut dst = xserv_cuda::allocator::cached_alloc(buffer.len())?;
dst.copy_from_device(buffer)?;
Ok(Storage::cuda(dst, *device))
}

23
crates/xserv-tokenizer/src/bpe.rs
View File

@@ -51,8 +51,6 @@ impl Tokenizer {
let tj: TokenizerJson = serde_json::from_str(&data)
.unwrap_or_else(|e| panic!("failed to parse tokenizer.json: {e}"));
let byte_fallback = tj.model.byte_fallback;
// Build encoder: token bytes → ID
// All HF tokenizers use GPT-2 byte-to-unicode mapping for vocab keys.
let mut encoder = HashMap::new();
@@ -170,16 +168,25 @@ impl Tokenizer {
}
// Fall back to per-byte encoding
let word_bytes: Vec<u8> = word.bytes().collect();
let mut token_ids: Vec<u32> = word_bytes.iter().map(|&b| {
let mut token_ids: Vec<u32> = word_bytes.iter().filter_map(|&b| {
if let Some(&id) = self.encoder.get(&vec![b]) {
id
Some(id)
} else if self.byte_fallback {
let hex_token = format!("<0x{:02X}>", b);
*self.special_tokens.get(&hex_token).unwrap_or_else(|| {
panic!("byte 0x{b:02X} not in vocab and no fallback token {hex_token}")
})
if let Some(&id) = self.special_tokens.get(&hex_token) {
Some(id)
} else if let Some(&id) = self.encoder.get(hex_token.as_bytes()) {
Some(id)
} else if let Some(&unk_id) = self.special_tokens.get("<unk>") {
eprintln!("warning: byte 0x{b:02X} not in vocab, using <unk> token");
Some(unk_id)
} else {
eprintln!("warning: byte 0x{b:02X} not in vocab and no fallback token, using token 0");
Some(0)
}
} else {
panic!("byte {b} (0x{b:02X}) not in vocab")
eprintln!("warning: byte {b} (0x{b:02X}) not in vocab, skipping");
None
}
}).collect();