pytorch - 💡(How to fix) Fix [cuDNN][SDPA] head_dim limit stuck at 128 on sm_120 (RTX 5090), blocking head_dim=256 models [1 comments, 2 participants]

Summary

On consumer Blackwell (sm_120, RTX 5090), check_cudnn_tensor_shapes caps head_dim_limit at 128 even though check_cudnn_hardware_support already whitelists the hardware in the [sm_80, sm_121] range. This silently rejects cuDNN attention for any head_dim=256 model (Gemma 3, Qwen3 ≥14B, Llama 3.1 70B, etc.) and forces a fall-through to the MATH backend with its O(seq²) fp32 softmax — OOMs at long context.

Adjacent prior art: PR #172621 (Jan 2026, merged) did the same kind of lift for datacenter Blackwell (sm_10.0/10.3 DeepSeek head_dim=192). The workstation/consumer sm_120 line was simply not included.

Empirical reproducer (RTX 5090, torch 2.8+cu129, cuDNN 9.10.02)

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

def try_backend(backend, head_dim, seq=256, heads=8, bs=1,
                dtype=torch.bfloat16):
    q = torch.randn(bs, heads, seq, head_dim, device="cuda", dtype=dtype)
    k = torch.randn(bs, heads, seq, head_dim, device="cuda", dtype=dtype)
    v = torch.randn(bs, heads, seq, head_dim, device="cuda", dtype=dtype)
    try:
        with sdpa_kernel([backend]):
            F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return "ACCEPT"
    except RuntimeError:
        return "REJECT"

for be in [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION,
           SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH]:
    print(be, {hd: try_backend(be, hd) for hd in [128, 192, 256]})

cuDNN error at rejection: "head_dim should be no more than 128" — exactly the warning from sdp_utils.cpp:522.

Source analysis

sdp_utils.cpp:504-525 (release/2.11):

auto head_dim_limit = 128;
// Hopper: head_dim<=256 support with cuDNN >= 9.10.0
if (cudnn_version >= 91000) {
  auto dprops = at::cuda::getCurrentDeviceProperties();
  if (dprops->major == 9 && !dprops->minor) {       // only sm_90
    head_dim_limit = 256;
  }
}
// (sm_10.x DeepSeek head_dim=192 case — irrelevant for sm_120)

sm_120 never enters the head_dim_limit = 256 branch, falls through with = 128. Since check_cudnn_hardware_support at lines 635-656 already accepts sm_120 (it's inside the [sm_80, sm_121] range), this is purely the shape check that's out of sync.

Proposed fix

   auto head_dim_limit = 128;
-  // Hopper: head_dim<=256 support with cuDNN >= 9.10.0
+  // Hopper (sm_90) and Blackwell consumer (sm_120, RTX 5090):
+  // head_dim<=256 support with cuDNN >= 9.10.0.
   if (cudnn_version >= 91000) {
     auto dprops = at::cuda::getCurrentDeviceProperties();
-    if (dprops->major == 9 && !dprops->minor) {
+    if ((dprops->major == 9 && !dprops->minor) ||
+        (dprops->major == 12 && dprops->minor == 0)) {
       head_dim_limit = 256;
     }
   }

Open questions for maintainers

1. Does cuDNN 9.10+ actually handle head_dim=256 kernels on sm_120? The pytorch gate is currently the only blocker per my reading, but if cuDNN rejects shape-wise independently on sm_120, lifting the pytorch gate would move the error from a readable warning to an opaque kernel-launch failure.

2. Is this a one-line patch you'd accept a PR for, or is there a broader Blackwell-shape-support effort in flight that we should wait on?

3. Any preferred test location? test_transformers.py seems to be where similar arch-gated checks live.

First-time pytorch contributor disclosure

This is my first issue on pytorch core. Apologies in advance for any style deviations. Happy to open a PR if the direction is right.

🤖 Drafted with Claude Code (Claude Opus 4.7), reviewed and posted by me.

cc @csarofeen @ptrblck @eqy @nWEIdia @msaroufim @jerryzh168 @tinglvv @drisspg @liangel-02 @howardzhang-cv