From 6cc151097edc3d082e3486078804faf5f9471e65 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 25 May 2026 01:08:38 +0000 Subject: [PATCH] Revert D2 multi-CTA attempts - keeping per-head launch approach (works correctly) --- dsv4/kernels/attention/fmha.py | 46 +++----------------- tests/unit/test_d2_multicta.py | 77 ---------------------------------- 2 files changed, 5 insertions(+), 118 deletions(-) delete mode 100644 tests/unit/test_d2_multicta.py diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index e0cb3555..f42693d4 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, num_query_heads=1, batch_size=1): + def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None, normalize=True): self.head_dim = head_dim self.s_k = s_k self.n_kv_tiles = s_k // 128 @@ -43,10 +43,6 @@ 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]) @@ -107,17 +103,7 @@ class FmhaKernel: @cute.jit def __call__(self, q, k, v, c, stream, lse=None): self.q_dtype = q.element_type; self.o_dtype = c.element_type; self.c_dtype = self.o_dtype - # For multi-CTA, q has shape (n_h, T, hd, 1). Layout depends on inner (T, hd) dims only. - # LayoutEnum.from_tensor needs 2D/3D tensors. For 4D Q, compute layout manually. - if const_expr(self.num_ctas > 1): - # q shape: (n_h, T, hd, 1). The inner (T, hd, 1) dims determine the MMA major mode. - # We know Q is MN-major (row-major in the attention M×K layout), so a_major = MN. - # This is always true for our FMHA kernel's Q layout. - self.a_major = cute.nvgpu.OperandMajorMode.MN - self.c_layout = LayoutEnum.ROW_MAJOR # O is always row-major - else: - self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() - self.c_layout = LayoutEnum.from_tensor(c) + self.a_major = LayoutEnum.from_tensor(q).mma_major_mode() self.b_major = LayoutEnum.from_tensor(k).mma_major_mode() v_fmha = cute.make_tensor( v.iterator, @@ -143,7 +129,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,self.num_ctas,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,1,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): @@ -186,32 +172,10 @@ 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) - # D2: Multi-CTA grid. Use block_idx_y to select Q and O for this CTA's head. - _bidx, head_cta_idx, _bidz = cute.arch.block_idx() # grid=(1, num_ctas, 1) - - # Print mQ shape at trace time to understand mode structure - print(f"D2 DEBUG: mQ shape={cute.shape(mQ)}, mK shape={cute.shape(mK)}") - - # Q: if num_ctas > 1, mQ has a head dimension. Use flat_divide for runtime coordinate. - # K/V: shared (MQA), always coordinate 0. - if const_expr(self.num_ctas > 1): - # flat_divide creates (tile_M, tile_K, rest...) layout. - # For mQ=(n_h, T, hd, 1) with tiler=(128, hd): - # tC_gQ shape = (128, hd, n_h, T/128, 1, 1) - # head_cta_idx indexes mode 2 (n_h) - q_tiler = cute.slice_(self.qk_mma_tiler, (None, 0, None)) # (128, hd) - tC_gQ = cute.flat_divide(mQ, q_tiler) - print(f"D2 DEBUG: tC_gQ shape={cute.shape(tC_gQ)}") - gQ = tC_gQ[None, None, head_cta_idx, None, None, None] - c_tiler = cute.slice_(self.pv_mma_tiler, (None, None, 0)) # (128, hd) - tC_gC = cute.flat_divide(mC, c_tiler) - print(f"D2 DEBUG: tC_gC shape={cute.shape(tC_gC)}") - gC = tC_gC[None, None, head_cta_idx, None, None] - else: - gQ = cute.local_tile(mQ,cute.slice_(self.qk_mma_tiler,(None,0,None)),(None,None,None)) - gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) + 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)) 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 deleted file mode 100644 index bc049ec0..00000000 --- a/tests/unit/test_d2_multicta.py +++ /dev/null @@ -1,77 +0,0 @@ -""" -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()