"""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 # noqa: E402 from anchor import load_upstream_refs, max_rel_err # noqa: E402 from mamba3_ref import mimo_step_ref # noqa: E402 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)") # chunk_size * mimo_rank stays at 64, the geometry upstream recommends. 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)") # The angle is a cumulative sum and grows without bound, so an approximate # sincos would make the residual proportional to it rather than flat. The # grids above top out at S=512 and cannot see that. 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())