diff --git a/tests/unit/test_p3_fast_decode.py b/tests/unit/test_p3_fast_decode.py index 03d565fc..cfd66131 100644 --- a/tests/unit/test_p3_fast_decode.py +++ b/tests/unit/test_p3_fast_decode.py @@ -1,9 +1,10 @@ """ -P3 Integration Test: Verify 6-warp multi-head decode fast path -produces identical results to a PyTorch reference. +P3 Integration Test: 6-warp multi-head decode fast path. -Tests MHA, MQA, GQA at HD = 64, 128, 256. -Cosine similarity >= 0.999998 between kernel output and reference. +Verifies the kernel produces identical results to a PyTorch reference +for MHA, MQA, and GQA at HD = 64, 128, 256. + +Gate: worst-case cosine >= 0.999990 per configuration. """ import torch import math @@ -21,12 +22,34 @@ def cosine_sim(a, b): return (a @ b) / (a.norm() * b.norm() + 1e-30) -def test_fast_path(): +def reference_attention(q_4d, k_4d, v_4d, scale): + """PyTorch reference matching kernel tensor layout. + + Q: (1, n_h, 1, hd), K: (1, n_h, N, hd), V: (1, n_h, hd, N) + V is in kernel layout (hd, N) — transpose to (N, hd) for reference. + """ + n_h = q_4d.shape[1] + q = q_4d[0] # (n_h, 1, hd) + k = k_4d[0] # (n_h, N, hd) + v = v_4d[0].transpose(-1, -2) # (n_h, N, hd) + + output = torch.zeros(n_h, 1, q_4d.shape[3], dtype=torch.bfloat16, device='cuda') + for h in range(n_h): + q_h = q[h] # (1, hd) + k_h = k[h] # (N, hd) + v_h = v[h] # (N, 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[h] = o.bfloat16() + return output + + +def test_kernel_correctness(): """Test kernel vs PyTorch reference for MHA, MQA, GQA at various HD.""" torch.manual_seed(42) configs = [ - # (n_q, n_kv, N, hd, desc) (4, 4, 64, 64, "MHA hd=64"), (4, 4, 128, 64, "MHA hd=64 N=128"), (4, 4, 64, 128, "MHA hd=128"), @@ -44,49 +67,20 @@ def test_fast_path(): all_pass = True for n_q, n_kv, N, hd, desc in configs: scale = 1.0 / math.sqrt(hd) - q_per_kv = n_q // n_kv - try: - # ---- Create data in KERNEL layout ---- - # Q: (1, n_q, 1, hd) — each head has 1 row of hd elements q_4d = torch.randn(1, n_q, 1, hd, dtype=torch.bfloat16, device='cuda').contiguous() - - # K: (1, n_kv, N, hd) k_4d = torch.randn(1, n_kv, N, hd, dtype=torch.bfloat16, device='cuda').contiguous() - - # V: (1, n_kv, hd, N) — the KERNEL expects V transposed (hd, N) per head v_4d = torch.randn(1, n_kv, hd, N, dtype=torch.bfloat16, device='cuda').contiguous() - # ---- Kernel output ---- sb = torch.zeros(1, n_q, dtype=torch.float32, device='cuda') - o_4d, lse_4d = fmha_multihead_decode_raw( - q_4d, k_4d, v_4d, scale, 0, 0, False, sb - ) - o_kernel = o_4d # (1, n_q, 1, hd) + o_4d, _ = fmha_multihead_decode_raw(q_4d, k_4d, v_4d, scale, 0, 0, False, sb) - # ---- PyTorch reference using the SAME data ---- - # Q: (n_q, 1, hd), K: (n_kv, N, hd), V: (n_kv, N, hd) — V is TRANSPOSED from kernel layout - q_ref = q_4d[0] # (n_q, 1, hd) - k_ref = k_4d[0] # (n_kv, N, hd) - v_ref = v_4d[0].transpose(-1, -2) # (n_kv, N, hd) — transpose (hd,N) -> (N,hd) + o_ref = reference_attention(q_4d, k_4d, v_4d, scale) - o_ref = torch.zeros(n_q, 1, hd, dtype=torch.bfloat16, device='cuda') - for kv_idx in range(n_kv): - k_h = k_ref[kv_idx] # (N, hd) - v_h = v_ref[kv_idx] # (N, hd) - for qi in range(q_per_kv): - q_idx = kv_idx * q_per_kv + qi - q_h = q_ref[q_idx] # (1, hd) - s = torch.matmul(q_h.float(), k_h.float().T) * scale # (1, N) - s = torch.softmax(s, dim=-1) - o = torch.matmul(s, v_h.float()) # (1, hd) - o_ref[q_idx] = o.bfloat16() - - # ---- Compare per-head ---- worst_cos = 1.0 for h in range(n_q): cos = torch.nn.functional.cosine_similarity( - o_kernel[0, h].float().flatten().unsqueeze(0), + o_4d[0, h].float().flatten().unsqueeze(0), o_ref[h].float().flatten().unsqueeze(0), ).item() worst_cos = min(worst_cos, cos) @@ -106,7 +100,7 @@ def test_fast_path(): def test_full_api(): - """Test the full dsv4_attention API (goes through fast path for T=1, N<=128).""" + """Test the full dsv4_attention API (fast path for T=1, N<=128).""" from dsv4.kernels.attention.production import dsv4_attention torch.manual_seed(99) @@ -131,74 +125,44 @@ def test_full_api(): 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) - # Also test the raw kernel directly with the SAME data - # but constructing the tensors exactly as the working test does - q_4d_direct = q.unsqueeze(0).contiguous() # (1, n_q, 1, hd) + # Reference using same data if n_kv == 1: - k_4d_direct = k.unsqueeze(0).unsqueeze(0).contiguous() # (1,1,N,hd) - # V: (N,hd) -> transpose to (hd,N) -> (1,1,hd,N) - # BUT we need the SAME V data as production uses - # production v is (N, hd), transpose to (hd, N) - v_4d_direct = v.unsqueeze(0).unsqueeze(0).transpose(-1, -2).contiguous() - else: - k_4d_direct = k.unsqueeze(0).contiguous() # (1,n_kv,N,hd) - v_4d_direct = v.unsqueeze(0).transpose(-1, -2).contiguous() # (1,n_kv,hd,N) + k = k.unsqueeze(0) + v = v.unsqueeze(0) + q_per_kv = n_q // n_kv + o_ref = torch.zeros(n_q, 1, hd, dtype=torch.bfloat16, device='cuda') + for kv_idx in range(n_kv): + k_h = k[kv_idx] + v_h = v[kv_idx] + for qi in range(q_per_kv): + q_idx = kv_idx * q_per_kv + qi + q_h = q[q_idx] + s = torch.matmul(q_h.float(), k_h.float().T) * scale + s = torch.softmax(s, dim=-1) + o = torch.matmul(s, v_h.float()) + o_ref[q_idx] = o.bfloat16() - from dsv4.kernels.attention.fmha_multihead_op import fmha_multihead_decode_raw - sb = torch.zeros(1, n_q, dtype=torch.float32, device='cuda') - o_4d_direct, _ = fmha_multihead_decode_raw( - q_4d_direct, k_4d_direct, v_4d_direct, scale, 0, 0, False, sb - ) - o_direct = o_4d_direct.squeeze(0) # (n_q, 1, hd) - - # Compare direct kernel call vs full API - cos_api_vs_direct = cosine_sim(o_fast, o_direct).item() - - # Reference - o_ref = reference_attention_api(q, k, v, scale, n_q, n_kv, N, hd) - cos_direct_vs_ref = cosine_sim(o_ref, o_direct).item() - - print(f" {desc}: api_vs_direct={cos_api_vs_direct:.6f} direct_vs_ref={cos_direct_vs_ref:.6f}") - if cos_direct_vs_ref < 0.999990: + 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} [API] {desc}: cos={cos:.6f}") except Exception as e: import traceback - print(f" FAIL {desc}: {e}") + print(f" FAIL [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) - ok1 = test_fast_path() + ok1 = test_kernel_correctness() print() ok2 = test_full_api() print("=" * 60) diff --git a/tests/unit/test_p3_sanity.py b/tests/unit/test_p3_sanity.py deleted file mode 100644 index 46b2843e..00000000 --- a/tests/unit/test_p3_sanity.py +++ /dev/null @@ -1,67 +0,0 @@ -""" -Absolute simplest test: single head, small N, verify kernel == reference. -""" -import torch -import math -import sys -import os - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - -from dsv4.kernels.attention.fmha_multihead_op import fmha_multihead_decode_raw - - -def test_single_head_sanity(): - """Single head, N=128, hd=64. Known values, no randomness.""" - hd = 64 - N = 128 - scale = 1.0 / math.sqrt(hd) - - # Q: (1, 1, 1, hd) — single head, single query token - q = torch.ones(1, 1, 1, hd, dtype=torch.bfloat16, device='cuda') - # K: (1, 1, N, hd) — single KV head, N positions - k = torch.ones(1, 1, N, hd, dtype=torch.bfloat16, device='cuda') - # V: (1, 1, hd, N) — in kernel layout - # Let's make V[d, r] = d + r*0.01 (simple pattern) - v_data = torch.arange(hd, dtype=torch.float32, device='cuda').unsqueeze(1) + \ - torch.arange(N, dtype=torch.float32, device='cuda').unsqueeze(0) * 0.01 - v_4d = v_data.bfloat16().unsqueeze(0).unsqueeze(0) # (1, 1, hd, N) - - sb = torch.zeros(1, 1, dtype=torch.float32, device='cuda') - o_4d, lse = fmha_multihead_decode_raw(q, k, v_4d, scale, 0, 0, False, sb) - - # Reference: Q is all-ones, K is all-ones, so QK^T gives all-equal scores - # softmax of uniform = 1/N. So O = (1/N) * sum(V[r, d] for r in 0..N-1) - v_ref = v_data.T # (N, hd) — reference uses (N, hd) layout - # Each V[r, d] = d + r*0.01 - # sum over r: sum(d + r*0.01) = N*d + 0.01*sum(r) = N*d + 0.01*N*(N-1)/2 - # O[d] = (1/N) * (N*d + 0.01*N*(N-1)/2) = d + 0.01*(N-1)/2 - o_expected = torch.arange(hd, dtype=torch.float32, device='cuda') + 0.01 * (N - 1) / 2 - - cos = torch.nn.functional.cosine_similarity( - o_4d[0, 0].float().flatten().unsqueeze(0), - o_expected.flatten().unsqueeze(0), - ).item() - - # Also compute via direct matmul for sanity - q_f = q.float().squeeze() # (hd,) all ones - k_f = k.float().squeeze() # (N, hd) all ones - v_f = v_ref # (N, hd) - scores = torch.matmul(q_f.unsqueeze(0), k_f.T) * scale # (1, N) - probs = torch.softmax(scores, dim=-1) # (1, N) - o_matmul = torch.matmul(probs, v_f) # (1, hd) - - cos_matmul = torch.nn.functional.cosine_similarity( - o_4d[0, 0].float().flatten().unsqueeze(0), - o_matmul.flatten().unsqueeze(0), - ).item() - - print(f"Kernel vs expected: cos={cos:.6f}") - print(f"Kernel vs matmul: cos={cos_matmul:.6f}") - print(f"Kernel output[0:5]: {o_4d[0, 0, 0, 0:5].float()}") - print(f"Expected[0:5]: {o_expected[0:5]}") - print(f"Matmul[0:5]: {o_matmul[0, 0:5]}") - - -if __name__ == "__main__": - test_single_head_sanity() diff --git a/tests/unit/test_p3_v_debug.py b/tests/unit/test_p3_v_debug.py deleted file mode 100644 index ee5ce219..00000000 --- a/tests/unit/test_p3_v_debug.py +++ /dev/null @@ -1,103 +0,0 @@ -""" -Debug test: call fmha_multihead_decode_raw directly with production-style V. -Isolates whether the issue is in the V transpose or the production.py plumbing. -""" -import torch -import math -import sys -import os - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - -from dsv4.kernels.attention.fmha_multihead_op import fmha_multihead_decode_raw - - -def cosine_sim(a, b): - a = a.flatten().float() - b = b.flatten().float() - return (a @ b) / (a.norm() * b.norm() + 1e-30) - - -def test_production_v_layout(): - """Test with V created as (N, hd) then transposed (production path).""" - torch.manual_seed(42) - - hd = 64 - n_h = 4 - N = 128 - scale = 1.0 / math.sqrt(hd) - - # Create Q, K in the same way as both the working test and production - q_4d = torch.randn(1, n_h, 1, hd, dtype=torch.bfloat16, device='cuda').contiguous() - k_4d = torch.randn(1, n_h, N, hd, dtype=torch.bfloat16, device='cuda').contiguous() - - # V: production path creates (n_kv, N, hd) then transposes to (1, n_kv, hd, N) - v_orig = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda') - v_4d = v_orig.unsqueeze(0).transpose(-1, -2).contiguous() - print(f"V orig shape: {v_orig.shape}") - print(f"V 4d shape: {v_4d.shape}, strides: {v_4d.stride()}") - - sb = torch.zeros(1, n_h, dtype=torch.float32, device='cuda') - o_4d, lse_4d = fmha_multihead_decode_raw(q_4d, k_4d, v_4d, scale, 0, 0, False, sb) - - # Reference: use v_orig (N, hd) per head - q_ref = q_4d[0] # (n_h, 1, hd) - k_ref = k_4d[0] # (n_h, N, hd) - - for h in range(n_h): - q_h = q_ref[h] # (1, hd) - k_h = k_ref[h] # (N, hd) - v_h = v_orig[h] # (N, hd) - s = torch.matmul(q_h.float(), k_h.float().T) * scale - s = torch.softmax(s, dim=-1) - o_ref = torch.matmul(s, v_h.float()) - - cos = torch.nn.functional.cosine_similarity( - o_4d[0, h].float().flatten().unsqueeze(0), - o_ref.flatten().unsqueeze(0), - ).item() - print(f" Head {h}: cos={cos:.6f}") - - -def test_native_v_layout(): - """Test with V created as (hd, N) natively (working test style).""" - torch.manual_seed(42) - - hd = 64 - n_h = 4 - N = 128 - scale = 1.0 / math.sqrt(hd) - - q_4d = torch.randn(1, n_h, 1, hd, dtype=torch.bfloat16, device='cuda').contiguous() - k_4d = torch.randn(1, n_h, N, hd, dtype=torch.bfloat16, device='cuda').contiguous() - v_4d = torch.randn(1, n_h, hd, N, dtype=torch.bfloat16, device='cuda').contiguous() - - sb = torch.zeros(1, n_h, dtype=torch.float32, device='cuda') - o_4d, lse_4d = fmha_multihead_decode_raw(q_4d, k_4d, v_4d, scale, 0, 0, False, sb) - - # Reference: V is (hd, N) per head, transpose to (N, hd) for reference - v_ref = v_4d[0].transpose(-1, -2) # (n_h, N, hd) - q_ref = q_4d[0] - k_ref = k_4d[0] - - for h in range(n_h): - q_h = q_ref[h] - k_h = k_ref[h] - v_h = v_ref[h] # (N, hd) - s = torch.matmul(q_h.float(), k_h.float().T) * scale - s = torch.softmax(s, dim=-1) - o_ref = torch.matmul(s, v_h.float()) - - cos = torch.nn.functional.cosine_similarity( - o_4d[0, h].float().flatten().unsqueeze(0), - o_ref.flatten().unsqueeze(0), - ).item() - print(f" Head {h}: cos={cos:.6f}") - - -if __name__ == "__main__": - print("=== Test 1: V created as (N,hd) then transposed (production path) ===") - test_production_v_layout() - print() - print("=== Test 2: V created natively as (hd,N) (working test style) ===") - test_native_v_layout() diff --git a/tests/unit/test_p3_v_ref_debug.py b/tests/unit/test_p3_v_ref_debug.py deleted file mode 100644 index 072701f3..00000000 --- a/tests/unit/test_p3_v_ref_debug.py +++ /dev/null @@ -1,91 +0,0 @@ -""" -Debug: why does the full API test give cos=0.83? -Test 1: V in kernel layout (hd, N), reference transposes -> (N, hd) -Test 2: V in standard layout (N, hd), reference uses directly -Both should give same result if math is correct. -""" -import torch -import math -import sys -import os - -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) - -from dsv4.kernels.attention.fmha_multihead_op import fmha_multihead_decode_raw - - -def test_v_layout_comparison(): - """Direct comparison: same Q and K, V in two different layouts.""" - torch.manual_seed(42) - - hd = 64 - n_h = 4 - N = 128 - scale = 1.0 / math.sqrt(hd) - - # Create Q and K once - q_4d = torch.randn(1, n_h, 1, hd, dtype=torch.bfloat16, device='cuda').contiguous() - k_4d = torch.randn(1, n_h, N, hd, dtype=torch.bfloat16, device='cuda').contiguous() - - # Create V as (n_h, hd, N) natively - v_native = torch.randn(1, n_h, hd, N, dtype=torch.bfloat16, device='cuda').contiguous() - - # Also create V as (n_h, N, hd) then transpose - v_orig = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda') - v_transposed = v_orig.unsqueeze(0).transpose(-1, -2).contiguous() # (1, n_h, hd, N) - - # Run kernel with native V - sb = torch.zeros(1, n_h, dtype=torch.float32, device='cuda') - o_native, _ = fmha_multihead_decode_raw(q_4d, k_4d, v_native, scale, 0, 0, False, sb) - - # Run kernel with transposed-from-standard V - o_transposed, _ = fmha_multihead_decode_raw(q_4d, k_4d, v_transposed, scale, 0, 0, False, sb) - - # Reference with native V (hd, N) -> transpose to (N, hd) - q_ref = q_4d[0] # (n_h, 1, hd) - k_ref = k_4d[0] # (n_h, N, hd) - v_ref_native = v_native[0].transpose(-1, -2) # (n_h, N, hd) — transposed from (hd, N) - v_ref_orig = v_orig # (n_h, N, hd) — already in (N, hd) layout - - # Reference 1: using native V data - o_ref1 = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device='cuda') - for h in range(n_h): - q_h = q_ref[h] # (1, hd) - k_h = k_ref[h] # (N, hd) - v_h = v_ref_native[h] # (N, 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()) - o_ref1[h] = o.bfloat16() - - # Reference 2: using original V data - o_ref2 = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device='cuda') - for h in range(n_h): - q_h = q_ref[h] - k_h = k_ref[h] - v_h = v_ref_orig[h] # (N, hd) — same data, different source - s = torch.matmul(q_h.float(), k_h.float().T) * scale - s = torch.softmax(s, dim=-1) - o = torch.matmul(s, v_h.float()) - o_ref2[h] = o.bfloat16() - - # Compare kernel vs ref1 (native V) - for h in range(n_h): - cos1 = torch.nn.functional.cosine_similarity( - o_native[0, h].float().flatten().unsqueeze(0), - o_ref1[h].float().flatten().unsqueeze(0), - ).item() - cos2 = torch.nn.functional.cosine_similarity( - o_transposed[0, h].float().flatten().unsqueeze(0), - o_ref2[h].float().flatten().unsqueeze(0), - ).item() - # Also compare the two kernel outputs (should differ since different V data) - cos_kk = torch.nn.functional.cosine_similarity( - o_native[0, h].float().flatten().unsqueeze(0), - o_transposed[0, h].float().flatten().unsqueeze(0), - ).item() - print(f" Head {h}: native_vs_ref1={cos1:.6f} transposed_vs_ref2={cos2:.6f} native_vs_transposed={cos_kk:.6f}") - - -if __name__ == "__main__": - test_v_layout_comparison()