Update verify.py

verify.py

@@ -32,7 +32,6 @@ model.config.use_cuda_kernel = False
 with torch.no_grad():
     out_f = model(ids).logits
 model.config.use_cuda_kernel = True
-
 diff = (out_k.float() - out_f.float()).abs()
 print(f"   max abs diff kernel vs fallback: {diff.max().item():.4e}")
 print(f"   mean abs diff:                   {diff.mean().item():.4e}")
@@ -41,7 +40,38 @@ print("   kernel top-1 next id:", out_k[0, -1].argmax().item(),
       "| fallback:", out_f[0, -1].argmax().item())
 
 print("=" * 70)
-print("4) Backward pass (gradients flow through kernel)")
+print("4) Cached vs uncached parity (prefill + stateful decode vs full recompute)")
+# The stateful path always runs on the PyTorch implementation; compare against a
+# full PyTorch-fallback forward so both sides share the same numerics (fp32
+# internal, sequential over time) -> should match (near-)exactly.
+SPLIT = 4  # decode the last SPLIT tokens one at a time
+assert ids.shape[1] > SPLIT
+model.config.use_cuda_kernel = False
+with torch.no_grad():
+    full = model(ids).logits
+    pre = model(ids[:, :-SPLIT], use_cache=True)      # prefill -> state
+    st, chunks = pre.state, [pre.logits]
+    for t in range(ids.shape[1] - SPLIT, ids.shape[1]):
+        o = model(ids[:, t:t + 1], state=st)          # passing state implies use_cache
+        st, chunks = o.state, chunks + [o.logits]
+    inc = torch.cat(chunks, dim=1)
+model.config.use_cuda_kernel = True
+diff_c = (inc.float() - full.float()).abs()
+print(f"   max abs diff cached vs uncached: {diff_c.max().item():.4e}")
+print(f"   mean abs diff:                   {diff_c.mean().item():.4e}")
+assert torch.allclose(inc.float(), full.float(), atol=1e-3, rtol=0), \
+    "cached (prefill+decode) logits diverge from full recompute"
+# also check against the kernel forward with a looser bf16 tolerance
+diff_ck = (inc.float() - out_k.float()).abs()
+print(f"   max abs diff cached vs kernel:   {diff_ck.max().item():.4e}")
+ax, wkv, fx = st[0]
+print("   state[0] shapes:",
+      "att_x_prev", tuple(ax.shape), ax.dtype, "|",
+      "wkv", tuple(wkv.shape), wkv.dtype, "|",
+      "ffn_x_prev", tuple(fx.shape), fx.dtype)
+
+print("=" * 70)
+print("5) Backward pass (gradients flow through kernel)")
 model.train()
 ids2 = tok("Backward test sentence for gradient check.", return_tensors="pt").input_ids.cuda()
 out = model(ids2, labels=ids2)
@@ -57,12 +87,29 @@ print(f"   blocks.6.att.w1 (decay-lora) grad norm: {g_w1.float().norm().item():.
 n_with_grad = sum(1 for p in model.parameters() if p.grad is not None and p.grad.abs().sum() > 0)
 n_total = sum(1 for _ in model.parameters())
 print(f"   params with non-zero grad: {n_with_grad}/{n_total}")
+model.zero_grad(set_to_none=True)
 
 print("=" * 70)
-print("5) generate()")
+print("6) generate() through the state cache")
 model.eval()
 with torch.no_grad():
     gen = model.generate(ids, max_new_tokens=10, do_sample=False)
 print("   generated:", tok.decode(gen[0]))
+
+# generate() must equal a manual RNN-mode greedy loop (same state-cache path)
+with torch.no_grad():
+    o = model(ids, use_cache=True)
+    cur = o.logits[:, -1].argmax(-1, keepdim=True)
+    st, toks = o.state, [cur]
+    for _ in range(9):
+        o = model(cur, state=st)
+        st = o.state
+        cur = o.logits[:, -1].argmax(-1, keepdim=True)
+        toks.append(cur)
+manual = torch.cat(toks, dim=1)
+match = bool((gen[:, ids.shape[1]:] == manual).all())
+print("   generate == manual stateful greedy:", match)
+assert match, "generate() output diverges from manual stateful greedy decode"
+
 print("=" * 70)
-print("ALL CHECKS DONE")
+print("ALL CHECKS DONE")