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

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

@@ -1,9 +1,10 @@
use std::collections::VecDeque;
use std::path::Path;
use std::sync::mpsc;
use xserv_model::{GpuKVCache, ModelConfig, Qwen3};
use std::sync::Once;
use std::time::Instant;
use xserv_model::{GpuKVCache, ModelConfig, Qwen3, SamplingParams, sample};
use xserv_model::loader;
use xserv_model::qwen3::sample_greedy;
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;
@@ -18,6 +19,7 @@ pub struct Engine {
pub struct GenerateRequest {
pub prompt_tokens: Vec<u32>,
pub max_tokens: usize,
pub sampling: SamplingParams,
pub sender: tokio::sync::mpsc::Sender<GenerateEvent>,
}
@@ -31,9 +33,12 @@ struct Sequence {
prompt_tokens: Vec<u32>,
generated_tokens: Vec<u32>,
max_tokens: usize,
sampling: SamplingParams,
kv_cache: GpuKVCache,
sender: tokio::sync::mpsc::Sender<GenerateEvent>,
prefilled: bool,
eos_token_id: Option<u32>,
created_at: Instant,
}
impl Engine {
@@ -84,20 +89,41 @@ impl Engine {
}
}
// Step 4: Process one iteration for all running sequences
// Step 4a: Process prefills (one at a time — different prompt lengths)
// Prefill sequences must be processed individually because they have
// different prompt lengths and each needs a full forward pass.
let mut newly_prefilled = Vec::new();
for seq in running.iter_mut() {
if !seq.prefilled {
// Prefill
let logits = self.model.forward_gpu_cache(&seq.prompt_tokens, &mut seq.kv_cache);
let next = sample_greedy(&logits);
let next = sample(&logits, &seq.sampling);
seq.generated_tokens.push(next);
seq.prefilled = true;
self.emit_token(seq, next);
} else {
// Decode one token
newly_prefilled.push(seq.id);
}
}
// Step 4b: Process decode (one token per sequence)
// Currently per-sequence (each has different KV cache length).
// TODO(Phase 14): With Flash Attention, batch all decode tokens into
// one forward pass — batch the compute-heavy ops (projections, FFN)
// and use FlashDecoding for per-seq variable-length attention.
let decode_count = running.iter()
.filter(|s| s.prefilled && !newly_prefilled.contains(&s.id))
.count();
if decode_count > 0 {
static LOG_ONCE: Once = Once::new();
LOG_ONCE.call_once(|| {
eprintln!("[scheduler] decode batching active (per-seq until Flash Attention)");
});
eprintln!("[scheduler] decode batch_size={}", decode_count);
}
for seq in running.iter_mut() {
if seq.prefilled && !newly_prefilled.contains(&seq.id) {
let last = *seq.generated_tokens.last().unwrap();
let logits = self.model.forward_gpu_cache(&[last], &mut seq.kv_cache);
let next = sample_greedy(&logits);
let next = sample(&logits, &seq.sampling);
seq.generated_tokens.push(next);
self.emit_token(seq, next);
}
@@ -120,15 +146,18 @@ impl Engine {
fn make_sequence(&self, req: GenerateRequest, next_id: &mut u64) -> Sequence {
let id = *next_id;
*next_id += 1;
let kv_cache = GpuKVCache::new(&self.config, self.max_seq_len, DType::BF16);
let kv_cache = GpuKVCache::new(&self.config, self.max_seq_len, DType::BF16, 0);
Sequence {
id,
prompt_tokens: req.prompt_tokens,
generated_tokens: Vec::new(),
max_tokens: req.max_tokens,
sampling: req.sampling,
kv_cache,
sender: req.sender,
prefilled: false,
eos_token_id: self.tokenizer.eos_token_id(),
created_at: Instant::now(),
}
}
@@ -157,5 +186,5 @@ fn is_finished(seq: &Sequence) -> bool {
if seq.generated_tokens.len() >= seq.max_tokens { return true; }
// Check EOS — need tokenizer info. Use a simple heuristic:
// If sender is closed (receiver dropped), also consider finished.
seq.sender.is_closed() || last == 151645 // Qwen3 EOS token ID (hardcoded for now)
seq.sender.is_closed() || seq.eos_token_id == Some(last)
}