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

D2: add per-head launch test

Browse Source
This commit is contained in:
biondizzle
2026-05-24 22:48:22 +00:00
parent 9b476d87f9
commit d563c93fc5
2 changed files with 146 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
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,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):

138
tests/unit/test_d2_perhead.py Normal file
View File

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