inferencelayer 0.2.9

//! End-to-end decode benchmark for the CUDA arm: `cuda-llm-bench <model-dir> [prompt-len] [ngen]`.
//!
//! Prints the generated token ids (so the output can be diffed against the WGSL engine for a
//! correctness gate) and the sustained decode tok/s — the number that goes head to head with
//! vLLM's measured 155.9 on the same card.
use anyhow::Result;
use inferencelayer::cuda_llm::CudaQwen3;

fn main() -> Result<()> {
    let a: Vec<String> = std::env::args().collect();
    if a.len() < 2 {
        eprintln!("usage: cuda-llm-bench <model-dir> [prompt-tokens=16] [ngen=128]");
        std::process::exit(2);
    }
    let dir = std::path::PathBuf::from(&a[1]);
    let plen: usize = a.get(2).and_then(|s| s.parse().ok()).unwrap_or(16);
    let ngen: usize = a.get(3).and_then(|s| s.parse().ok()).unwrap_or(128);

    // Resolve the prompt FIRST: max_seq (KV cache, scores, token log) is sized from it, and
    // sizing from the argv default while actually running a longer prompt walks off the end.
    let prompt: Vec<u32> = match a.get(4) {
        Some(f) => std::fs::read_to_string(f)?
            .trim()
            .split(',')
            .filter(|s| !s.is_empty())
            .map(|s| s.trim().parse::<u32>().unwrap())
            .collect(),
        None => (0..plen).map(|i| (1000 + i as u32) % 1000 + 100).collect(),
    };
    let plen = prompt.len();

    let t0 = std::time::Instant::now();
    let mut m = CudaQwen3::load(&dir, plen + ngen + 16)?;
    m.sync()?;
    println!("loaded in {:.1}s  cfg={:?}", t0.elapsed().as_secs_f64(), m.cfg());

    // Deterministic synthetic prompt (token ids only — this binary is about decode throughput,
    // not tokenization). Prefill runs as sequential decode: irrelevant to the metric, one path.
    let tp = std::time::Instant::now();
    for (i, &t) in prompt.iter().enumerate() {
        m.seed(t, i)?;
        m.step(0, 0)?;
    }
    m.sync()?;
    if std::env::var("VERIFY_GEMM").is_ok() && m.cfg().num_experts == 0 {
        m.init_batch(1, plen + ngen + 16)?;
        m.verify_dequant()?;
        let (d, scale) = m.verify_gemm()?;
        println!("== GEMM vs GEMV: max|diff| = {d:.4}, max|value| = {scale:.4}  (rel {:.4}) => {}",
                 d / scale.max(1e-6), if d / scale.max(1e-6) < 0.05 { "OK" } else { "WRONG" });
        return Ok(());
    }

    // Which arm is RIGHT, not merely different. The single-stream q4 path (q4 GEMV, f32
    // intermediates) is the reference: it is the one that has been correct since before any of this
    // work. Seed all three the same way — one token, same position — and compare greedy ids.
    if std::env::var("VERIFY_ARM").is_ok() && m.cfg().num_experts == 0 {
        let n = 24usize;
        let seed = prompt[0];
        m.seed(seed, 0)?;
        let mut single = Vec::new();
        for _ in 0..n {
            m.step(0, 0)?;
            m.sync()?;
            single.push(m.last_token()?);
        }
        // step-by-step, recording the top-2 margin so a divergence can be classified
        let mut run = |eng: &mut inferencelayer::cuda_llm::CudaQwen3, q4: bool|
         -> anyhow::Result<(Vec<u32>, Vec<(u32, f32, f32)>)> {
            eng.set_q4_arm(q4);
            eng.seed_batch(&[seed, seed], 0)?;
            let mut marg = Vec::new();
            for _ in 0..n {
                eng.step_batch()?;
                eng.sync()?;
                let lg = eng.batch_logits(0)?;
                let mut best = (0usize, f32::MIN);
                let mut second = f32::MIN;
                for (i, &v) in lg.iter().enumerate() {
                    if v > best.1 { second = best.1; best = (i, v); }
                    else if v > second { second = v; }
                }
                marg.push((best.0 as u32, best.1, best.1 - second));
            }
            let toks = eng.batch_tokens(0, 1, n + 1)?;
            Ok((toks, marg))
        };
        m.init_batch(2, n + 8)?;
        let (bf16, mb) = run(&mut m, false)?;
        let (q4, mq) = run(&mut m, true)?;
        if let Some(k) = (0..n.min(mb.len()).min(mq.len())).find(|&i| mb[i].0 != mq[i].0) {
            println!("== FIRST DIVERGENCE at step {k}");
            println!("   bf16 arm: top1 id {} logit {:.5}, margin over top2 {:.5}", mb[k].0, mb[k].1, mb[k].2);
            println!("   q4   arm: top1 id {} logit {:.5}, margin over top2 {:.5}", mq[k].0, mq[k].1, mq[k].2);
            let verdict = if mb[k].2 < 0.05 || mq[k].2 < 0.05 { "NEAR-TIE — rounding, not a defect" }
                          else { "WIDE MARGIN — real defect in the arm" };
            println!("   => {verdict}");
        } else {
            println!("== no divergence in the sampled ids");
        }
        let agree = |a: &[u32], b: &[u32]| a.iter().zip(b).take_while(|(x, y)| x == y).count();
        println!("== single-stream (reference): {:?}", &single[..12.min(single.len())]);
        println!("== batched bf16 mirrors:      {:?}", &bf16[..12.min(bf16.len())]);
        println!("== batched q4 arm:            {:?}", &q4[..12.min(q4.len())]);
        println!("   bf16 agrees with reference for {}/{} tokens", agree(&single, &bf16), n);
        println!("   q4   agrees with reference for {}/{} tokens", agree(&single, &q4), n);
        return Ok(());
    }

    // DETERMINISM: bit-identical logits across repeated runs, not merely "same argmax".
    //
    // This is a real differentiator — SGLang sells deterministic inference as a paid feature — but
    // it only counts if it is measured. The engine is built for it: the q4 GEMM's split-K reduction
    // is a fixed-order second pass rather than f32 atomicAdd, because atomics reorder and a shifted
    // last bit can flip an argmax. This checks the claim end to end.
    if std::env::var("VERIFY_DET").is_ok() && m.cfg().num_experts == 0 {
        m.init_batch(2, plen + 64)?;
        let a = m.prefill_logits(&prompt)?;
        let b = m.prefill_logits(&prompt)?;
        let c = m.prefill_logits(&prompt)?;
        let bitsame = |x: &[f32], y: &[f32]| x.iter().zip(y).all(|(p, q)| p.to_bits() == q.to_bits());
        let ndiff = a.iter().zip(&b).filter(|(p, q)| p.to_bits() != q.to_bits()).count();
        println!("== PREFILL DETERMINISM over {} logits, 3 runs:", a.len());
        println!("   run1 vs run2: {}", if bitsame(&a, &b) { "BIT-IDENTICAL" } else { "DIFFERS" });
        println!("   run1 vs run3: {}", if bitsame(&a, &c) { "BIT-IDENTICAL" } else { "DIFFERS" });
        if ndiff > 0 { println!("   {ndiff} logits differ"); }
        // and decode: same ids from repeated batched runs
        let mut runs = Vec::new();
        for _ in 0..3 {
            m.seed_batch(&[prompt[0], prompt[0]], 0)?;
            for _ in 0..24 { m.step_batch()?; }
            m.sync()?;
            runs.push(m.batch_tokens(0, 1, 25)?);
        }
        let same = runs[0] == runs[1] && runs[0] == runs[2];
        println!("== DECODE DETERMINISM, 3 runs x 24 tokens: {}",
                 if same { "IDENTICAL" } else { "DIFFERS" });
        println!("   {:?}", &runs[0][..8.min(runs[0].len())]);
        return Ok(());
    }

    // Seq-vs-batch numeric bisect: ONE decode step through both paths from IDENTICAL state
    // (same slot-0 prefill KV, same input token, same position), logits diffed. The sequential
    // path's output has never matched HF while the batch path now does; this finds where.
    if std::env::var("VERIFY_SEQ").is_ok() && m.cfg().num_experts == 0 {
        // SLOTS/BMAX mirror the server's geometry — the seq failure reproduces only there.
        let slots: usize = std::env::var("SLOTS").ok().and_then(|v| v.parse().ok()).unwrap_or(2);
        let bmax2: usize = std::env::var("BMAX").ok().and_then(|v| v.parse().ok()).unwrap_or(plen + 64);
        m.init_batch(slots, bmax2)?;
        m.prefill_slot(&prompt, Some(0))?;   // eager, graph-free, HF-validated path
        m.sync()?;
        let first = m.last_token()?;
        // Path B first: sequential step from (first, plen) over the prefill's KV.
        m.seed_seq_head(first, plen)?;
        m.step(0, 0)?;
        m.sync()?;
        let lb = m.seq_logits()?;
        // Path A: batch step at width 1 from the same state (rewrites KV[plen] with its own
        // rope of the same input, so each path is self-consistent).
        m.set_width(1)?;
        m.occupy_slot(0, first, plen)?;
        m.step_batch()?;
        m.sync()?;
        let la = m.batch_logits(0)?;
        let top = |l: &[f32]| {
            let mut b = (0usize, f32::MIN); let mut s2 = f32::MIN;
            for (i, &v) in l.iter().enumerate() {
                if v > b.1 { s2 = b.1; b = (i, v); } else if v > s2 { s2 = v; }
            }
            (b.0, b.1, b.1 - s2)
        };
        let (ai, av, am) = top(&la);
        let (bi, bv, bm) = top(&lb);
        let mut md = 0f32; let mut mi = 0usize;
        for i in 0..la.len().min(lb.len()) {
            let d = (la[i] - lb[i]).abs();
            if d > md { md = d; mi = i; }
        }
        println!("== SEQ vs BATCH one step from identical state (input={first}, pos={plen})");
        println!("   batch: top1 id {ai} logit {av:.4} margin {am:.4}");
        println!("   seq:   top1 id {bi} logit {bv:.4} margin {bm:.4}");
        println!("   max|logit diff| = {md:.4} at id {mi}  (batch {:.4} vs seq {:.4})", la[mi], lb[mi]);

        // Level 2: the CAPTURED seq graph, one replay, from the same re-seeded state — diffed
        // against the eager step above. The step kernels just proved numerically equivalent, so
        // if the replay diverges the disease is in capture/replay, not math.
        m.seed_seq_head(first, plen)?;
        m.capture()?;
        m.seed_seq_head(first, plen)?;
        m.run_graph(1)?;
        m.sync()?;
        let lg = m.seq_logits()?;
        let (gi, gv, gm) = top(&lg);
        let mut md2 = 0f32; let mut mi2 = 0usize;
        for i in 0..lb.len().min(lg.len()) {
            let d = (lb[i] - lg[i]).abs();
            if d > md2 { md2 = d; mi2 = i; }
        }
        println!("   graph-replay: top1 id {gi} logit {gv:.4} margin {gm:.4}");
        println!("   max|eager-seq vs replay diff| = {md2:.4} at id {mi2}");

        // Level 3: the SERVER'S exact chain — GRAPHED prefill (cold then warm replay), NO
        // explicit seed (rely on the prefill's leftover d_token/d_pos, as the server does),
        // then one seq graph replay. If THIS diverges, the graphed prefill leaves different
        // head state than the eager prefill.
        m.prefill_slot_graphed(&prompt, 0)?;   // warm (cold capture happened via eager? no — first graphed call captures)
        m.sync()?;
        m.prefill_slot_graphed(&prompt, 0)?;   // guaranteed replay
        m.sync()?;
        let f2 = m.last_token()?;
        println!("   graphed-prefill leftover: d_token={f2} (eager gave {first})");
        // The server's exact admission tail: async slot bookkeeping + the burst readback.
        m.occupy_slot_async(0, plen)?;
        let firsts = m.first_tokens()?;
        m.set_slot_token_async(0, firsts[0])?;
        m.run_graph(1)?;
        m.sync()?;
        let t3 = m.last_token()?;
        println!("   seq step after graphed prefill + slot_ctl/first_tokens: token={t3} (want 11)");

        // Level 4: the LAST delta vs the server — captured batch width-graphs coexisting with
        // the seq graph — plus an 8-token seq loop exactly like the server's cadence.
        m.set_width(1)?;
        m.capture_batch()?;
        m.prefill_slot_graphed(&prompt, 0)?;
        m.sync()?;
        let mut toks = Vec::new();
        toks.push(m.last_token()?);
        for _ in 0..8 {
            m.run_graph(1)?;
            m.sync()?;
            toks.push(m.last_token()?);
        }
        println!("   with batch graph resident, seq loop: {toks:?} (expect [12095, 11/13, ...coherent])");
        return Ok(());
    }

    if std::env::var("Q4GEMM").is_ok() && m.cfg().num_experts == 0 {
        m.init_batch(1, 64)?;
        // Sweep the widths that matter: dequant cost is per weight tile and amortises over ROWS,
        // so q4-direct and bf16-mirror swap places somewhere between decode and prefill widths.
        for t in [8usize, 16, 32, 64, 128, 256, 2000] {
            println!("--- M = {t}");
            m.verify_q4_gemm(t, 200)?;
        }
        return Ok(());
    }

    if std::env::var("GEMM_BENCH").is_ok() && m.cfg().num_experts == 0 {
        m.init_batch(1, 64)?;
        m.bench_gemm(plen.max(2000), 100)?;
        return Ok(());
    }

    // Flash attention vs the cuBLAS score-matrix path, on real weights, all 151936 logits.
    // An argmax-only check would pass on a kernel that is subtly wrong everywhere else.
    if std::env::var("VERIFY_FA").is_ok() && m.cfg().num_experts == 0 {
        m.init_batch(1, plen + 32)?;
        m.set_fa(false);
        let a = m.prefill_logits(&prompt)?;
        m.set_fa(true);
        let b = m.prefill_logits(&prompt)?;
        let (mut d, mut sc) = (0f32, 0f32);
        for (x, y) in a.iter().zip(b.iter()) {
            d = d.max((x - y).abs());
            sc = sc.max(x.abs());
        }
        let am = |v: &[f32]| v.iter().enumerate().max_by(|p, q| p.1.total_cmp(q.1)).unwrap().0;
        let rel = d / sc.max(1e-6);
        println!("== FLASH vs SCORE-MATRIX: max|diff| = {d:.4}, max|logit| = {sc:.4} (rel {rel:.5})");
        println!("   argmax gemm = {}, argmax flash = {}  => {}", am(&a), am(&b),
                 if am(&a) == am(&b) && rel < 0.02 { "OK" } else { "WRONG" });
        return Ok(());
    }

    // batched prefill (dense): needs the f16 mirrors, so init_batch(1, ..) first
    if std::env::var("PREFILL").is_ok() && m.cfg().num_experts == 0 {
        m.init_batch(1, plen + ngen + 16)?;
        m.prefill(&prompt)?;            // warm
        m.sync()?;
        let tp2 = std::time::Instant::now();
        m.prefill(&prompt)?;
        m.sync()?;
        let dt = tp2.elapsed().as_secs_f64();
        println!("== BATCHED PREFILL: {plen} tokens in {dt:.4}s = {:.0} tok/s", plen as f64 / dt);
        // same work, one submission instead of ~340
        m.capture_prefill(&prompt)?;
        m.run_prefill_graph(&prompt)?;
        m.sync()?;
        let tg = std::time::Instant::now();
        m.run_prefill_graph(&prompt)?;
        m.sync()?;
        let dg = tg.elapsed().as_secs_f64();
        println!("== GRAPH PREFILL:   {plen} tokens in {dg:.4}s = {:.0} tok/s", plen as f64 / dg);
        // Correctness of prefill itself: its greedy token must equal what the sequential path
        // samples at the same position. Different => the batched prefill math is wrong.
        let pf_tok = m.last_token()?;
        let mut seq_tok = 0u32;
        for (i, &t) in prompt.iter().enumerate() {
            m.seed(t, i)?;
            m.step(0, 0)?;
            if i + 1 == prompt.len() { m.sync()?; seq_tok = m.last_token()?; }
        }
        println!("   prefill token = {pf_tok}, sequential token = {seq_tok}  => {}",
                 if pf_tok == seq_tok { "MATCH" } else { "MISMATCH (prefill math wrong)" });
    }
    let prefill_end = plen;
    println!("prefill({plen} tok, sequential) {:.2}s", tp.elapsed().as_secs_f64());

    // ---- eager path (one launch per kernel) ----
    // BATCH_ONLY skips both single-stream phases. They issue ~35k launches before the batched
    // section starts, which puts the M=64 step out of reach of any sane `ncu --launch-skip`.
    let batch_only = std::env::var("BATCH_ONLY").is_ok();
    m.seed(m.tokens(prefill_end - 1, prefill_end)?[0], prefill_end)?;
    m.step(0, 0)?;
    m.sync()?;
    let te = std::time::Instant::now();
    if !batch_only { for _ in 0..ngen { m.step(0, 0)?; } }
    m.sync()?;
    let eager = ngen as f64 / te.elapsed().as_secs_f64();

    // ---- graph path (one replay per token) ----
    m.seed(m.tokens(prefill_end - 1, prefill_end)?[0], prefill_end)?;
    m.capture()?;
    m.run_graph(2)?;
    m.sync()?;
    let start = prefill_end;
    m.seed(m.tokens(prefill_end - 1, prefill_end)?[0], start)?;
    let tg = std::time::Instant::now();
    m.run_graph(if batch_only { 1 } else { ngen })?;
    m.sync()?;
    let dt = tg.elapsed().as_secs_f64();
    let ids = m.tokens(start, start + ngen)?;

    println!("first 24 generated ids: {:?}", &ids[..ids.len().min(24)]);
    // full list for a token-for-token diff against vLLM. Not under BATCH_ONLY: only one token
    // was actually generated, and overwriting the file would feed a correctness probe garbage.
    if !batch_only {
        std::fs::write("/root/ours_ids.txt",
            ids.iter().map(|t| t.to_string()).collect::<Vec<_>>().join(","))?;
    }
    println!("== EAGER decode: {eager:.1} tok/s");
    println!("\n== DECODE: {ngen} tokens in {dt:.3}s = {:.1} tok/s", ngen as f64 / dt);
    println!("   vLLM reference on the same H100 (Qwen3-30B-A3B BF16): 155.9 tok/s");

    // ---- concurrency: M sequences decoding in lockstep ----
    if let Ok(ms) = std::env::var("BATCH") {
        let bm: usize = ms.parse().unwrap_or(4);
        if m.cfg().num_experts != 0 {
            println!("\n== CONCURRENCY: skipped (batched path is the dense arm)");
            return Ok(());
        }
        // BMAX overrides the per-slot KV stride: the server allocates 4096 and measured slower
        // GPU steps than this bench at 256 — this knob exists to test exactly that variable.
        let bmax = std::env::var("BMAX").ok().and_then(|v| v.parse().ok()).unwrap_or(256usize);
        m.init_batch(bm, bmax)?;
        let seeds: Vec<u32> = (0..bm).map(|i| prompt[i % prompt.len()]).collect();
        m.seed_batch(&seeds, 0)?;
        for _ in 0..8 { m.step_batch()?; }
        m.sync()?;
        m.capture_batch()?;
        m.run_graph_batch(2)?;
        m.sync()?;
        let tb = std::time::Instant::now();
        if std::env::var("SYNC_EACH").is_ok() {
            // The server's cadence: one replay, one sync, one 256-byte readback per step. Exists
            // to answer whether the through-HTTP decode deficit is the protocol or the engine.
            for _ in 0..ngen {
                m.run_graph_batch(1)?;
                m.sync()?;
                let _ = m.batch_last_tokens()?;
            }
        } else {
            m.run_graph_batch(ngen)?;
            m.sync()?;
        }
        let dt2 = tb.elapsed().as_secs_f64();
        println!("\n== CONCURRENCY M={bm}: {:.1} tok/s per seq, AGGREGATE {:.1} tok/s",
                 ngen as f64 / dt2, (bm * ngen) as f64 / dt2);
        println!("   seq0 ids[8..20]: {:?}", m.batch_tokens(0, 8, 20)?);
    }
    Ok(())
}