add inference benchmarking

README.md

@@ -176,6 +176,7 @@ I've published a number of guides that might contain helpful information, most r
 │   ├── chat_eval.py                # Chat model: eval tasks
 │   ├── chat_rl.py                  # Chat model: reinforcement learning
 │   ├── chat_sft.py                 # Chat model: train SFT
+│   ├── infer_bench.py              # Inference: latency/throughput/VRAM bench
 │   ├── tok_eval.py                 # Tokenizer: evaluate compression rate
 │   └── tok_train.py                # Tokenizer: train it
 ├── tasks

nanochat/common.py

@@ -277,3 +277,52 @@ def get_peak_flops(device_name: str) -> float:
     # Unknown GPU - return inf so MFU shows as 0% rather than a wrong guess
     logger.warning(f"Peak flops undefined for: {device_name}, MFU will show as 0%")
     return float('inf')
+
+def get_peak_bandwidth(device_name: str) -> float:
+    """Peak HBM/GDDR memory bandwidth in bytes/sec. The decode phase of inference
+    is memory-bandwidth-bound, so this is the roofline for tokens/sec (see MBU)."""
+    name = device_name.lower()
+
+    # Table order matters: more specific patterns first.
+    _PEAK_BANDWIDTH_TABLE = (
+        # NVIDIA Blackwell (HBM3e)
+        (["gb200"], 8.0e12),
+        (["grace blackwell"], 8.0e12),
+        (["b200"], 8.0e12),
+        (["b100"], 8.0e12),
+        # NVIDIA Hopper
+        (["h200"], 4.8e12),
+        (["h100", "nvl"], 3.9e12),
+        (["h100", "pcie"], 2.0e12),
+        (["h100"], 3.35e12), # SXM
+        (["h800", "pcie"], 2.0e12),
+        (["h800"], 3.35e12), # SXM
+        # NVIDIA Ampere data center (A100 80GB; the 40GB variant is 1.6e12)
+        (["a100"], 2.0e12),
+        (["a800"], 2.0e12),
+        (["a40"], 696e9),
+        (["a30"], 933e9),
+        # NVIDIA Ada data center
+        (["l40s"], 864e9),
+        (["l40-s"], 864e9),
+        (["l40 s"], 864e9),
+        (["l4"], 300e9),
+        # AMD CDNA accelerators
+        (["mi355"], 8.0e12),
+        (["mi325"], 6.0e12),
+        (["mi300x"], 5.3e12),
+        (["mi300a"], 5.3e12),
+        (["mi250x"], 3.28e12),
+        (["mi250"], 3.28e12),
+        # Consumer RTX
+        (["5090"], 1.79e12),
+        (["4090"], 1.01e12),
+        (["3090"], 936e9),
+    )
+    for patterns, bandwidth in _PEAK_BANDWIDTH_TABLE:
+        if all(p in name for p in patterns):
+            return bandwidth
+
+    # Unknown GPU - return inf so MBU shows as 0% rather than a wrong guess
+    logger.warning(f"Peak bandwidth undefined for: {device_name}, MBU will show as 0%")
+    return float('inf')

scripts/infer_bench.py

