diff --git a/tests/unit/test_p3_ctypes_minimal.py b/tests/unit/test_p3_ctypes_minimal.py new file mode 100644 index 00000000..3e3bd7cd --- /dev/null +++ b/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()