109 lines
3.5 KiB
Python
109 lines
3.5 KiB
Python
"""
|
|
P6 Integration Test: One-way TMEM→regs→SMEM→GMEM epilogue.
|
|
|
|
Verifies the epilogue refactoring:
|
|
1. TMEM → registers (tcgen05.ld, warp-collective)
|
|
2. epilogue_op in registers (normalize, FP4 hook)
|
|
3. Registers → SMEM (row-major)
|
|
4. SMEM → GMEM (direct write)
|
|
|
|
Gate: worst-case cosine >= 0.999990 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_one_way_epilogue():
|
|
"""Test the one-way TMEM→regs→SMEM→GMEM epilogue path."""
|
|
from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw
|
|
from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw
|
|
|
|
torch.manual_seed(42)
|
|
configs = [
|
|
# Single-tile (N<=128)
|
|
(4, 4, 64, 64, "MHA hd=64", False),
|
|
(4, 4, 128, 128, "MHA hd=128", False),
|
|
(4, 4, 64, 256, "MHA hd=256", False),
|
|
(4, 1, 64, 64, "MQA hd=64", False),
|
|
(128, 1, 64, 64, "MQA Pro hd=64", False),
|
|
(8, 2, 64, 64, "GQA hd=64", False),
|
|
# Multi-tile (N>128)
|
|
(4, 4, 256, 64, "MHA hd=64 N=256 (2 tiles)", True),
|
|
(4, 4, 512, 64, "MHA hd=64 N=512 (4 tiles)", True),
|
|
(4, 1, 256, 128, "MQA hd=128 N=256 (2 tiles)", True),
|
|
]
|
|
|
|
all_pass = True
|
|
for n_q, n_kv, N, hd, desc, multitile 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()
|
|
|
|
if multitile:
|
|
o_4d, _ = fmha_multitile_decode_raw(q_4d, k_4d, v_4d, scale)
|
|
else:
|
|
sb = torch.zeros(1, n_q, dtype=torch.float32, device='cuda')
|
|
o_4d, _ = fmha_multitile_decode_raw(q_4d, k_4d, v_4d, scale)
|
|
|
|
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.999994 else "FAIL"
|
|
if worst_cos < 0.999994:
|
|
all_pass = False
|
|
print(f" {status} {desc}: worst_cos={worst_cos:.6f}")
|
|
|
|
return all_pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("P6 Integration Test: One-way TMEM→regs→SMEM→GMEM epilogue")
|
|
print("=" * 60)
|
|
|
|
if test_one_way_epilogue():
|
|
print("ALL PASS")
|
|
else:
|
|
print("SOME FAILED")
|
|
sys.exit(1)
|