xserv/crates/xserv-model/src/bin/bench-gpt-oss.rs

use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;

use xserv_distributed::{TpContext, UniqueId, get_unique_id};
use xserv_model::{loader, GptOss, ModelConfig, PagedKVCache, BLOCK_SIZE};
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-gpt-oss <model-dir> [--max-tokens N] [--tp N]");
        std::process::exit(1);
    }

    let model_dir = PathBuf::from(&args[1]);
    let max_tokens: usize = get_arg(&args, "--max-tokens").unwrap_or(32);
    let world: usize = get_arg(&args, "--tp").unwrap_or(2);

    let config = ModelConfig::from_file(&model_dir.join("config.json"));
    let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));

    eprintln!(
        "gpt-oss-20b: layers={}, hidden={}, heads={}/{} kv, experts={}, top_k={}, vocab={}",
        config.num_layers(), config.hidden(), config.num_heads(),
        config.num_kv_heads(), config.num_experts(), config.experts_per_token(),
        config.vocab_size
    );
    eprintln!("TP world={world}, max_tokens={max_tokens}");

    let max_seq_len: usize = 2048;
    let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;

    // TP setup
    let uid = get_unique_id();
    let local_kv = config.num_kv_heads() / world;

    // Spawn worker threads for ranks 1..world
    let mut worker_handles = Vec::new();
    let mut worker_txs = Vec::new();
    for rank in 1..world {
        let (tx, rx) = std::sync::mpsc::channel::<WorkerCmd>();
        let (ack_tx, ack_rx) = std::sync::mpsc::channel::<()>();
        let cfg = config.clone();
        let md = model_dir.clone();
        let uid_copy = uid;
        worker_handles.push((
            std::thread::spawn(move || {
                worker_loop(rank, world, uid_copy, md, cfg, max_seq_len, rx, ack_tx);
            }),
            ack_rx,
        ));
        worker_txs.push(tx);
    }

    // Rank 0 setup
    xserv_cuda::device::set_device(0).unwrap();
    let tp0 = Arc::new(TpContext::init(0, world, uid, 0));
    eprintln!("[rank 0] Loading weights...");
    let weights = loader::load_model_dir(&model_dir, Device::Cpu);
    eprintln!("[rank 0] Loaded {} tensors, building model...", weights.len());
    let model = GptOss::from_weights_tp(config.clone(), weights, 0, world, 0, Some(tp0));
    let total_blocks = max_blocks_per_seq + 64;
    let mut cache = PagedKVCache::new_tp(
        &config, local_kv, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, 0,
    );
    eprintln!("[rank 0] Ready.");

    // Prompt
    let prompt_arg = get_arg::<String>(&args, "--prompt");
    let prompt = prompt_arg.as_deref().unwrap_or("What is the meaning of life?");
    let token_ids = tokenizer.encode(prompt);
    eprintln!("Prompt ({} tokens): {prompt}", token_ids.len());

    // Register sequence
    let slot = 0;
    cache.register_sequence(slot).unwrap();
    broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Register(slot));

    // Teacher-forced diagnostic: prefill (prompt + forced ids) in one shot and
    // report, for each forced position, whether xserv's argmax == the forced
    // (oracle) next token. Removes free-running compounding so it isolates
    // whether per-position logits agree with the llama.cpp trajectory.
    if let Some(forced) = get_arg::<String>(&args, "--forced") {
        let forced_ids: Vec<u32> = forced.split(',').filter_map(|s| s.trim().parse().ok()).collect();
        let mut seq = token_ids.clone();
        seq.extend_from_slice(&forced_ids);
        // Workers must run the same prefill in lockstep (TP AllReduces match up).
        broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Prefill { tokens: seq.clone(), slot });
        let logits = model.forward_prefill_paged(&seq, slot, &mut cache);
        wait_workers(&worker_handles);
        let logits_cpu = logits.to_device(Device::Cpu);
        let vocab = logits.shape()[1];
        let data = logits_cpu.as_slice::<half::bf16>();
        let plen = token_ids.len();
        let mut matches = 0usize;
        let mut total = 0usize;
        // position i predicts seq[i+1]; we check the forced region
        for i in (plen - 1)..(seq.len() - 1) {
            let row = &data[i * vocab..(i + 1) * vocab];
            let argmax = row.iter().enumerate()
                .max_by(|a, b| a.1.to_f32().partial_cmp(&b.1.to_f32()).unwrap())
                .map(|(j, _)| j as u32).unwrap();
            let expected = seq[i + 1];
            let ok = argmax == expected;
            if ok { matches += 1; }
            total += 1;
            eprintln!("pos {i}: xserv_argmax={argmax} oracle={expected} {}", if ok {"OK"} else {"DIFF"});
        }
        eprintln!("\nTeacher-forced top-1 agreement: {matches}/{total} = {:.1}%",
            100.0 * matches as f64 / total as f64);
        broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Shutdown);
        for (h, _) in worker_handles { h.join().unwrap(); }
        return;
    }

    // Teacher-forced DECODE diagnostic: prefill the prompt, then walk the oracle
    // trajectory through the autoregressive decode path (NOT prefill), recording
    // per-position top-1 agreement bucketed by position. Localizes long-context
    // decode degradation (which prefill teacher-forcing cannot see).
    if let Some(forced) = get_arg::<String>(&args, "--forced-decode") {
        let forced_ids: Vec<u32> = forced.split(',').filter_map(|s| s.trim().parse().ok()).collect();
        broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Prefill { tokens: token_ids.clone(), slot });
        let logits = model.forward_prefill_paged(&token_ids, slot, &mut cache);
        wait_workers(&worker_handles);
        let mut pred = sample_greedy_last(&logits); // prediction for forced[0]
        let bucket = 50usize;
        let mut buckets: Vec<(usize, usize)> = Vec::new();
        let (mut matches, mut total) = (0usize, 0usize);
        for (i, &f) in forced_ids.iter().enumerate() {
            let ok = pred == f;
            matches += ok as usize;
            total += 1;
            let b = i / bucket;
            if buckets.len() <= b { buckets.push((0, 0)); }
            buckets[b].0 += ok as usize;
            buckets[b].1 += 1;
            // Teacher-force: feed the oracle token through the decode path.
            let pos = cache.seq_len(slot);
            broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Decode {
                tokens: vec![f], positions: vec![pos], slots: vec![slot],
            });
            let logits = model.forward_decode_paged(&[f], &[pos], &[slot], &mut cache);
            wait_workers(&worker_handles);
            pred = sample_greedy_last(&logits);
        }
        eprintln!("Teacher-forced DECODE agreement: {matches}/{total} = {:.1}%",
            100.0 * matches as f64 / total as f64);
        for (b, (m, t)) in buckets.iter().enumerate() {
            eprintln!("  pos[{:>4}..{:<4}]: {m:>3}/{t:<3} = {:.0}%",
                b * bucket, b * bucket + t, 100.0 * (*m as f64) / (*t as f64));
        }
        broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Shutdown);
        for (h, _) in worker_handles { h.join().unwrap(); }
        return;
    }

    // Prefill
    let t0 = Instant::now();
    broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Prefill {
        tokens: token_ids.clone(), slot,
    });
    let logits = model.forward_prefill_paged(&token_ids, slot, &mut cache);
    wait_workers(&worker_handles);
    let ttft = t0.elapsed();

    let mut next = sample_greedy_last(&logits);
    let mut output_tokens = vec![next];

    eprintln!("TTFT: {:.1}ms", ttft.as_secs_f64() * 1000.0);
    print!("{prompt}");

    // Decode
    let decode_start = Instant::now();
    for _ in 1..max_tokens {
        let text = tokenizer.decode(&[next]);
        print!("{text}");

        if tokenizer.eos_token_id() == Some(next) { break; }

        let pos = cache.seq_len(slot);
        broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Decode {
            tokens: vec![next], positions: vec![pos], slots: vec![slot],
        });
        let logits = model.forward_decode_paged(&[next], &[pos], &[slot], &mut cache);
        wait_workers(&worker_handles);

        next = sample_greedy_last(&logits);
        output_tokens.push(next);
    }
    let decode_elapsed = decode_start.elapsed();
    println!();

    let gen_tokens = output_tokens.len();
    let full_text = tokenizer.decode(&output_tokens);
    eprintln!("\nGenerated text: {full_text}");
    eprintln!("Token IDs: {:?}", &output_tokens[..output_tokens.len().min(20)]);
    let tpot = if gen_tokens > 1 {
        decode_elapsed.as_secs_f64() * 1000.0 / (gen_tokens - 1) as f64
    } else { 0.0 };
    let tok_s = if gen_tokens > 1 {
        (gen_tokens - 1) as f64 / decode_elapsed.as_secs_f64()
    } else { 0.0 };

    eprintln!("\n--- Performance ---");
    eprintln!("Generated: {} tokens", gen_tokens);
    eprintln!("TTFT: {:.1}ms", ttft.as_secs_f64() * 1000.0);
    eprintln!("TPOT: {:.1}ms", tpot);
    eprintln!("Throughput: {:.1} tok/s", tok_s);

    // Cleanup
    broadcast_cmd(&worker_txs, &worker_handles, WorkerCmd::Shutdown);
    for (h, _) in worker_handles {
        h.join().unwrap();
    }
}

