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

D1: Fix test - remove gP param

Browse Source
This commit is contained in:
biondizzle
2026-05-24 03:41:44 +00:00
parent 3b6f19043f
commit 5d0f34f0e7
1 changed files with 7 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
tests/unit/test_d1_smem128.py
View File

@@ -1,4 +1,4 @@
"""D1: Test SMEM-P at hd=128 via gP→TMA path."""
"""D1: Test SMEM-P at hd=128 (current kernel, no gP)."""
import torch, math
import cutlass.cute as cute
import cutlass.torch as ct
@@ -6,7 +6,7 @@ import cuda.bindings.driver as cuda
from dsv4.kernels.attention.fmha import FmhaKernel
def test_smem_p_gP(hd, n_kv=128):
def test_smem_p(hd, n_kv=128):
m = 128
torch.manual_seed(42)
q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
@@ -20,7 +20,6 @@ def test_smem_p_gP(hd, n_kv=128):
attn_exp = torch.exp(qf @ kf.T * scale - attn_max)
attn_sum = attn_exp.sum(dim=-1, keepdim=True)
ref_unnorm = attn_exp @ v.float()
ref_lse = (torch.log(attn_sum.squeeze(-1)) + attn_max.squeeze(-1))[0].item()
lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
kernel = FmhaKernel(head_dim=hd, s_k=n_kv, use_smem_p=True)
@@ -29,26 +28,20 @@ def test_smem_p_gP(hd, n_kv=128):
v_tile = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1)
c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
# gP: global buffer for P matrix (128 x s_k), partitioned by QK C-fragment
gP = torch.zeros(m, n_kv, 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_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile))
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))
mGP = ct.from_dlpack(gP).mark_layout_dynamic(leading_dim=ct.get_leading_dim(gP))
print(f'hd={hd} SMEM-P (gP→TMA): Compiling...', flush=True)
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE, mGP)
compiled(mQ, mK, mV, mC, stream, mLSE, mGP)
print(f'hd={hd} SMEM-P: Compiling...', flush=True)
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE)
compiled(mQ, mK, mV, mC, stream, mLSE)
torch.cuda.synchronize()
out = c_tile[:, :, 0].float()
cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0)).item()
kernel_lse = lse_tensor[0, 0, 0].item()
lse_err = abs(kernel_lse - ref_lse)
print(f'hd={hd} SMEM-P (gP→TMA): cos={cos:.6f} lse_err={lse_err:.6f} {"PASS" if cos >= 0.99 else "FAIL"}')
print(f'hd={hd} SMEM-P: cos={cos:.6f} {"PASS" if cos >= 0.99 else "FAIL"}')
if cos < 0.99:
print(f' out[0,:4]={out[0,:4].tolist()}')
print(f' ref[0,:4]={ref_unnorm[0,:4].tolist()}')
@@ -56,5 +49,4 @@ def test_smem_p_gP(hd, n_kv=128):
if __name__ == '__main__':
print("=== SMEM-P via gP→TMA ===\n")
test_smem_p_gP(128)
test_smem_p(128)