diff --git a/DEPLOYMENT.md b/DEPLOYMENT.md index e71a738..10728fe 100644 --- a/DEPLOYMENT.md +++ b/DEPLOYMENT.md @@ -126,7 +126,7 @@ Then install the attention and kernel dependencies: git -C deploy/tmp/cutlass checkout dcf215af # build only this machine's arch (the default is a 5-arch fat binary — # sm_80..120a — which multiplies compile time ~5x); auto-detected: - export JOYOMNI_OPS_CUDA_ARCHS=$(python -c "import torch; cc = torch.cuda.get_device_capability(0); print(f'{cc[0]}{cc[1]}a' if cc[0] >= 10 else f'{cc[0]}{cc[1]}')") + export JOYOMNI_OPS_CUDA_ARCHS=$(python -c "import torch; cc = torch.cuda.get_device_capability(0); print(f'{cc[0]}{cc[1]}a' if cc[0] >= 9 else f'{cc[0]}{cc[1]}')") echo "building joyomni_ops for sm_$JOYOMNI_OPS_CUDA_ARCHS" JOYOMNI_OPS_CUTLASS_DIR=$(pwd)/deploy/tmp/cutlass \ python -m pip install --no-build-isolation ./deploy/joyomni_ops diff --git a/deploy/joyomni_ops/setup.py b/deploy/joyomni_ops/setup.py index 2b6d083..df31a5b 100644 --- a/deploy/joyomni_ops/setup.py +++ b/deploy/joyomni_ops/setup.py @@ -9,7 +9,7 @@ - fp8_scaled_mm : FP8 per-token x per-channel scaled GEMM (cutlass) GPU arch coverage is chosen from the local nvcc version: - - always: sm_80, sm_89, sm_90 + - always: sm_80, sm_89, sm_90a - CUDA >= 12.8: also sm_100a (B200) and sm_120a (RTX PRO 6000 / RTX 5090) So building on a CUDA 12.8+ toolchain automatically yields Blackwell support. """ @@ -57,7 +57,12 @@ def _gencodes(): flags = [ "-gencode=arch=compute_80,code=sm_80", "-gencode=arch=compute_89,code=sm_89", - "-gencode=arch=compute_90,code=sm_90", + # Hopper must be sm_90a: fp8_scaled_mm selects cutlass' + # KernelTmaWarpSpecializedPingpongFP8FastAccum, whose WGMMA path is + # arch-conditional. Built as plain sm_90 it compiles and loads, then + # aborts every launch ("Arch conditional MMA instruction used without + # targeting appropriate compute capability"). + "-gencode=arch=compute_90a,code=sm_90a", ] if (mj, mn) >= (12, 8): flags += [