| """05_benchmark.py -- reproduce the perf numbers in the README. |
| |
| Runs the SISO and MIMO forward + decode benchmarks at target dims. Reports |
| warm median latency + p95, and a summary table matching the README. |
| |
| Run: |
| python 05_benchmark.py [--full] |
| |
| --full: also runs seqlen=512 and seqlen=1024 (adds ~5 minutes) |
| |
| Note: MIMO does not use torch.compile due to a known regression -- see the |
| project's ticket 61. SISO uses torch.compile at batch=1 for optimal single-request |
| latency. |
| """ |
|
|
| import argparse |
| import statistics |
| import time |
| import torch |
| from _loader import load_kernel |
|
|
|
|
| def _sync(): |
| try: |
| torch.neuron.synchronize() |
| except Exception: |
| pass |
|
|
|
|
| def bench_forward(mixer, batch, seqlen, warmup=3, iters=10): |
| u = torch.randn(batch, seqlen, 1024, device="neuron") |
| for _ in range(warmup): |
| y, _ = mixer(u) |
| _sync() |
| times = [] |
| for _ in range(iters): |
| t0 = time.perf_counter() |
| y, _ = mixer(u) |
| _sync() |
| times.append((time.perf_counter() - t0) * 1000) |
| return statistics.median(times), sorted(times)[int(len(times) * 0.95)] |
|
|
|
|
| def bench_decode(mixer, batch, prefill_len, n_steps=32, warmup=3): |
| u_pre = torch.randn(batch, prefill_len, 1024, device="neuron") |
| _, cache = mixer(u_pre) |
| _sync() |
| tokens = [torch.randn(batch, 1, 1024, device="neuron") for _ in range(warmup + n_steps)] |
| for tok in tokens[:warmup]: |
| y, cache = mixer.step(tok, cache) |
| _sync() |
| times = [] |
| for tok in tokens[warmup:]: |
| t0 = time.perf_counter() |
| y, cache = mixer.step(tok, cache) |
| _sync() |
| times.append((time.perf_counter() - t0) * 1000) |
| return statistics.median(times), sorted(times)[int(len(times) * 0.95)] |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--full", action="store_true", help="Include seqlen 512 and 1024") |
| args = ap.parse_args() |
|
|
| print("=" * 78) |
| print("mamba3-neuron-kernels benchmark (target dims)") |
| print("=" * 78) |
| print("Setup: d_model=1024, d_state=128, headdim=64, nheads=32, batch=1") |
| print("Env: TORCH_NEURONX_ENABLE_CONCATENATION=1 recommended") |
|
|
| mamba3 = load_kernel() |
|
|
| seqlens = [64, 256, 512, 1024] if args.full else [64, 256] |
|
|
| |
| print("\n### SISO forward (NKI SSD + torch.compile) ###") |
| print(f"{'seqlen':>7s} {'median (ms)':>12s} {'p95 (ms)':>10s}") |
| print("-" * 32) |
| for seqlen in seqlens: |
| torch.manual_seed(0) |
| m = mamba3.NeuronMamba3Mixer( |
| d_model=1024, d_state=128, headdim=64, chunk_size=64, mimo_rank=1, |
| use_nki_ssd=True, |
| ).to("neuron") |
| m_compiled = torch.compile(m, backend="neuron") |
| try: |
| med, p95 = bench_forward(m_compiled, 1, seqlen) |
| print(f"{seqlen:>7d} {med:>10.2f} {p95:>8.2f}") |
| except Exception as e: |
| print(f"{seqlen:>7d} FAILED: {type(e).__name__}: {str(e)[:40]}") |
| del m, m_compiled |
|
|
| |
| print("\n### MIMO forward (NKI SSD, eager -- torch.compile disabled per ticket 61) ###") |
| print(f"{'seqlen':>7s} {'median (ms)':>12s} {'p95 (ms)':>10s}") |
| print("-" * 32) |
| for seqlen in seqlens: |
| torch.manual_seed(0) |
| m = mamba3.NeuronMamba3Mixer( |
| d_model=1024, d_state=128, headdim=64, chunk_size=16, mimo_rank=4, |
| use_nki_ssd=True, |
| ).to("neuron") |
| try: |
| med, p95 = bench_forward(m, 1, seqlen) |
| print(f"{seqlen:>7d} {med:>10.2f} {p95:>8.2f}") |
| except Exception as e: |
| print(f"{seqlen:>7d} FAILED: {type(e).__name__}: {str(e)[:40]}") |
| del m |
|
|
| |
| print("\n### Decode step (single-token, prefill_len=64) ###") |
| print(f"{'mode':>6s} {'median (ms)':>12s} {'p95 (ms)':>10s}") |
| print("-" * 32) |
| for mode in ["SISO", "MIMO"]: |
| torch.manual_seed(0) |
| if mode == "SISO": |
| m = mamba3.NeuronMamba3Mixer( |
| d_model=1024, d_state=128, headdim=64, chunk_size=64, mimo_rank=1, |
| use_nki_ssd=True, |
| ).to("neuron") |
| else: |
| m = mamba3.NeuronMamba3Mixer( |
| d_model=1024, d_state=128, headdim=64, chunk_size=16, mimo_rank=4, |
| use_nki_ssd=True, |
| ).to("neuron") |
| med, p95 = bench_decode(m, 1, prefill_len=64, n_steps=32) |
| print(f"{mode:>6s} {med:>10.2f} {p95:>8.2f}") |
| del m |
|
|
| print("\n" + "=" * 78) |
| print("Done.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
