biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

D2: add multi-CTA grid with block_idx_y for Q/O head indexing

Browse Source
This commit is contained in:
biondizzle
2026-05-24 23:27:38 +00:00
parent 335e310c79
commit 4c79e5533e
2 changed files with 97 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
dsv4/kernels/attention/fmha.py
View File

@@ -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)

77
tests/unit/test_d2_multicta.py Normal file
View File

@@ -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()