| """Checks for the fake registrations, the backward shared-memory guard, and the |
| backward's workspace cost. |
| |
| run.cmd test_fixes.py |
| """ |
|
|
| 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 / "tests")) |
|
|
| import load_local as mamba3 |
| from test_mamba3 import build |
|
|
| DEV = "cuda" |
| fails = [] |
|
|
|
|
| def check(label, ok, detail=""): |
| print(f" {label:<52} {'ok' if ok else 'FAIL'} {detail}") |
| if not ok: |
| fails.append(label) |
|
|
|
|
| def case(B, S, H, G, P, N, R, Na, seed=0, grad=False): |
| c = build(B, S, H, G, P, N, R, Na, DEV, seed=seed) |
| if grad: |
| for k in ("q", "k", "v", "z", "q_bias", "k_bias", "mimo_v", "mimo_o", |
| "mimo_z", "D"): |
| c[k] = c[k].detach().requires_grad_(True) |
| return c |
|
|
|
|
| def fwd(c, C): |
| return mamba3.forward( |
| c["q"], c["k"], c["v"], c["q_bias"], c["k_bias"], c["mimo_v"], c["mimo_o"], |
| c["angles"], c["adt"], c["dt"], c["trap"], z=c["z"], mimo_z=c["mimo_z"], |
| D=c["D"], chunk_size=C) |
|
|
|
|
| print(f"torch {torch.__version__} {torch.cuda.get_device_name(0)}\n") |
|
|
| print("1. fake implementations") |
| import _m3pkg |
| registered = getattr(_m3pkg, "FAKE_OPS", ()) |
| check("every op carries a meta kernel", len(registered) == 4, str(registered)) |
|
|
| meta = torch.device("meta") |
| try: |
| with torch._subclasses.FakeTensorMode(): |
| c = case(1, 32, 4, 1, 32, 64, 4, 32) |
| y = fwd(c, 16) |
| shape_ok = tuple(y.shape) == (1, 32, 4, 32) |
| check("forward traces under FakeTensorMode", shape_ok, str(tuple(y.shape))) |
| except Exception as exc: |
| check("forward traces under FakeTensorMode", False, f"{type(exc).__name__}: {exc}") |
|
|
| print("\n2. torch.compile(fullgraph=True)") |
| try: |
| c = case(1, 32, 4, 1, 32, 64, 4, 32) |
| eager = fwd(c, 16) |
| compiled = torch.compile(lambda: fwd(c, 16), fullgraph=True)() |
| |
| |
| d = float((compiled - eager).abs().max() / eager.abs().max()) |
| check("compiles fullgraph and matches eager", d < 1e-5, f"max rel {d:.2e}") |
| except Exception as exc: |
| check("compiles fullgraph and matches eager", False, f"{type(exc).__name__}: {exc}") |
|
|
| print("\n3. backward shared-memory guard (no streamed path above C*R = 64)") |
| for C, R, P, want in [(16, 4, 64, True), (32, 4, 64, False), (16, 4, 128, False)]: |
| c = case(1, 64, 4, 1, P, 128, R, 64, grad=True) |
| try: |
| fwd(c, C).sum().backward() |
| got, why = True, "ran" |
| except RuntimeError as exc: |
| msg = str(exc) |
| got = False |
| why = "guarded" if "shared memory" in msg else f"UNGUARDED: {msg[:70]}" |
| check(f"C={C} R={R} P={P} C*R={C*R}", got == want or why == "guarded", why) |
|
|
| print("\n4. backward peak memory (H=32 P=64 N=128 R=4 C=16)") |
| for S in (512, 1024, 2048, 4096): |
| c = case(1, S, 32, 1, 64, 128, 4, 64, grad=True) |
| torch.cuda.empty_cache() |
| torch.cuda.reset_peak_memory_stats() |
| y = fwd(c, 16) |
| fwd_peak = torch.cuda.max_memory_allocated() / 2**20 |
| y.sum().backward() |
| torch.cuda.synchronize() |
| tot_peak = torch.cuda.max_memory_allocated() / 2**20 |
| print(f" S={S:<6} forward {fwd_peak:7.0f} MB forward+backward {tot_peak:7.0f} MB") |
| del c, y |
|
|
| print("\n" + ("PASS" if not fails else f"FAIL ({len(fails)}): " + "; ".join(fails))) |
| sys.exit(0 if not fails else 1) |
|
|
