File size: 2,670 Bytes
d5a6d53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """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 # noqa: E402
from mamba3_ref import mimo_step_ref # noqa: E402
GRID = [
# S H G P N R
(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}")
# fused-norm variant
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())
|