Skip to content

Update core.py Fix for CUDA <12.8 #202

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
umerkay wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Update core.py Fix for CUDA <12.8 #202

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6090cb4
Update core.py
umerkay Jul 2, 2025
85d52fd
More efficient: Get CUDA version once at module load
umerkay Jul 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 14 additions & 6 deletions sageattention/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,8 @@

import subprocess
import re
def get_cuda_version():
# Get CUDA version once at module load
def _detect_cuda_version():
try:
output = subprocess.check_output(['nvcc', '--version']).decode()
match = re.search(r'release (\d+)\.(\d+)', output)
Expand All @@ -66,6 +67,13 @@ def get_cuda_version():
print("Failed to get CUDA version:", e)
return None, None

# Store it in a module-level variable
_CUDA_VERSION = _detect_cuda_version()

# Now the function just returns the cached value
def get_cuda_version():
return _CUDA_VERSION

def get_cuda_arch_versions():
cuda_archs = []
for i in range(torch.cuda.device_count()):
Expand Down Expand Up @@ -684,10 +692,10 @@ def sageattn_qk_int8_pv_fp8_cuda(
assert q.device == k.device == v.device, "All tensors must be on the same device."
assert q.dtype == k.dtype == v.dtype, "All tensors must have the same dtype."

# cuda_major_version, cuda_minor_version = get_cuda_version()
# if(cuda_major_version, cuda_minor_version) < (12, 8) and pv_accum_dtype == 'fp32+fp16':
# warnings.warn("cuda version < 12.8, change pv_accum_dtype to 'fp32+fp32'")
# pv_accum_dtype = 'fp32+fp32'
cuda_major_version, cuda_minor_version = get_cuda_version()
if(cuda_major_version, cuda_minor_version) < (12, 8) and pv_accum_dtype == 'fp32+fp16':
warnings.warn("cuda version < 12.8, change pv_accum_dtype to 'fp32+fp32'")
pv_accum_dtype = 'fp32+fp32'

# FIXME(DefTruth): make sage attention work compatible with distributed
# env, for example, xDiT which launch by torchrun. Without this workaround,
Expand Down Expand Up @@ -930,4 +938,4 @@ def sageattn_qk_int8_pv_fp8_cuda_sm90(
if return_lse:
return o, lse / 1.44269504 + lse_correction * sm_scale if smooth_k else lse / 1.44269504
else:
return o
return o