Clean test suite for production attention (1/2/4 segments, multi-head)

tests/unit/test_production.py

@@ -4,112 +4,90 @@ import math
 from dsv4.kernels.attention.production import dsv4_attention
 
 
-def test_production_basic():
-    """Test basic single-head attention."""
+def test_single_head_128():
+    """Single head, 1 KV segment (N=128)."""
     torch.manual_seed(42)
-    hd = 64
-    n_h = 1
-    T = 128
-    N = 128
-
-    q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
-    k = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda')
-    v = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda')
-
-    # PyTorch reference (un-normalized)
-    qf = q[0].float()
-    kf = k.float()
-    vf = v.float()
-    scale = 1.0 / math.sqrt(hd)
-    attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0]
-    attn_exp = torch.exp(qf @ kf.T * scale - attn_max)
-    attn_sum = attn_exp.sum(dim=-1, keepdim=True)
-    ref_unnorm = attn_exp @ vf
-    ref_norm = (attn_exp / attn_sum) @ vf
-
-    out = dsv4_attention(q, k, v)
-
-    cos_unnorm = torch.nn.functional.cosine_similarity(
-        out.flatten().unsqueeze(0), ref_unnorm.unsqueeze(0).flatten().unsqueeze(0)
-    ).item()
-    cos_norm = torch.nn.functional.cosine_similarity(
-        out.flatten().unsqueeze(0), ref_norm.unsqueeze(0).flatten().unsqueeze(0)
-    ).item()
-    print(f"  hd={hd}, n_h={n_h}, N={N}: cos_unnorm {cos_unnorm:.6f}  cos_norm {cos_norm:.6f}")
-
-
-def test_production_multi_head():
-    """Test multi-head attention (per-head launch)."""
-    torch.manual_seed(42)
-    hd = 64
-    n_h = 4
-    T = 128
-    N = 256
-
+    hd = 64; n_h = 1; T = 128; N = 128
     q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
     k = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
     v = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
 
-    # PyTorch reference
+    out = dsv4_attention(q, k, v)
+
+    qf = q[0].float(); kf = k[0].float(); vf = v[0].float()
+    scale = 1.0 / math.sqrt(hd)
+    attn = qf @ kf.T * scale
+    ref = torch.softmax(attn, dim=-1) @ vf
+
+    cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.unsqueeze(0).flatten().unsqueeze(0)).item()
+    print(f"  hd={hd}, n_h={n_h}, N={N}: cos {cos:.6f}  {'PASS' if cos >= 0.99 else 'FAIL'}")
+
+
+def test_single_head_256():
+    """Single head, 2 KV segments (N=256)."""
+    torch.manual_seed(42)
+    hd = 64; n_h = 1; T = 128; N = 256
+    q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
+    k = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
+    v = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
+
+    out = dsv4_attention(q, k, v)
+
+    qf = q[0].float(); kf = k[0].float(); vf = v[0].float()
+    scale = 1.0 / math.sqrt(hd)
+    attn = qf @ kf.T * scale
+    ref = torch.softmax(attn, dim=-1) @ vf
+
+    cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.unsqueeze(0).flatten().unsqueeze(0)).item()
+    print(f"  hd={hd}, n_h={n_h}, N={N}: cos {cos:.6f}  {'PASS' if cos >= 0.99 else 'FAIL'}")
+
+
+def test_single_head_512():
+    """Single head, 4 KV segments (N=512)."""
+    torch.manual_seed(42)
+    hd = 64; n_h = 1; T = 128; N = 512
+    q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
+    k = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
+    v = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
+
+    out = dsv4_attention(q, k, v)
+
+    qf = q[0].float(); kf = k[0].float(); vf = v[0].float()
+    scale = 1.0 / math.sqrt(hd)
+    attn = qf @ kf.T * scale
+    ref = torch.softmax(attn, dim=-1) @ vf
+
+    cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.unsqueeze(0).flatten().unsqueeze(0)).item()
+    print(f"  hd={hd}, n_h={n_h}, N={N}: cos {cos:.6f}  {'PASS' if cos >= 0.99 else 'FAIL'}")
+
+
+def test_multi_head():
+    """Multi-head, 2 KV segments."""
+    torch.manual_seed(42)
+    hd = 64; n_h = 4; T = 128; N = 256
+    q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
+    k = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
+    v = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
+
+    out = dsv4_attention(q, k, v)
+
     scale = 1.0 / math.sqrt(hd)
     ref = torch.zeros_like(q)
     for h in range(n_h):
-        qf = q[h].float()
-        kf = k[h].float()
-        vf = v[h].float()
+        qf = q[h].float(); kf = k[h].float(); vf = v[h].float()
         attn = qf @ kf.T * scale
-        attn_max = attn.max(dim=-1, keepdim=True)[0]
-        attn_exp = torch.exp(attn - attn_max)
-        attn_sum = attn_exp.sum(dim=-1, keepdim=True)
-        ref[h] = ((attn_exp / attn_sum) @ vf).bfloat16()
+        ref[h] = (torch.softmax(attn, dim=-1) @ vf).bfloat16()
 
-    out = dsv4_attention(q, k, v)
-
-    cos = torch.nn.functional.cosine_similarity(
-        out.flatten().unsqueeze(0), ref.float().flatten().unsqueeze(0)
-    ).item()
-    status = "PASS" if cos >= 0.99 else "FAIL"
-    print(f"  hd={hd}, n_h={n_h}, N={N}: cos {cos:.6f}  {status}")
-
-
-def test_production_multi_kv():
-    """Test multi-KV-tile with Python KV merge."""
-    torch.manual_seed(42)
-    hd = 64
-    n_h = 1
-    T = 128
-    N = 256  # 2 KV segments
-
-    q = torch.randn(n_h, T, hd, dtype=torch.bfloat16, device='cuda')
-    k = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
-    v = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
-
-    # PyTorch reference
-    scale = 1.0 / math.sqrt(hd)
-    qf = q[0].float()
-    kf = k[0].float()
-    vf = v[0].float()
-    attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0]
-    attn_exp = torch.exp(qf @ kf.T * scale - attn_max)
-    attn_sum = attn_exp.sum(dim=-1, keepdim=True)
-    ref_norm = (attn_exp / attn_sum) @ vf
-    ref_unnorm = attn_exp @ vf
-
-    out = dsv4_attention(q, k, v)
-
-    cos_unnorm = torch.nn.functional.cosine_similarity(
-        out.flatten().unsqueeze(0), ref_unnorm.unsqueeze(0).flatten().unsqueeze(0)
-    ).item()
-    cos_norm = torch.nn.functional.cosine_similarity(
-        out.flatten().unsqueeze(0), ref_norm.unsqueeze(0).flatten().unsqueeze(0)
-    ).item()
-    print(f"  hd={hd}, n_h={n_h}, N={N}: cos_unnorm {cos_unnorm:.6f}  cos_norm {cos_norm:.6f}")
+    cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.float().flatten().unsqueeze(0)).item()
+    print(f"  hd={hd}, n_h={n_h}, N={N}: cos {cos:.6f}  {'PASS' if cos >= 0.99 else 'FAIL'}")
 
 
 def test():
-    print("=== Production DSV4 Attention Wrapper ===\n")
-    test_production_basic()
-    test_production_multi_kv()
+    print("=== Production DSV4 Attention ===\n")
+    test_single_head_128()
+    test_single_head_256()
+    test_single_head_512()
+    test_multi_head()
 
 
 if __name__ == '__main__':