// --- Worker infrastructure ---

#[derive(Clone)]
enum WorkerCmd {
    Register(usize),
    Prefill { tokens: Vec<u32>, slot: usize },
    Decode { tokens: Vec<u32>, positions: Vec<usize>, slots: Vec<usize> },
    Shutdown,
}

fn worker_loop(
    rank: usize,
    world: usize,
    uid: UniqueId,
    model_dir: PathBuf,
    config: ModelConfig,
    max_seq_len: usize,
    rx: std::sync::mpsc::Receiver<WorkerCmd>,
    ack_tx: std::sync::mpsc::Sender<()>,
) {
    xserv_cuda::device::set_device(rank as u32).unwrap();
    let tp = Arc::new(TpContext::init(rank, world, uid, rank as u32));
    eprintln!("[rank {rank}] Loading weights...");
    let weights = loader::load_model_dir(&model_dir, Device::Cpu);
    let model = GptOss::from_weights_tp(config.clone(), weights, rank, world, rank as u32, Some(tp));
    let local_kv = config.num_kv_heads() / world;
    let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
    let total_blocks = max_blocks_per_seq + 64;
    let mut cache = PagedKVCache::new_tp(
        &config, local_kv, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, rank as u32,
    );
    eprintln!("[rank {rank}] Ready.");
    ack_tx.send(()).unwrap();

    while let Ok(cmd) = rx.recv() {
        match cmd {
            WorkerCmd::Register(slot) => {
                let _ = cache.register_sequence(slot);
            }
            WorkerCmd::Prefill { tokens, slot } => {
                let _ = model.forward_prefill_paged(&tokens, slot, &mut cache);
            }
            WorkerCmd::Decode { tokens, positions, slots } => {
                let _ = model.forward_decode_paged(&tokens, &positions, &slots, &mut cache);
            }
            WorkerCmd::Shutdown => break,
        }
        ack_tx.send(()).unwrap();
    }
}

