P5: minimal Python multi-tile test
This commit is contained in:
@@ -130,5 +130,9 @@ def fmha_multitile_decode_raw(
|
||||
ctypes.c_float(scale),
|
||||
)
|
||||
if ret != 0:
|
||||
# Check CUDA error state
|
||||
err = torch.cuda.current_device()
|
||||
raise RuntimeError(f"Multi-tile kernel failed: {ret}")
|
||||
# Synchronize to catch async errors
|
||||
torch.cuda.synchronize()
|
||||
return o, lse
|
||||
|
||||
25
tests/unit/test_p5_python_minimal.py
Normal file
25
tests/unit/test_p5_python_minimal.py
Normal file
@@ -0,0 +1,25 @@
|
||||
"""Minimal multi-tile test via Python."""
|
||||
import torch, sys, math
|
||||
sys.path.insert(0, '/root/dsv4-nvfp4-workspace/kernel')
|
||||
|
||||
from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw
|
||||
|
||||
torch.manual_seed(42)
|
||||
hd = 64
|
||||
N = 256
|
||||
scale = 1.0 / math.sqrt(hd)
|
||||
|
||||
q = torch.randn(1, 1, 1, hd, dtype=torch.bfloat16, device='cuda').contiguous()
|
||||
k = torch.randn(1, 1, N, hd, dtype=torch.bfloat16, device='cuda').contiguous()
|
||||
v = torch.randn(1, 1, hd, N, dtype=torch.bfloat16, device='cuda').contiguous()
|
||||
|
||||
print(f'q align: {q.data_ptr() % 128}, k align: {k.data_ptr() % 128}, v align: {v.data_ptr() % 128}')
|
||||
print(f'q shape: {q.shape}, k shape: {k.shape}, v shape: {v.shape}')
|
||||
|
||||
try:
|
||||
o, lse = fmha_multitile_decode_raw(q, k, v, scale)
|
||||
print(f'Output[0,0,0,:5]: {o[0,0,0,:5].float()}')
|
||||
print(f'LSE: {lse[0,0,0].item():.4f}')
|
||||
print('SUCCESS')
|
||||
except Exception as e:
|
||||
print(f'FAILED: {e}')
|
||||
Reference in New Issue
Block a user