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

fix: comprehensive review + 14 bug fixes + Phase 12/14 overhaul

Browse Source 
Strict code review identified 30+ issues across correctness, performance,
and architecture. This commit addresses 14 of them with verified fixes,
restructures Phase 12 for honest continuous batching, and updates Phase 14
to target FA2 (RTX 5090 SM120 lacks TMEM required by FA4).

Bug fixes:
- FIX-01: Global cuBLAS handle (thread-local singleton, was per-call)
- FIX-02: Remove 19 unnecessary cudaDeviceSynchronize calls from kernels
- FIX-03: Qwen3 ChatML template (was plain text concatenation)
- FIX-04: EOS token from tokenizer (was hardcoded 151645)
- FIX-05: Storage tracks actual GPU device ordinal (was always Cuda(0))
- FIX-06: unsqueeze stride preserves contiguous layout
- FIX-08: CudaDeviceProp replaced with heap buffer (was UB-prone padding)
- FIX-09: Tokenizer byte_fallback to <0xNN> tokens (was panic)

Feature additions:
- FIX-10: SSE streaming (/v1/chat/completions, OpenAI-compatible)
- FIX-11: Correct usage statistics (prompt/completion/total tokens)
- FIX-13: Temperature / top-k / top-p sampling with SamplingParams

Performance improvements:
- FIX-07: Caching allocator wired up (thread-local pool, pooled flag)
- FIX-12: KV cache staging buffers (zero-alloc get_kv_len via borrow_raw)
- FIX-14: GPU strided copy kernel (eliminates contiguous() CPU round-trip)

Architecture:
- Phase 12 engine restructured: prefill/decode separation, honest TODO
  for batched GPU forward (requires Flash Attention)
- Phase 14 updated: FA2 for SM120 (FA4 requires TMEM, absent on 5090)
- Qwen3-7B → Qwen3-8B typo fixed across all docs (36 layers, hidden 4096)

Validated on dash5 (8x RTX 5090):
- 52/52 API prompts pass (EN/CN/code), SSE streaming verified
- Logits match HF transformers 9/10 top-1, 4.0/5 avg top-5 overlap
- 8 concurrent requests: 5.99x scheduling speedup (batch_size=4)
- Throughput: 10.3 tok/s (serial), 30% of HF baseline

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-22 17:53:28 +08:00
parent d8493bd70f
commit ee68d3565d
38 changed files with 3012 additions and 259 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/xserv-tensor/src/lib.rs
View File

@@ -6,4 +6,4 @@ pub mod tensor;
pub use dtype::{DType, TensorDType};
pub use shape::Dims;
pub use storage::{Device, Storage};
pub use tensor::Tensor;
pub use tensor::{register_gpu_contiguous, Tensor};

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

