Add microgpt: from-scratch GPT with RoPE, KV cache, sampling, tests and benchmark

scripts/benchmark.py

@@ -0,0 +1,44 @@
+"""Measure decoding throughput with and without the KV cache.
+
+    python scripts/benchmark.py            # prints a markdown table
+"""
+
+from __future__ import annotations
+
+import time
+
+import torch
+
+from microgpt.model import GPT, GPTConfig
+
+
+def time_generate(model: GPT, prompt: torch.Tensor, n: int, use_cache: bool) -> float:
+    t0 = time.perf_counter()
+    model.generate(prompt, n, temperature=0, use_cache=use_cache)
+    return time.perf_counter() - t0
+
+
+def main() -> None:
+    torch.manual_seed(0)
+    torch.set_num_threads(max(1, torch.get_num_threads()))
+    cfg = GPTConfig(vocab_size=96, block_size=512, n_layer=6, n_head=6, n_embd=192)
+    model = GPT(cfg).eval()
+    prompt = torch.randint(0, cfg.vocab_size, (1, 16))
+    print(f"model: {model.num_params() / 1e6:.2f}M params, {torch.get_num_threads()} CPU threads\n")
+    for use_cache in (False, True):  # warm up allocator / thread pool before timing
+        time_generate(model, prompt, 32, use_cache)
+    print("| new tokens | no cache (tok/s) | KV cache (tok/s) | speedup | identical output |")
+    print("|-----------:|-----------------:|-----------------:|--------:|:-----------------:|")
+    for n in (32, 64, 128, 256, 480):
+        slow = time_generate(model, prompt, n, use_cache=False)
+        fast = time_generate(model, prompt, n, use_cache=True)
+        same = torch.equal(
+            model.generate(prompt, n, temperature=0, use_cache=False),
+            model.generate(prompt, n, temperature=0, use_cache=True),
+        )
+        verdict = "yes" if same else "NO"
+        print(f"| {n} | {n / slow:.1f} | {n / fast:.1f} | {slow / fast:.1f}x | {verdict} |")
+
+
+if __name__ == "__main__":
+    main()