91 lines
2.4 KiB
Python
91 lines
2.4 KiB
Python
"""
|
|
Test: Compile FMHA SM100 kernel with nvcc directly.
|
|
|
|
Step 1: Try to compile the .cuh to check for C++ errors.
|
|
Step 2: If that works, try torch.utils.cpp_extension JIT.
|
|
"""
|
|
import subprocess
|
|
import sys
|
|
import os
|
|
|
|
def get_repo_root():
|
|
"""Find repo root from this file's location."""
|
|
d = os.path.dirname(os.path.abspath(__file__))
|
|
# Go up until we find dsv4/ directory
|
|
while d != '/':
|
|
if os.path.exists(os.path.join(d, 'dsv4')):
|
|
return d
|
|
d = os.path.dirname(d)
|
|
return None
|
|
|
|
REPO = get_repo_root()
|
|
print(f"Repo root: {REPO}")
|
|
|
|
CUTLASS = "/root/cutlass"
|
|
|
|
# Step 1: Try nvcc compile (just syntax check)
|
|
print("\n" + "=" * 60)
|
|
print("Step 1: nvcc syntax check")
|
|
print("=" * 60)
|
|
|
|
nvcc_cmd = [
|
|
"/usr/local/cuda-13.2/bin/nvcc",
|
|
"--std=c++20",
|
|
"-gencode=arch=compute_100a,code=sm_100a",
|
|
f"-I{CUTLASS}/include",
|
|
f"-I{REPO}",
|
|
"-DCUTE_ARCH_TCGEN05_MMA_ENABLED",
|
|
"-DCUTE_ARCH_TCGEN05_TMEM_ENABLED",
|
|
"-DCUTE_ARCH_TCGEN05_F16F32_MMA_ENABLED",
|
|
"-c",
|
|
f"{REPO}/dsv4/kernels/attention/fmha_sm100.cuh",
|
|
"--x", "cu",
|
|
"-o", "/tmp/fmha_sm100_test.o",
|
|
"--ptxas-options=-v",
|
|
"--expt-relaxed-constexpr",
|
|
]
|
|
|
|
print(f"nvcc command: {' '.join(nvcc_cmd)}")
|
|
result = subprocess.run(nvcc_cmd, capture_output=True, text=True, timeout=120)
|
|
print(f"Exit code: {result.returncode}")
|
|
if result.stdout:
|
|
print(f"STDOUT:\n{result.stdout[-1000:]}")
|
|
if result.stderr:
|
|
print(f"STDERR:\n{result.stderr[-2000:]}")
|
|
|
|
if result.returncode != 0:
|
|
print("\n❌ nvcc compilation FAILED — fix errors above before proceeding")
|
|
sys.exit(1)
|
|
|
|
print("\n✅ nvcc compilation PASSED!")
|
|
|
|
# Step 2: JIT compile with torch
|
|
print("\n" + "=" * 60)
|
|
print("Step 2: torch.utils.cpp_extension JIT")
|
|
print("=" * 60)
|
|
|
|
try:
|
|
from torch.utils.cpp_extension import load
|
|
|
|
module = load(
|
|
name="fmha_sm100",
|
|
sources=[f"{REPO}/dsv4/kernels/attention/fmha_sm100_launch.cu"],
|
|
extra_cuda_cflags=[
|
|
"-gencode=arch=compute_100a,code=sm_100a",
|
|
f"-I{REPO}",
|
|
f"-I{CUTLASS}/include",
|
|
],
|
|
extra_cflags=[
|
|
f"-I/usr/local/cuda-13.2/include",
|
|
],
|
|
verbose=True,
|
|
)
|
|
print("\n✅ JIT compilation PASSED!")
|
|
print(f"Module: {module}")
|
|
print(f"fmha_decode: {module.fmha_decode}")
|
|
except Exception as e:
|
|
print(f"\n❌ JIT compilation FAILED: {e}")
|
|
sys.exit(1)
|
|
|
|
print("\n✅ ALL COMPILATION TESTS PASSED!")
|