biondizzle
adcf3e04ab
- fmha_multihead_capi.cu: pure C API wrapper, no ATen/pybind11 deps - fmha_multihead_op.py: nvcc precompile + ctypes load (sm_100a) - Removed fmha_multihead_launch.cu (ATen approach didn't work) - Updated test to call kernel directly via ctypes API
134 lines
4.5 KiB
Python
134 lines
4.5 KiB
Python
"""
|
|
P3 Integration Test: Verify 6-warp multi-head decode fast path
|
|
produces identical results to the CuTeDSL slow path.
|
|
|
|
Tests:
|
|
1. MHA (n_q == n_kv), MQA (n_kv == 1), GQA (n_q > n_kv)
|
|
2. HD = 64, 128, 256
|
|
3. Single KV segment (N <= 128), T = 1
|
|
4. Cosine similarity >= 0.999998 between fast and slow paths
|
|
5. Launch count: fast path = 1 kernel, 0 cudaDeviceSynchronize
|
|
"""
|
|
import torch
|
|
import math
|
|
import sys
|
|
import os
|
|
|
|
# Ensure dsv4 is importable
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from dsv4.kernels.attention.production import dsv4_attention, _run_fmha_segmented
|
|
|
|
|
|
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, k, v, scale):
|
|
"""Pure PyTorch reference: (n_q, T, hd) x (n_kv, N, hd) -> (n_q, T, hd)"""
|
|
n_q, T, hd = q.shape
|
|
n_kv = k.shape[0] if k.dim() == 3 else 1
|
|
N = k.shape[-2] if k.dim() == 3 else k.shape[0]
|
|
q_per_kv = n_q // n_kv
|
|
|
|
if k.dim() == 2:
|
|
k = k.unsqueeze(0)
|
|
if v.dim() == 2:
|
|
v = v.unsqueeze(0)
|
|
|
|
output = torch.zeros(n_q, T, hd, dtype=torch.bfloat16, device='cuda')
|
|
for kv_idx in range(n_kv):
|
|
k_h = k[kv_idx] # (N, hd)
|
|
v_h = v[kv_idx] # (N, hd)
|
|
for qi in range(q_per_kv):
|
|
q_idx = kv_idx * q_per_kv + qi
|
|
q_h = q[q_idx] # (T, hd) — T=1
|
|
# S = q @ k^T / sqrt(hd)
|
|
s = torch.matmul(q_h.float(), k_h.float().T) * scale # (1, N)
|
|
s = torch.softmax(s, dim=-1)
|
|
o = torch.matmul(s, v_h.float()) # (1, hd)
|
|
output[q_idx] = o.bfloat16()
|
|
return output
|
|
|
|
|
|
def test_fast_path_matches_reference():
|
|
"""Test that the 6-warp fast path matches PyTorch reference."""
|
|
torch.manual_seed(42)
|
|
|
|
configs = [
|
|
# (n_q, n_kv, N, hd, desc)
|
|
(8, 8, 64, 64, "MHA hd=64"),
|
|
(8, 8, 128, 64, "MHA hd=64 N=128"),
|
|
(8, 8, 64, 128, "MHA hd=128"),
|
|
(8, 8, 64, 256, "MHA hd=256"),
|
|
(8, 1, 64, 64, "MQA hd=64"),
|
|
(8, 1, 128, 64, "MQA hd=64 N=128"),
|
|
(8, 1, 64, 128, "MQA hd=128"),
|
|
(128, 1, 64, 64, "MQA Pro hd=64"),
|
|
(128, 1, 64, 128, "MQA Pro hd=128"),
|
|
(8, 2, 64, 64, "GQA hd=64"),
|
|
(8, 4, 64, 128, "GQA hd=128"),
|
|
]
|
|
|
|
all_pass = True
|
|
for n_q, n_kv, N, hd, desc in configs:
|
|
T = 1
|
|
scale = 1.0 / math.sqrt(hd)
|
|
|
|
q = torch.randn(n_q, T, hd, dtype=torch.bfloat16, device='cuda')
|
|
if n_kv == 1:
|
|
k = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda')
|
|
v = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda')
|
|
else:
|
|
k = torch.randn(n_kv, N, hd, dtype=torch.bfloat16, device='cuda')
|
|
v = torch.randn(n_kv, N, hd, dtype=torch.bfloat16, device='cuda')
|
|
|
|
try:
|
|
from dsv4.kernels.attention.fmha_multihead_op import fmha_multihead_decode_raw
|
|
|
|
# Prepare tensors in the shape the kernel expects:
|
|
# Q: (1, n_q, 1, hd) BF16
|
|
# K: (1, n_kv, N, hd) BF16
|
|
# V: (1, n_kv, hd, N) BF16 (transposed!)
|
|
if n_kv == 1:
|
|
q_4d = q.unsqueeze(0).contiguous()
|
|
k_4d = k.unsqueeze(0).unsqueeze(0).contiguous()
|
|
v_4d = v.unsqueeze(0).unsqueeze(0).transpose(-1, -2).contiguous()
|
|
else:
|
|
q_4d = q.unsqueeze(0).contiguous()
|
|
k_4d = k.unsqueeze(0).contiguous()
|
|
v_4d = v.unsqueeze(0).transpose(-1, -2).contiguous()
|
|
|
|
sb = torch.zeros(1, n_q, dtype=torch.float32, device='cuda')
|
|
o_4d, lse_4d = fmha_multihead_decode_raw(
|
|
q_4d, k_4d, v_4d, scale, 0, 0, False, sb
|
|
)
|
|
o_fast = o_4d.squeeze(0) # (n_q, 1, hd)
|
|
o_ref = reference_attention(q, k, v, scale)
|
|
cos = cosine_sim(o_ref, o_fast).item()
|
|
status = "PASS" if cos >= 0.999998 else "FAIL"
|
|
if status == "FAIL":
|
|
all_pass = False
|
|
print(f" {status} {desc}: cos={cos:.6f}")
|
|
except Exception as e:
|
|
import traceback
|
|
print(f" FAIL {desc}: {e}")
|
|
traceback.print_exc()
|
|
all_pass = False
|
|
|
|
return all_pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("P3 Integration Test: 6-warp decode fast path vs reference")
|
|
print("=" * 60)
|
|
ok = test_fast_path_matches_reference()
|
|
print("=" * 60)
|
|
if ok:
|
|
print("ALL PASS")
|
|
else:
|
|
print("SOME FAILED")
|
|
sys.exit(0 if ok else 1)
|