233 lines
8.9 KiB
Rust
233 lines
8.9 KiB
Rust
use std::path::PathBuf;
|
|
use std::time::Instant;
|
|
use xserv_model::qwen3::sample_greedy;
|
|
use xserv_model::{DecodeGraphState, GpuKVCache, ModelConfig, Qwen3, loader};
|
|
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] [--cuda-graph]");
|
|
std::process::exit(1);
|
|
}
|
|
let model_dir = PathBuf::from(&args[1]);
|
|
let gen_tokens: usize = args
|
|
.iter()
|
|
.position(|a| a == "--gen-tokens")
|
|
.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();
|
|
|
|
let config = ModelConfig::from_file(&model_dir.join("config.json"));
|
|
eprintln!("Loading Qwen3-8B weights...");
|
|
let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
|
|
eprintln!("Loaded {} tensors", weights.len());
|
|
let model = Qwen3::from_weights(config.clone(), weights);
|
|
let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
|
|
|
|
// Warmup
|
|
{
|
|
let ids = tokenizer.encode("warmup");
|
|
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![
|
|
"The capital of France is",
|
|
"Once upon a time in a land far away",
|
|
"Hello, how are you doing today",
|
|
"In a shocking finding, scientists discovered a",
|
|
"The weather today is sunny, so I decided to",
|
|
"Alan Turing was a British mathematician who",
|
|
"The best way to learn programming is",
|
|
"Artificial intelligence will change the world because",
|
|
"The history of the internet began in the",
|
|
"A good morning routine starts with",
|
|
"The stock market crashed because investors",
|
|
"Deep learning is a subset of machine learning that",
|
|
"The president of the United States announced",
|
|
"In the year 2050, humans will",
|
|
"The secret to happiness is",
|
|
"When I was a child, I used to",
|
|
"The most important scientific discovery of the century",
|
|
"Climate change is caused by",
|
|
"The recipe for chocolate cake requires",
|
|
"In conclusion, the evidence suggests that",
|
|
"The cat sat on the mat and",
|
|
"According to recent studies, exercise can",
|
|
"The first step in solving any problem is",
|
|
"Technology has transformed the way we",
|
|
"The novel begins with the protagonist",
|
|
"Education is the most powerful weapon",
|
|
"The ocean covers more than seventy percent of",
|
|
"Last night I had a dream about",
|
|
"The company announced its quarterly earnings",
|
|
"Music has the power to",
|
|
"The difference between success and failure is",
|
|
"In the beginning, there was nothing but",
|
|
"The doctor told me that I should",
|
|
"Python is a popular programming language because",
|
|
"The ancient Romans built roads that",
|
|
"A balanced diet should include",
|
|
"The movie received mixed reviews from critics",
|
|
"Space exploration has led to many",
|
|
"The teacher asked the students to",
|
|
"Global warming is one of the most",
|
|
"The bridge collapsed due to structural",
|
|
"Quantum computing promises to revolutionize",
|
|
"The new policy will affect millions of",
|
|
"During the winter months, it is important to",
|
|
"The human brain contains approximately",
|
|
"Democracy depends on the active participation of",
|
|
"The train arrived at the station exactly",
|
|
"Researchers at MIT have developed a new",
|
|
"The smartphone has become an essential part of",
|
|
"After careful consideration, the committee decided to",
|
|
];
|
|
|
|
println!("[");
|
|
for (i, prompt) in prompts.iter().enumerate() {
|
|
let input_ids = tokenizer.encode(prompt);
|
|
let input_len = input_ids.len();
|
|
|
|
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);
|
|
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 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;
|
|
}
|
|
}
|
|
|
|
let num_generated = generated.len();
|
|
let generated_text = tokenizer.decode(&generated);
|
|
let tbt_us = if !token_times.is_empty() {
|
|
token_times.iter().sum::<u128>() / token_times.len() as u128
|
|
} else {
|
|
0
|
|
};
|
|
let total_gen_us: u128 = ttft_us + token_times.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('\\', "\\\\")
|
|
.replace('"', "\\\"")
|
|
.replace('\n', "\\n")
|
|
.replace('\r', "\\r")
|
|
.replace('\t', "\\t");
|
|
let gen_ids_str: Vec<String> = generated.iter().map(|id| id.to_string()).collect();
|
|
|
|
print!(" {{\"prompt\": \"{}\", ", prompt.replace('"', "\\\""));
|
|
print!("\"input_len\": {input_len}, ");
|
|
print!("\"num_generated\": {num_generated}, ");
|
|
print!("\"generated_ids\": [{}], ", gen_ids_str.join(", "));
|
|
print!("\"generated_text\": \"{gen_text_escaped}\", ");
|
|
print!("\"ttft_us\": {ttft_us}, ");
|
|
print!("\"tbt_us\": {tbt_us}, ");
|
|
print!("\"tpot_us\": {tpot_us}}}");
|
|
if i < prompts.len() - 1 {
|
|
println!(",");
|
|
} else {
|
|
println!();
|
|
}
|
|
|
|
let display_text = generated_text.replace('\n', " ");
|
|
let truncated: String = display_text.chars().take(60).collect();
|
|
eprintln!(
|
|
"[{}/{}] input={input_len}tok gen={num_generated}tok ttft={:.1}ms tbt={:.1}ms | {}",
|
|
i + 1,
|
|
prompts.len(),
|
|
ttft_us as f64 / 1000.0,
|
|
tbt_us as f64 / 1000.0,
|
|
truncated
|
|
);
|
|
}
|
|
println!("]");
|
|
}
|