diff --git a/dsv4/kernels/attention/fmha_multitile_op.py b/dsv4/kernels/attention/fmha_multitile_op.py index b246b108..9a6ff221 100644 --- a/dsv4/kernels/attention/fmha_multitile_op.py +++ b/dsv4/kernels/attention/fmha_multitile_op.py @@ -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 diff --git a/tests/unit/test_p5_python_minimal.py b/tests/unit/test_p5_python_minimal.py new file mode 100644 index 00000000..9033be0a --- /dev/null +++ b/tests/unit/test_p5_python_minimal.py @@ -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}')