diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index f42693d4..bd07340f 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,10 @@ 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) + # D2: Multi-CTA grid. Each CTA handles one (head, batch) pair. + self.num_query_heads = num_query_heads + self.batch_size = 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 +133,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): @@ -172,10 +176,22 @@ class FmhaKernel: _p_swizzle = cute.make_layout(((1,1),1,(1,1),1)) sP = smem.allocate_tensor(element_type=self.q_dtype,layout=_p_layout,byte_alignment=128,swizzle=_p_swizzle) - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) + # D2: Multi-CTA grid. Use block_idx_y to select Q and O for this CTA's head. + head_cta_idx = cute.arch.block_idx(dim=1) # block_idx_y + + # Q: if num_ctas > 1, mQ has a head dimension. local_tile indexes into it. + # K/V: shared (MQA), always coordinate 0. + # For single-CTA (num_ctas=1), head_cta_idx=0 and coordinates are the same as before. + if const_expr(self.num_ctas > 1): + gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(head_cta_idx,None,None)) + else: + gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) + if const_expr(self.num_ctas > 1): + gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(head_cta_idx,None,None)) + else: + gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) n_kv_tiles = cute.size(gK, mode=[3]) qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) diff --git a/tests/unit/test_d2_multicta.py b/tests/unit/test_d2_multicta.py new file mode 100644 index 00000000..bc049ec0 --- /dev/null +++ b/tests/unit/test_d2_multicta.py @@ -0,0 +1,77 @@ +""" +D2: Multi-CTA grid test. + +Tests the multi-CTA FMHA kernel with Q shape (n_h, T, hd, 1). +Each CTA (indexed by block_idx_y) handles one query head. +K/V are shared (MQA) — all CTAs load the same K/V. +""" +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_multi_cta(hd=64, n_h=2, s_k=128): + T = 128 + torch.manual_seed(42) + + # Q: (n_h, T, hd, 1) — head dimension outermost + q = torch.randn(n_h, T, hd, 1, dtype=torch.bfloat16, device='cuda') + # K/V: (s_k, hd, 1) — shared KV (no head dim) + k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda') + v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda') + # O: (n_h, T, hd, 1) + o = torch.zeros(n_h, T, hd, 1, dtype=torch.bfloat16, device='cuda') + + # FP32 reference (un-normalized) + qf = q[:, :, :, 0].float() # (n_h, T, hd) + kf = k[:, :, 0].float() # (s_k, hd) + vf = v.float() # (s_k, hd) + scale = 1.0 / math.sqrt(hd) + + ref_unnorm = torch.zeros(n_h, T, hd, dtype=torch.float32, device='cuda') + for h in range(n_h): + attn = qf[h] @ kf.T * scale + attn_max = attn.max(dim=-1, keepdim=True)[0] + attn_exp = torch.exp(attn - attn_max) + ref_unnorm[h] = attn_exp @ vf + + lse_tensor = torch.zeros(T, 1, 1, dtype=torch.float32, device='cuda') + + kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False, + num_query_heads=n_h, batch_size=1) + pv_n_tile = kernel.pv_n_tile + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + # Compile with Q having head dimension + v_tile = v[:, 0:pv_n_tile].contiguous() + v_kernel = v_tile.unsqueeze(-1) + c_tile = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') + + mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) + mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) + mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel)) + mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile)) + mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor)) + + print(f' Compiling (hd={hd}, n_h={n_h})...', flush=True) + compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) + compiled(mQ, mK, mV, mC, stream, mLSE) + torch.cuda.synchronize() + + # Check output + out = o[:, :, :, 0].float() # (n_h, T, hd) + cos = torch.nn.functional.cosine_similarity( + out.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) + ).item() + print(f' hd={hd}, n_h={n_h}: cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}') + + +def test(): + print("=== D2: Multi-CTA Grid ===\n") + test_multi_cta(64, 2) + + +if __name__ == '__main__': + test()