xserv/crates/xserv-model/src/bin/bench-gpt2.rs

use std::path::PathBuf;
use std::time::Instant;
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: bench-gpt2 <model-dir> [--gen-tokens N] [--no-cache]");
        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_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.clone(), weights);
    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));

    // Warmup
    {
        let ids = tokenizer.encode("warmup");
        let _ = model.forward(&ids);
    }

    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",
        "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 (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)
        };

        let num_generated = generated_ids.len();
        let generated_text = tokenizer.decode(&generated_ids);

        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('\\', "\\\\")
            .replace('"', "\\\"")
            .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('"', "\\\""));
        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!(); }

        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,
            &generated_text.replace('\n', " ")[..generated_text.len().min(60)]
        );
    }
    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(),
        xserv_tensor::DType::F32, 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)
}