diff --git a/tests/unit/test_p3_fast_decode.py b/tests/unit/test_p3_fast_decode.py index a6ae33fd..b9c72682 100644 --- a/tests/unit/test_p3_fast_decode.py +++ b/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)