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

phase 9: KV cache + autoregressive generation

Browse Source 
- KVCache: per-layer, per-head storage with append + reconstruct
- forward_with_cache: prefill (full prompt) + decode (single token) modes
- Fixed data layout bug: per-head vectors avoid cross-head interleaving
- CLI updated to use KV cache by default
- bench-gpt2 supports --no-cache flag for comparison

Benchmark results (50 prompts × 20 tokens):
- KV cache vs no-cache: 50/50 bit-identical (cache is correct)
- 18x speedup: TTFT 400→24ms, TBT 407→22ms, throughput 2.5→44 tok/s
- vs HF transformers: 40/50 match (10 are FP divergence, avg logit gap 0.20)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-21 23:39:41 +08:00
parent cb12250ef0
commit 64084d3489
7 changed files with 395 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

107
crates/xserv-model/src/bin/bench-gpt2.rs
View File

@@ -1,6 +1,6 @@
use std::path::PathBuf;
use std::time::Instant;
use xserv_model::gpt2::sample_greedy;
use xserv_model::gpt2::{sample_greedy, KVCache};
use xserv_model::{loader, GPT2, ModelConfig};
use xserv_tensor::Device;
use xserv_tokenizer::Tokenizer;
@@ -8,7 +8,7 @@ use xserv_tokenizer::Tokenizer;
fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() < 2 {
eprintln!("Usage: bench-gpt2 <model-dir> [--gen-tokens N]");
eprintln!("Usage: bench-gpt2 <model-dir> [--gen-tokens N] [--no-cache]");
std::process::exit(1);
}
let model_dir = PathBuf::from(&args[1]);
@@ -18,12 +18,13 @@ fn main() {
.and_then(|i| args.get(i + 1))
.and_then(|s| s.parse().ok())
.unwrap_or(20);
let use_cache = !args.iter().any(|a| a == "--no-cache");
xserv_cuda::device::set_device(0).unwrap();
let config = ModelConfig::from_file(&model_dir.join("config.json"));
let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
let model = GPT2::from_weights(config, weights);
let model = GPT2::from_weights(config.clone(), weights);
let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
// Warmup
@@ -32,7 +33,9 @@ fn main() {
let _ = model.forward(&ids);
}
let prompts = vec![
eprintln!("mode: {}", if use_cache { "KV cache" } else { "no cache" });
let prompts: Vec<&str> = vec![
"The capital of France is",
"Once upon a time in a land far away",
"Hello, how are you doing today",
@@ -85,44 +88,25 @@ fn main() {
"After careful consideration, the committee decided to",
];
// JSON output
println!("[");
for (i, prompt) in prompts.iter().enumerate() {
let input_ids = tokenizer.encode(prompt);
let input_len = input_ids.len();
let mut all_ids = input_ids.clone();
// TTFT: time for first forward pass (prefill)
let t0 = Instant::now();
let logits = model.forward(&all_ids);
let first_token = sample_greedy(&logits);
let ttft_us = t0.elapsed().as_micros();
all_ids.push(first_token);
let (generated_ids, ttft_us, token_times_us) = if use_cache {
generate_with_cache(&model, &config, &tokenizer, &input_ids, gen_tokens)
} else {
generate_no_cache(&model, &tokenizer, &input_ids, gen_tokens)
};
// Generate remaining tokens, measure each
let mut token_times_us = Vec::new();
for _ in 1..gen_tokens {
let t_start = Instant::now();
let logits = model.forward(&all_ids);
let next = sample_greedy(&logits);
let elapsed = t_start.elapsed().as_micros();
token_times_us.push(elapsed);
all_ids.push(next);
if tokenizer.eos_token_id() == Some(next) {
break;
}
}
let generated_ids: Vec<u32> = all_ids[input_len..].to_vec();
let generated_text = tokenizer.decode(&generated_ids);
let num_generated = generated_ids.len();
let generated_text = tokenizer.decode(&generated_ids);
let total_gen_us: u128 = ttft_us + token_times_us.iter().sum::<u128>();
let tpot_us = if num_generated > 0 { total_gen_us / num_generated as u128 } else { 0 };
let tbt_us = if !token_times_us.is_empty() {
token_times_us.iter().sum::<u128>() / token_times_us.len() as u128
} else { 0 };
let total_gen_us: u128 = ttft_us + token_times_us.iter().sum::<u128>();
let tpot_us = if num_generated > 0 { total_gen_us / num_generated as u128 } else { 0 };
let gen_text_escaped = generated_text
.replace('\\', "\\\\")
@@ -130,7 +114,6 @@ fn main() {
.replace('\n', "\\n")
.replace('\r', "\\r")
.replace('\t', "\\t");
let gen_ids_str: Vec<String> = generated_ids.iter().map(|id| id.to_string()).collect();
print!(" {{\"prompt\": \"{}\", ", prompt.replace('"', "\\\""));
@@ -153,3 +136,63 @@ fn main() {
}
println!("]");
}
fn generate_with_cache(
model: &GPT2, config: &ModelConfig, tokenizer: &Tokenizer,
input_ids: &[u32], gen_tokens: usize,
) -> (Vec<u32>, u128, Vec<u128>) {
let mut cache = KVCache::new(
config.num_layers(), config.num_heads(), config.head_dim(),
Device::Cuda(0),
);
// Prefill
let t0 = Instant::now();
let logits = model.forward_with_cache(input_ids, &mut cache);
let first_token = sample_greedy(&logits);
let ttft_us = t0.elapsed().as_micros();
let mut generated = vec![first_token];
let mut token_times = Vec::new();
// Decode
for _ in 1..gen_tokens {
let last = *generated.last().unwrap();
let t_start = Instant::now();
let logits = model.forward_with_cache(&[last], &mut cache);
let next = sample_greedy(&logits);
token_times.push(t_start.elapsed().as_micros());
generated.push(next);
if tokenizer.eos_token_id() == Some(next) { break; }
}
(generated, ttft_us, token_times)
}
fn generate_no_cache(
model: &GPT2, tokenizer: &Tokenizer,
input_ids: &[u32], gen_tokens: usize,
) -> (Vec<u32>, u128, Vec<u128>) {
let mut all_ids = input_ids.to_vec();
let t0 = Instant::now();
let logits = model.forward(&all_ids);
let first_token = sample_greedy(&logits);
let ttft_us = t0.elapsed().as_micros();
all_ids.push(first_token);
let mut generated = vec![first_token];
let mut token_times = Vec::new();
for _ in 1..gen_tokens {
let t_start = Instant::now();
let logits = model.forward(&all_ids);
let next = sample_greedy(&logits);
token_times.push(t_start.elapsed().as_micros());
all_ids.push(next);
generated.push(next);
if tokenizer.eos_token_id() == Some(next) { break; }
}
(generated, ttft_us, token_times)
}

54
crates/xserv-model/src/bin/xserv-cli.rs
View File

@@ -1,21 +1,20 @@
use std::io::{self, Write};
use std::path::PathBuf;
use xserv_model::{GPT2, ModelConfig};
use xserv_model::loader;
use xserv_model::gpt2::sample_greedy;
use xserv_tokenizer::Tokenizer;
use xserv_model::gpt2::{sample_greedy, KVCache};
use xserv_model::{loader, GPT2, ModelConfig};
use xserv_tensor::Device;
use xserv_tokenizer::Tokenizer;
fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() < 2 {
eprintln!("Usage: xserv-cli <model-dir> [--max-tokens N]");
eprintln!(" model-dir: path to HF model directory (containing model.safetensors, config.json, tokenizer.json)");
std::process::exit(1);
}
let model_dir = PathBuf::from(&args[1]);
let max_tokens: usize = args.iter()
let max_tokens: usize = args
.iter()
.position(|a| a == "--max-tokens")
.and_then(|i| args.get(i + 1))
.and_then(|s| s.parse().ok())
@@ -25,26 +24,24 @@ fn main() {
let info = xserv_cuda::device::device_info(0).unwrap();
eprintln!("GPU: {} ({} MB free)", info.name, info.free_memory / 1024 / 1024);
// Load config
let config = ModelConfig::from_file(&model_dir.join("config.json"));
eprintln!("Model: {:?}, layers={}, hidden={}, heads={}, vocab={}",
config.model_type, config.num_layers(), config.hidden(),
config.num_heads(), config.vocab_size);
eprintln!(
"Model: {:?}, layers={}, hidden={}, heads={}, vocab={}",
config.model_type,
config.num_layers(),
config.hidden(),
config.num_heads(),
config.vocab_size
);
// Load weights
eprintln!("Loading weights...");
let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
eprintln!("Loaded {} tensors", weights.len());
// GPT-2 uses weight names without "model." prefix
let model = GPT2::from_weights(config, weights);
// Load tokenizer
let model = GPT2::from_weights(config.clone(), weights);
let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
eprintln!("Tokenizer loaded (vocab_size={})", tokenizer.vocab_size());
eprintln!("Ready.\n");
eprintln!("Ready (KV cache enabled).\n");
// Interactive loop
loop {
print!("xserv> ");
io::stdout().flush().unwrap();
@@ -56,22 +53,27 @@ fn main() {
if input.is_empty() { continue; }
if input == "quit" || input == "exit" { break; }
let mut token_ids = tokenizer.encode(input);
let token_ids = tokenizer.encode(input);
let mut cache = KVCache::new(
config.num_layers(), config.num_heads(), config.head_dim(),
Device::Cuda(0),
);
// Prefill
let logits = model.forward_with_cache(&token_ids, &mut cache);
let mut next = sample_greedy(&logits);
print!("{input}");
io::stdout().flush().unwrap();
for _ in 0..max_tokens {
let logits = model.forward(&token_ids);
let next = sample_greedy(&logits);
token_ids.push(next);
let text = tokenizer.decode(&[next]);
print!("{text}");
io::stdout().flush().unwrap();
if tokenizer.eos_token_id() == Some(next) {
break;
}
if tokenizer.eos_token_id() == Some(next) { break; }
let logits = model.forward_with_cache(&[next], &mut cache);
next = sample_greedy(&logits);
}
println!();
}