model: paged KV cache with CPU swap pool, decode graph, qwen3 updates
- paged_kv_cache: new block-paged KV cache; adds a pinned-host swap pool with
a second BlockAllocator, per-sequence Location {Gpu,Cpu}, and lossless
swap_out/swap_in (block-granular D2H/H2D) for vLLM-style preemption.
bytes_per_block helper exposes per-block cost for VRAM-based sizing.
- decode_graph: CUDA-graph decode path.
- qwen3/gpt2/kv_cache: paged prefill/decode forward + related updates.
- tokenizer/bins: BPE updates, new xserv-chat CLI, bench-qwen3 tweaks.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -1,14 +1,14 @@
|
||||
use std::path::PathBuf;
|
||||
use std::time::Instant;
|
||||
use xserv_model::qwen3::sample_greedy;
|
||||
use xserv_model::{loader, GpuKVCache, ModelConfig, Qwen3};
|
||||
use xserv_model::{loader, DecodeGraphState, GpuKVCache, ModelConfig, Qwen3};
|
||||
use xserv_tensor::{DType, Device};
|
||||
use xserv_tokenizer::Tokenizer;
|
||||
|
||||
fn main() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
if args.len() < 2 {
|
||||
eprintln!("Usage: bench-qwen3 <model-dir> [--gen-tokens N]");
|
||||
eprintln!("Usage: bench-qwen3 <model-dir> [--gen-tokens N] [--cuda-graph]");
|
||||
std::process::exit(1);
|
||||
}
|
||||
let model_dir = PathBuf::from(&args[1]);
|
||||
@@ -18,6 +18,7 @@ fn main() {
|
||||
.and_then(|i| args.get(i + 1))
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(20);
|
||||
let use_cuda_graph = args.iter().any(|a| a == "--cuda-graph");
|
||||
|
||||
xserv_cuda::device::set_device(0).unwrap();
|
||||
|
||||
@@ -34,6 +35,18 @@ fn main() {
|
||||
let mut cache = GpuKVCache::new(&config, 256, DType::BF16, 0);
|
||||
let _ = model.forward_gpu_cache(&ids, &mut cache);
|
||||
}
|
||||
|
||||
// CUDA Graph setup
|
||||
let layer_ptrs = model.layer_weight_ptrs();
|
||||
let (norm_w, lm_head, embed, cos, sin) = model.graph_capture_ptrs();
|
||||
let mut decode_graph = if use_cuda_graph {
|
||||
eprintln!("CUDA Graph mode enabled");
|
||||
Some(DecodeGraphState::new(&config))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let mut graph_captured = false;
|
||||
|
||||
eprintln!("Warmup done. Running benchmark...");
|
||||
|
||||
let prompts: Vec<&str> = vec![
|
||||
@@ -96,6 +109,12 @@ fn main() {
|
||||
|
||||
let mut cache = GpuKVCache::new(&config, 256, DType::BF16, 0);
|
||||
|
||||
// Reset graph state for new prompt
|
||||
graph_captured = false;
|
||||
if let Some(ref mut g) = decode_graph {
|
||||
g.invalidate();
|
||||
}
|
||||
|
||||
// Prefill
|
||||
let t0 = Instant::now();
|
||||
let logits = model.forward_gpu_cache(&input_ids, &mut cache);
|
||||
@@ -109,8 +128,35 @@ fn main() {
|
||||
for _ in 1..gen_tokens {
|
||||
let last = *generated.last().unwrap();
|
||||
let t_start = Instant::now();
|
||||
let logits = model.forward_gpu_cache(&[last], &mut cache);
|
||||
let next = sample_greedy(&logits);
|
||||
|
||||
let next = if let Some(ref mut graph) = decode_graph {
|
||||
if !graph_captured {
|
||||
// First decode token: run ungraphed, then capture
|
||||
let logits = model.forward_gpu_cache(&[last], &mut cache);
|
||||
graph_captured = true;
|
||||
graph.capture(&layer_ptrs, norm_w, lm_head, embed, cos, sin);
|
||||
sample_greedy(&logits)
|
||||
} else {
|
||||
// Replay captured graphs
|
||||
let pos = cache.seq_len() as u32;
|
||||
graph.execute(last, pos, &mut cache, &layer_ptrs, embed, config.vocab_size as i32, config.hidden() as i32);
|
||||
cache.advance_seq_len(1);
|
||||
// Read logits from graph buffer
|
||||
let vocab_size = config.vocab_size;
|
||||
let mut logits_bytes = vec![0u8; vocab_size * 2];
|
||||
graph.logits_buffer().copy_to_host(&mut logits_bytes).unwrap();
|
||||
let logits_data: &[half::bf16] = unsafe {
|
||||
std::slice::from_raw_parts(logits_bytes.as_ptr() as *const half::bf16, vocab_size)
|
||||
};
|
||||
logits_data.iter().enumerate()
|
||||
.max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
|
||||
.map(|(idx, _)| idx as u32).unwrap()
|
||||
}
|
||||
} else {
|
||||
let logits = model.forward_gpu_cache(&[last], &mut cache);
|
||||
sample_greedy(&logits)
|
||||
};
|
||||
|
||||
token_times.push(t_start.elapsed().as_micros());
|
||||
generated.push(next);
|
||||
if tokenizer.eos_token_id() == Some(next) { break; }
|
||||
|
||||
419
crates/xserv-model/src/bin/xserv-chat.rs
Normal file
419
crates/xserv-model/src/bin/xserv-chat.rs
Normal file
@@ -0,0 +1,419 @@
|
||||
use std::io::{self, IsTerminal, Write};
|
||||
use std::path::PathBuf;
|
||||
|
||||
use xserv_model::{loader, sample, ModelConfig, PagedKVCache, Qwen3, SamplingParams, BLOCK_SIZE};
|
||||
use xserv_tensor::{DType, Device};
|
||||
use xserv_tokenizer::Tokenizer;
|
||||
|
||||
const SLOT: usize = 0;
|
||||
|
||||
struct CliOptions {
|
||||
model_dir: PathBuf,
|
||||
max_tokens: usize,
|
||||
max_seq_len: usize,
|
||||
sampling: SamplingParams,
|
||||
system_prompt: Option<String>,
|
||||
enable_thinking: bool,
|
||||
color: bool,
|
||||
}
|
||||
|
||||
enum Finish {
|
||||
Stop { token_id: u32 },
|
||||
Length,
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let opts = parse_args();
|
||||
|
||||
xserv_cuda::device::set_device(0).unwrap();
|
||||
let info = xserv_cuda::device::device_info(0).unwrap();
|
||||
eprintln!(
|
||||
"GPU: {} ({} MB free)",
|
||||
info.name,
|
||||
info.free_memory / 1024 / 1024
|
||||
);
|
||||
|
||||
let config = ModelConfig::from_file(&opts.model_dir.join("config.json"));
|
||||
let model_type = config.model_type.as_deref().unwrap_or("unknown");
|
||||
if !model_type.contains("qwen") {
|
||||
eprintln!("xserv-chat currently supports Qwen-style ChatML models only; got model_type={model_type}");
|
||||
std::process::exit(2);
|
||||
}
|
||||
|
||||
let max_seq_len = opts.max_seq_len.min(config.max_seq_len()).max(1);
|
||||
eprintln!(
|
||||
"Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}, max_seq_len={}",
|
||||
config.num_layers(),
|
||||
config.hidden(),
|
||||
config.num_heads(),
|
||||
config.num_kv_heads(),
|
||||
config.vocab_size,
|
||||
max_seq_len
|
||||
);
|
||||
|
||||
eprintln!("Loading weights...");
|
||||
let weights = loader::load_model_dir(&opts.model_dir, Device::Cuda(0));
|
||||
eprintln!("Loaded {} tensors", weights.len());
|
||||
let model = Qwen3::from_weights(config.clone(), weights);
|
||||
|
||||
let tokenizer = Tokenizer::from_file(&opts.model_dir.join("tokenizer.json"));
|
||||
let mut cache = new_paged_cache(&config, max_seq_len);
|
||||
cache.register_sequence(SLOT).expect("register chat slot");
|
||||
let use_color = opts.color && io::stdout().is_terminal();
|
||||
|
||||
eprintln!("Ready (paged KV cache, persistent chat slot).");
|
||||
eprintln!("Commands: /exit, /quit, /clear\n");
|
||||
|
||||
loop {
|
||||
print!("user> ");
|
||||
io::stdout().flush().unwrap();
|
||||
|
||||
let mut input = String::new();
|
||||
if io::stdin().read_line(&mut input).unwrap() == 0 {
|
||||
break;
|
||||
}
|
||||
let input = input.trim();
|
||||
if input.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
match input {
|
||||
"/exit" | "/quit" | "exit" | "quit" => break,
|
||||
"/clear" => {
|
||||
cache.free_sequence(SLOT);
|
||||
cache.register_sequence(SLOT).expect("register chat slot");
|
||||
eprintln!("history and KV cache cleared");
|
||||
continue;
|
||||
}
|
||||
"/help" => {
|
||||
print_help();
|
||||
continue;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
let include_system = cache.seq_len(SLOT) == 0;
|
||||
let prompt = build_turn_prompt(
|
||||
opts.system_prompt.as_deref(),
|
||||
include_system,
|
||||
input,
|
||||
opts.enable_thinking,
|
||||
);
|
||||
let prompt_tokens = tokenizer.encode(&prompt);
|
||||
if prompt_tokens.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let used = cache.seq_len(SLOT);
|
||||
let remaining = max_seq_len.saturating_sub(used);
|
||||
if prompt_tokens.len() >= remaining {
|
||||
eprintln!(
|
||||
"context full: {used}/{max_seq_len} tokens used, new turn needs {} tokens; use /clear",
|
||||
prompt_tokens.len()
|
||||
);
|
||||
continue;
|
||||
}
|
||||
let max_new_tokens = opts.max_tokens.min(remaining - prompt_tokens.len());
|
||||
|
||||
print!("assistant> ");
|
||||
io::stdout().flush().unwrap();
|
||||
let finish = generate_with_paged_cache(
|
||||
&model,
|
||||
&mut cache,
|
||||
&tokenizer,
|
||||
&prompt_tokens,
|
||||
&opts.sampling,
|
||||
max_new_tokens,
|
||||
use_color,
|
||||
);
|
||||
match finish {
|
||||
Finish::Stop { token_id } => {
|
||||
append_after_stop(&model, &mut cache, &tokenizer, max_seq_len, token_id);
|
||||
}
|
||||
Finish::Length => {
|
||||
append_text_to_cache(&model, &mut cache, &tokenizer, max_seq_len, "<|im_end|>\n");
|
||||
}
|
||||
}
|
||||
println!();
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_args() -> CliOptions {
|
||||
let args: Vec<String> = std::env::args().skip(1).collect();
|
||||
if args.is_empty() || args.iter().any(|a| a == "--help" || a == "-h") {
|
||||
print_usage_and_exit(0);
|
||||
}
|
||||
|
||||
let mut model_dir = None;
|
||||
let mut max_tokens = 256usize;
|
||||
let mut max_seq_len = 2048usize;
|
||||
let mut temperature = 0.0f32;
|
||||
let mut top_k = 0usize;
|
||||
let mut top_p = 1.0f32;
|
||||
let mut system_prompt = None;
|
||||
let mut enable_thinking = false;
|
||||
let mut color = true;
|
||||
|
||||
let mut i = 0;
|
||||
while i < args.len() {
|
||||
match args[i].as_str() {
|
||||
"-m" | "--model" => {
|
||||
i += 1;
|
||||
model_dir = args.get(i).map(PathBuf::from);
|
||||
}
|
||||
"--max-tokens" => {
|
||||
i += 1;
|
||||
max_tokens = parse_value(&args, i, "--max-tokens");
|
||||
}
|
||||
"--max-seq-len" => {
|
||||
i += 1;
|
||||
max_seq_len = parse_value(&args, i, "--max-seq-len");
|
||||
}
|
||||
"--temperature" => {
|
||||
i += 1;
|
||||
temperature = parse_value(&args, i, "--temperature");
|
||||
}
|
||||
"--top-k" => {
|
||||
i += 1;
|
||||
top_k = parse_value(&args, i, "--top-k");
|
||||
}
|
||||
"--top-p" => {
|
||||
i += 1;
|
||||
top_p = parse_value(&args, i, "--top-p");
|
||||
}
|
||||
"--system" => {
|
||||
i += 1;
|
||||
system_prompt = args.get(i).cloned();
|
||||
if system_prompt.is_none() {
|
||||
eprintln!("missing value for --system");
|
||||
std::process::exit(2);
|
||||
}
|
||||
}
|
||||
"--think" => {
|
||||
enable_thinking = true;
|
||||
}
|
||||
"--no-color" => {
|
||||
color = false;
|
||||
}
|
||||
arg if arg.starts_with('-') => {
|
||||
eprintln!("unknown option: {arg}");
|
||||
print_usage_and_exit(2);
|
||||
}
|
||||
arg => {
|
||||
if model_dir.is_some() {
|
||||
eprintln!("unexpected extra argument: {arg}");
|
||||
print_usage_and_exit(2);
|
||||
}
|
||||
model_dir = Some(PathBuf::from(arg));
|
||||
}
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
|
||||
CliOptions {
|
||||
model_dir: model_dir.unwrap_or_else(|| {
|
||||
eprintln!("missing model directory");
|
||||
print_usage_and_exit(2);
|
||||
}),
|
||||
max_tokens: max_tokens.max(1),
|
||||
max_seq_len: max_seq_len.max(1),
|
||||
sampling: SamplingParams {
|
||||
temperature,
|
||||
top_k,
|
||||
top_p,
|
||||
},
|
||||
system_prompt,
|
||||
enable_thinking,
|
||||
color,
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_value<T: std::str::FromStr>(args: &[String], i: usize, name: &str) -> T {
|
||||
args.get(i).and_then(|s| s.parse().ok()).unwrap_or_else(|| {
|
||||
eprintln!("invalid or missing value for {name}");
|
||||
std::process::exit(2);
|
||||
})
|
||||
}
|
||||
|
||||
fn print_usage_and_exit(code: i32) -> ! {
|
||||
eprintln!(
|
||||
"Usage: xserv-chat <model-dir> [options]\n\
|
||||
\n\
|
||||
Options:\n\
|
||||
\t-m, --model DIR Model directory\n\
|
||||
\t--max-tokens N Max generated tokens per turn (default: 256)\n\
|
||||
\t--max-seq-len N Persistent KV context length (default: 2048)\n\
|
||||
\t--temperature F Sampling temperature, 0 = greedy (default: 0)\n\
|
||||
\t--top-k N Top-k sampling, 0 = disabled (default: 0)\n\
|
||||
\t--top-p F Top-p sampling (default: 1.0)\n\
|
||||
\t--system TEXT System prompt for the first turn after start or /clear\n\
|
||||
\t--think Let Qwen3 emit thinking; rendered gray on terminals\n\
|
||||
\t--no-color Disable ANSI color for thinking output\n\
|
||||
\t-h, --help Show this help"
|
||||
);
|
||||
std::process::exit(code);
|
||||
}
|
||||
|
||||
fn print_help() {
|
||||
eprintln!("Commands:");
|
||||
eprintln!(" /clear clear chat history and free/recreate the paged KV slot");
|
||||
eprintln!(" /exit quit");
|
||||
eprintln!(" /quit quit");
|
||||
}
|
||||
|
||||
fn new_paged_cache(config: &ModelConfig, max_seq_len: usize) -> PagedKVCache {
|
||||
let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
||||
let total_blocks = (max_blocks_per_seq + 1).max(2);
|
||||
// Single-slot interactive CLI: no swap pool (cpu_total_blocks = 0).
|
||||
PagedKVCache::new(config, total_blocks, 0, 1, max_blocks_per_seq, DType::BF16, 0)
|
||||
}
|
||||
|
||||
fn build_turn_prompt(
|
||||
system: Option<&str>,
|
||||
include_system: bool,
|
||||
user_input: &str,
|
||||
enable_thinking: bool,
|
||||
) -> String {
|
||||
let mut prompt = String::new();
|
||||
if include_system {
|
||||
if let Some(system) = system {
|
||||
if !system.trim().is_empty() {
|
||||
prompt.push_str("<|im_start|>system\n");
|
||||
prompt.push_str(system.trim());
|
||||
prompt.push_str("<|im_end|>\n");
|
||||
}
|
||||
}
|
||||
}
|
||||
prompt.push_str("<|im_start|>user\n");
|
||||
prompt.push_str(user_input);
|
||||
prompt.push_str("<|im_end|>\n");
|
||||
prompt.push_str("<|im_start|>assistant\n");
|
||||
if !enable_thinking {
|
||||
prompt.push_str("<think>\n\n</think>\n\n");
|
||||
}
|
||||
prompt
|
||||
}
|
||||
|
||||
fn generate_with_paged_cache(
|
||||
model: &Qwen3,
|
||||
cache: &mut PagedKVCache,
|
||||
tokenizer: &Tokenizer,
|
||||
prompt_tokens: &[u32],
|
||||
sampling: &SamplingParams,
|
||||
max_tokens: usize,
|
||||
use_color: bool,
|
||||
) -> Finish {
|
||||
let logits = model.forward_prefill_paged(prompt_tokens, SLOT, cache);
|
||||
let mut next = sample(&logits, sampling);
|
||||
let mut decode_buffer = Vec::new();
|
||||
let mut in_thinking = false;
|
||||
|
||||
for _ in 0..max_tokens {
|
||||
let position = cache.seq_len(SLOT);
|
||||
let logits = model.forward_decode_paged(&[next], &[position], &[SLOT], cache);
|
||||
if is_stop_token(tokenizer, next) {
|
||||
print_stream_text(
|
||||
&tokenizer.flush_decode_stream(&mut decode_buffer),
|
||||
in_thinking,
|
||||
use_color,
|
||||
);
|
||||
io::stdout().flush().unwrap();
|
||||
return Finish::Stop { token_id: next };
|
||||
}
|
||||
|
||||
print_generated_token(
|
||||
tokenizer,
|
||||
next,
|
||||
&mut decode_buffer,
|
||||
&mut in_thinking,
|
||||
use_color,
|
||||
);
|
||||
io::stdout().flush().unwrap();
|
||||
next = sample(&logits, sampling);
|
||||
}
|
||||
|
||||
print_stream_text(
|
||||
&tokenizer.flush_decode_stream(&mut decode_buffer),
|
||||
in_thinking,
|
||||
use_color,
|
||||
);
|
||||
io::stdout().flush().unwrap();
|
||||
Finish::Length
|
||||
}
|
||||
|
||||
fn append_after_stop(
|
||||
model: &Qwen3,
|
||||
cache: &mut PagedKVCache,
|
||||
tokenizer: &Tokenizer,
|
||||
max_seq_len: usize,
|
||||
stop_token_id: u32,
|
||||
) {
|
||||
if tokenizer.special_token_id("<|im_end|>") == Some(stop_token_id) {
|
||||
append_text_to_cache(model, cache, tokenizer, max_seq_len, "\n");
|
||||
}
|
||||
}
|
||||
|
||||
fn append_text_to_cache(
|
||||
model: &Qwen3,
|
||||
cache: &mut PagedKVCache,
|
||||
tokenizer: &Tokenizer,
|
||||
max_seq_len: usize,
|
||||
text: &str,
|
||||
) {
|
||||
let tokens = tokenizer.encode(text);
|
||||
if tokens.is_empty() || cache.seq_len(SLOT) + tokens.len() > max_seq_len {
|
||||
return;
|
||||
}
|
||||
let _ = model.forward_prefill_paged(&tokens, SLOT, cache);
|
||||
}
|
||||
|
||||
fn print_generated_token(
|
||||
tokenizer: &Tokenizer,
|
||||
token_id: u32,
|
||||
decode_buffer: &mut Vec<u8>,
|
||||
in_thinking: &mut bool,
|
||||
use_color: bool,
|
||||
) {
|
||||
if tokenizer.special_token_id("<think>") == Some(token_id) {
|
||||
print_stream_text(
|
||||
&tokenizer.flush_decode_stream(decode_buffer),
|
||||
*in_thinking,
|
||||
use_color,
|
||||
);
|
||||
*in_thinking = true;
|
||||
print_stream_text("<think>", true, use_color);
|
||||
return;
|
||||
}
|
||||
|
||||
if tokenizer.special_token_id("</think>") == Some(token_id) {
|
||||
print_stream_text(
|
||||
&tokenizer.flush_decode_stream(decode_buffer),
|
||||
*in_thinking,
|
||||
use_color,
|
||||
);
|
||||
print_stream_text("</think>", true, use_color);
|
||||
*in_thinking = false;
|
||||
return;
|
||||
}
|
||||
|
||||
let text = tokenizer.decode_token_stream(token_id, decode_buffer);
|
||||
print_stream_text(&text, *in_thinking, use_color);
|
||||
}
|
||||
|
||||
fn print_stream_text(text: &str, in_thinking: bool, use_color: bool) {
|
||||
if text.is_empty() {
|
||||
return;
|
||||
}
|
||||
if in_thinking && use_color {
|
||||
print!("\x1b[90m{text}\x1b[0m");
|
||||
} else {
|
||||
print!("{text}");
|
||||
}
|
||||
}
|
||||
|
||||
fn is_stop_token(tokenizer: &Tokenizer, token_id: u32) -> bool {
|
||||
tokenizer.eos_token_id() == Some(token_id)
|
||||
|| tokenizer.special_token_id("<|im_end|>") == Some(token_id)
|
||||
|| tokenizer.special_token_id("<|endoftext|>") == Some(token_id)
|
||||
|| tokenizer.special_token_id("<|end_of_text|>") == Some(token_id)
|
||||
}
|
||||
Reference in New Issue
Block a user