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

Test reference FMHA with proper API

Browse Source
This commit is contained in:
biondizzle
2026-05-22 21:12:49 +00:00
parent a800a83d5c
commit 4fc264e034
1 changed files with 35 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

80
tests/test_reference_fmha.py
View File

@@ -7,49 +7,43 @@ import torch
import math
import cutlass
import cutlass.cute as cute
import cutlass.torch as ct
import cuda.bindings.driver as cuda
from cute.blackwell.kernel.attention.fmha.fmha import FMHA, FusedMask, FusedMaskScale, FMHA_OperandMajorMode
from cute.blackwell.kernel.attention.fmha.fmha import (
BlackwellFusedMultiHeadAttentionForward,
FusedMask, FusedMaskScale, FMHA_OperandMajorMode
)
from cutlass.utils import LayoutEnum
def test_reference_fmha():
def test_reference():
HEAD_DIM = 64
for n in [128, 256, 512]:
torch.manual_seed(42)
m = 128
q = torch.randn(m, HEAD_DIM, 1, dtype=torch.bfloat16, device='cuda')
k = torch.randn(n, HEAD_DIM, 1, dtype=torch.bfloat16, device='cuda')
v = torch.randn(n, HEAD_DIM, dtype=torch.bfloat16, device='cuda')
c = torch.zeros(m, HEAD_DIM, 1, dtype=torch.bfloat16, device='cuda')
# Reference
qf = q[:, :, 0].float()
kf = k[:, :, 0].float()
batch = 1
q = torch.randn(batch, 1, m, HEAD_DIM, dtype=torch.bfloat16, device='cuda')
k = torch.randn(batch, 1, n, HEAD_DIM, dtype=torch.bfloat16, device='cuda')
v = torch.randn(batch, 1, n, HEAD_DIM, dtype=torch.bfloat16, device='cuda')
c = torch.zeros(batch, 1, m, HEAD_DIM, dtype=torch.bfloat16, device='cuda')
# Reference: PyTorch softmax attention
qf = q[0, 0].float()
kf = k[0, 0].float()
vf = v[0, 0].float()
scale = 1.0 / math.sqrt(HEAD_DIM)
attn = qf @ kf.T * scale
attn = torch.softmax(attn, dim=-1)
ref = attn @ v.float()
# CUTLASS reference FMHA
from cutlass.utils import LayoutEnum
import cutlass.torch as ct
q_tensor = q
k_tensor = k
v_tensor = v.unsqueeze(-1) # Add batch dim
c_tensor = c
q_maj = LayoutEnum.ROW_MAJOR # K major
k_maj = LayoutEnum.ROW_MAJOR
v_maj = LayoutEnum.ROW_MAJOR
o_maj = LayoutEnum.ROW_MAJOR
# Build the FMHA kernel
kernel = FMHA(
q_major_mode=FMHA_OperandMajorMode.from_LayoutEnum(q_maj),
k_major_mode=FMHA_OperandMajorMode.from_LayoutEnum(k_maj),
v_major_mode=FMHA_OperandMajorMode.from_LayoutEnum(v_maj),
o_major_mode=FMHA_OperandMajorMode.from_LayoutEnum(o_maj),
ref = attn @ vf
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
kernel = BlackwellFusedMultiHeadAttentionForward(
q_major_mode=FMHA_OperandMajorMode.K,
k_major_mode=FMHA_OperandMajorMode.K,
v_major_mode=FMHA_OperandMajorMode.MN,
o_major_mode=FMHA_OperandMajorMode.K,
q_head_dim=HEAD_DIM,
kv_head_dim=HEAD_DIM,
num_q_heads=1,
@@ -62,19 +56,13 @@ def test_reference_fmha():
epilogue_dtype=cutlass.Float32,
use_2cta_instrs=False,
)
# Set dimensions
kernel.set_dims(m, n)
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
print(f'n={n}: Compiling reference FMHA...', flush=True)
try:
compiled = cute.compile(kernel, q_tensor, k_tensor, v_tensor, c_tensor, stream)
compiled(q_tensor, k_tensor, v_tensor, c_tensor, stream)
result = kernel.run(q, k, v, c, stream)
torch.cuda.synchronize()
out = c[:, :, 0].float()
out = c[0, 0].float()
cos = torch.nn.functional.cosine_similarity(
out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)
).item()
@@ -84,7 +72,9 @@ def test_reference_fmha():
print(f' out[0,:4]={out[0,:4].tolist()}')
print(f' ref[0,:4]={ref[0,:4].tolist()}')
except Exception as e:
print(f'Reference FMHA n={n}: FAILED with error: {e}')
import traceback
print(f'Reference FMHA n={n}: FAILED')
traceback.print_exc()
if __name__ == '__main__':
test_reference_fmha()
test_reference()