inferencelayer 0.2.3

//! `lfm2-profile-stage` — per-kernel GPU-time profile of one serving batch step on a layer
//! shard: where does T_stage actually go? Prints the top dispatches and an aggregate by kind.
//!
//! ```text
//! lfm2-profile-stage --model <dir> --layers a:b --cols 16 --pos 64 [--wide] [--slots N]
//! --slots N packs the `cols` columns into N sequences (PREFILL shape: contiguous per-slot
//! runs — the DN chunk kernel's real geometry); default = cols distinct slots (decode shape).
//! ```
use anyhow::{Context, Result};
use inferencelayer::{BatchCol, GpuCtx, Lfm2Gpu, Weights};

fn main() -> Result<()> {
    let mut model = None;
    let mut layers = None;
    let mut cols_n = 16usize;
    let mut pos = 64u32;
    let mut wide = false;
    let mut spec = false;
    let mut logits = false;
    let mut gap_probe = false;
    let mut full = false;
    let mut slots: Option<usize> = None;
    let mut args = std::env::args().skip(1);
    while let Some(a) = args.next() {
        match a.as_str() {
            "--model" => model = args.next(),
            "--layers" => layers = args.next(),
            "--cols" => cols_n = args.next().context("--cols")?.parse()?,
            "--pos" => pos = args.next().context("--pos")?.parse()?,
            "--wide" => wide = true,
            "--spec" => spec = true,
            "--logits" => logits = true,
            "--gap-probe" => gap_probe = true,
            // `--full`: load the WHOLE model through the ordinary loader — the only way to
            // profile non-Qwen arches (GLM-OCR), whose loaders `load_shard` refuses.
            "--full" => full = true,
            "--slots" => slots = Some(args.next().context("--slots value")?.parse()?),
            other => anyhow::bail!("unknown argument {other}"),
        }
    }
    let model = model.context("--model required")?;
    let ctx = GpuCtx::new()?;
    eprintln!("backend: {} ts={}", ctx.backend, ctx.timestamps);
    let (w, a, b) = if full {
        let w = Weights::load(&ctx, std::path::Path::new(&model))?;
        let n = w.cfg.n_layers;
        (w, 0, n)
    } else {
        let (a, b) = layers
            .context("--layers a:b required (or --full)")?
            .split_once(':')
            .map(|(x, y)| (x.parse::<usize>().unwrap(), y.parse::<usize>().unwrap()))
            .context("--layers a:b")?;
        (
            Weights::load_shard(&ctx, std::path::Path::new(&model), a, b)?,
            a,
            b,
        )
    };
    let gpu = Lfm2Gpu::new(&ctx, w);
    // `--spec` profiles the SPECULATIVE plan — the one production's verify actually runs. Without
    // it this tool profiles `make_batch_plan`, which omits the DeltaNet snapshot machinery
    // (`dn_snaps`, per-column parent-snapshot loads for rollback) that the spec plan dispatches for
    // free-rollback on partial acceptance. Profiling the plain plan and comparing its kernel-time
    // sum against a production verify therefore charges the spec plan's snapshot work to "dead
    // time" — an attribution error that invents a fusion program out of kernels the profiler simply
    // never ran. See `bench/P7_profile_the_verify_is_gap_bound.md`.
    anyhow::ensure!(!(wide && spec), "--wide and --spec are different plans");
    let bp = if wide {
        gpu.make_batch_plan_wide(&ctx, cols_n)
    } else if spec {
        gpu.make_batch_plan_spec(&ctx, cols_n)
    } else {
        gpu.make_batch_plan(&ctx, cols_n)
    };
    // Decode-realistic columns: cols_n DISTINCT sequences, one token each at `pos`; block tables
    // pre-pointed at disjoint pool ranges.
    let n_slots = slots.unwrap_or(cols_n).max(1).min(cols_n);
    let per_slot = cols_n / n_slots;
    let blocks_per = (pos as usize + per_slot + 2).div_ceil(16);
    let mut cols = Vec::new();
    for s in 0..n_slots {
        let row = (s + 1) as u32;
        let blocks: Vec<u32> = (0..blocks_per as u32)
            .map(|j| (s * blocks_per) as u32 + j)
            .collect();
        gpu.write_btab_row(&ctx, row, &blocks);
        gpu.zero_dn_slot(&ctx, row as usize);
    }
    for i in 0..cols_n {
        // slots=N packs contiguous per-slot runs (prefill shape: successive positions of one
        // sequence); the default (N = cols) is the decode shape (one column per sequence).
        let s = i / per_slot.max(1);
        let row = (s.min(n_slots - 1) + 1) as u32;
        cols.push(BatchCol::text(
            1000 + i as u32,
            pos + (i % per_slot.max(1)) as u32,
            row,
            // Production's speculative verify needs a logit for EVERY column -- that is how each
            // draft's argmax is checked for acceptance -- and on this checkpoint the LM head is a
            // 248320-row matrix. Profiling with need_logit:false silently omits it, so the kernel-time
            // sum comes out short and the shortfall gets misread as inter-dispatch dead time.
            logits,
        ));
    }
    // Warm (also fills KV up to pos so attention spans are realistic).
    for p in 0..pos {
        let warm: Vec<BatchCol> = cols
            .iter()
            .map(|c| BatchCol::text(c.token, p, c.btrow, c.need_logit))
            .collect();
        let hidden_zero = vec![0f32; warm.len() * gpu.w.cfg.hidden];
        let hidden_in = (!gpu.w.cfg.stage_first).then_some(hidden_zero.as_slice());
        let _ = gpu.batch_stage_step(&ctx, &bp, &warm, hidden_in)?;
    }
    let reps = 10;
    let mut agg: Vec<(u32, u32, f64, &'static str)> = Vec::new();
    let mut wall_sum = 0.0;
    for _ in 0..reps {
        let (steps, wall) = gpu.profile_batch_step(&ctx, &bp, &cols)?;
        if agg.is_empty() {
            agg = steps
                .iter()
                .map(|(_, gx, gy, _, tag)| (*gx, *gy, 0.0, *tag))
                .collect();
        }
        for (i, (_, _, _, us, _)) in steps.iter().enumerate() {
            agg[i].2 += us;
        }
        wall_sum += wall;
    }
    let total: f64 = agg.iter().map(|(_, _, t, _)| t / reps as f64).sum();
    println!(
        "== batch step profile: {cols_n} cols, pos {pos}, layers {a}:{b}, wide={wide} spec={spec} logits={logits} =="
    );
    println!(
        "GPU busy {:.0} µs/step  wall {:.0} µs  → {:.1} tok/s/stage at this width",
        total,
        wall_sum / reps as f64,
        cols_n as f64 * 1e6 / total
    );
    println!(
        "  (this wall is NOT production's: one compute pass PER DISPATCH here vs ONE pass total in \
         batch_stage_step. Use --gap-probe for anything about gaps.)"
    );

    // `--gap-probe`: encode the plan 1× and 2× inside ONE command buffer and ONE compute pass. The
    // dispatch count doubles; submit, poll and readback happen once either way, so they cancel in
    // the difference:
    //     wall(2×) − wall(1×) = kernel time + inter-dispatch cost, for one plan pass.
    // Subtract the kernel-time sum above and the inter-dispatch dead time is left alone. This is
    // the measurement that decides whether the ~47% of production's verify that is NOT kernel
    // execution is reclaimable by fusing dispatches, or is just submit/poll latency that fusion
    // cannot touch. See bench/P7_profile_the_verify_is_gap_bound.md.
    if gap_probe {
        let probe_reps = 10;
        let mut w1 = f64::MAX;
        let mut w2 = f64::MAX;
        let (mut h1, mut h2) = (f64::MAX, f64::MAX);
        let head = logits && gpu.w.cfg.stage_last;
        for _ in 0..probe_reps {
            // MIN, not mean: contention only ever makes a sample look slower.
            w1 = w1.min(gpu.time_batch_step_repeated(&ctx, &bp, 1, false)?);
            w2 = w2.min(gpu.time_batch_step_repeated(&ctx, &bp, 2, false)?);
            if head {
                h1 = h1.min(gpu.time_batch_step_repeated(&ctx, &bp, 1, true)?);
                h2 = h2.min(gpu.time_batch_step_repeated(&ctx, &bp, 2, true)?);
            }
        }
        let (w1_us, w2_us) = (w1 * 1e6, w2 * 1e6);
        let marginal = w2_us - w1_us; // one extra plan pass, with poll/submit cancelled out
        let fixed = w1_us - marginal; // what one pass costs BEYOND its own dispatches = submit+poll
        let gaps = marginal - total; // dispatch-attributable time that is not kernel execution
        println!("-- gap probe (1× vs 2× plan in ONE pass; poll/submit cancel in the delta) --");
        println!("  wall 1×            {w1_us:>9.0} µs");
        println!("  wall 2×            {w2_us:>9.0} µs");
        println!("  marginal pass      {marginal:>9.0} µs   = kernels + inter-dispatch cost");
        println!("  kernel-time sum    {total:>9.0} µs   (profiled above; excludes the LM head)");
        println!(
            "  ⇒ inter-dispatch   {gaps:>9.0} µs   ({:.0}% of the marginal pass)",
            100.0 * gaps / marginal.max(1e-9)
        );
        println!("  ⇒ fixed submit+poll{fixed:>9.0} µs   (cancels out of the delta)");
        if gaps > 0.25 * marginal {
            println!(
                "  VERDICT: dead time is PER-DISPATCH ⇒ fusing dispatches reclaims it. Fusion is on."
            );
        } else {
            println!(
                "  VERDICT: dead time is NOT per-dispatch ⇒ it is submit/poll latency. Fusion buys \
                 nothing; do not write a fused kernel."
            );
        }

        // The LM head lives OUTSIDE bp.plan, so profile_batch_step never times it and the kernel
        // sum above omits it entirely. On a 248320-token vocab that omission is the size of the
        // whole mystery. Same subtraction, one more term: the marginal pass WITH the head, minus
        // the marginal pass without, is the head — with submit/poll cancelling out of both.
        if head {
            let (h1_us, h2_us) = (h1 * 1e6, h2 * 1e6);
            let marginal_head = h2_us - h1_us;
            let lm = marginal_head - marginal;
            let bytes = gpu.w.cfg.vocab as f64 * gpu.w.cfg.hidden as f64;
            println!("-- LM head (vocab {}) --", gpu.w.cfg.vocab);
            println!("  marginal pass +head{marginal_head:>9.0} µs");
            println!(
                "  ⇒ LM head          {lm:>9.0} µs   ({:.0}% of a pass WITH the head)",
                100.0 * lm / marginal_head.max(1e-9)
            );
            println!(
                "     ↳ {:.0} GB/s if the head is f16 ({:.0} MB) | {:.0} GB/s if q4_0 ({:.0} MB)",
                bytes * 2.0 / (lm * 1e-6) / 1e9,
                bytes * 2.0 / 1e6,
                bytes * 0.5625 / (lm * 1e-6) / 1e9,
                bytes * 0.5625 / 1e6,
            );
        }
    }
    // Per-kernel aggregation (the actionable view): total µs and share per tag.
    let mut by_tag: std::collections::HashMap<&str, (f64, usize)> =
        std::collections::HashMap::new();
    for (_, _, t, tag) in &agg {
        let e = by_tag.entry(tag).or_insert((0.0, 0));
        e.0 += t / reps as f64;
        e.1 += 1;
    }
    let mut tags: Vec<_> = by_tag.into_iter().collect();
    tags.sort_by(|a, b| b.1.0.partial_cmp(&a.1.0).unwrap());
    println!("-- by kernel --");
    for (tag, (t, n)) in &tags {
        println!(
            "  {tag:<28} n={n:<4} {:>9.1} µs ({:>4.1}%)  avg {:>7.1} µs",
            t,
            100.0 * t / total,
            t / *n as f64
        );
    }
    let mut idx: Vec<usize> = (0..agg.len()).collect();
    idx.sort_by(|&x, &y| agg[y].2.partial_cmp(&agg[x].2).unwrap());
    println!("-- top single dispatches --");
    for &i in idx.iter().take(8) {
        let (gx, gy, t, tag) = agg[i];
        println!(
            "  step[{i:>3}] {tag:<26} gx={gx:<6} gy={gy:<3} {:>9.1} µs ({:>4.1}%)",
            t / reps as f64,
            100.0 * t / reps as f64 / total
        );
    }
    Ok(())
}