78 lines
3.3 KiB
Python
78 lines
3.3 KiB
Python
"""
|
|
D1.5 Debug: NO-OP TMEM round-trip test.
|
|
Tests s_k=256 with rescale_factor=1.0 (NO-OP).
|
|
If the round-trip itself is broken, even NO-OP will corrupt O.
|
|
If it produces the same (wrong) result as without rescale, the round-trip works.
|
|
"""
|
|
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 reference_attention(q, k, v, scale):
|
|
qf = q.float(); kf = k.float()
|
|
attn = qf @ kf.T * scale
|
|
attn_max = attn.max(dim=-1, keepdim=True)[0]
|
|
attn_exp = torch.exp(attn - attn_max)
|
|
return attn_exp @ v.float()
|
|
|
|
|
|
def run_test(s_k, debug_noop=False, label=""):
|
|
hd = 64; m = 128; scale = 1.0 / math.sqrt(hd)
|
|
torch.manual_seed(42)
|
|
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
q = torch.randn(m, hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda')
|
|
v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
|
|
|
|
kernel = FmhaKernel(head_dim=hd, s_k=s_k, use_smem_p=False, normalize=False, debug_noop_rescale=debug_noop)
|
|
pv_n_tile = kernel.pv_n_tile
|
|
c = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
|
|
lse = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
rs = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda')
|
|
|
|
v_t = v[:, 0:pv_n_tile].contiguous().unsqueeze(-1)
|
|
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_t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_t))
|
|
mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c))
|
|
mLSE = ct.from_dlpack(lse).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse))
|
|
mRS = ct.from_dlpack(rs).mark_layout_dynamic(leading_dim=ct.get_leading_dim(rs))
|
|
|
|
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE, row_sums=mRS)
|
|
compiled(mQ, mK, mV, mC, stream, mLSE, row_sums=mRS)
|
|
torch.cuda.synchronize()
|
|
|
|
ref = reference_attention(q[:, :, 0], k[:, :, 0], v, scale)
|
|
out = c[:, :, 0].float()
|
|
cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item()
|
|
print(f'{label}: cos={cos:.6f} {"PASS" if cos >= 0.999 else "FAIL"}', flush=True)
|
|
return out
|
|
|
|
|
|
def test():
|
|
# Test 1: s_k=128 baseline
|
|
run_test(128, label="s_k=128 baseline")
|
|
|
|
# Test 2: s_k=256 WITH rescale (should be correct if TMEM round-trip works)
|
|
run_test(256, label="s_k=256 with rescale")
|
|
|
|
# Test 3: s_k=256 NOOP rescale (acc_scale=1.0)
|
|
# This should produce the same result as s_k=256 WITHOUT any rescale
|
|
# (which is: O = P[0]*V[0] + P[1]*V[1] with no O rescale — mathematically wrong
|
|
# but should be stable if TMEM round-trip doesn't corrupt)
|
|
out_noop = run_test(256, debug_noop=True, label="s_k=256 NOOP rescale")
|
|
|
|
# Test 4: s_k=256 WITHOUT any rescale (old code path)
|
|
# Compare with NOOP to see if TMEM round-trip itself corrupts
|
|
# We can't easily disable the rescale in the current code,
|
|
# but NOOP rescale with factor=1.0 is equivalent to a successful round-trip
|
|
# followed by multiply-by-1. If the output matches a "no-rescale" baseline,
|
|
# the TMEM round-trip is working correctly.
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test()
|