OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2026-01-15 12:15:51 +00:00
Code Issues Actions 52 Packages Projects Releases Wiki Activity

[cuDNN][SDPA] Use-same prefer-cuDNN settings for Blackwell and Blackwell Ultra (#170800)

Browse Source 
GB300 can run GB200 SDPA kernels

Pull Request resolved: https://github.com/pytorch/pytorch/pull/170800
Approved by: https://github.com/Skylion007
This commit is contained in:
Eddie Yan
2025-12-31 17:35:59 +00:00
committed by PyTorch MergeBot
parent 766be25a17
commit 178ebac3a9
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
aten/src/ATen/native/transformers/cuda/sdp_utils.cpp
View File

@@ -84,7 +84,8 @@ bool check_prefer_cudnn_attention() {
try {
auto dprops = at::cuda::getCurrentDeviceProperties();
auto major = dprops->major;
return (major == 9 || major == 10) && !dprops->minor;
auto minor = dprops->minor;
return (major == 9 || major == 10) && (!minor || minor == 3);
} catch ([[maybe_unused]] c10::Error const& e) {
#ifdef DEBUG
TORCH_WARN("check_prefer_cudnn_attention() caught exception ", e.what());