| """Check the independent reference against the upstream Mamba-3 reference.""" |
|
|
| import argparse |
| import sys |
|
|
| import torch |
|
|
| sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent)) |
| from anchor import build_case, load_upstream_refs, max_rel_err |
| from mamba3_ref import mimo_step_ref |
|
|
| GRID = [ |
| |
| (128, 8, 1, 32, 64, 4), |
| (128, 8, 1, 32, 64, 1), |
| (256, 4, 1, 64, 128, 4), |
| (256, 8, 2, 32, 64, 2), |
| (64, 8, 1, 64, 128, 8), |
| ] |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--mamba-tests", required=True) |
| ap.add_argument("--device", default="cpu") |
| args = ap.parse_args() |
| up = load_upstream_refs(args.mamba_tests) |
|
|
| print(f"torch {torch.__version__} device={args.device}") |
| print("mimo_step_ref (ours) vs mamba3_MIMO_step_ref (upstream):") |
| worst = 0.0 |
| for S, H, G, P, N, R in GRID: |
| B, C = 1, 16 |
| c = build_case(B, S, H, G, P, N, R, C, args.device, torch.float32, seed=0) |
|
|
| ours, _ = 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"], |
| ) |
| ref, _ = up.mamba3_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"], |
| ) |
| rel = max_rel_err(ours, ref) |
| worst = max(worst, rel) |
| print(f" S={S:4d} H={H:2d} G={G} P={P:3d} N={N:3d} R={R} max_rel={rel:.3e}") |
|
|
| |
| c = build_case(1, 128, 8, 1, 32, 64, 4, 16, args.device, torch.float32, seed=3) |
| w = torch.randn((8, 32), device=args.device, dtype=torch.float32) |
| ours, _ = 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"], |
| fused_norm=True, outproj_norm_weight=w, |
| ) |
| ref, _ = up.mamba3_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"], |
| fused_norm=True, outproj_norm_weight=w, |
| ) |
| rel = max_rel_err(ours, ref) |
| worst = max(worst, rel) |
| print(f" fused_norm max_rel={rel:.3e}") |
|
|
| print(f"\nworst: {worst:.3e}") |
| ok = worst < 1e-5 |
| print("PASS" if ok else "FAIL") |
| return 0 if ok else 1 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
