P3: add full API integration test

tests/unit/test_p3_fast_decode.py

@@ -105,10 +105,80 @@ def test_fast_path():
     return all_pass
 
 
+def test_full_api():
+    """Test the full dsv4_attention API (goes through fast path for T=1, N<=128)."""
+    from dsv4.kernels.attention.production import dsv4_attention
+
+    torch.manual_seed(99)
+
+    configs = [
+        (8, 8, 128, 64, "MHA hd=64"),
+        (8, 8, 128, 128, "MHA hd=128"),
+        (8, 1, 128, 64, "MQA hd=64"),
+        (8, 1, 128, 128, "MQA hd=128"),
+        (8, 2, 128, 64, "GQA hd=64"),
+    ]
+
+    all_pass = True
+    for n_q, n_kv, N, hd, desc in configs:
+        scale = 1.0 / math.sqrt(hd)
+        try:
+            q = torch.randn(n_q, 1, hd, dtype=torch.bfloat16, device='cuda')
+            if n_kv == 1:
+                k = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda')
+                v = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda')
+            else:
+                k = torch.randn(n_kv, N, hd, dtype=torch.bfloat16, device='cuda')
+                v = torch.randn(n_kv, N, hd, dtype=torch.bfloat16, device='cuda')
+
+            # Full API call (should use fast path)
+            o_fast = dsv4_attention(q, k, v, scale=scale)
+
+            # Reference
+            o_ref = reference_attention_api(q, k, v, scale, n_q, n_kv, N, hd)
+
+            cos = cosine_sim(o_ref, o_fast).item()
+            status = "PASS" if cos >= 0.999990 else "FAIL"
+            if status == "FAIL":
+                all_pass = False
+            print(f"  {status} [full API] {desc}: cos={cos:.6f}")
+        except Exception as e:
+            import traceback
+            print(f"  FAIL [full API] {desc}: {e}")
+            traceback.print_exc()
+            all_pass = False
+
+    return all_pass
+
+
+def reference_attention_api(q, k, v, scale, n_q, n_kv, N, hd):
+    """Reference that matches dsv4_attention input format."""
+    q_per_kv = n_q // n_kv
+    if k.dim() == 2:
+        k = k.unsqueeze(0)
+    if v.dim() == 2:
+        v = v.unsqueeze(0)
+    output = torch.zeros(n_q, 1, hd, dtype=torch.bfloat16, device='cuda')
+    for kv_idx in range(n_kv):
+        k_h = k[kv_idx]  # (N, hd)
+        v_h = v[kv_idx]  # (N, hd)
+        for qi in range(q_per_kv):
+            q_idx = kv_idx * q_per_kv + qi
+            q_h = q[q_idx]  # (1, hd)
+            s = torch.matmul(q_h.float(), k_h.float().T) * scale
+            s = torch.softmax(s, dim=-1)
+            o = torch.matmul(s, v_h.float())
+            output[q_idx] = o.bfloat16()
+    return output
+
+
 if __name__ == "__main__":
     print("P3 Integration Test: 6-warp decode fast path")
     print("=" * 60)
-    ok = test_fast_path()
+    ok1 = test_fast_path()
+    print()
+    ok2 = test_full_api()
     print("=" * 60)
+    ok = ok1 and ok2
     print("ALL PASS" if ok else "SOME FAILED")
     sys.exit(0 if ok else 1)