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

1
crates/xserv-model/Cargo.toml
View File

@@ -13,3 +13,4 @@ smallvec.workspace = true
serde.workspace = true
serde_json.workspace = true
safetensors.workspace = true
rand.workspace = true

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

@@ -31,7 +31,7 @@ fn main() {
// Warmup
{
let ids = tokenizer.encode("warmup");
let mut cache = GpuKVCache::new(&config, 256, DType::BF16);
let mut cache = GpuKVCache::new(&config, 256, DType::BF16, 0);
let _ = model.forward_gpu_cache(&ids, &mut cache);
}
eprintln!("Warmup done. Running benchmark...");
@@ -94,7 +94,7 @@ fn main() {
let input_ids = tokenizer.encode(prompt);
let input_len = input_ids.len();
let mut cache = GpuKVCache::new(&config, 256, DType::BF16);
let mut cache = GpuKVCache::new(&config, 256, DType::BF16, 0);
// Prefill
let t0 = Instant::now();

1
crates/xserv-model/src/gpt2.rs
View File

@@ -116,6 +116,7 @@ fn tensor_from_raw_bytes(bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor
impl GPT2 {
pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
crate::init_kernels();
let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
};

53
crates/xserv-model/src/kv_cache.rs
View File

@@ -9,16 +9,22 @@ pub struct GpuKVCache {
// Layout: [num_kv_heads, max_seq_len, head_dim] — contiguous per head
k_bufs: Vec<GpuBuffer>,
v_bufs: Vec<GpuBuffer>,
// Per layer: pre-allocated staging buffers for get_kv_len output.
// Size: num_kv_heads * max_seq_len * head_dim * elem_size (max possible output).
// Avoids cudaMalloc/cudaFree on every get_kv_len call.
k_staging: Vec<GpuBuffer>,
v_staging: Vec<GpuBuffer>,
seq_len: usize,
max_seq_len: usize,
num_kv_heads: usize,
head_dim: usize,
elem_size: usize,
dtype: DType,
device: u32,
}
impl GpuKVCache {
pub fn new(config: &ModelConfig, max_seq_len: usize, dtype: DType) -> Self {
pub fn new(config: &ModelConfig, max_seq_len: usize, dtype: DType, device: u32) -> Self {
let num_layers = config.num_layers();
let num_kv_heads = config.num_kv_heads();
let head_dim = config.head_dim();
@@ -27,6 +33,8 @@ impl GpuKVCache {
let mut k_bufs = Vec::with_capacity(num_layers);
let mut v_bufs = Vec::with_capacity(num_layers);
let mut k_staging = Vec::with_capacity(num_layers);
let mut v_staging = Vec::with_capacity(num_layers);
for _ in 0..num_layers {
let mut k = GpuBuffer::alloc(buf_size).expect("alloc KV cache K");
let mut v = GpuBuffer::alloc(buf_size).expect("alloc KV cache V");
@@ -34,9 +42,11 @@ impl GpuKVCache {
v.zero().unwrap();
k_bufs.push(k);
v_bufs.push(v);
k_staging.push(GpuBuffer::alloc(buf_size).expect("alloc KV staging K"));
v_staging.push(GpuBuffer::alloc(buf_size).expect("alloc KV staging V"));
}
Self { k_bufs, v_bufs, seq_len: 0, max_seq_len, num_kv_heads, head_dim, elem_size, dtype }
Self { k_bufs, v_bufs, k_staging, v_staging, seq_len: 0, max_seq_len, num_kv_heads, head_dim, elem_size, dtype, device }
}
pub fn seq_len(&self) -> usize { self.seq_len }
@@ -69,45 +79,58 @@ impl GpuKVCache {
}
/// Get K/V cache tensors for a layer up to `seq_len` tokens: [1, num_kv_heads, seq_len, head_dim]
pub fn get_kv(&self, layer: usize) -> (Tensor, Tensor) {
pub fn get_kv(&mut self, layer: usize) -> (Tensor, Tensor) {
let sl = self.seq_len;
self.get_kv_len(layer, sl)
}
pub fn get_kv_len(&self, layer: usize, sl: usize) -> (Tensor, Tensor) {
pub fn get_kv_len(&mut self, layer: usize, sl: usize) -> (Tensor, Tensor) {
let hd = self.head_dim;
let nh = self.num_kv_heads;
let es = self.elem_size;
let max_s = self.max_seq_len;
// Allocate output tensors [1, nh, sl, hd]
// Copy each head's valid portion into pre-allocated staging buffers.
// Split borrows: staging (mut) vs cache (shared) are separate struct fields,
// so the borrow checker allows simultaneous &mut staging + &cache.
let out_size = nh * sl * hd * es;
let mut k_out = GpuBuffer::alloc(out_size).expect("alloc k_out");
let mut v_out = GpuBuffer::alloc(out_size).expect("alloc v_out");
// Copy each head's valid portion
let k_stg = &mut self.k_staging[layer];
let k_buf = &self.k_bufs[layer];
let v_stg = &mut self.v_staging[layer];
let v_buf = &self.v_bufs[layer];
for h in 0..nh {
let src_off = (h * max_s) * hd * es;
let dst_off = (h * sl) * hd * es;
let count = sl * hd * es;
k_out.copy_from_device_at(&self.k_bufs[layer], src_off, dst_off, count).unwrap();
v_out.copy_from_device_at(&self.v_bufs[layer], src_off, dst_off, count).unwrap();
k_stg.copy_from_device_at(k_buf, src_off, dst_off, count).unwrap();
v_stg.copy_from_device_at(v_buf, src_off, dst_off, count).unwrap();
}
// Grab raw pointers before dropping the mutable borrows
let k_ptr = k_stg.as_mut_ptr();
let v_ptr = v_stg.as_mut_ptr();
// Create Tensors that borrow from the staging buffers (no cudaMalloc/cudaFree).
// Safety: staging buffers are owned by GpuKVCache and outlive the returned Tensors
// in practice (Tensors are consumed within the same forward pass before the next
// get_kv_len call overwrites the staging buffer).
let shape = &[1usize, nh, sl, hd];
let k = unsafe { tensor_from_gpu_buffer(k_out, shape, self.dtype) };
let v = unsafe { tensor_from_gpu_buffer(v_out, shape, self.dtype) };
let k = unsafe {
tensor_from_gpu_buffer(GpuBuffer::borrow_raw(k_ptr, out_size), shape, self.dtype, self.device)
};
let v = unsafe {
tensor_from_gpu_buffer(GpuBuffer::borrow_raw(v_ptr, out_size), shape, self.dtype, self.device)
};
(k, v)
}
}
/// Create a Tensor from a GpuBuffer (takes ownership).
unsafe fn tensor_from_gpu_buffer(buf: GpuBuffer, shape: &[usize], dtype: DType) -> Tensor {
unsafe fn tensor_from_gpu_buffer(buf: GpuBuffer, shape: &[usize], dtype: DType, device: u32) -> Tensor {
use xserv_tensor::storage::Storage;
use xserv_tensor::shape::contiguous_strides;
use smallvec::SmallVec;
let storage = Storage::cuda(buf);
let storage = Storage::cuda(buf, device);
Tensor::from_storage(
storage,
SmallVec::from_slice(shape),

8
crates/xserv-model/src/lib.rs
View File

@@ -3,8 +3,16 @@ pub mod gpt2;
pub mod kv_cache;
pub mod loader;
pub mod qwen3;
pub mod sampling;
pub use config::ModelConfig;
pub use gpt2::{GPT2, KVCache};
pub use kv_cache::GpuKVCache;
pub use qwen3::Qwen3;
pub use sampling::{SamplingParams, sample};
/// Initialize GPU kernel hooks. Called automatically by model constructors,
/// but safe to call multiple times (idempotent via OnceLock).
pub fn init_kernels() {
xserv_kernels::init();
}

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

@@ -32,6 +32,7 @@ struct Qwen3Block {
impl Qwen3 {
pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
crate::init_kernels();
let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
};

120
crates/xserv-model/src/sampling.rs Normal file
View File

@@ -0,0 +1,120 @@
use half::bf16;
use rand::Rng;
use xserv_tensor::{DType, Device, Tensor};
pub struct SamplingParams {
pub temperature: f32,
pub top_k: usize,
pub top_p: f32,
}
impl Default for SamplingParams {
fn default() -> Self {
Self { temperature: 0.0, top_k: 0, top_p: 1.0 }
}
}
/// Sample a token from logits with shape [seq_len, vocab_size].
/// Uses the last position's logits. Handles both F32 and BF16 dtypes.
pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
assert_eq!(logits.ndim(), 2);
let vocab_size = logits.shape()[1];
let seq_len = logits.shape()[0];
let logits_cpu = logits.to_device(Device::Cpu);
// Extract last row as f32
let last_row: Vec<f32> = match logits.dtype() {
DType::F32 => {
let data = logits_cpu.as_slice::<f32>();
data[(seq_len - 1) * vocab_size..seq_len * vocab_size].to_vec()
}
DType::BF16 => {
let data = logits_cpu.as_slice::<bf16>();
data[(seq_len - 1) * vocab_size..seq_len * vocab_size]
.iter()
.map(|v| v.to_f32())
.collect()
}
_ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()),
};
// Greedy
if params.temperature == 0.0 {
return argmax(&last_row);
}
// Apply temperature
let mut logits_f32: Vec<f32> = last_row.iter().map(|v| v / params.temperature).collect();
// Top-k filtering
if params.top_k > 0 && params.top_k < vocab_size {
let mut indices: Vec<usize> = (0..vocab_size).collect();
indices.select_nth_unstable_by(params.top_k, |&a, &b| {
logits_f32[b].partial_cmp(&logits_f32[a]).unwrap()
});
// Everything after top_k should be masked
for &i in &indices[params.top_k..] {
logits_f32[i] = f32::NEG_INFINITY;
}
}
// Top-p (nucleus) filtering
if params.top_p < 1.0 {
// Sort indices by descending logit value
let mut indices: Vec<usize> = (0..vocab_size).collect();
indices.sort_unstable_by(|&a, &b| logits_f32[b].partial_cmp(&logits_f32[a]).unwrap());
// Compute softmax probabilities for the sorted order
let max_val = logits_f32[indices[0]];
let sorted_probs: Vec<f32> = indices
.iter()
.map(|&i| (logits_f32[i] - max_val).exp())
.collect();
let sum: f32 = sorted_probs.iter().sum();
let sorted_probs: Vec<f32> = sorted_probs.iter().map(|v| v / sum).collect();
// Cumulative sum, find cutoff
let mut cumsum = 0.0f32;
let mut cutoff = indices.len();
for (rank, &prob) in sorted_probs.iter().enumerate() {
cumsum += prob;
if cumsum > params.top_p {
cutoff = rank + 1; // keep at least this many
break;
}
}
// Mask everything beyond cutoff
for &i in &indices[cutoff..] {
logits_f32[i] = f32::NEG_INFINITY;
}
}
// Softmax
let max_val = logits_f32.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let exps: Vec<f32> = logits_f32.iter().map(|v| (v - max_val).exp()).collect();
let sum: f32 = exps.iter().sum();
let probs: Vec<f32> = exps.iter().map(|v| v / sum).collect();
// Weighted random sampling
let mut rng = rand::thread_rng();
let r: f32 = rng.r#gen();
let mut cumsum = 0.0f32;
for (i, &p) in probs.iter().enumerate() {
cumsum += p;
if cumsum > r {
return i as u32;
}
}
// Fallback (rounding edge case)
(vocab_size - 1) as u32
}
fn argmax(data: &[f32]) -> u32 {
data.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
.map(|(i, _)| i as u32)
.unwrap()
}