| """Full validation for the Mamba-3 MIMO kernels.""" |
|
|
| 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 |
| from mamba3_ref import mimo_step_ref |
|
|
| VAL = ["q", "k", "v", "z", "q_bias", "k_bias", "mimo_v", "mimo_o", "mimo_z", "D"] |
|
|
|
|
| def build(B, S, H, G, P, N, R, Na, dev, seed, dtype=torch.float32): |
| torch.manual_seed(seed) |
| f = lambda *s: torch.randn(*s, device=dev) |
| dt = F.softplus(-3.0 + f(B, H, S)) |
| c = { |
| "q": f(B, S, R, G, N), "k": f(B, S, R, G, N), |
| "v": f(B, S, H, P), "z": f(B, S, H, P), |
| "q_bias": f(H, R, N), "k_bias": f(H, R, N), |
| "mimo_v": torch.rand(H, R, P, device=dev) / R, |
| "mimo_o": torch.rand(H, R, P, device=dev) / R, |
| "mimo_z": torch.rand(H, R, P, device=dev) / R, |
| "D": f(H), "angles": torch.rand(B, S, H, Na, device=dev), |
| "dt": dt, "adt": -F.softplus(f(B, H, S)).clamp(max=-1e-4) * dt, |
| "trap": torch.rand(B, H, S, device=dev) * 0.5, |
| } |
| if dtype is not torch.float32: |
| c = {k: (t.to(dtype).contiguous() if k in VAL else t) for k, t in c.items()} |
| return c |
|
|
|
|
| def run_decode(c, B, S, H, P, N, R, Na, dev, dtype=torch.float32): |
| st = mamba3.DecodeState(B, H, P, N, R, Na, device=dev) |
| ys = [] |
| for t in range(S): |
| ys.append(st.step( |
| c["q"][:, t].contiguous(), c["k"][:, t].contiguous(), c["v"][:, t].contiguous(), |
| c["q_bias"], c["k_bias"], c["mimo_v"], c["mimo_o"], |
| c["angles"][:, t].contiguous(), c["adt"][:, :, t], c["dt"][:, :, t], |
| c["trap"][:, :, t], z=c["z"][:, t].contiguous(), |
| mimo_z=c["mimo_z"], D=c["D"])) |
| return torch.stack(ys, dim=1), st |
|
|
|
|
| def run_fwd(c, C, dtype=torch.float32, norm_w=None): |
| 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, norm_weight=norm_w) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--mamba-tests", required=True) |
| args = ap.parse_args() |
| up = load_upstream_refs(args.mamba_tests) |
| dev = "cuda" |
| fails = [] |
|
|
| def check(label, rel, tol): |
| ok = rel < tol |
| print(f" {label:<52} {rel:.3e} {'ok' if ok else 'FAIL'}") |
| if not ok: |
| fails.append(label) |
|
|
| print(f"torch {torch.__version__} {torch.cuda.get_device_name(0)}\n") |
|
|
| print("1. decode step vs independent PyTorch reference (float32)") |
| for B, S, H, G, P, N, R, Na in [(1, 24, 32, 1, 64, 128, 4, 32), |
| (2, 16, 8, 1, 32, 64, 4, 16), |
| (1, 16, 8, 2, 32, 64, 2, 32), |
| (1, 16, 8, 1, 64, 128, 1, 64), |
| (1, 12, 4, 1, 128, 64, 8, 16)]: |
| c = build(B, S, H, G, P, N, R, Na, dev, seed=0) |
| got, _ = run_decode(c, B, S, H, P, N, R, Na, dev) |
| ref, _ = mimo_step_ref(c["q"], c["k"], c["v"], c["adt"], c["dt"], c["trap"], |
| c["q_bias"], c["k_bias"], c["angles"], c["mimo_v"], |
| c["mimo_o"], D=c["D"], z=c["z"], mimo_z=c["mimo_z"]) |
| check(f"B={B} H={H} G={G} P={P} N={N} R={R} Na={Na}", max_rel_err(got, ref), 2e-5) |
|
|
| print("\n2. chunked forward vs upstream mamba3_MIMO_chunk_ref (float32)") |
| for S, H, G, P, N, R, C in [(128, 8, 1, 32, 64, 4, 16), (256, 32, 1, 64, 128, 4, 16), |
| (128, 8, 2, 32, 64, 2, 32), (128, 8, 1, 64, 128, 1, 64), |
| (512, 4, 1, 64, 128, 4, 16)]: |
| Na, B = N // 2, 1 |
| c = build(B, S, H, G, P, N, R, Na, dev, seed=0) |
| ang = mamba3.cumulative_angles(c["angles"], c["dt"]) |
| dacs, dacsr = mamba3.chunk_decay(c["adt"], C) |
| ref, _, _ = up.mamba3_MIMO_chunk_ref( |
| c["q"], c["k"], c["v"], c["q_bias"], c["k_bias"], c["mimo_v"], c["mimo_o"], |
| c["z"], c["mimo_z"], ang, dacs, dacsr, c["dt"], c["trap"], c["D"], |
| chunk_size=C, rotary_dim_divisor=2, dtype=torch.float32, rotate_pairwise=True) |
| check(f"S={S} H={H} G={G} P={P} N={N} R={R} C={C}", |
| max_rel_err(run_fwd(c, C), ref), 5e-5) |
|
|
| print("\n3. chunked forward vs decode kernel (partial rotation, float32)") |
| for S, H, G, P, N, R, Na, C in [(64, 8, 1, 64, 128, 4, 32, 16), |
| (64, 8, 1, 64, 128, 4, 16, 16), |
| (96, 4, 2, 32, 64, 2, 16, 32)]: |
| B = 1 |
| c = build(B, S, H, G, P, N, R, Na, dev, seed=1) |
| dec, _ = run_decode(c, B, S, H, P, N, R, Na, dev) |
| check(f"S={S} H={H} G={G} P={P} N={N} R={R} Na={Na} C={C}", |
| max_rel_err(run_fwd(c, C), dec), 5e-5) |
|
|
| print("\n4. sequence length not a multiple of chunk_size (forward vs decode)") |
| |
| for S, C, R in [(100, 16, 4), (17, 16, 4), (63, 32, 2), (200, 64, 1), (33, 16, 4)]: |
| B, H, G, P, N, Na = 1, 8, 1, 32, 64, 32 |
| c = build(B, S, H, G, P, N, R, Na, dev, seed=2) |
| dec, _ = run_decode(c, B, S, H, P, N, R, Na, dev) |
| check(f"S={S} C={C} R={R} ({S % C} left over)", |
| max_rel_err(run_fwd(c, C), dec), 5e-5) |
|
|
| print("\n5. fused head-wise RMSNorm (forward vs upstream reference)") |
| for S, H, P, N, R, C in [(128, 8, 32, 64, 4, 16), (256, 8, 64, 128, 4, 16)]: |
| B, G, Na = 1, 1, N // 2 |
| c = build(B, S, H, G, P, N, R, Na, dev, seed=3) |
| w = torch.randn(H, P, device=dev) |
| ang = mamba3.cumulative_angles(c["angles"], c["dt"]) |
| dacs, dacsr = mamba3.chunk_decay(c["adt"], C) |
| ref, _, _ = up.mamba3_MIMO_chunk_ref( |
| c["q"], c["k"], c["v"], c["q_bias"], c["k_bias"], c["mimo_v"], c["mimo_o"], |
| c["z"], c["mimo_z"], ang, dacs, dacsr, c["dt"], c["trap"], c["D"], |
| chunk_size=C, rotary_dim_divisor=2, dtype=torch.float32, rotate_pairwise=True, |
| fused_norm=True, outproj_norm_weight=w) |
| check(f"S={S} H={H} P={P} N={N} R={R} C={C}", |
| max_rel_err(run_fwd(c, C, norm_w=w), ref), 5e-5) |
|
|
| print("\n6. bfloat16 against the float32 kernel (~8 mantissa bits)") |
| for S, H, P, N, R, Na, C in [(256, 32, 64, 128, 4, 64, 16), (2048, 8, 64, 128, 4, 64, 16)]: |
| B, G = 1, 1 |
| c32 = build(B, S, H, G, P, N, R, Na, dev, seed=0) |
| c16 = build(B, S, H, G, P, N, R, Na, dev, seed=0, dtype=torch.bfloat16) |
| check(f"forward S={S} H={H} P={P} N={N} R={R}", |
| max_rel_err(run_fwd(c16, C).float(), run_fwd(c32, C)), 3e-2) |
| for S, H, P, N, R, Na in [(16, 8, 64, 128, 4, 32)]: |
| B, G = 1, 1 |
| c32 = build(B, S, H, G, P, N, R, Na, dev, seed=0) |
| c16 = build(B, S, H, G, P, N, R, Na, dev, seed=0, dtype=torch.bfloat16) |
| d16, _ = run_decode(c16, B, S, H, P, N, R, Na, dev) |
| d32, _ = run_decode(c32, B, S, H, P, N, R, Na, dev) |
| check(f"decode S={S} H={H} P={P} N={N} R={R}", |
| max_rel_err(d16.float(), d32), 3e-2) |
|
|
| print("\n7. long sequences (the rotation angle accumulates over the sequence)") |
| |
| |
| |
| for S in [2048, 8192, 32768]: |
| B, H, G, P, N, R, C = 1, 4, 1, 64, 128, 4, 16 |
| Na = N // 2 |
| c = build(B, S, H, G, P, N, R, Na, dev, seed=0) |
| ang = mamba3.cumulative_angles(c["angles"], c["dt"]) |
| dacs, dacsr = mamba3.chunk_decay(c["adt"], C) |
| ref, _, _ = up.mamba3_MIMO_chunk_ref( |
| c["q"], c["k"], c["v"], c["q_bias"], c["k_bias"], c["mimo_v"], c["mimo_o"], |
| c["z"], c["mimo_z"], ang, dacs, dacsr, c["dt"], c["trap"], c["D"], |
| chunk_size=C, rotary_dim_divisor=2, dtype=torch.float32, rotate_pairwise=True) |
| check(f"S={S} (max angle {ang.abs().max().item():.0f} rad)", |
| max_rel_err(run_fwd(c, C), ref), 5e-6) |
| del c, ref, ang |
| torch.cuda.empty_cache() |
|
|
| print("\n" + ("PASS" if not fails else f"FAIL ({len(fails)}): " + "; ".join(fails))) |
| return 0 if not fails else 1 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
