biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

debug: compare api vs direct kernel vs reference

Browse Source
This commit is contained in:
biondizzle
2026-05-30 08:23:43 +00:00
parent 78e6d58b85
commit a3c5f817e1
1 changed files with 28 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
tests/unit/test_p3_fast_decode.py
View File

@@ -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