""" 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 # ----------------------------------------------------------------------------- # Measurement # (the architecture-side cost accounting - FLOPs, KV cache bytes - lives on the # GPT model itself: estimate_decode_flops, estimate_prefill_flops, kv_bytes_per_token, kv_read_bytes) 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 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 = model.kv_bytes_per_token() context_mid = prompt_len + args.decode_tokens // 2 # representative decode context kv_read = model.kv_read_bytes(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": model.estimate_decode_flops(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 * model.estimate_prefill_flops(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 * model.estimate_decode_flops(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()