Gumiho #3: speculative-decoding throughput/latency benchmark (standard JSON)

benchmarks/apple_gpu/README.md

@@ -1,12 +1,31 @@
 # Apple GPU benchmarks
 
-Two benchmark drivers live here:
+Benchmark drivers live here:
 
 | Driver | Coverage |
 |---|---|
 | [`benchmark_fusion.py`](benchmark_fusion.py) | Phase 8.4.x MSL fusion sweep — `matmul → softmax`, SwiGLU MLP block. Fused-vs-sequential pairing. |
+| [`benchmark_encode_session.py`](benchmark_encode_session.py) | R2 command-buffer batching — a dependent bmm chain per-op (one sync/op) vs. batched into one command buffer. |
+| [`benchmark_gumiho.py`](benchmark_gumiho.py) | Gumiho speculative decoding — algorithmic tokens-per-target-pass + resident vs per-op serial draft wall-clock. |
 | [`benchmark_ga_ebm.py`](benchmark_ga_ebm.py) | GA + EBM end-to-end stack walk **plus workload mode**. 17 GA primitives + **9 native EBM primitives** (including `ebm_partition_exact`) + Python-reference comparison rows + **4 workload rows** (GA feature pipeline + EBT-tiny refinement, each in apple_gpu + python_ref variants), plus opt-in `--ebt-sweep`. |
 
+## Gumiho speculative-decoding benchmark
+
+[`benchmark_gumiho.py`](benchmark_gumiho.py) emits the standard JSON schema with
+two op groups:
+
+- `op="gumiho_decode"` — **tokens committed per target forward pass** for
+  `vanilla` (1.0), `speculative_untrained`, and `speculative_trained` (rows
+  carry `tokens_per_step`, `speedup_vs_vanilla`, `mean_accepted_length`,
+  `tokens_per_sec`). This is the model-size-independent speculative win. On a
+  recent M-series run: untrained ~1.86×, trained ~4.33× tokens/pass; trained
+  also wins ~2.3× on wall-clock throughput because 4.3 tokens/step amortizes the
+  verify.
+- `op="gumiho_serial_draft"` — `resident` (one command buffer per token) vs.
+  `per_op` serial-draft wall-clock; `speedup_vs_per_op` ~2.6× from removing the
+  per-op CPU↔GPU sync. Headline numbers drift across machines/toolchains — the
+  schema is the stable contract.
+
 ## GA + EBM benchmark — what it walks
 
 For every primitive the driver exercises the full stack:

benchmarks/apple_gpu/benchmark_gumiho.py

