🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/172108

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit f21f5be with merge base bfed04b ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@liangel-02 Would it be OK to use the repro you shared as a test here or is that still considered private?

go for it :)

@pytorchmergebot rebase

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

Successfully rebased wipvarlen onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout wipvarlen && git pull --rebase)

Sure just FYI need to revisit this after we upgrade to cuDNN 9.18+, as otherwise these tests are expected to fail

@claude can you do a review of this PR? Also it looks like we pass through strides for all tensors but dO is there a possibility the stride info for dO is wrong or are we guaranteed it will match q based off of construction?

Claude finished @drisspg's task in 1m 25s —— View job

PR Review: cuDNN varlen SDPA Fixes

• Read and understand the changed files
• Review aten/src/ATen/native/cudnn/MHA.cpp changes
• Review aten/src/ATen/native/transformers/cuda/attention_backward.cu changes
• Review torch/nn/attention/varlen.py changes
• Investigate dO stride handling question
• Provide comprehensive review feedback

Summary

This PR fixes several issues with the cuDNN variable-length (varlen) SDPA backward pass:

1. Separate ragged offsets for dO and output tensors - Previously, dO was incorrectly using RAG_O_OFF_ instead of its own ragged offset
2. Correct stride handling for gradient tensors (dQ, dK, dV) - Previously using strides from input tensors (q, k, v) instead of the output gradient tensors
3. Fix softmaxstats/logsumexp layout - cuDNN expects T,H,1 but torch/FA convention is H,T - now properly handled via unsqueeze/transpose
4. Enable cuDNN path on Hopper (SM 9.x) and Blackwell (SM 10.x) - _should_use_cudnn now returns True for these architectures
5. Use explicit kwargs in Python API calls - Improves readability and reduces positional argument errors

Detailed Review

aten/src/ATen/native/cudnn/MHA.cpp

1. New UIDs for separate ragged offsets (lines 420-429)
The addition of RAG_DQ_OFF, RAG_DK_OFF, RAG_DV_OFF, and RAG_DO_OFF is the right approach. This allows the gradient tensors to have their own ragged offset computations independent of the input tensors.

2. dO stride handling (lines 1296-1306)
This is the key fix for @drisspg's question:

auto do_strides = dO.strides();
auto DO_ = mha_graph->tensor(fe::graph::Tensor_attributes()
                                 .set_ragged_offset(RAG_DO_OFF_)  // Changed from RAG_O_OFF_
                                 .set_uid(DO)
                                 .set_name("DO")
                                 .set_dim({b, h_q, s_q, d_v})
                                 .set_stride(
                                     {INT_MAX,
                                      do_strides[strideidx0],
                                      do_strides[strideidx1],
                                      do_strides[strideidx2]}));

To answer your specific question: Yes, this PR now correctly passes through strides for dO. Previously at line 1259 (old code), DO_ was using RAG_O_OFF_ instead of its own ragged offset. The PR adds RAG_DO_OFF_ and also captures do_strides = dO.strides() to set the correct stride information for the dO tensor.

3. Gradient tensor strides (lines 1308-1335)
The change from using input tensor strides (q_strides, k_strides, v_strides) to output gradient tensor strides (dq_strides, dk_strides, dv_strides) is correct. The gradients may have different memory layouts than their corresponding inputs, and cuDNN needs accurate stride information.

4. softmaxstats layout fix (lines 1524-1532, 1811-1818)
The fix for H,T vs T,H,1 layout is important for compatibility with the Flash Attention convention:

// cuDNN wants T, H, 1, but torch/FA convention is H, T
softmaxstats = at::empty({h_q, q.size(0)}, q.options().dtype(kFloat));
// ...
if (softmaxstats.dim() == 2) {
    softmaxstats_ = softmaxstats.unsqueeze(-1).transpose(0, 1);
}

aten/src/ATen/native/transformers/cuda/attention_backward.cu

The removal of max_seqlen_batch_q and max_seqlen_batch_k in favor of using the passed-in max_q and max_k parameters is correct. For nested tensors, query.size(2) doesn't represent the max sequence length - it's the head dimension in THD format.