fn broadcast_cmd(
    txs: &[std::sync::mpsc::Sender<WorkerCmd>],
    _handles: &[(std::thread::JoinHandle<()>, std::sync::mpsc::Receiver<()>)],
    cmd: WorkerCmd,
) {
    for tx in txs {
        tx.send(cmd.clone()).unwrap();
    }
}

fn wait_workers(handles: &[(std::thread::JoinHandle<()>, std::sync::mpsc::Receiver<()>)]) {
    for (_, rx) in handles {
        rx.recv().unwrap();
    }
}

fn sample_greedy_last(logits: &xserv_tensor::Tensor) -> u32 {
    use half::bf16;
    assert_eq!(logits.ndim(), 2);
    let logits_cpu = logits.to_device(Device::Cpu);
    let vocab_size = logits.shape()[1];
    let seq_len = logits.shape()[0];
    let data = logits_cpu.as_slice::<bf16>();
    let last = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];


    last.iter().enumerate()
        .max_by(|a, b| {
            let af = a.1.to_f32();
            let bf = b.1.to_f32();
            af.partial_cmp(&bf).unwrap_or(std::cmp::Ordering::Equal)
        })
        .map(|(i, _)| i as u32).unwrap()
}

fn get_arg<T: std::str::FromStr>(args: &[String], flag: &str) -> Option<T> {
    args.iter()
        .position(|a| a == flag)
        .and_then(|i| args.get(i + 1))
        .and_then(|s| s.parse().ok())
}