test: minimal ctypes debug test for P3

tests/unit/test_p3_ctypes_minimal.py

@@ -0,0 +1,130 @@
+"""
+Minimal ctypes test: exact same setup as standalone test_fmha_6warp_multihead_hd64.cu
+Uses raw CUDA memory, not PyTorch tensors, to isolate kernel correctness.
+"""
+import torch
+import ctypes
+import math
+import os
+import sys
+import subprocess
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from dsv4.kernels.attention.fmha_multihead_op import _find_nvcc, _ensure_built, BUILD_DIR, SO_NAME
+
+def f32_to_bf16_bits(f):
+    """Convert float to BF16 bit pattern (uint16)."""
+    import struct
+    u = struct.unpack('I', struct.pack('f', f))[0]
+    return (u >> 16) & 0xFFFF
+
+def bf16_bits_to_f32(h):
+    """Convert BF16 bit pattern (uint16) to float."""
+    import struct
+    u = h << 16
+    return struct.unpack('f', struct.pack('I', u))[0]
+
+def test_minimal():
+    # Build the .so
+    lib = _ensure_built()
+
+    hd = 64
+    n_h = 4
+    N = 128  # SK
+    batch = 1
+    scale = 1.0 / math.sqrt(hd)
+
+    # Create data on GPU using PyTorch (easier than raw CUDA malloc for setup)
+    torch.manual_seed(42)
+
+    # Q: (batch, n_h, 1, hd) — each head has 1 row of hd elements
+    q_data = torch.randn(n_h, hd, dtype=torch.bfloat16, device='cuda')
+    q_4d = q_data.unsqueeze(0).unsqueeze(2).contiguous()  # (1, n_h, 1, hd)
+
+    # K: (batch, n_h, N, hd) — each head has N rows of hd elements
+    k_data = torch.randn(n_h, N, hd, dtype=torch.bfloat16, device='cuda')
+    k_4d = k_data.unsqueeze(0).contiguous()  # (1, n_h, N, hd)
+
+    # V: (batch, n_h, hd, N) — transposed
+    v_data = torch.randn(n_h, hd, N, dtype=torch.bfloat16, device='cuda')
+    v_4d = v_data.unsqueeze(0).contiguous()  # (1, n_h, hd, N)
+
+    # Output
+    o_4d = torch.zeros(1, n_h, 1, hd, dtype=torch.bfloat16, device='cuda')
+    lse_4d = torch.zeros(1, n_h, 1, dtype=torch.float32, device='cuda')
+
+    # Strides
+    q_hs = q_4d.stride(1)  # hd
+    q_bs = q_4d.stride(0)  # n_h * 1 * hd = n_h * hd
+    k_hs = k_4d.stride(1)  # N * hd
+    k_bs = k_4d.stride(0)  # n_h * N * hd
+    v_hs = v_4d.stride(1)  # hd * N
+    v_bs = v_4d.stride(0)  # n_h * hd * N
+    o_hs = o_4d.stride(1)
+    o_bs = o_4d.stride(0)
+    lse_hs = lse_4d.stride(1)
+    lse_bs = lse_4d.stride(0)
+
+    print(f"Q shape: {q_4d.shape}, strides: {q_4d.stride()}")
+    print(f"K shape: {k_4d.shape}, strides: {k_4d.stride()}")
+    print(f"V shape: {v_4d.shape}, strides: {v_4d.stride()}")
+    print(f"O shape: {o_4d.shape}, strides: {o_4d.stride()}")
+    print(f"LSE shape: {lse_4d.shape}, strides: {lse_4d.stride()}")
+    print(f"q_hs={q_hs}, q_bs={q_bs}, k_hs={k_hs}, k_bs={k_bs}")
+    print(f"v_hs={v_hs}, v_bs={v_bs}, o_hs={o_hs}, o_bs={o_bs}")
+    print(f"lse_hs={lse_hs}, lse_bs={lse_bs}")
+
+    ret = lib.fmha_multihead_decode_launch(
+        ctypes.c_void_p(q_4d.data_ptr()),
+        ctypes.c_void_p(k_4d.data_ptr()),
+        ctypes.c_void_p(v_4d.data_ptr()),
+        ctypes.c_void_p(o_4d.data_ptr()),
+        ctypes.c_void_p(lse_4d.data_ptr()),
+        ctypes.c_int(batch),
+        ctypes.c_int(n_h),
+        ctypes.c_int(n_h),  # n_kv = n_h for MHA
+        ctypes.c_int(N),
+        ctypes.c_int(hd),
+        ctypes.c_int(q_hs),
+        ctypes.c_int(q_bs),
+        ctypes.c_int(k_hs),
+        ctypes.c_int(k_bs),
+        ctypes.c_int(v_hs),
+        ctypes.c_int(v_bs),
+        ctypes.c_int(o_hs),
+        ctypes.c_int(o_bs),
+        ctypes.c_int(lse_hs),
+        ctypes.c_int(lse_bs),
+        ctypes.c_float(scale),
+    )
+    print(f"Kernel return: {ret}")
+
+    # Reference: pure PyTorch
+    o_ref = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device='cuda')
+    for h in range(n_h):
+        q_h = q_data[h:h+1]  # (1, hd)
+        k_h = k_data[h]      # (N, hd)
+        v_h = v_data[h].T    # (N, hd) — V is (hd, N), transpose to (N, 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[h] = o.bfloat16()
+
+    # Compare
+    o_kernel = o_4d.squeeze(0).squeeze(1)  # (n_h, hd)
+    o_ref_flat = o_ref.squeeze(1)  # (n_h, hd)
+
+    for h in range(n_h):
+        cos = torch.nn.functional.cosine_similarity(
+            o_kernel[h].float().unsqueeze(0),
+            o_ref_flat[h].float().unsqueeze(0)
+        ).item()
+        print(f"  Head {h}: cos={cos:.6f}")
+
+    torch.cuda.synchronize()
+    print("Done")
+
+
+if __name__ == "__main__":
+    test_minimal()