diff --git a/crates/xserv-cuda/src/device.rs b/crates/xserv-cuda/src/device.rs index a1cd8a6..5cb60bf 100644 --- a/crates/xserv-cuda/src/device.rs +++ b/crates/xserv-cuda/src/device.rs @@ -46,7 +46,8 @@ pub fn current_device() -> Result { pub fn device_info(device: u32) -> Result { // 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) })?; diff --git a/crates/xserv-kernels/src/embedding.rs b/crates/xserv-kernels/src/embedding.rs index 7b56cc9..a992271 100644 --- a/crates/xserv-kernels/src/embedding.rs +++ b/crates/xserv-kernels/src/embedding.rs @@ -26,7 +26,7 @@ pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor { num_tokens * std::mem::size_of::(), ) }; - 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()); diff --git a/crates/xserv-kernels/src/gemm.rs b/crates/xserv-kernels/src/gemm.rs index afe39f7..68949e6 100644 --- a/crates/xserv-kernels/src/gemm.rs +++ b/crates/xserv-kernels/src/gemm.rs @@ -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 = 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: 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 } diff --git a/crates/xserv-kernels/src/rope.rs b/crates/xserv-kernels/src/rope.rs index f2b3965..830a061 100644 --- a/crates/xserv-kernels/src/rope.rs +++ b/crates/xserv-kernels/src/rope.rs @@ -58,7 +58,7 @@ pub fn rope_inplace(x: &Tensor, cache: &RopeCache, positions: &[u32]) { num_tokens * std::mem::size_of::(), ) }; - 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 { diff --git a/crates/xserv-model/src/qwen3.rs b/crates/xserv-model/src/qwen3.rs index e1e32c9..c98279f 100644 --- a/crates/xserv-model/src/qwen3.rs +++ b/crates/xserv-model/src/qwen3.rs @@ -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]); diff --git a/crates/xserv-server/src/api.rs b/crates/xserv-server/src/api.rs index 3edfc77..7dcf069 100644 --- a/crates/xserv-server/src/api.rs +++ b/crates/xserv-server/src/api.rs @@ -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, ) -> 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, req: ChatRequest) -> Json { +async fn chat_non_stream(state: Arc, 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, req: ChatRequest) -> Json= 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::(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, req: ChatRequest) -> Json, req: ChatRequest, -) -> Sse>> { +) -> 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::(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( diff --git a/crates/xserv-server/src/engine.rs b/crates/xserv-server/src/engine.rs index 69712d9..98f4001 100644 --- a/crates/xserv-server/src/engine.rs +++ b/crates/xserv-server/src/engine.rs @@ -34,7 +34,7 @@ struct Sequence { generated_tokens: Vec, max_tokens: usize, sampling: SamplingParams, - kv_cache: GpuKVCache, + kv_cache: Option, sender: tokio::sync::mpsc::Sender, prefilled: bool, eos_token_id: Option, @@ -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) { let mut waiting: VecDeque = 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 = 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 = 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(), }); diff --git a/crates/xserv-server/src/main.rs b/crates/xserv-server/src/main.rs index 72bd506..bf1605c 100644 --- a/crates/xserv-server/src/main.rs +++ b/crates/xserv-server/src/main.rs @@ -10,13 +10,14 @@ pub struct AppState { pub model_name: String, pub engine_sender: Mutex>, pub engine_tokenizer: Mutex, + pub max_seq_len: usize, } #[tokio::main] async fn main() { let args: Vec = std::env::args().collect(); if args.len() < 2 { - eprintln!("Usage: xserv-server [--port PORT] [--max-batch N]"); + eprintln!("Usage: xserv-server [--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(); } diff --git a/crates/xserv-tensor/src/storage.rs b/crates/xserv-tensor/src/storage.rs index a854ce7..eb5aaf0 100644 --- a/crates/xserv-tensor/src/storage.rs +++ b/crates/xserv-tensor/src/storage.rs @@ -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)) } diff --git a/crates/xserv-tokenizer/src/bpe.rs b/crates/xserv-tokenizer/src/bpe.rs index 08b6ca9..fa65c3e 100644 --- a/crates/xserv-tokenizer/src/bpe.rs +++ b/crates/xserv-tokenizer/src/bpe.rs @@ -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 = word.bytes().collect(); - let mut token_ids: Vec = word_bytes.iter().map(|&b| { + let mut token_ids: Vec = 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("") { + eprintln!("warning: byte 0x{b:02X} not in vocab, using 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(); diff --git a/csrc/attention/causal_mask.cu b/csrc/attention/causal_mask.cu index e6a60a3..e01da7a 100644 --- a/csrc/attention/causal_mask.cu +++ b/csrc/attention/causal_mask.cu @@ -27,8 +27,7 @@ __global__ void causal_mask_bf16( int col = blockIdx.x * blockDim.x + threadIdx.x; if (col < cols && col > row + offset) { - // BF16 doesn't have proper -inf literal, use a very large negative - scores[batch_idx * rows * cols + row * cols + col] = __float2bfloat16(-1e9f); + scores[batch_idx * rows * cols + row * cols + col] = __float2bfloat16(-INFINITY); } } diff --git a/csrc/normalization/layernorm.cu b/csrc/normalization/layernorm.cu index f4ad2d2..dcae11c 100644 --- a/csrc/normalization/layernorm.cu +++ b/csrc/normalization/layernorm.cu @@ -14,27 +14,34 @@ __global__ void layernorm_f32( const float* x_row = x + row * hidden_size; float* out_row = out + row * hidden_size; - // Welford online: compute mean and variance in one pass + // Pass 1: compute mean float local_sum = 0.0f; - float local_sum_sq = 0.0f; for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) { - float v = x_row[i]; - local_sum += v; - local_sum_sq += v * v; + local_sum += x_row[i]; } local_sum = block_reduce_sum(local_sum); - local_sum_sq = block_reduce_sum(local_sum_sq); __shared__ float s_mean, s_inv_std; if (threadIdx.x == 0) { - float mean = local_sum / hidden_size; - float var = local_sum_sq / hidden_size - mean * mean; - s_mean = mean; - s_inv_std = rsqrtf(var + eps); + s_mean = local_sum / hidden_size; } __syncthreads(); float mean = s_mean; + + // Pass 2: compute variance = sum((x - mean)^2) / N + float local_var = 0.0f; + for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) { + float d = x_row[i] - mean; + local_var += d * d; + } + local_var = block_reduce_sum(local_var); + + if (threadIdx.x == 0) { + s_inv_std = rsqrtf(local_var / hidden_size + eps); + } + __syncthreads(); + float inv_std = s_inv_std; for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) { out_row[i] = gamma[i] * (x_row[i] - mean) * inv_std + beta[i]; @@ -52,26 +59,34 @@ __global__ void layernorm_bf16( const __nv_bfloat16* x_row = x + row * hidden_size; __nv_bfloat16* out_row = out + row * hidden_size; + // Pass 1: compute mean float local_sum = 0.0f; - float local_sum_sq = 0.0f; for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) { - float v = __bfloat162float(x_row[i]); - local_sum += v; - local_sum_sq += v * v; + local_sum += __bfloat162float(x_row[i]); } local_sum = block_reduce_sum(local_sum); - local_sum_sq = block_reduce_sum(local_sum_sq); __shared__ float s_mean, s_inv_std; if (threadIdx.x == 0) { - float mean = local_sum / hidden_size; - float var = local_sum_sq / hidden_size - mean * mean; - s_mean = mean; - s_inv_std = rsqrtf(var + eps); + s_mean = local_sum / hidden_size; } __syncthreads(); float mean = s_mean; + + // Pass 2: compute variance = sum((x - mean)^2) / N + float local_var = 0.0f; + for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) { + float d = __bfloat162float(x_row[i]) - mean; + local_var += d * d; + } + local_var = block_reduce_sum(local_var); + + if (threadIdx.x == 0) { + s_inv_std = rsqrtf(local_var / hidden_size + eps); + } + __syncthreads(); + float inv_std = s_inv_std; for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) { float v = __bfloat162float(x_row[i]); @@ -86,6 +101,7 @@ extern "C" { void launch_layernorm_f32(const void* x, const void* gamma, const void* beta, void* out, int rows, int hidden_size, float eps, void* stream) { int block = (hidden_size < 1024) ? hidden_size : 1024; + if (block < 32) block = 32; layernorm_f32<<>>( (const float*)x, (const float*)gamma, (const float*)beta, (float*)out, hidden_size, eps); @@ -94,6 +110,7 @@ void launch_layernorm_f32(const void* x, const void* gamma, const void* beta, void launch_layernorm_bf16(const void* x, const void* gamma, const void* beta, void* out, int rows, int hidden_size, float eps, void* stream) { int block = (hidden_size < 1024) ? hidden_size : 1024; + if (block < 32) block = 32; layernorm_bf16<<>>( (const __nv_bfloat16*)x, (const __nv_bfloat16*)gamma, (const __nv_bfloat16*)beta, (__nv_bfloat16*)out, hidden_size, eps); diff --git a/csrc/normalization/rmsnorm.cu b/csrc/normalization/rmsnorm.cu index cc8c700..480520a 100644 --- a/csrc/normalization/rmsnorm.cu +++ b/csrc/normalization/rmsnorm.cu @@ -108,6 +108,7 @@ extern "C" { void launch_rmsnorm_f32(const void* x, const void* gamma, void* out, int rows, int hidden_size, float eps, void* stream) { int block = (hidden_size < 1024) ? hidden_size : 1024; + if (block < 32) block = 32; rmsnorm_f32<<>>( (const float*)x, (const float*)gamma, (float*)out, hidden_size, eps); } @@ -115,6 +116,7 @@ void launch_rmsnorm_f32(const void* x, const void* gamma, void* out, void launch_rmsnorm_bf16(const void* x, const void* gamma, void* out, int rows, int hidden_size, float eps, void* stream) { int block = (hidden_size < 1024) ? hidden_size : 1024; + if (block < 32) block = 32; rmsnorm_bf16<<>>( (const __nv_bfloat16*)x, (const __nv_bfloat16*)gamma, (__nv_bfloat16*)out, hidden_size, eps); @@ -124,6 +126,7 @@ void launch_add_rmsnorm_bf16(const void* x, const void* residual, const void* ga void* normed_out, void* sum_out, int rows, int hidden_size, float eps, void* stream) { int block = (hidden_size < 1024) ? hidden_size : 1024; + if (block < 32) block = 32; add_rmsnorm_bf16<<>>( (const __nv_bfloat16*)x, (const __nv_bfloat16*)residual, (const __nv_bfloat16*)gamma, diff --git a/docs/TO-BE-FIXED.md b/docs/TO-BE-FIXED.md index 5f42ffd..c2fc668 100644 --- a/docs/TO-BE-FIXED.md +++ b/docs/TO-BE-FIXED.md @@ -1,287 +1,214 @@ -# xserv — To Be Fixed +# xserv — To Be Fixed (2026-05-23 审查更新) -> 由最严格审查产出的修复清单。每项修复有明确验收标准，禁止 reward hacking。 +> 由全面审查产出的修复清单。每项修复有明确验收标准。 > 优先级: P0 (阻塞可用性) > P1 (严重bug/性能) > P2 (重要改进) > P3 (设计债务) --- -## FIX-01: 全局 cuBLAS handle，消除 per-call 创建 [P0-性能] +## 第一批：P0 — 阻塞可用性 -**问题**: `gemm.rs` 中每次 `matmul` / `batched_matmul` 调用都 `cublasCreate_v2` + `cublasDestroy_v2`。Qwen3-8B 一次 forward 约 168 次 matmul，每次创建/销毁 handle 耗费数毫秒。 +### FIX-01: 全局 cuBLAS handle [P0-性能] ❌未修 + +**问题**: `gemm.rs` 中 `matmul` (line 146) 和 `batched_matmul` (line 224) 每次调用都 `CublasContext::new()` 创建+销毁 handle。Qwen3-8B 一次 forward ~252 次 matmul。 **修复要求**: -- 使用 thread-local 或全局单例 cuBLAS handle -- handle 生命周期覆盖整个进程，不在 matmul 内创建/销毁 -- `CublasContext` 支持 `set_stream` 切换 stream +- 使用 thread-local 单例 cuBLAS handle +- handle 生命周期覆盖整个进程 +- `matmul` / `batched_matmul` 函数体内不再有 `CublasContext::new()` **验收标准**: -1. `grep -rn "cublasCreate_v2" crates/xserv-kernels/src/gemm.rs` 只出现 1 次（初始化处） -2. `matmul` 和 `batched_matmul` 函数体内不再有 `CublasContext::new()` -3. 编译通过，现有 gemm_test 全部通过 +1. `grep -n "CublasContext::new" crates/xserv-kernels/src/gemm.rs` 只出现 1 次（thread_local 初始化处） +2. 编译通过，现有 gemm_test 全部通过 --- -## FIX-02: 移除不必要的 cudaDeviceSynchronize [P0-性能] +### FIX-16: EOS token 泄漏到 API 响应 [P0-功能] ❌新发现 -**问题**: 几乎每个 kernel wrapper 结尾都有 `xserv_cuda::device::synchronize()`（即 `cudaDeviceSynchronize`），完全杀死 GPU pipeline。 +**问题**: `engine.rs:218` 中 `emit_token` 先发 `GenerateEvent::Token { text: "<|im_end|>" }` 再发 `Done`。`api.rs:110-111` 把所有 Token text 拼到 content 里，导致最终响应包含 `<|im_end|>` 乱码。 **修复要求**: -- 删除所有 kernel wrapper 中的 `device::synchronize()` 调用 -- 仅在需要读回 GPU 数据到 CPU 时同步（如 `sample_greedy`, `to_device(Cpu)`, benchmark） -- 在 `Tensor::to_device(Cpu)` 路径中已有隐式同步（`cudaMemcpy` 是同步的），不需要额外 sync -- 如果 kernel 使用 null stream（默认 stream），`cudaMemcpy` 会隐式等待默认 stream 上的所有操作 +- `emit_token` 中，当 token 是 EOS 时，不发送 Token event（或发送空 text），直接发 Done +- 或者: API 层收到 Done 时丢弃最后一个 token 的 text（如果 finish_reason == "stop"） **验收标准**: -1. `grep -rn "device::synchronize" crates/xserv-kernels/src/` 返回 0 行 -2. `grep -rn "device::synchronize" crates/xserv-model/src/` 只出现在 benchmark binary 中，不在 forward path 中 -3. 编译通过，现有测试全部通过 -4. 模型推理结果与修复前 bit-exact 一致（greedy decode 相同 prompt 产生相同 token 序列） - ---- - -## FIX-03: 修复 Chat Template [P0-功能] - -**问题**: `api.rs` 的 `build_prompt` 只是简单拼接文本，没有 ChatML special tokens。Qwen3 模型收到的 prompt 没有对话结构。 - -**修复要求**: -- 生成符合 Qwen3 ChatML 格式的 prompt： - ``` - <|im_start|>system\n{content}<|im_end|>\n<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n - ``` -- 如果没有 system message，跳过 system 部分 -- 如果有多轮 assistant/user 交替，按顺序生成 -- 结尾始终是 `<|im_start|>assistant\n`（让模型生成 assistant 回复） - -**验收标准**: -1. 单元测试: 给定 `[{role: "user", content: "Hello"}]`，生成的 prompt 字符串包含 `<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\n` -2. 单元测试: 给定 system + user + assistant + user 四条消息，格式正确 +1. 发送请求，响应 content 不包含 `<|im_end|>` 或其他 special token 文本 +2. streaming 模式下最后一个 content chunk 不是 EOS 文本 3. 编译通过 --- -## FIX-04: 修复 `is_finished` 硬编码 EOS [P0-功能] +### FIX-17: max_seq_len 硬编码 256 [P0-功能] ❌新发现 -**问题**: `engine.rs:160` 硬编码 `last == 151645` 作为 EOS 判断。 +**问题**: `engine.rs:53` 硬编码 `let max_seq_len = 256`，超过就 KV cache panic。 **修复要求**: -- `Sequence` struct 增加 `eos_token_id: Option` 字段 -- 在 `make_sequence` 中从 tokenizer 获取 EOS token ID -- `is_finished` 使用该字段判断 +- `Engine::load` 接受 `max_seq_len` 参数（或从 config 读取，上限为 config.max_seq_len()） +- `main.rs` 中通过命令行参数或环境变量传入，默认值改为 2048 +- 同步更新 RoPE cache 上限（当前 `qwen3.rs:45` 限制 8192，应与 max_seq_len 一致） **验收标准**: -1. `grep -rn "151645" crates/xserv-server/` 返回 0 行 -2. `is_finished` 函数不包含任何硬编码 token ID +1. `grep -n "let max_seq_len = 256" crates/xserv-server/` 返回 0 行 +2. 启动 server 时 `--max-seq-len 4096` 可用 3. 编译通过 --- -## FIX-05: 修复 `Storage::device()` 丢失设备信息 [P1-Bug] +### FIX-18: max_tokens 无上限校验 [P0-功能] ❌新发现 -**问题**: `storage.rs:43` 对所有 GPU storage 返回 `Device::Cuda(0)`，不追踪实际设备。 +**问题**: API 不校验 `max_tokens`，客户端可发 `max_tokens: 1000000` 导致 KV cache panic。 **修复要求**: -- `StorageInner::Cuda` 增加 `device: u32` 字段 -- `Storage::cuda()` 接受 device 参数，或从 `GpuBuffer` 推断 -- `Storage::device()` 返回实际设备 -- 所有创建 `Storage::cuda()` 的调用点更新 +- `api.rs` 中 clamp `max_tokens` 到 `engine.max_seq_len - prompt_tokens.len()` +- 如果 prompt 已超过 max_seq_len，返回 400 错误 **验收标准**: -1. 创建一个 `Device::Cuda(3)` 的 tensor，`tensor.device()` 返回 `Device::Cuda(3)` +1. 发送 `max_tokens: 999999`，不 panic，正常生成到 seq_len 上限 +2. 发送超长 prompt（> max_seq_len），返回 HTTP 400 +3. 编译通过 + +--- + +## 第二批：P1 — 严重 bug/性能 + +### FIX-07: 使用 CachingAllocator [P1-性能] ❌未修 + +**问题**: `CachingAllocator` 已实现（`allocator.rs`）但从未使用。所有 GPU 分配直接 `cudaMalloc`。 + +**修复要求**: +- `Tensor::empty` 对 GPU device 使用 `cached_alloc` 而非 `GpuBuffer::alloc` +- `GpuBuffer::Drop` 调用 `cached_dealloc` 归还到池（而非 `cudaFree`） +- 或者更简单：在 `GpuBuffer::alloc` 内部接入 caching allocator（全局透明替换） + +**验收标准**: +1. 连续运行 10 次 decode step，`cudaMalloc` 调用次数应显著低于总分配次数 +2. 编译通过，现有测试通过 +3. 推理结果与修复前一致 + +--- + +### FIX-08: CudaDeviceProp FFI 安全性 [P1-Bug] ❌未修 + +**问题**: `ffi.rs:31` 用 `_pad: [u8; 4096]` 猜测 `cudaDeviceProp` struct 大小，CUDA 12.9 可能更大。 + +**修复要求**: +- 增大 pad 到 `[u8; 8192]` 或使用 `cudaDeviceGetAttribute` 替代 name 查询 +- 可参考 `device.rs` 中已有的 `cudaDeviceGetAttribute` 用法 + +**验收标准**: +1. `device_info()` 返回正确的 device name +2. 编译通过 + +--- + +### FIX-09: Tokenizer byte_fallback panic [P1-Bug] ❌未修 + +**问题**: `bpe.rs:176-182` 中 Qwen3 tokenizer 遇到不在 vocab 的单字节时 panic。 + +**修复要求**: +- 当 `byte_fallback == true` 且单字节不在 vocab 时，查找 `<0xNN>` 格式 token +- 如果 `<0xNN>` 也不存在，返回 unk_token_id（而非 panic） + +**验收标准**: +1. 包含所有 256 个字节值的字符串可以 encode 不 panic +2. 编译通过 + +--- + +### FIX-19: 因果掩码 -1e9 应改为 -inf [P1-Bug] ❌新发现 + +**问题**: `csrc/attention/causal_mask.cu:31` 用 `-1e9f` 代替 `-inf`，注释说 "BF16 没有 -inf" 但这是错误的。 + +**修复要求**: +- BF16 路径改为 `__float2bfloat16(-INFINITY)` +- F32 路径改为 `-INFINITY`（如果还没有的话） + +**验收标准**: +1. causal mask 中被遮蔽的值为 `-inf`（而非 `-1e9`） +2. 编译通过，attention test 通过 + +--- + +### FIX-20: LayerNorm 数值稳定性 [P1-Bug] ❌新发现 + +**问题**: `csrc/normalization/layernorm.cu:19-25` 注释写 "Welford online" 但实际用 `E[x²] - E[x]²`，大均值小方差时会灾难性抵消。 + +**修复要求**: +- 改为真正的 two-pass 或 Welford online 算法 +- pass 1: 求 mean; pass 2: 求 variance = E[(x-mean)²] + +**验收标准**: +1. 对 mean=1e6, std=1e-3 的输入，layernorm 输出与 PyTorch 一致（relative error < 1e-3） 2. 编译通过，现有测试通过 --- -## FIX-06: 修复 `unsqueeze` stride 计算 [P1-Bug] +### FIX-21: LayerNorm/RMSNorm 最小 block size [P1-Bug] ❌新发现 -**问题**: `tensor.rs:128` 中 unsqueeze 的 stride 计算错误。对 `[3,4]` strides `[4,1]` 做 `unsqueeze(0)` 得到 strides `[4,4,1]`，而正确应为 `[12,4,1]`。虽然 size-1 维度的 stride 不影响寻址，但导致 `is_contiguous()` 误判为 false，触发不必要的 copy。 +**问题**: `layernorm.cu:88` 和 `rmsnorm.cu` 对 hidden_size < 32 的输入会崩溃（block_reduce 需要至少一个完整 warp）。 **修复要求**: -- size-1 维度的 stride 应设为 `shape[dim+1] * strides[dim+1]`（如果 dim 不是最后一维），使其满足 contiguous 条件 -- 或者更简单: unsqueeze 后如果原 tensor 是 contiguous 的，直接重算 contiguous strides +- launch 时 `block = max(min(hidden_size, 1024), 32)` **验收标准**: -1. 单元测试: `[3,4]` contiguous tensor 做 `unsqueeze(0)` 后 `is_contiguous()` 返回 true -2. 单元测试: `[3,4]` contiguous tensor 做 `unsqueeze(1)` 后 `is_contiguous()` 返回 true -3. 单元测试: `[3,4]` contiguous tensor 做 `unsqueeze(2)` 后 `is_contiguous()` 返回 true -4. 编译通过，现有测试通过 +1. hidden_size=16 的 layernorm/rmsnorm 不崩溃 +2. 编译通过 --- -## FIX-07: 使用 Caching Allocator [P1-性能] +## 第三批：P2 — 重要改进 -**问题**: `CachingAllocator` 已实现但从未使用。所有 GPU 分配直接 `cudaMalloc`。 +### FIX-22: Engine dummy KV cache 分配 [P2-性能] ❌新发现 + +**问题**: `engine.rs:142-148` 每次 batched decode 用 `std::mem::replace` 创建 dummy `GpuKVCache::new(..., 1, ...)` 来绕过 borrow checker，每步分配 `num_layers * 2` 个 GPU buffer。 **修复要求**: -- 创建一个全局或 thread-local `CachingAllocator` 实例 -- `Tensor::zeros` 等分配路径通过 caching allocator -- 或者至少: `GpuKVCache::get_kv_len` 中的临时 buffer 分配通过 caching allocator（这是最热的分配路径） -- `GpuBuffer::Drop` 需要与 allocator 配合（return to pool 而非 cudaFree） +- 将 `running` 从 `Vec` 改为存储方式让 KV cache 可以独立借出 +- 或使用 `Option` + `.take()` / `.insert()` 避免 dummy 分配 **验收标准**: -1. 在 decode loop 中连续调用 `get_kv_len` 100 次，`AllocStats.cuda_malloc_count` < 10（大部分命中 cache） -2. 编译通过，现有测试通过 +1. batched decode 路径不再分配 dummy KV cache +2. 编译通过，功能不变 --- -## FIX-08: 修复 `CudaDeviceProp` FFI 安全性 [P1-Bug] +### FIX-23: RoPE cache 硬限 8192 [P2-功能] ❌新发现 -**问题**: `ffi.rs:31` 使用 `_pad: [u8; 4096]` 假设 cudaDeviceProp 总大小。CUDA 12.9 的实际结构可能更大。 +**问题**: `qwen3.rs:45` `config.max_seq_len().min(8192)` 人为截断。 **修复要求**: -- 删除 `CudaDeviceProp` struct（或仅保留 name 字段所需的最小 struct） -- 如果只需要 name: 分配一个足够大的 buffer（如 `[u8; 8192]`）并直接读取 name offset（前 256 bytes） -- 或者更安全: 使用 `cudaDeviceGetAttribute` + 单独的 name 查询 API（`device.rs` 已经用 getAttribute 查其他属性了，只差 name） +- 去掉 `.min(8192)`，或改为与 engine 的 max_seq_len 一致 +- 确保 RoPE cache 覆盖实际使用的 max_seq_len **验收标准**: -1. 不再有 `CudaDeviceProp` struct，或 padding 大小基于 `std::mem::size_of` 动态确定 -2. `device_info()` 仍能返回正确的 device name -3. 编译通过，现有测试通过 +1. RoPE cache 长度 >= engine max_seq_len +2. 编译通过 --- -## FIX-09: 修复 Tokenizer byte_fallback panic [P1-Bug] +### FIX-15: GPT-2 消除 CPU round-trip [P3-性能] ❌未修 -**问题**: `bpe.rs:173-176` 中 Qwen3 tokenizer 遇到不在 vocab 的单字节时 panic。 - -**修复要求**: -- 当 `byte_fallback == true` 且单字节不在 vocab 时，查找 `<0xNN>` 格式的 special token -- 如果 `<0xNN>` 也不存在，才 panic（带有明确的错误信息） - -**验收标准**: -1. 使用 Qwen3 tokenizer encode 包含所有 256 个字节值的字符串不 panic -2. encode 后 decode 回来的字节序列与原始一致 -3. 编译通过 +**问题**: GPT-2 `split_qkv`、`merge_heads`、`add_bias` 全在 CPU 做。优先级低（GPT-2 不是主力模型）。 --- -## FIX-10: 实现 SSE Streaming [P2-功能] +## 修复依赖图和执行顺序 -**问题**: API 只支持阻塞式响应，不支持 SSE streaming。 +``` +第一批 P0 (可并行): + FIX-01 (cuBLAS handle) ← 独立 + FIX-16 (EOS 泄漏) ← 独立 + FIX-17 (max_seq_len) ← 独立，FIX-23 依赖此 + FIX-18 (max_tokens 校验) ← 依赖 FIX-17（需要知道 max_seq_len） -**修复要求**: -- `ChatRequest` 增加 `stream: Option` 字段 -- 当 `stream == true` 时，返回 `text/event-stream` content type -- 每生成一个 token 发送一个 SSE event，格式与 OpenAI 兼容: - ``` - data: {"id":"chatcmpl-xxx","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"token"},"finish_reason":null}]} - ``` -- 最后发送 `data: [DONE]` -- 非 streaming 模式行为不变 +第二批 P1 (可并行): + FIX-07 (caching allocator) ← 独立 + FIX-08 (CudaDeviceProp) ← 独立 + FIX-09 (byte_fallback) ← 独立 + FIX-19 (causal mask -inf) ← 独立 + FIX-20 (layernorm 稳定性) ← 独立 + FIX-21 (min block size) ← 独立 -**验收标准**: -1. `curl` 请求 `stream: true` 能看到逐行 SSE 输出 -2. 每行 SSE data 是合法 JSON，包含 `choices[0].delta.content` -3. 最后一行是 `data: [DONE]` -4. 非 streaming 请求仍正常工作 -5. 编译通过 - ---- - -## FIX-11: 修复 Usage 统计 [P2-功能] - -**问题**: API 返回的 usage 全是 0。 - -**修复要求**: -- 追踪 prompt token 数量和 completion token 数量 -- 在 non-streaming 响应中返回正确的 usage -- 在 streaming 最后一个 chunk（或 `[DONE]` 前）可选择性包含 usage - -**验收标准**: -1. 发送一个 non-streaming 请求，`usage.prompt_tokens` > 0，`usage.completion_tokens` > 0 -2. `usage.total_tokens == usage.prompt_tokens + usage.completion_tokens` -3. 编译通过 - ---- - -## FIX-12: `GpuKVCache::get_kv_len` 避免重复分配 [P2-性能] - -**问题**: 每次调用 `get_kv_len` 都 `GpuBuffer::alloc` 新内存，decode 循环中每步每层一次。 - -**修复要求**: -- 方案 A: 返回 view/slice 到已有的预分配 buffer（零分配），需要构造 Tensor 时使用正确的 strides 指向 padded buffer -- 方案 B: 在 GpuKVCache 中预分配 output buffer，get_kv_len 做 D2D copy 到固定 buffer（每层 2 个 output buffer） -- 方案 A 更优但实现复杂度更高 - -**验收标准**: -1. 连续调用 `get_kv_len` 100 次，`cudaMalloc` 调用次数 <= 2（初始分配） -2. 返回的 tensor 数据正确（与修改前 bit-exact） -3. 编译通过，现有测试通过 - ---- - -## FIX-13: 实现 Sampling Strategies [P2-功能] - -**问题**: 只有 greedy sampling，没有 temperature / top-k / top-p。 - -**修复要求**: -- 实现 `SamplingParams { temperature, top_k, top_p }` struct -- temperature: `logits = logits / temperature` 后 softmax 后按概率采样 -- top_k: 保留 top-k logits，其余置 -inf -- top_p: 按概率降序累加到 >= p 后截断 -- greedy 作为 `temperature = 0` 或独立模式 -- `GenerateRequest` 接收 sampling params -- API 层解析 temperature / top_k / top_p 参数 - -**验收标准**: -1. temperature=0.0 与 greedy 结果一致 -2. temperature=1.0 多次生成同一 prompt 产生不同结果 -3. top_k=1 与 greedy 结果一致 -4. 编译通过 - ---- - -## FIX-14: GPU Tensor contiguous() 用 GPU kernel [P2-性能] - -**问题**: `tensor.rs:148` 中非 contiguous GPU tensor 做 contiguous 需要 GPU→CPU→CPU copy→CPU→GPU。 - -**修复要求**: -- 实现一个通用的 strided copy GPU kernel（或至少对常见的 transpose 情况有 kernel） -- `contiguous()` 对 GPU tensor 直接在 GPU 上完成 - -**验收标准**: -1. 对一个 GPU 上的 transposed tensor 调用 `contiguous()`，不触发任何 `cudaMemcpy` H2D/D2H -2. 结果与 CPU 实现 bit-exact -3. 编译通过，现有测试通过 - ---- - -## FIX-15: GPT-2 消除 CPU round-trip (split_qkv, merge_heads, add_bias) [P3-性能] - -**问题**: GPT-2 的 `split_qkv`, `merge_heads`, `add_bias` 全在 CPU 上做。 - -**修复要求**: -- `add_bias`: 实现 broadcast-add GPU kernel（[S,N] + [N] → [S,N]） -- `split_qkv`: 实现 GPU kernel 将 [S, 3H] 分成 Q/K/V 并 reshape 为 [1, heads, S, D] -- `merge_heads`: 复用已有的 `merge_heads_gpu` kernel（目前只有 BF16 版本，需要 F32 版本） - -**验收标准**: -1. GPT-2 forward path 中 `grep -n "to_device(Device::Cpu)"` 只出现在 `sample_greedy` 中 -2. 推理结果与修复前一致（greedy decode bit-exact） -3. 编译通过，现有测试通过 - ---- - -## 修复优先级排序 - -**第一批 (必须先做，其他依赖它们)**: -1. FIX-01: 全局 cuBLAS handle -2. FIX-02: 移除 device sync -3. FIX-03: Chat template -4. FIX-04: is_finished EOS - -**第二批 (重要 bug 修复)**: -5. FIX-05: Storage device tracking -6. FIX-06: unsqueeze stride -7. FIX-08: CudaDeviceProp -8. FIX-09: byte_fallback panic - -**第三批 (功能完善)**: -9. FIX-10: SSE streaming -10. FIX-11: Usage stats -11. FIX-13: Sampling strategies - -**第四批 (性能优化)**: -12. FIX-07: Caching allocator -13. FIX-12: KV cache alloc -14. FIX-14: GPU contiguous -15. FIX-15: GPT-2 CPU round-trip +第三批 P2: + FIX-22 (dummy KV cache) ← 独立 + FIX-23 (RoPE cache) ← 依赖 FIX-17 +```