81 lines
2.7 KiB
Python
81 lines
2.7 KiB
Python
"""Test the CUTLASS reference Blackwell FMHA on the B200.
|
|
Does it actually work multi-tile?"""
|
|
import sys
|
|
sys.path.insert(0, '/root/cutlass/examples/python/CuTeDSL')
|
|
|
|
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 (
|
|
BlackwellFusedMultiHeadAttentionForward,
|
|
FusedMask, FusedMaskScale, FMHA_OperandMajorMode
|
|
)
|
|
from cutlass.utils import LayoutEnum
|
|
|
|
def test_reference():
|
|
HEAD_DIM = 64
|
|
for n in [128, 256, 512]:
|
|
torch.manual_seed(42)
|
|
m = 128
|
|
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 @ 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,
|
|
num_kv_heads=1,
|
|
q_dtype=cutlass.BFloat16,
|
|
k_dtype=cutlass.BFloat16,
|
|
v_dtype=cutlass.BFloat16,
|
|
o_dtype=cutlass.BFloat16,
|
|
acc_dtype=cutlass.Float32,
|
|
epilogue_dtype=cutlass.Float32,
|
|
use_2cta_instrs=False,
|
|
)
|
|
|
|
print(f'n={n}: Compiling reference FMHA...', flush=True)
|
|
try:
|
|
result = kernel.run(q, k, v, c, stream)
|
|
torch.cuda.synchronize()
|
|
|
|
out = c[0, 0].float()
|
|
cos = torch.nn.functional.cosine_similarity(
|
|
out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)
|
|
).item()
|
|
n_tiles = n // 128
|
|
print(f'Reference FMHA n={n} ({n_tiles} tiles): 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[0,:4].tolist()}')
|
|
except Exception as e:
|
|
import traceback
|
|
print(f'Reference FMHA n={n}: FAILED')
|
|
traceback.print_exc()
|
|
|
|
if __name__ == '__main__':
|
|
test_reference()
|