From 6a7726e7642da038f31ad771b1a60a6f3ccada37 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 16:58:24 +0000 Subject: [PATCH] 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. --- tests/unit/test_p6_tma_epilogue.py | 247 +++++++++++++++++++++++++++++ 1 file changed, 247 insertions(+) create mode 100644 tests/unit/test_p6_tma_epilogue.py diff --git a/tests/unit/test_p6_tma_epilogue.py b/tests/unit/test_p6_tma_epilogue.py new file mode 100644 index 00000000..57c03310 --- /dev/null +++ b/tests/unit/test_p6_tma_epilogue.py @@ -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)