| """Streamed-intra path (C*R > 64) against the upstream reference and small chunks.""" |
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(ROOT)) |
| sys.path.insert(0, str(ROOT / "ref")) |
| import load_local as mamba3 |
| from anchor import load_upstream_refs, max_rel_err |
|
|
| ap = argparse.ArgumentParser(); ap.add_argument("--mamba-tests", required=True) |
| up = load_upstream_refs(ap.parse_args().mamba_tests) |
| dev = "cuda" |
| worst, fails = 0.0, [] |
|
|
| |
| GRID = [(256, 8, 1, 64, 128, 4, 64), (256, 8, 1, 32, 64, 4, 32), |
| (512, 4, 1, 64, 128, 4, 64), (128, 8, 2, 32, 64, 2, 64), |
| (256, 8, 1, 64, 128, 1, 128), (192, 4, 1, 64, 128, 8, 32)] |
|
|
| for S, H, G, P, N, R, C in GRID: |
| B, Na = 1, N // 2 |
| torch.manual_seed(0) |
| f = lambda *s: torch.randn(*s, device=dev) |
| dt = F.softplus(-3.0 + f(B, H, S)) |
| adt = -F.softplus(f(B, H, S)).clamp(max=-1e-4) * dt |
| trap = torch.rand(B, H, S, device=dev) * 0.5 |
| raw = torch.rand(B, S, H, Na, device=dev) |
| q, k, v, z = f(B, S, R, G, N), f(B, S, R, G, N), f(B, S, H, P), f(B, S, H, P) |
| qb, kb = f(H, R, N), f(H, R, N) |
| mv, mo, mz = (torch.rand(H, R, P, device=dev) / R for _ in range(3)) |
| D = f(H) |
|
|
| got = mamba3.forward(q, k, v, qb, kb, mv, mo, raw, adt, dt, trap, |
| z=z, mimo_z=mz, D=D, chunk_size=C) |
| ang = mamba3.cumulative_angles(raw, dt) |
| dcs, dcsr = mamba3.chunk_decay(adt, C) |
| ref, _, _ = up.mamba3_MIMO_chunk_ref(q, k, v, qb, kb, mv, mo, z, mz, ang, dcs, dcsr, |
| dt, trap, D, chunk_size=C, rotary_dim_divisor=2, |
| dtype=torch.float32, rotate_pairwise=True) |
| e = max_rel_err(got, ref) |
| worst = max(worst, e) |
| ok = e < 5e-5 |
| if not ok: |
| fails.append(f"S{S}C{C}R{R}") |
| print(f" S={S:4d} H={H:2d} G={G} P={P:3d} N={N:3d} R={R} C={C:3d} CR={C*R:3d}" |
| f" {e:.3e} {'ok' if ok else 'FAIL'}") |
|
|
| print(f"\nworst {worst:.3e}") |
| print("PASS" if not fails else "FAIL: " + ", ".join(fails)) |
| sys.exit(0 if not fails else 1) |
|
|
