""" Numerical parity test: CUDA causal TTT vs the torch reference. Validates `src/models/ttt_cuda` against `causal_block_fast_weight_swish_glu` from `src/models/ttt.py`, at the REAL vision-expert shapes: head_dim=64, d_h=64, num_heads=12, 256 image tokens, chunk_size=256, with the global VLM-context pre-update (ctx present), muon_update_steps=0. Checks: 1. output max-abs-err vs torch reference (bf16 tolerance) 2. updated fast-weight (w0/w1/w2) max-abs-err 3. causality: perturbing query token p must not change outputs < p (only meaningful when chunk_size < L; tested with a small chunk_size too) Run (needs a GPU; timing is NOT measured here so a shared card is fine): TORCHDYNAMO_DISABLE=1 PYTHONPATH="$PWD" \ /mnt/afs-h200/yuyangcheng/venvs/fla_triton32/bin/python -m src.models.test_ttt_cuda """ import os import sys import torch # eager reference (compile disabled keeps it simple + matches training env) os.environ.setdefault("TORCHDYNAMO_DISABLE", "1") try: torch._dynamo.config.suppress_errors = True except Exception: pass from src.models.ttt import causal_block_fast_weight_swish_glu from src.models import ttt_cuda def _mk(B, L, d, dh, T_ctx, dtype, device, seed=0): g = torch.Generator(device=device).manual_seed(seed) rk = lambda *s: torch.randn(*s, generator=g, device=device, dtype=dtype) w0 = rk(B, d, dh) w1 = rk(B, dh, d) w2 = rk(B, d, dh) q = rk(B, L, d) k = rk(B, L, d) v = rk(B, L, d) # lr is fp32 and positive (softplus output in the real path) lr0 = torch.rand(B, L, 1, generator=g, device=device, dtype=torch.float32) * 0.02 + 0.001 lr1 = torch.rand(B, L, 1, generator=g, device=device, dtype=torch.float32) * 0.02 + 0.001 lr2 = torch.rand(B, L, 1, generator=g, device=device, dtype=torch.float32) * 0.02 + 0.001 ck = rk(B, T_ctx, d) cv = rk(B, T_ctx, d) cl0 = torch.rand(B, T_ctx, 1, generator=g, device=device, dtype=torch.float32) * 0.02 + 0.001 cl1 = torch.rand(B, T_ctx, 1, generator=g, device=device, dtype=torch.float32) * 0.02 + 0.001 cl2 = torch.rand(B, T_ctx, 1, generator=g, device=device, dtype=torch.float32) * 0.02 + 0.001 return dict(w0=w0, w1=w1, w2=w2, q=q, k=k, v=v, lr0=lr0, lr1=lr1, lr2=lr2, ck=ck, cv=cv, cl0=cl0, cl1=cl1, cl2=cl2) def _run_ref(t, chunk_size): return causal_block_fast_weight_swish_glu( t["w0"].clone(), t["w1"].clone(), t["w2"].clone(), t["q"], t["k"], t["v"], t["lr0"], t["lr1"], t["lr2"], chunk_size=chunk_size, muon_update_steps=0, vlm_k=t["ck"], vlm_v=t["cv"], vlm_lr0=t["cl0"], vlm_lr1=t["cl1"], vlm_lr2=t["cl2"], ) def _run_cuda(t, chunk_size): return ttt_cuda.causal_ttt_forward( t["w0"].clone(), t["w1"].clone(), t["w2"].clone(), t["q"], t["k"], t["v"], t["lr0"], t["lr1"], t["lr2"], chunk_size=chunk_size, vlm_k=t["ck"], vlm_v=t["cv"], vlm_lr0=t["cl0"], vlm_lr1=t["cl1"], vlm_lr2=t["cl2"], ) def _err(a, b): return (a.float() - b.float()).abs().max().item() def _rel_err(a, b): # max abs err normalized by the reference's max abs magnitude denom = a.float().abs().max().item() + 1e-9 return (a.float() - b.float()).abs().max().item() / denom def main(): if not torch.cuda.is_available(): print("SKIP: no CUDA device") return 0 dev = "cuda" # bf16: the real training/inference dtype; fp32 also tested for tighter bound. results = [] # ---- fp32: strict absolute/relative parity (the math must be correct) ---- for chunk_size in (256, 64): # 256 = real config (1 chunk); 64 = multi-chunk causal t = _mk(B=24, L=256, d=64, dh=64, T_ctx=32, dtype=torch.float32, device=dev, seed=1) o_ref, w0r, w1r, w2r = _run_ref(t, chunk_size) try: o_c, w0c, w1c, w2c = _run_cuda(t, chunk_size) except RuntimeError as e: print(f"FAIL build/run (fp32, cs={chunk_size}): {e}") return 1 eo = _rel_err(o_ref, o_c) ew = max(_err(w0r, w0c), _err(w1r, w1c), _err(w2r, w2c)) tol = 2e-3 ok = eo < tol and ew < tol print(f"[{'OK ' if ok else 'BAD'}] fp32 cs={chunk_size:3d} | " f"out_relerr={eo:.2e} w_err={ew:.2e} (tol={tol:.0e})") results.append(ok) # ---- bf16: the reference is ITSELF only bf16-accurate, so an absolute tol # is meaningless. Fair criterion: CUDA-bf16 must be no less accurate than # torch-bf16, both measured against the fp32 ground truth. Allow 1.5x slack. for chunk_size in (256, 64): # same bf16 inputs for all three; fp32 truth = those inputs upcast. t16 = _mk(B=24, L=256, d=64, dh=64, T_ctx=32, dtype=torch.bfloat16, device=dev, seed=1) t32 = {kk: (vv.float() if torch.is_tensor(vv) and vv.dtype == torch.bfloat16 else vv) for kk, vv in t16.items()} o_truth, w0t, w1t, w2t = _run_ref(t32, chunk_size) # fp32 ground truth o_tref, w0tr, w1tr, w2tr = _run_ref(t16, chunk_size) # torch bf16 o_cu, w0cu, w1cu, w2cu = _run_cuda(t16, chunk_size) # cuda bf16 # output relative error vs fp32 truth ref_oerr = _rel_err(o_truth, o_tref) cu_oerr = _rel_err(o_truth, o_cu) # weight error vs fp32 truth ref_werr = max(_err(w0t, w0tr), _err(w1t, w1tr), _err(w2t, w2tr)) cu_werr = max(_err(w0t, w0cu), _err(w1t, w1cu), _err(w2t, w2cu)) ok = cu_oerr <= ref_oerr * 1.5 + 1e-6 and cu_werr <= ref_werr * 1.5 + 1e-6 print(f"[{'OK ' if ok else 'BAD'}] bf16 cs={chunk_size:3d} | " f"out: cuda={cu_oerr:.2e} vs torch={ref_oerr:.2e} | " f"w: cuda={cu_werr:.2e} vs torch={ref_werr:.2e} (cuda<=1.5x torch)") results.append(ok) # causality (fp32, small chunk so it's a real test): perturb q token p, # outputs at positions < p (in earlier chunks) must be unchanged. t = _mk(B=4, L=128, d=64, dh=64, T_ctx=16, dtype=torch.float32, device=dev, seed=2) cs = 32 o0, *_ = _run_cuda(t, cs) p = 100 t2 = {kk: (vv.clone() if torch.is_tensor(vv) else vv) for kk, vv in t.items()} t2["q"][:, p, :] += 1.0 t2["k"][:, p, :] += 1.0 t2["v"][:, p, :] += 1.0 o1, *_ = _run_cuda(t2, cs) # positions strictly before p's chunk start must be bitwise-ish unchanged chunk_start = (p // cs) * cs delta_before = (o0[:, :chunk_start] - o1[:, :chunk_start]).abs().max().item() delta_after = (o0[:, p:] - o1[:, p:]).abs().max().item() causal_ok = delta_before < 1e-4 and delta_after > 1e-3 print(f"[{'OK ' if causal_ok else 'BAD'}] causality | " f"before_p_delta={delta_before:.2e} (want~0) after_p_delta={delta_after:.2e} (want>0)") results.append(causal_ok) allok = all(results) print("\n=== PARITY", "PASS ===" if allok else "FAIL ===") return 0 if allok else 1 if __name__ == "__main__": sys.exit(main())