File size: 4,533 Bytes
e95c403 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """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]
# ---- SISO Forward ----
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
# ---- MIMO Forward ----
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
# ---- Decode step ----
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()
|