torch/nn/attention/varlen.py

1. Enable cuDNN for Hopper/Blackwell (lines 29-35)

def _should_use_cudnn(device_index: int) -> bool:
    major_cap = torch.cuda.get_device_capability(device_index)[0]
    if major_cap == 9 or major_cap == 10:
        return True
    return False

2. Additional checks before using cuDNN (lines 67-69)

head_dim_cudnn_ok = query.shape[-1] % 8 == 0 and value.shape[-1] % 8 == 0
window_size_cudnn_ok = window_size[0] == -1 and window_size[1] == -1
use_cudnn = window_size_cudnn_ok and query.is_cuda and _should_use_cudnn(query.device.index) and head_dim_cudnn_ok

These checks are important - cuDNN has restrictions on head dimensions (must be divisible by 8) and doesn't support sliding window attention.

3. Explicit kwargs in backward call (lines 314-331)
The switch to explicit keyword arguments is good for maintainability.

Potential Issues / Suggestions

1. Line 301 redundant assignment: In _varlen_attn_backward, line 303 assigns use_cudnn and then line 306 immediately reassigns it with the full condition. The first assignment on line 303 is dead code:

use_cudnn = query.is_cuda and _should_use_cudnn(query.device.index)  # Line 303 - dead code
head_dim_cudnn_ok = query.shape[-1] % 8 == 0 and value.shape[-1] % 8 == 0
window_size_cudnn_ok = window_size[0] == -1 and window_size[1] == -1
use_cudnn = window_size_cudnn_ok and query.is_cuda and _should_use_cudnn(query.device.index) and head_dim_cudnn_ok  # Line 306 - overwrites

Fix this →

2. philox_offset reuse: In the backward Python call, philox_offset=rng_state is passed with a comment "should be unused". While this works since dropout is disabled, it might be cleaner to pass a separate empty tensor or document this more explicitly in the function signature.

3. Review comment from @Skylion007: The use of NamedTuple for AuxRequest is noted as deprecated. Consider using the class-based approach instead.

Answer to @drisspg's Question

Also it looks like we pass through strides for all tensors but dO is there a possibility the stride info for dO is wrong or are we guaranteed it will match q based off of construction?

This PR fixes exactly this issue. Before this PR, dO was:

1. Using RAG_O_OFF_ (output's ragged offset) instead of its own RAG_DO_OFF_
2. The stride information was implicitly derived

After this PR:

• dO gets its own ragged offset (RAG_DO_OFF_) computed from cum_seqlen_q.mul(dO_.stride(-3))
• dO explicitly captures and uses do_strides = dO.strides()

The gradient output tensor's strides are not guaranteed to match q's layout. In fact, there's already handling at lines 1821-1824 that may call permute_to_matching_layout(o, dO_) if the innermost stride isn't 1. This ensures the layout is compatible with cuDNN's requirements.

Overall, this is a solid bug fix PR that addresses real issues with stride/offset handling in the cuDNN varlen backward pass. The changes are well-structured and address the core problem systematically.
| Branch: wipvarlen-53282b5c

#172108 (comment)

can we still enable the tests in this PR in test_varlen_attention.py?

@liangel-02 sorry if I am misunderstanding the question but I thought this PR is enabling the tests in that it will exercise cuDNN. The issue that remains is we also need #174310 for the cuDNN backend fix that is required

@eqy oh i meant can we have a test case where we set _should_use_cudnn = True in test_varlen_attention.py? but i can also add that in a separate PR once the bwd is fixed so whatever's easier !

@liangel-02 Did you mean torch/nn/attention/varlen.py
I updated that to

@lru_cache(maxsize=8)
def _should_use_cudnn(device_index: int) -> bool:
    """Cache device capability check to avoid repeated CUDA calls."""
    major_cap = torch.cuda.get_device_capability(device_index)[0]
    if major_cap == 9 or major_cap == 10:
        return True
    return False

Or should this be parametrized in the test as well to sweep both options?

@eqy yeah i think it'd be good to parametrize in the test as well so that the tests can run for both backends

Did we pull in that version yet?

Not done yet, I'm not sure this is fixed as of 9.20