@@ -0,0 +1,297 @@
+"""
+Inference benchmark: measure latency, throughput, memory and bandwidth
+utilization of a trained checkpoint, sweeping over the decode batch size.
+
+This is an "eval" in its own right: intelligence metrics (CORE etc.) say
+nothing about what a model costs to run. Architecture choices (GQA, sliding
+windows, ...) show up here, so improvements can be judged on both axes.
+
+Background: inference has two very different regimes.
+- Prefill processes the whole prompt in parallel: big matmuls, compute-bound.
+- Decode generates one token at a time: every step re-reads all weights and
+  the KV cache to do a tiny amount of math, so it is memory-bandwidth-bound
+  and batching is nearly free until compute saturates.
+Sweeping the batch size traces out the latency <-> throughput tradeoff curve
+between these regimes. MBU (model bandwidth utilization) is the decode
+counterpart of training MFU: achieved bytes/sec over the peak bandwidth of
+the GPU. It measures how far the implementation is from the physical ceiling.
+
+Output: a human-readable card and table, and then the very last line of stdout
+is a single compact JSON document with all of the same data, so that scripts
+can consume the benchmark without parsing the pretty formatting:
+
+    result = json.loads(subprocess.run([...], capture_output=True, text=True).stdout.splitlines()[-1])
+
+Examples:
+
+    # benchmark a base model checkpoint on one GPU
+    python -m scripts.infer_bench -i base -g d12
+
+    # benchmark the SFT model, custom sweep
+    python -m scripts.infer_bench -i sft --batch-sizes 1,4,16,64 --decode-tokens 512
+
+    # machine-readable: grab the last line
+    python -m scripts.infer_bench -i base -g d12 | tail -1 | jq .sweep
+"""
+
+import argparse
+import json
+import time
+import torch
+
+from nanochat.common import compute_init, compute_cleanup, autodetect_device_type, get_peak_bandwidth, get_peak_flops
+from nanochat.checkpoint_manager import load_model
+from nanochat.engine import Engine
+
+# -----------------------------------------------------------------------------
+# Static accounting: what the architecture implies about inference cost
+
+def kv_bytes_per_token(model):
+    """Bytes to *store* one token of KV cache, per row (all layers)."""
+    config = model.config
+    head_dim = config.n_embd // config.n_head
+    kv_dtype_bytes = 2 # KV cache is kept in bf16
+    return config.n_layer * 2 * config.n_kv_head * head_dim * kv_dtype_bytes
+
+def kv_read_bytes(model, context_len):
+    """Bytes of KV cache *read* by one decode step at a given context length, per row.
+    Sliding window layers only attend to (and read) the last `window` tokens."""
+    config = model.config
+    head_dim = config.n_embd // config.n_head
+    kv_dtype_bytes = 2
+    total = 0
+    for window, _ in model.window_sizes:
+        total += 2 * config.n_kv_head * head_dim * kv_dtype_bytes * min(context_len, window)
+    return total
+
+def weight_bytes(model):
+    """Bytes of parameters as stored (each decode step reads all of them)."""
+    return sum(p.numel() * p.element_size() for p in model.parameters())
+
+def matmul_param_count(model):
+    """Number of matmul (weight) parameters. Backed out of estimate_flops() so that
+    the exclusion list of non-matmul params (embeddings, scalars, ...) lives in one place."""
+    config = model.config
+    h = config.n_head
+    q = config.n_embd // config.n_head
+    t = config.sequence_len
+    train_attn_flops = sum(12 * h * q * min(window, t) for window, _ in model.window_sizes)
+    matmul_params = (model.estimate_flops() - train_attn_flops) // 6
+    return matmul_params
+
+def decode_flops_per_token(model, context_len):
+    """Forward FLOPs to decode one token at a given context length: 2 FLOPs per
+    matmul param, plus attention reads of min(context, window) per layer."""
+    config = model.config
+    h = config.n_head
+    q = config.n_embd // config.n_head
+    attn_flops = sum(4 * h * q * min(context_len, window) for window, _ in model.window_sizes)
+    return 2 * matmul_param_count(model) + attn_flops
+
+def prefill_flops(model, prompt_len):
+    """Forward FLOPs to prefill a prompt: causal, so token t attends to min(t, window)."""
+    config = model.config
+    h = config.n_head
+    q = config.n_embd // config.n_head
+    attn_flops = 0
+    for window, _ in model.window_sizes:
+        w = min(window, prompt_len)
+        attended = w * (w + 1) // 2 + (prompt_len - w) * w # ramp up to w, then flat
+        attn_flops += 4 * h * q * attended
+    return 2 * matmul_param_count(model) * prompt_len + attn_flops
+
+# -----------------------------------------------------------------------------
+# Measurement
+
+def bench_generate(engine, prompt_tokens, batch_size, decode_tokens, temperature):
+    """Run one timed generation. Returns dict of measurements."""
+    device = engine.model.get_device()
+    torch.cuda.reset_peak_memory_stats(device)
+    torch.cuda.synchronize(device)
+
+    generator = engine.generate(prompt_tokens, num_samples=batch_size,
+                                max_tokens=decode_tokens, temperature=temperature)
+    # The first next() runs the batch=1 prefill, the KV cache replication to
+    # batch_size rows, and samples the first token: that is the TTFT.
+    t_start = time.perf_counter()
+    next(generator)
+    torch.cuda.synchronize(device)
+    ttft = time.perf_counter() - t_start
+
+    # Every subsequent next() is one decode step for all rows.
+    step_times = []
+    while True:
+        t0 = time.perf_counter()
+        try:
+            next(generator)
+        except StopIteration:
+            break
+        torch.cuda.synchronize(device)
+        step_times.append(time.perf_counter() - t0)
+
+    peak_vram = torch.cuda.max_memory_allocated(device)
+    return dict(ttft=ttft, step_times=step_times, peak_vram=peak_vram)
+
+def build_prompt(tokenizer, num_tokens):
+    """A natural-language prompt of exactly num_tokens tokens (random ids would
+    do for speed, but a real prompt keeps argmax decoding from degenerating)."""
+    paragraph = ("The history of science is the study of the development of science, "
+                 "including both the natural and social sciences. Science is a body of "
+                 "empirical, theoretical, and practical knowledge about the natural world. ")
+    text = paragraph * (num_tokens // 10) # more than enough tokens
+    tokens = tokenizer.encode(text, prepend="<|bos|>")
+    assert len(tokens) >= num_tokens, "prompt text too short, increase the repetition"
+    return tokens[:num_tokens]
+
+# -----------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(description="Inference benchmark")
+    parser.add_argument("-i", "--source", type=str, default="base", help="Checkpoint source: base|mid|sft")
+    parser.add_argument("-g", "--model-tag", type=str, default=None, help="Model tag to load")
+    parser.add_argument("-s", "--step", type=int, default=None, help="Step to load (default = last)")
+    parser.add_argument("--prompt-tokens", type=int, default=2048, help="Prompt length for prefill")
+    parser.add_argument("--decode-tokens", type=int, default=256, help="Tokens to generate per row")
+    parser.add_argument("--batch-sizes", type=str, default="1,8,32,128", help="Comma-separated decode batch sizes")
+    parser.add_argument("-t", "--temperature", type=float, default=0.0)
+    args = parser.parse_args()
+
+    device_type = autodetect_device_type()
+    assert device_type == "cuda", "infer_bench currently assumes a CUDA GPU (for timing and VRAM measurement)"
+    ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type)
+    assert ddp_world_size == 1, "infer_bench is a single GPU benchmark, run without torchrun"
+
+    model, tokenizer, meta = load_model(args.source, device, phase="eval", model_tag=args.model_tag, step=args.step)
+    config = model.config
+    engine = Engine(model, tokenizer)
+
+    # Clamp the prompt so prompt + decode fits in the training context
+    max_prompt = config.sequence_len - args.decode_tokens
+    prompt_len = min(args.prompt_tokens, max_prompt)
+    if prompt_len < args.prompt_tokens:
+        print(f"note: clamping prompt to {prompt_len} tokens so prompt+decode fits sequence_len={config.sequence_len}")
+    prompt_tokens = build_prompt(tokenizer, prompt_len)
+
+    # ------------------------------------------------------------------------
+    # Static card: inference cost implied by the architecture, before measuring
+    device_name = torch.cuda.get_device_name(device)
+    peak_bw = get_peak_bandwidth(device_name)
+    peak_flops = get_peak_flops(device_name)
+    total_vram = torch.cuda.get_device_properties(device).total_memory
+    w_bytes = weight_bytes(model)
+    num_params = sum(p.numel() for p in model.parameters())
+    kv_store = kv_bytes_per_token(model)
+    context_mid = prompt_len + args.decode_tokens // 2 # representative decode context
+    kv_read = kv_read_bytes(model, context_mid)
+    # tokens/sec ceiling at batch 1: every step must at least re-read weights + KV
+    ceiling_bs1 = peak_bw / (w_bytes + kv_read)
+    # how many rows of full-context KV fit next to the weights
+    max_rows = int((total_vram - w_bytes) / (kv_store * config.sequence_len))
+
+    print("=" * 100)
+    print(f"Model: {args.source} {meta.get('model_tag', '')} (step {meta['step']}) | "
+          f"depth {config.n_layer}, dim {config.n_embd}, heads {config.n_head}, kv heads {config.n_kv_head} (GQA)")
+    print(f"GPU: {device_name} | peak bandwidth {peak_bw/1e12:.2f} TB/s | peak compute {peak_flops/1e12:.0f} TFLOPS | VRAM {total_vram/2**30:.0f} GiB")
+    print("-" * 100)
+    dtype_counts = {}
+    for p in model.parameters():
+        dtype_name = str(p.dtype).replace("torch.", "")
+        dtype_counts[dtype_name] = dtype_counts.get(dtype_name, 0) + p.numel()
+    param_dtypes = ", ".join(f"{n:,} {dtype_name}" for dtype_name, n in sorted(dtype_counts.items()))
+    print(f"Parameters: {num_params:,} ({param_dtypes}) | weight bytes as stored: {w_bytes/2**20:.0f} MiB")
+    print(f"KV cache: {kv_store:,} bytes/token stored | {kv_read:,} bytes read/step at context {context_mid} "
+          f"(window pattern {config.window_pattern})")
+    print(f"Theoretical decode ceiling at batch 1: {ceiling_bs1:,.0f} tok/s | "
+          f"max ~{max_rows:,} full-context rows in VRAM")
+    print("=" * 100)
+
+    # Everything printed above also goes into the final JSON line for scripts
+    payload = {
+        "source": args.source,
+        "step": meta["step"],
+        "model_config": meta["model_config"],
+        "gpu": device_name,
+        # None (not Infinity) for unknown GPUs, so the last line stays valid JSON
+        "peak_bandwidth_bytes_per_sec": peak_bw if peak_bw != float("inf") else None,
+        "total_vram_bytes": total_vram,
+        "num_params": num_params,
+        "param_dtypes": dtype_counts,
+        "weight_bytes": w_bytes,
+        "kv_bytes_per_token": kv_store,
+        "kv_read_bytes_per_step": kv_read,
+        "context_mid": context_mid,
+        "peak_flops_per_sec": peak_flops if peak_flops != float("inf") else None,
+        "decode_flops_per_token": decode_flops_per_token(model, context_mid),
+        "ceiling_bs1_tok_per_sec": round(ceiling_bs1, 1) if ceiling_bs1 != float("inf") else None,
+        "max_full_context_rows": max_rows,
+        "prompt_tokens": prompt_len,
+        "decode_tokens": args.decode_tokens,
+        "temperature": args.temperature,
+        "sweep": [],
+    }
+
+    # ------------------------------------------------------------------------
+    # Prefill measurement: batch 1, a single decode step, so TTFT ~= prefill time.
+    # Prefill is compute-bound, so MFU (not MBU) is its distance from the roofline.
+    bench_generate(engine, prompt_tokens, 1, 2, args.temperature) # warmup
+    prefill_result = bench_generate(engine, prompt_tokens, 1, 2, args.temperature)
+    prefill_time = prefill_result["ttft"]
+    prefill_mfu = 100 * prefill_flops(model, prompt_len) / prefill_time / peak_flops
+    prefill_tok_per_sec = prompt_len / prefill_time
+    print(f"Prefill (batch 1, {prompt_len} tokens): {prefill_tok_per_sec:,.0f} tok/s | MFU {prefill_mfu:.1f}%")
+    payload["prefill"] = {
+        "tok_per_sec": round(prefill_tok_per_sec, 1),
+        "mfu_percent": round(prefill_mfu, 2),
+        "time_sec": round(prefill_time, 6),
+    }
+
+    # ------------------------------------------------------------------------
+    # Measured sweep over batch sizes. Decode reads all weights + KV every step:
+    # MBU is the distance from the bandwidth roofline (binds at small batch),
+    # MFU the distance from the compute roofline (binds at large batch).
+    batch_sizes = [int(b) for b in args.batch_sizes.split(",")]
+    header = f"{'batch':>6} {'TTFT ms':>9} {'TPOT ms':>9} {'tok/s':>10} {'MBU %':>7} {'MFU %':>7} {'VRAM GiB':>9} {'steps':>6}"
+    print(header)
+    print("-" * len(header))
+    for batch_size in batch_sizes:
+        # warmup (cublas autotune, allocator warm, attention kernels)
+        bench_generate(engine, prompt_tokens, batch_size, 8, args.temperature)
+        # timed run
+        result = bench_generate(engine, prompt_tokens, batch_size, args.decode_tokens, args.temperature)
+        step_times = result["step_times"]
+        num_steps = len(step_times)
+        if num_steps == 0:
+            print(f"{batch_size:>6}  all rows terminated during warmup?! skipping")
+            continue
+        tpot = sorted(step_times)[num_steps // 2] # median decode step time
+        tok_per_sec = batch_size * num_steps / sum(step_times)
+        # MBU: bytes each decode step must move, over what the GPU can move
+        bytes_per_step = w_bytes + batch_size * kv_read
+        mbu = 100 * (bytes_per_step / tpot) / peak_bw
+        # MFU: FLOPs each decode step must do, over what the GPU can do
+        flops_per_step = batch_size * decode_flops_per_token(model, context_mid)
+        mfu = 100 * (flops_per_step / tpot) / peak_flops
+        vram_gib = result["peak_vram"] / 2**30
+        note = "" if num_steps == args.decode_tokens - 1 else f" (early stop @ {num_steps})"
+        print(f"{batch_size:>6} {result['ttft']*1e3:>9.1f} {tpot*1e3:>9.2f} {tok_per_sec:>10,.0f} "
+              f"{mbu:>7.1f} {mfu:>7.2f} {vram_gib:>9.2f} {num_steps:>6}{note}")
+        payload["sweep"].append({
+            "batch_size": batch_size,
+            "ttft_sec": round(result["ttft"], 6), # microsecond resolution, plenty for wall clock
+            "tpot_sec": round(tpot, 6),
+            "tok_per_sec": round(tok_per_sec, 1),
+            "mbu_percent": round(mbu, 2),
+            "mfu_percent": round(mfu, 4), # decode MFU is tiny at small batch, keep the signal
+
+            "peak_vram_bytes": result["peak_vram"],
+            "decode_steps": num_steps,
+        })
+
+    # The last line of stdout is the machine-readable version of the whole run
+    print("-" * len(header))
+    print(json.dumps(payload))
+
+    compute_cleanup()
+
+if __name__ == "__main__":
+    main()