@@ -3,7 +3,7 @@ use xserv_cuda::{GpuBuffer, Result as CudaResult};
enum StorageInner {
Cpu { data: Vec<u8> },
Cuda { buffer: GpuBuffer },
Cuda { buffer: GpuBuffer, device: u32 },
}
/// Reference-counted storage for tensor data. Multiple tensors can share
@@ -31,21 +31,21 @@ impl Storage {
Self(Arc::new(StorageInner::Cpu { data }))
}
pub fn cuda(buffer: GpuBuffer) -> Self {
Self(Arc::new(StorageInner::Cuda { buffer }))
pub fn cuda(buffer: GpuBuffer, device: u32) -> Self {
Self(Arc::new(StorageInner::Cuda { buffer, device }))
}
pub fn device(&self) -> Device {
match self.0.as_ref() {
StorageInner::Cpu { .. } => Device::Cpu,
StorageInner::Cuda { .. } => Device::Cuda(0),
StorageInner::Cuda { device, .. } => Device::Cuda(*device),
}
}
pub fn len_bytes(&self) -> usize {
match self.0.as_ref() {
StorageInner::Cpu { data } => data.len(),
StorageInner::Cuda { buffer } => buffer.len(),
StorageInner::Cuda { buffer, .. } => buffer.len(),
}
}
@@ -59,7 +59,7 @@ impl Storage {
pub fn gpu_buffer(&self) -> &GpuBuffer {
match self.0.as_ref() {
StorageInner::Cuda { buffer } => buffer,
StorageInner::Cuda { buffer, .. } => buffer,
StorageInner::Cpu { .. } => panic!("cannot access CPU storage as GPU buffer"),
}
}
@@ -71,11 +71,11 @@ impl Storage {
return Ok(self.clone());
}
match (current, target) {
(Device::Cpu, Device::Cuda(_dev)) => {
(Device::Cpu, Device::Cuda(dev)) => {
let cpu_data = self.as_cpu_bytes();
let mut buf = GpuBuffer::alloc(cpu_data.len())?;
buf.copy_from_host(cpu_data)?;
Ok(Storage::cuda(buf))
Ok(Storage::cuda(buf, dev))
}
(Device::Cuda(_), Device::Cpu) => {
let gpu_buf = self.gpu_buffer();
@@ -83,11 +83,11 @@ impl Storage {
gpu_buf.copy_to_host(&mut data)?;
Ok(Storage::cpu(data))
}
(Device::Cuda(_), Device::Cuda(_)) => {
(Device::Cuda(_), Device::Cuda(dev)) => {
let src = self.gpu_buffer();
let mut dst = GpuBuffer::alloc(src.len())?;
dst.copy_from_device(src)?;
Ok(Storage::cuda(dst))
Ok(Storage::cuda(dst, dev))
}
_ => unreachable!(),
}
@@ -97,10 +97,10 @@ impl Storage {
pub fn deep_copy(&self) -> CudaResult<Self> {
match self.0.as_ref() {
StorageInner::Cpu { data } => Ok(Storage::cpu(data.clone())),
StorageInner::Cuda { buffer } => {
StorageInner::Cuda { buffer, device } => {
let mut dst = GpuBuffer::alloc(buffer.len())?;
dst.copy_from_device(buffer)?;
Ok(Storage::cuda(dst))
Ok(Storage::cuda(dst, *device))
}
}
}
@@ -109,10 +109,10 @@ impl Storage {
pub fn zeros(len_bytes: usize, device: Device) -> CudaResult<Self> {
match device {
Device::Cpu => Ok(Storage::cpu(vec![0u8; len_bytes])),
Device::Cuda(_) => {
let mut buf = GpuBuffer::alloc(len_bytes)?;
Device::Cuda(dev) => {
let mut buf = xserv_cuda::allocator::cached_alloc(len_bytes)?;
buf.zero()?;
Ok(Storage::cuda(buf))
Ok(Storage::cuda(buf, dev))
}
}
}

87
crates/xserv-tensor/src/tensor.rs
View File

@@ -1,7 +1,21 @@
use std::sync::OnceLock;
use crate::dtype::{DType, TensorDType};
use crate::shape::{self, Dims};
use crate::storage::{Device, Storage};
/// Global hook for GPU strided-to-contiguous copy.
/// Set by `xserv-kernels` (or any crate that provides a GPU kernel) via
/// `register_gpu_contiguous`. When set, `contiguous()` on a non-contiguous
/// GPU tensor calls this instead of doing a CPU round-trip.
static GPU_CONTIGUOUS_FN: OnceLock<fn(&Tensor) -> Tensor> = OnceLock::new();
/// Register a function that makes a non-contiguous GPU tensor contiguous.
/// Intended to be called once by the kernel crate at startup.
pub fn register_gpu_contiguous(f: fn(&Tensor) -> Tensor) {
let _ = GPU_CONTIGUOUS_FN.set(f);
}
/// Multi-dimensional array with CPU or GPU storage.
///
/// Tensors support view semantics: transpose, slice, etc. share
@@ -123,10 +137,15 @@ impl Tensor {
pub fn unsqueeze(&self, dim: usize) -> Self {
assert!(dim <= self.ndim());
let mut new_shape = self.shape.clone();
let mut new_strides = self.strides.clone();
new_shape.insert(dim, 1);
let stride_val = if dim < self.strides.len() { self.strides[dim] } else { 1 };
new_strides.insert(dim, stride_val);
let new_strides = if self.is_contiguous() {
shape::contiguous_strides(&new_shape)
} else {
let mut s = self.strides.clone();
let stride_val = if dim < self.strides.len() { self.strides[dim] } else { 1 };
s.insert(dim, stride_val);
s
};
Self {
storage: self.storage.clone(),
shape: new_shape,
@@ -142,9 +161,12 @@ impl Tensor {
if self.is_contiguous() {
return self.clone();
}
// For GPU tensors: round-trip through CPU (correct but slow).
// TODO: write a GPU contiguous-copy kernel for performance.
// For GPU tensors: use the registered GPU kernel if available,
// otherwise fall back to CPU round-trip.
if matches!(self.device(), Device::Cuda(_)) {
if let Some(gpu_fn) = GPU_CONTIGUOUS_FN.get() {
return gpu_fn(self);
}
let cpu = self.to_device(Device::Cpu);
let contig = cpu.contiguous();
return contig.to_device(self.device());
@@ -237,3 +259,58 @@ impl std::fmt::Debug for Tensor {
)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn contiguous_2d() -> Tensor {
Tensor::from_slice(&[1.0f32; 12], &[3, 4])
}
#[test]
fn unsqueeze_dim0_contiguous() {
let t = contiguous_2d();
let u = t.unsqueeze(0);
assert_eq!(u.shape(), &[1, 3, 4]);
assert!(u.is_contiguous());
assert_eq!(u.strides(), &[12, 4, 1]);
}
#[test]
fn unsqueeze_dim1_contiguous() {
let t = contiguous_2d();
let u = t.unsqueeze(1);
assert_eq!(u.shape(), &[3, 1, 4]);
assert!(u.is_contiguous());
assert_eq!(u.strides(), &[4, 4, 1]);
}
#[test]
fn unsqueeze_dim2_contiguous() {
let t = contiguous_2d();
let u = t.unsqueeze(2);
assert_eq!(u.shape(), &[3, 4, 1]);
assert!(u.is_contiguous());
assert_eq!(u.strides(), &[4, 1, 1]);
}
#[test]
fn unsqueeze_noncontiguous() {
// Transpose makes [3,4] into [4,3] with strides [1,4] (non-contiguous)
let t = contiguous_2d().transpose(0, 1);
assert!(!t.is_contiguous());
let u = t.unsqueeze(0);
assert_eq!(u.shape(), &[1, 4, 3]);
// Non-contiguous path: stride_val copied from strides[0]=1
assert_eq!(u.strides(), &[1, 1, 4]);
}
#[test]
fn unsqueeze_squeeze_roundtrip() {
let t = contiguous_2d();
let u = t.unsqueeze(1).squeeze(1);
assert_eq!(u.shape(), t.shape());
assert!(u.is_contiguous());
}
}