P6: Add integration test for TMA store epilogue
test_p6_tma_epilogue.py: Tests direct GMEM path, TMA store path, and parity between both. Gate: cos >= 0.999998.
This commit is contained in:
247
tests/unit/test_p6_tma_epilogue.py
Normal file
247
tests/unit/test_p6_tma_epilogue.py
Normal file
@@ -0,0 +1,247 @@
|
||||
"""
|
||||
P6 Integration Test: One-way TMEM→regs→SMEM→TMA store epilogue.
|
||||
|
||||
Tests both the direct GMEM write fallback (tma_o=nullptr) and the
|
||||
proper TMA store pipeline. Verifies the epilogue refactoring hasn't
|
||||
regressed numerics and that the TMA store path produces identical results.
|
||||
|
||||
Gate: worst-case cosine >= 0.999998 per configuration.
|
||||
"""
|
||||
import torch
|
||||
import math
|
||||
import sys
|
||||
import os
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
|
||||
def cosine_sim(a, b):
|
||||
a = a.flatten().float()
|
||||
b = b.flatten().float()
|
||||
return (a @ b) / (a.norm() * b.norm() + 1e-30)
|
||||
|
||||
|
||||
def reference_attention(q_4d, k_4d, v_4d, scale):
|
||||
"""PyTorch reference matching kernel tensor layout."""
|
||||
n_h = q_4d.shape[1]
|
||||
n_kv = k_4d.shape[1]
|
||||
N = k_4d.shape[2]
|
||||
q_per_kv = n_h // n_kv
|
||||
q = q_4d[0]
|
||||
k = k_4d[0]
|
||||
v = v_4d[0].transpose(-1, -2)
|
||||
|
||||
output = torch.zeros(n_h, 1, q_4d.shape[3], dtype=torch.bfloat16, device='cuda')
|
||||
for h in range(n_h):
|
||||
kv_idx = h // q_per_kv
|
||||
q_h = q[h]
|
||||
k_h = k[kv_idx]
|
||||
v_h = v[kv_idx]
|
||||
s = torch.matmul(q_h.float(), k_h.float().T) * scale
|
||||
s = torch.softmax(s, dim=-1)
|
||||
o = torch.matmul(s, v_h.float())
|
||||
output[h] = o.bfloat16()
|
||||
return output
|
||||
|
||||
|
||||
def test_direct_gmem_path():
|
||||
"""Test the direct GMEM write path (tma_o=nullptr, backward compatible)."""
|
||||
from dsv4.kernels.attention.fmha_multihead_op import fmha_multihead_decode_raw
|
||||
|
||||
torch.manual_seed(42)
|
||||
configs = [
|
||||
(4, 4, 64, 64, "MHA hd=64"),
|
||||
(4, 4, 128, 128, "MHA hd=128"),
|
||||
(4, 4, 64, 256, "MHA hd=256"),
|
||||
(4, 1, 64, 64, "MQA hd=64"),
|
||||
(128, 1, 64, 64, "MQA Pro hd=64"),
|
||||
]
|
||||
|
||||
all_pass = True
|
||||
for n_q, n_kv, N, hd, desc in configs:
|
||||
scale = 1.0 / math.sqrt(hd)
|
||||
q_4d = torch.randn(1, n_q, 1, hd, dtype=torch.bfloat16, device='cuda').contiguous()
|
||||
k_4d = torch.randn(1, n_kv, N, hd, dtype=torch.bfloat16, device='cuda').contiguous()
|
||||
v_4d = torch.randn(1, n_kv, hd, N, dtype=torch.bfloat16, device='cuda').contiguous()
|
||||
|
||||
sb = torch.zeros(1, n_q, dtype=torch.float32, device='cuda')
|
||||
o_4d, _ = fmha_multihead_decode_raw(q_4d, k_4d, v_4d, scale, 0, 0, False, sb)
|
||||
o_ref = reference_attention(q_4d, k_4d, v_4d, scale)
|
||||
|
||||
worst_cos = 1.0
|
||||
for h in range(n_q):
|
||||
cos = torch.nn.functional.cosine_similarity(
|
||||
o_4d[0, h].flatten().float().unsqueeze(0),
|
||||
o_ref[h].flatten().float().unsqueeze(0)
|
||||
).item()
|
||||
worst_cos = min(worst_cos, cos)
|
||||
|
||||
status = "PASS" if worst_cos >= 0.999998 else "FAIL"
|
||||
if worst_cos < 0.999998:
|
||||
all_pass = False
|
||||
print(f" {status} [direct] {desc}: worst_cos={worst_cos:.6f}")
|
||||
|
||||
return all_pass
|
||||
|
||||
|
||||
def test_tma_store_path():
|
||||
"""Test the TMA store epilogue path (tma_o set, proper async pipeline)."""
|
||||
from dsv4.kernels.attention.fmha_multihead_op import _get_lib
|
||||
import ctypes
|
||||
|
||||
lib = _get_lib()
|
||||
|
||||
torch.manual_seed(123)
|
||||
configs = [
|
||||
(4, 4, 64, 64, "MHA hd=64"),
|
||||
(4, 4, 128, 128, "MHA hd=128"),
|
||||
(4, 4, 64, 256, "MHA hd=256"),
|
||||
(4, 1, 64, 64, "MQA hd=64"),
|
||||
(128, 1, 64, 64, "MQA Pro hd=64"),
|
||||
]
|
||||
|
||||
all_pass = True
|
||||
for n_q, n_kv, N, hd, desc in configs:
|
||||
scale = 1.0 / math.sqrt(hd)
|
||||
q = torch.randn(1, n_q, 1, hd, dtype=torch.bfloat16, device='cuda').contiguous()
|
||||
k = torch.randn(1, n_kv, N, hd, dtype=torch.bfloat16, device='cuda').contiguous()
|
||||
v = torch.randn(1, n_kv, hd, N, dtype=torch.bfloat16, device='cuda').contiguous()
|
||||
|
||||
# GQA expansion
|
||||
q_per_kv = n_q // n_kv
|
||||
if n_kv < n_q:
|
||||
k = k.repeat_interleave(q_per_kv, dim=1)
|
||||
v = v.repeat_interleave(q_per_kv, dim=1)
|
||||
|
||||
# Pad N to 128
|
||||
if N < 128:
|
||||
pad = 128 - N
|
||||
k = torch.cat([k, torch.zeros(1, k.shape[1], pad, hd, dtype=torch.bfloat16, device='cuda')], dim=2)
|
||||
v = torch.cat([v, torch.zeros(1, v.shape[1], hd, pad, dtype=torch.bfloat16, device='cuda')], dim=3)
|
||||
N = 128
|
||||
k = k.contiguous()
|
||||
v = v.contiguous()
|
||||
q = q.contiguous()
|
||||
|
||||
o = torch.zeros(1, n_q, 1, hd, dtype=torch.bfloat16, device='cuda')
|
||||
lse = torch.zeros(1, n_q, 1, dtype=torch.float32, device='cuda')
|
||||
|
||||
# Call the TMA store variant
|
||||
ret = lib.fmha_multihead_decode_tma_launch(
|
||||
ctypes.c_void_p(q.data_ptr()),
|
||||
ctypes.c_void_p(k.data_ptr()),
|
||||
ctypes.c_void_p(v.data_ptr()),
|
||||
ctypes.c_void_p(o.data_ptr()),
|
||||
ctypes.c_void_p(lse.data_ptr()),
|
||||
ctypes.c_int(1), ctypes.c_int(n_q), ctypes.c_int(n_q), ctypes.c_int(N), ctypes.c_int(hd),
|
||||
ctypes.c_int(q.stride(1)), ctypes.c_int(q.stride(0)),
|
||||
ctypes.c_int(k.stride(1)), ctypes.c_int(k.stride(0)),
|
||||
ctypes.c_int(v.stride(1)), ctypes.c_int(v.stride(0)),
|
||||
ctypes.c_int(o.stride(1)), ctypes.c_int(o.stride(0)),
|
||||
ctypes.c_int(lse.stride(1)), ctypes.c_int(lse.stride(0)),
|
||||
ctypes.c_float(scale),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if ret != 0:
|
||||
print(f" FAIL [TMA] {desc}: kernel launch failed (ret={ret})")
|
||||
all_pass = False
|
||||
continue
|
||||
|
||||
# Compare with reference
|
||||
o_ref = reference_attention(
|
||||
q.reshape(1, n_q, 1, hd),
|
||||
k.reshape(1, n_q, N, hd),
|
||||
v.reshape(1, n_q, hd, N),
|
||||
scale
|
||||
)
|
||||
|
||||
worst_cos = 1.0
|
||||
for h in range(n_q):
|
||||
cos = torch.nn.functional.cosine_similarity(
|
||||
o[0, h].flatten().float().unsqueeze(0),
|
||||
o_ref[h].flatten().float().unsqueeze(0)
|
||||
).item()
|
||||
worst_cos = min(worst_cos, cos)
|
||||
|
||||
status = "PASS" if worst_cos >= 0.999998 else "FAIL"
|
||||
if worst_cos < 0.999998:
|
||||
all_pass = False
|
||||
print(f" {status} [TMA] {desc}: worst_cos={worst_cos:.6f}")
|
||||
|
||||
return all_pass
|
||||
|
||||
|
||||
def test_direct_vs_tma_parity():
|
||||
"""Verify direct and TMA paths produce identical results."""
|
||||
from dsv4.kernels.attention.fmha_multihead_op import fmha_multihead_decode_raw, _get_lib
|
||||
import ctypes
|
||||
|
||||
lib = _get_lib()
|
||||
torch.manual_seed(999)
|
||||
|
||||
configs = [
|
||||
(4, 64, 64, "hd=64"),
|
||||
(4, 128, 128, "hd=128"),
|
||||
(4, 64, 256, "hd=256"),
|
||||
]
|
||||
|
||||
all_pass = True
|
||||
for n_q, N, hd, desc in configs:
|
||||
scale = 1.0 / math.sqrt(hd)
|
||||
q = torch.randn(1, n_q, 1, hd, dtype=torch.bfloat16, device='cuda').contiguous()
|
||||
k = torch.randn(1, n_q, N, hd, dtype=torch.bfloat16, device='cuda').contiguous()
|
||||
v = torch.randn(1, n_q, hd, N, dtype=torch.bfloat16, device='cuda').contiguous()
|
||||
|
||||
# Direct path
|
||||
sb = torch.zeros(1, n_q, dtype=torch.float32, device='cuda')
|
||||
o_direct, _ = fmha_multihead_decode_raw(q, k, v, scale, 0, 0, False, sb)
|
||||
|
||||
# TMA path
|
||||
o_tma = torch.zeros_like(q)
|
||||
lse = torch.zeros(1, n_q, 1, dtype=torch.float32, device='cuda')
|
||||
ret = lib.fmha_multihead_decode_tma_launch(
|
||||
ctypes.c_void_p(q.data_ptr()),
|
||||
ctypes.c_void_p(k.data_ptr()),
|
||||
ctypes.c_void_p(v.data_ptr()),
|
||||
ctypes.c_void_p(o_tma.data_ptr()),
|
||||
ctypes.c_void_p(lse.data_ptr()),
|
||||
ctypes.c_int(1), ctypes.c_int(n_q), ctypes.c_int(n_q), ctypes.c_int(N), ctypes.c_int(hd),
|
||||
ctypes.c_int(q.stride(1)), ctypes.c_int(q.stride(0)),
|
||||
ctypes.c_int(k.stride(1)), ctypes.c_int(k.stride(0)),
|
||||
ctypes.c_int(v.stride(1)), ctypes.c_int(v.stride(0)),
|
||||
ctypes.c_int(o_tma.stride(1)), ctypes.c_int(o_tma.stride(0)),
|
||||
ctypes.c_int(lse.stride(1)), ctypes.c_int(lse.stride(0)),
|
||||
ctypes.c_float(scale),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if ret != 0:
|
||||
print(f" FAIL [parity] {desc}: TMA kernel failed")
|
||||
all_pass = False
|
||||
continue
|
||||
|
||||
# Direct vs TMA should be bit-identical (same math, same BF16 rounding)
|
||||
cos = cosine_sim(o_direct, o_tma)
|
||||
status = "PASS" if cos >= 0.999999 else "FAIL"
|
||||
if cos < 0.999999:
|
||||
all_pass = False
|
||||
print(f" {status} [parity] direct vs TMA {desc}: cos={cos:.8f}")
|
||||
|
||||
return all_pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("P6 Integration Test: One-way TMEM→regs→SMEM→TMA store epilogue")
|
||||
print("=" * 60)
|
||||
|
||||
p1 = test_direct_gmem_path()
|
||||
p2 = test_tma_store_path()
|
||||
p3 = test_direct_vs_tma_parity()
|
||||
|
||||
print()
|
||||
if p1 and p2 and p3:
|
||||
print("ALL PASS")
|
||||
else:
|
||||
print("SOME FAILED")
|
||||
sys.exit(1)
|
||||
Reference in New Issue
Block a user