From a3c5f817e11f02f1bbb807e3876e3b839570ad47 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 08:23:43 +0000 Subject: [PATCH] debug: compare api vs direct kernel vs reference --- tests/unit/test_p3_fast_decode.py | 33 ++++++++++++++++++++++++++----- 1 file changed, 28 insertions(+), 5 deletions(-) diff --git a/tests/unit/test_p3_fast_decode.py b/tests/unit/test_p3_fast_decode.py index b9c72682..03d565fc 100644 --- a/tests/unit/test_p3_fast_decode.py +++ b/tests/unit/test_p3_fast_decode.py @@ -134,17 +134,40 @@ def test_full_api(): # 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) + 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) + + 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() - cos = cosine_sim(o_ref, o_fast).item() - status = "PASS" if cos >= 0.999990 else "FAIL" - if status == "FAIL": + 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: 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}") + print(f" FAIL {desc}: {e}") traceback.print_exc() all_pass = False