From 2cc66bff68e7788888362b0dc05beb7de66453ac Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 22:26:10 +0000 Subject: [PATCH] D2: add initial multi-head test file --- tests/unit/test_d2_multihead.py | 118 ++++++++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 tests/unit/test_d2_multihead.py diff --git a/tests/unit/test_d2_multihead.py b/tests/unit/test_d2_multihead.py new file mode 100644 index 00000000..fecea619 --- /dev/null +++ b/tests/unit/test_d2_multihead.py @@ -0,0 +1,118 @@ +""" +FMHA D2: Multi-Query Grid with Head Packing. + +Start with n_h=1 (regression), then n_h=2, n_h=8, etc. +Uses s_k=128 (1 KV tile, no O rescale needed). + +Strategy B: Head as grid dimension. +Grid: (ceil_div(T, 128), num_query_heads, batch) +Each CTA handles one query head. +""" +import torch, math +import cutlass.cute as cute +import cutlass.torch as ct +import cuda.bindings.driver as cuda +from dsv4.kernels.attention.fmha import FmhaKernel + + +def test_multihead(hd=64, n_h=1, batch=1, T=128, s_k=128): + torch.manual_seed(42) + + # Q: (batch, n_h, T, hd) + q = torch.randn(batch, n_h, T, hd, dtype=torch.bfloat16, device='cuda') + # K/V: (batch, 1, s_k, hd) — MQA: shared KV + k = torch.randn(batch, 1, s_k, hd, dtype=torch.bfloat16, device='cuda') + v = torch.randn(batch, 1, s_k, hd, dtype=torch.bfloat16, device='cuda') + # O: (batch, n_h, T, hd) + o = torch.zeros(batch, n_h, T, hd, dtype=torch.bfloat16, device='cuda') + # LSE: (batch, n_h, T) + lse = torch.zeros(batch, n_h, T, dtype=torch.float32, device='cuda') + + # FP32 reference + qf = q.float() # (batch, n_h, T, hd) + kf = k[:, 0].float() # (batch, s_k, hd) — shared KV + vf = v[:, 0].float() # (batch, s_k, hd) + scale = 1.0 / math.sqrt(hd) + + ref_o = torch.zeros(batch, n_h, T, hd, dtype=torch.float32, device='cuda') + for b in range(batch): + for h in range(n_h): + attn = qf[b, h] @ kf[b].T * scale # (T, s_k) + attn_max = attn.max(dim=-1, keepdim=True)[0] + attn_exp = torch.exp(attn - attn_max) + attn_sum = attn_exp.sum(dim=-1, keepdim=True) + ref_o[b, h] = (attn_exp / attn_sum) @ vf[b] + + # For now, test with n_h=1 to verify the kernel works. + # D2 multi-head requires kernel changes (grid, TMA, etc.) + # This test will FAIL until D2 is implemented in the kernel. + + # Current kernel expects Q: (T, hd, 1), K: (s_k, hd, 1), V: (s_k, hd) + # For n_h=1, this is just a reshape. + if n_h == 1 and batch == 1: + q_kernel = q[0, 0].unsqueeze(-1) # (T, hd, 1) + k_kernel = k[0, 0].unsqueeze(-1) # (s_k, hd, 1) + v_kernel = v[0, 0] # (s_k, hd) + + kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False) + pv_n_tile = kernel.pv_n_tile + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + # For hd=64, n_pv_tiles=1, but we handle general case + n_pv_tiles = hd // pv_n_tile + o_kernel = torch.zeros(T, hd, dtype=torch.float32, device='cuda') + + # Compile + v_tile = v_kernel[:, 0:pv_n_tile].contiguous() + v_k = v_tile.unsqueeze(-1) + c_tile = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') + lse_t = torch.zeros(T, dtype=torch.float32, device='cuda') + + mQ = ct.from_dlpack(q_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_kernel)) + mK = ct.from_dlpack(k_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_kernel)) + mV = ct.from_dlpack(v_k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_k)) + mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) + mLSE = ct.from_dlpack(lse_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_t)) + + print(f' Compiling (hd={hd}, n_h={n_h}, T={T}, s_k={s_k})...', flush=True) + compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) + + for nt in range(n_pv_tiles): + v_start = nt * pv_n_tile + v_end = v_start + pv_n_tile + v_tile = v_kernel[:, v_start:v_end].contiguous() + v_k = v_tile.unsqueeze(-1) + c_tile = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') + lse_t.zero_() + + mQ = ct.from_dlpack(q_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_kernel)) + mK = ct.from_dlpack(k_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_kernel)) + mV = ct.from_dlpack(v_k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_k)) + mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) + mLSE = ct.from_dlpack(lse_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_t)) + + compiled(mQ, mK, mV, mC, stream, mLSE) + torch.cuda.synchronize() + + o_kernel[:, v_start:v_end] = c_tile[:, :, 0].float() + + cos = torch.nn.functional.cosine_similarity( + o_kernel.flatten().unsqueeze(0), ref_o[0, 0].flatten().unsqueeze(0) + ).item() + print(f' hd={hd}, n_h={n_h}, T={T}, s_k={s_k}: cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') + else: + print(f' n_h={n_h}, batch={batch}: SKIPPED (D2 multi-head not yet implemented)') + + +def test(): + print("=== D2: Multi-Query Grid ===\n") + + # Regression: n_h=1 (same as existing tests) + test_multihead(64, 1, 1, 128, 128) + + # n_h=2 (first multi-head test, will need kernel changes) + test_multihead(64, 2, 1, 128, 128) + + +if __name__ == '__main__': + test()