inferencelayer 0.2.9

//! Accept/reject evidence for the CUDA fast path, both arms, on one NVIDIA GPU:
//!
//!   PREFILL (compute-bound)  — cuBLAS GemmEx on tensor cores.  Target: vLLM parity.
//!   DECODE  (memory-bound)   — q4_0 GEMV.  Target: BEAT vLLM, via 4x fewer weight bytes.
//!
//! Reference points from the 2026-07-24 H100 head-to-head (`bench/CROSS_GPU_2026-07-24.md`):
//! vLLM BF16 compiled = 31 679 tok/s prefill, 155.9 tok/s decode; our WGSL arm = 199 / 28.5.
//! Requires `--features cudarc` + an NVIDIA GPU; a default build never compiles this file.
use anyhow::Result;
use inferencelayer::cublas_prefill::{CublasPrefill, GemvKernel};

/// Qwen3-30B-A3B: ~3.3B ACTIVE params/token (MoE, 8 of 128 experts).
const ACTIVE_PARAMS: f64 = 3.3e9;
/// q4_0 = 20 bytes per 32 weights.
const Q4_BYTES_PER_PARAM: f64 = 20.0 / 32.0;
/// vLLM serves this model BF16 = 2 bytes/param.
const BF16_BYTES_PER_PARAM: f64 = 2.0;
const VLLM_DECODE: f64 = 155.9;
const VLLM_PREFILL: f64 = 31679.0;

fn main() -> Result<()> {
    let g = match CublasPrefill::new() {
        Ok(g) => g,
        Err(e) => {
            eprintln!("cuBLAS unavailable ({e}) — needs an NVIDIA GPU; WGSL is the fallback path");
            std::process::exit(1);
        }
    };

    println!("== PREFILL arm: cuBLAS GemmEx (tensor cores), Qwen3 prefill shapes");
    println!("{:>18} | {:>14}", "shape (m,n,k)", "TFLOP/s");
    for (m, n, k) in [
        (2048usize, 2048usize, 2048usize),
        (2624, 2048, 2048),
        (4096, 4096, 4096),
    ] {
        println!(
            "{:>6},{:>5},{:>5} | {:>14.1}",
            m,
            n,
            k,
            g.bench_tflops(m, n, k, 50)?
        );
    }

    println!("\n== DECODE arm: q4_0 GEMV (the bandwidth win path)");
    let gemv = g.q4_gemv()?;
    println!(
        "{:>17} | {:>7} | {:>7} | {:>8} | {:>8} | {:>8} | {:>9}",
        "shape (rows,cols)", "v1 blk", "v2 warp", "v3 shX", "v4 f16u", "v5 h2bit", "best %peak"
    );
    let mut best_gbps = 0f64;
    let mut best_is_f16 = false;
    for (rows, cols) in [(32768usize, 8192usize), (65536, 8192), (32768, 16384)] {
        let v1 = gemv.bench_gbps(rows, cols, 20, GemvKernel::BlockPerRow)?;
        let v2 = gemv.bench_gbps(rows, cols, 20, GemvKernel::WarpPerRow)?;
        let v3 = gemv.bench_gbps(rows, cols, 20, GemvKernel::SharedX)?;
        let v4 = gemv.bench_gbps(rows, cols, 20, GemvKernel::F16ScalesUnroll)?;
        let v5 = gemv.bench_gbps(rows, cols, 20, GemvKernel::Half2BitTrick)?;
        let best = v1.max(v2).max(v3).max(v4).max(v5);
        if best > best_gbps {
            best_gbps = best;
            best_is_f16 = best == v4 || best == v5;
        }
        println!(
            "{:>8},{:>8} | {:>7.0} | {:>7.0} | {:>8.0} | {:>8.0} | {:>8.0} | {:>8.0}%",
            rows, cols, v1, v2, v3, v4, v5, 100.0 * best / 2000.0
        );
    }
    // v4/v5 carry f16 scales (18 B per 32 weights) — that is the production WGSL binding too.
    let q4_bpp = if best_is_f16 { 18.0 / 32.0 } else { Q4_BYTES_PER_PARAM };

    // Decode tok/s = achieved weight bandwidth / weight bytes per token.
    let ours = best_gbps * 1e9 / (ACTIVE_PARAMS * q4_bpp);
    let vllm_ceiling = best_gbps * 1e9 / (ACTIVE_PARAMS * BF16_BYTES_PER_PARAM);
    println!("\n== Projected single-stream decode for Qwen3-30B-A3B @ {best_gbps:.0} GB/s");
    println!(
        "  ours  (q4_0, {:.2} GB/tok): {:>8.0} tok/s",
        ACTIVE_PARAMS * q4_bpp / 1e9,
        ours
    );
    println!(
        "  BF16 ceiling ({:.2} GB/tok): {:>8.0} tok/s  [vLLM MEASURED {VLLM_DECODE}]",
        ACTIVE_PARAMS * BF16_BYTES_PER_PARAM / 1e9,
        vllm_ceiling
    );
    println!("  => vs vLLM measured decode: {:.1}x", ours / VLLM_DECODE);
    println!(
        "  (WGSL arm measured 28.5 tok/s => CUDA decode uplift {:.0}x)",
        ours / 28.5
    );
    println!("\n  vLLM prefill reference: {VLLM_PREFILL:.0} tok/s (ours WGSL: 199)");
    Ok(())
}