@@ -0,0 +1,191 @@
+"""apple_gpu Gumiho speculative-decoding benchmark.
+
+Two measurements, both emitted in the standard ``benchmark_gemm.py`` JSON schema
+(roofline-ingestible):
+
+1. **Algorithmic speculative speedup** (``op="gumiho_decode"``) — tokens
+   committed per target forward pass for vanilla autoregressive decode (1.0) vs.
+   speculative decode before and after distilling the draft. This is the real
+   speculative-decoding win and is **model-size independent**: at toy scale the
+   draft is no cheaper than the target, so wall-clock won't show it, but the
+   ``tokens_per_step`` ratio is exactly what scales to a real draft≪target gap.
+
+2. **Resident vs per-op serial draft** (``op="gumiho_serial_draft"``) — wall-clock
+   of the serial draft run GPU-resident (one command buffer per token) vs. the
+   per-op path (one dispatch + sync per op). This sync-reduction **is** a real
+   wall-clock win even at toy scale, because per-op CPU↔GPU sync dominates
+   small-batch decode.
+
+Best-effort: runs only on Darwin with Metal; elsewhere it writes an empty run
+list and exits 0.
+
+Usage:
+    python benchmarks/apple_gpu/benchmark_gumiho.py --reps 5 --max-new-tokens 24
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+_EX = Path(__file__).resolve().parents[2] / "examples" / "advanced" / "gumiho"
+if str(_EX) not in sys.path:
+    sys.path.insert(0, str(_EX))
+
+from gumiho.backend import NumpyBackend, make_backend  # noqa: E402
+from gumiho.config import tiny_config  # noqa: E402
+from gumiho.decode import run_multistep_decode  # noqa: E402
+from gumiho.model import SerialHead, TargetModel, make_weights  # noqa: E402
+from gumiho.resident import ResidentSerialDraft  # noqa: E402
+from gumiho.training import distill, trajectory_contexts  # noqa: E402
+
+
+def _median_ms(fn, reps: int) -> tuple[float, float]:
+    fn()  # warm up
+    s = []
+    for _ in range(reps):
+        t0 = time.perf_counter_ns()
+        fn()
+        s.append((time.perf_counter_ns() - t0) / 1e6)
+    return statistics.median(s), (statistics.stdev(s) if reps > 1 else 0.0)
+
+
+def _vanilla_decode(cfg, weights, prompts, n, backend):
+    """Plain autoregressive decode — one target forward per token."""
+    tgt = TargetModel(weights, cfg)
+    total = 0
+    for prompt in prompts:
+        seq = list(int(t) for t in prompt)
+        for _ in range(n):
+            ctx = np.asarray(seq[-cfg.context_len:], np.int64)
+            _h, logits = tgt.forward(backend, ctx)
+            seq.append(int(np.argmax(logits[-1])))
+            total += 1
+    return total
+
+
+def _device_name() -> str:
+    return "apple_silicon_metal" if sys.platform == "darwin" else "non-darwin-fallback"
+
+
+def _tessera_version() -> str:
+    try:
+        import importlib.metadata
+        return importlib.metadata.version("tessera")
+    except Exception:
+        return "dev"
+
+
+def _row(op, mode, shape, ms, stdev, reps, device, version, **extra):
+    row = {
+        "backend": "apple_gpu", "op": op, "shape": shape, "dtype": "f32",
+        "mode": mode, "reps": reps, "latency_ms": ms, "stdev_ms": stdev,
+        "tflops": 0.0, "memory_bw_gb_s": 0.0,
+        "device": device, "tessera_version": version,
+    }
+    row.update(extra)
+    return row
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--reps", type=int, default=5)
+    parser.add_argument("--prompts", type=int, default=6)
+    parser.add_argument("--max-new-tokens", type=int, default=24)
+    parser.add_argument("--train-steps", type=int, default=500)
+    parser.add_argument("--target", default="apple_gpu",
+                        choices=["apple_gpu", "apple_cpu", "numpy"])
+    parser.add_argument("--output", type=Path, default=None)
+    args = parser.parse_args(argv)
+
+    if sys.platform != "darwin" and args.target != "numpy":
+        if args.output is not None:
+            args.output.write_text(json.dumps(
+                {"runs": [], "skipped_apple_gpu": "non-Darwin host — no Metal"},
+                indent=2, sort_keys=True))
+        print("gumiho benchmark: skipping (non-Darwin host)", file=sys.stderr)
+        return 0
+
+    cfg = tiny_config()
+    device, version = _device_name(), _tessera_version()
+    shape = (f"v{cfg.vocab}_d{cfg.d_model}_s{cfg.serial_tokens}"
+             f"_p{cfg.parallel_heads}_t{args.max_new_tokens}")
+    rows: list[dict[str, Any]] = []
+
+    weights = make_weights(cfg, seed=0)
+    rng = np.random.default_rng(0)
+    prompts = rng.integers(0, cfg.vocab, size=(args.prompts, cfg.context_len),
+                           dtype=np.int64)
+    ctxs = trajectory_contexts(cfg, weights, prompts, args.max_new_tokens)
+    trained = distill(cfg, weights, contexts=ctxs, serial_steps=args.train_steps,
+                      parallel_steps=args.train_steps, lr=0.05, seed=0)
+
+    # ── 1. Algorithmic speculative speedup (tokens per target pass) ──────────
+    n = args.max_new_tokens
+    ref = NumpyBackend(eps=cfg.rmsnorm_eps)
+    vbackend = make_backend(args.target, eps=cfg.rmsnorm_eps)
+
+    vms, vsd = _median_ms(
+        lambda: _vanilla_decode(cfg, weights, prompts, n, vbackend), args.reps)
+    vtoks = args.prompts * n
+    rows.append(_row("gumiho_decode", "vanilla", shape, vms, vsd, args.reps,
+                     device, version, tokens_per_step=1.0, speedup_vs_vanilla=1.0,
+                     tokens_per_sec=vtoks / (vms / 1000.0) if vms else 0.0))
+
+    for mode, w in (("speculative_untrained", weights),
+                    ("speculative_trained", trained)):
+        holder: dict[str, Any] = {}
+
+        def _run(_w=w, _h=holder):
+            _h["m"] = run_multistep_decode(cfg, _w, prompts=prompts,
+                                           max_new_tokens=n, target=args.target,
+                                           trained=(_w is trained))
+        ms, sd = _median_ms(_run, args.reps)
+        m = holder["m"]
+        tps = m.tokens_per_step
+        tok_s = (m.tokens_generated / (ms / 1000.0)) if ms else 0.0
+        rows.append(_row("gumiho_decode", mode, shape, ms, sd, args.reps,
+                         device, version, tokens_per_step=tps,
+                         speedup_vs_vanilla=m.speedup_vs_vanilla,
+                         mean_accepted_length=m.mean_accepted_length,
+                         tokens_per_sec=tok_s))
+
+    # ── 2. Resident vs per-op serial draft (wall-clock sync reduction) ───────
+    tgt = TargetModel(trained, cfg)
+    last_hidden, _ = tgt.forward(ref, prompts[0].astype(np.int64))
+    h_init, root = last_hidden[-1], int(prompts[0][-1])
+
+    rd = ResidentSerialDraft(trained, cfg)
+    sh = SerialHead(trained, cfg)
+    op_backend = make_backend(args.target, eps=cfg.rmsnorm_eps)
+
+    rms, rsd = _median_ms(lambda: rd.generate(h_init, root), args.reps)
+    pms, psd = _median_ms(
+        lambda: sh.generate(op_backend, tgt, h_init, root), args.reps)
+    sshape = f"d{cfg.d_model}_layers{cfg.serial_layers}_steps{cfg.serial_tokens}"
+    speedup = pms / rms if rms else 0.0
+    rows.append(_row("gumiho_serial_draft", "resident", sshape, rms, rsd,
+                     args.reps, device, version, command_buffers=cfg.serial_tokens,
+                     speedup_vs_per_op=speedup, backend_path=rd._metal and "metal" or "numpy"))
+    rows.append(_row("gumiho_serial_draft", "per_op", sshape, pms, psd,
+                     args.reps, device, version,
+                     ops_per_step=rd.ops_per_step, speedup_vs_per_op=1.0))
+    rd.free()
+
+    payload = {"runs": rows}
+    out = json.dumps(payload, indent=2, sort_keys=True)
+    if args.output is not None:
+        args.output.write_text(out)
+    else:
+        print(out)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/unit/test_benchmark_gumiho.py

@@ -0,0 +1,65 @@
+"""Smoke test for benchmarks/apple_gpu/benchmark_gumiho.py.
+
+Runs the driver with tiny args and validates the standard JSON schema + that the
+expected op/mode rows are present. Keeps the perf driver honest as the example
+and runtime evolve (numbers are machine-dependent, so only structure is checked).
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import sys
+from pathlib import Path
+
+import pytest
+
+_BENCH = (Path(__file__).resolve().parents[2] / "benchmarks" / "apple_gpu"
+          / "benchmark_gumiho.py")
+_EX = Path(__file__).resolve().parents[2] / "examples" / "advanced" / "gumiho"
+
+
+@pytest.fixture(scope="module")
+def bench_main():
+    if not _BENCH.exists():
+        pytest.skip("benchmark_gumiho.py not present")
+    if str(_EX) not in sys.path:
+        sys.path.insert(0, str(_EX))
+    spec = importlib.util.spec_from_file_location("benchmark_gumiho", _BENCH)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod.main
+
+
+def test_benchmark_emits_schema(tmp_path, bench_main):
+    out = tmp_path / "gumiho.json"
+    rc = bench_main([
+        "--reps", "1", "--prompts", "2", "--max-new-tokens", "8",
+        "--train-steps", "120", "--target", "numpy", "--output", str(out),
+    ])
+    assert rc == 0
+    payload = json.loads(out.read_text())
+    runs = payload["runs"]
+    assert runs, "no benchmark rows emitted"
+
+    required = {"backend", "op", "shape", "dtype", "mode", "latency_ms",
+                "tflops", "memory_bw_gb_s", "device", "tessera_version"}
+    for row in runs:
+        assert required <= set(row), f"row missing schema keys: {row}"
+        assert row["backend"] == "apple_gpu"
+
+    ops = {(r["op"], r["mode"]) for r in runs}
+    assert ("gumiho_decode", "vanilla") in ops
+    assert ("gumiho_decode", "speculative_trained") in ops
+    assert ("gumiho_serial_draft", "resident") in ops
+    assert ("gumiho_serial_draft", "per_op") in ops
+
+
+def test_trained_beats_vanilla_tokens_per_pass(bench_main, tmp_path):
+    out = tmp_path / "g.json"
+    bench_main(["--reps", "1", "--prompts", "2", "--max-new-tokens", "8",
+                "--train-steps", "200", "--target", "numpy", "--output", str(out)])
+    runs = json.loads(out.read_text())["runs"]
+    by_mode = {r["mode"]: r for r in runs if r["op"] == "gumiho_decode"}
+    assert by_mode["vanilla"]["tokens_per_step"] == 1.0
+    assert by_mode["speculative_trained"]["tokens_per_step"] > 1.5