diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index f42693d4..a2ff1b4e 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -16,7 +16,7 @@ import math class FmhaKernel: - def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None, normalize=True): + def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None, normalize=True, num_query_heads=1, batch_size=1): self.head_dim = head_dim self.s_k = s_k self.n_kv_tiles = s_k // 128 @@ -43,6 +43,12 @@ class FmhaKernel: self.num_c_stage = 1 if head_dim > 256 else 2 # Reduce SMEM at hd=512 self.scale_softmax = scale_softmax if scale_softmax is not None else 1.0 / math.sqrt(self.head_dim) self.scale_softmax_log2 = self.scale_softmax * math.log2(math.e) + self.num_query_heads = num_query_heads + self.batch_size = batch_size + # D2: Multi-CTA grid. Each CTA handles one (head, batch) pair. + # Grid: (1, num_query_heads * batch_size, 1) + # Total CTAs = num_query_heads * batch_size + self.num_ctas = num_query_heads * batch_size def _setup(self, qk_mma, pv_mma): qk_ik = cute.size(qk_mma.shape_mnk, mode=[2]) @@ -129,7 +135,7 @@ class FmhaKernel: # CuTeDSL doesn't support None parameters in @cute.kernel. if const_expr(lse is None): lse = cute.make_tensor(c.iterator, cute.make_layout((1,), stride=(0,))) - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.p_smem_s,self.c_smem_s,self.epi_tile,lse).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) + self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.p_smem_s,self.c_smem_s,self.epi_tile,lse).launch(grid=(1,self.num_ctas,1),block=[self.threads_per_cta,1,1],stream=stream) @cute.kernel def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, p_smem_s, c_smem_s, epi_tile, mLSE): diff --git a/tests/unit/test_d2_perhead.py b/tests/unit/test_d2_perhead.py new file mode 100644 index 00000000..38abd1bb --- /dev/null +++ b/tests/unit/test_d2_perhead.py @@ -0,0 +1,138 @@ +""" +FMHA D2: Multi-Head via per-head kernel launch (simple approach). + +For DSV4 MQA, each query head shares the same K/V. +We launch the kernel once per (head, batch) pair. +""" +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 = torch.randn(batch, n_h, T, hd, dtype=torch.bfloat16, device='cuda') + k = torch.randn(batch, s_k, hd, dtype=torch.bfloat16, device='cuda') + v = torch.randn(batch, s_k, hd, dtype=torch.bfloat16, device='cuda') + o = torch.zeros(batch, n_h, T, hd, dtype=torch.bfloat16, device='cuda') + + # FP32 reference + qf = q.float() + kf = k.float() + vf = v.float() + scale = 1.0 / math.sqrt(hd) + + ref_o = torch.zeros_like(qf) + for b in range(batch): + for h in range(n_h): + attn = qf[b, h] @ kf[b].T * scale + 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] + + # Run kernel per (head, batch) + 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) + + # Compile once with first head's data + q0 = q[0, 0].unsqueeze(-1) # (T, hd, 1) + k0 = k[0].unsqueeze(-1) # (s_k, hd, 1) + v0_tile = v[0, :, 0:pv_n_tile].contiguous() + v0_k = v0_tile.unsqueeze(-1) + c0 = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') + lse0 = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda') + + mQ = ct.from_dlpack(q0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q0)) + mK = ct.from_dlpack(k0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k0)) + mV = ct.from_dlpack(v0_k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v0_k)) + mC = ct.from_dlpack(c0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c0)) + mLSE = ct.from_dlpack(lse0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse0)) + + print(f' Compiling (hd={hd}, n_h={n_h}, batch={batch}, T={T}, s_k={s_k})...', flush=True) + compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) + + for b in range(batch): + for h in range(n_h): + q_h = q[b, h].unsqueeze(-1) # (T, hd, 1) + k_b = k[b].unsqueeze(-1) # (s_k, hd, 1) + v_b = v[b] # (s_k, hd) + c_h = torch.zeros(T, hd, dtype=torch.bfloat16, device='cuda') + + # Run per PV tile + for nt in range(hd // pv_n_tile): + v_start = nt * pv_n_tile + v_end = v_start + pv_n_tile + v_tile = v_b[:, 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') + lse0.zero_() + + mQ = ct.from_dlpack(q_h).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_h)) + mK = ct.from_dlpack(k_b).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_b)) + 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(lse0).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse0)) + + compiled(mQ, mK, mV, mC, stream, mLSE) + c_h[:, v_start:v_end] = c_tile[:, :, 0] + + o[b, h] = c_h + + torch.cuda.synchronize() + + # Compare (normalized) + o_norm = o.float() + for b in range(batch): + for h in range(n_h): + qf_h = qf[b, h] + kf_b = kf[b] + attn = qf_h @ kf_b.T * scale + 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) + o_norm[b, h] = (attn_exp / attn_sum) @ vf[b] + + cos = torch.nn.functional.cosine_similarity( + o.flatten().unsqueeze(0), ref_o.flatten().unsqueeze(0) + ).item() + + # Wait, o is the kernel output (un-normalized), ref_o is normalized. Need to compare properly. + # Actually, the kernel with normalize=False outputs un-normalized O. + # For a fair comparison, let me compute un-normalized reference. + ref_unnorm = torch.zeros_like(ref_o) + for b in range(batch): + for h in range(n_h): + qf_h = qf[b, h] + kf_b = kf[b] + attn = qf_h @ kf_b.T * scale + attn_max = attn.max(dim=-1, keepdim=True)[0] + attn_exp = torch.exp(attn - attn_max) + ref_unnorm[b, h] = attn_exp @ vf[b] + + cos = torch.nn.functional.cosine_similarity( + o.flatten().float().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) + ).item() + + print(f' hd={hd}, n_h={n_h}, batch={batch}, T={T}, s_k={s_k}: cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') + + +def test(): + print("=== D2: Multi-Head (per-head launch) ===\n") + + # n_h=1 regression + test_multihead(64, 1, 1, 128, 128) + + # n_h=2 + test_multihead(64, 2, 1, 128, 128) + + # n_h=8, batch=2 + test_multihead(64, 8, 2, 128, 128) + + +if __name__ == '__main__': + test()