huggingface / huggingface/diffusers
cuDNN attention backward corrupts gradients: K/V transposed twice, Q once
- Dominant language
- Python
- Stars
- 34.5k
- Forks
- 7.3k
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 91
Description
# cuDNN attention backward pass corrupts gradients (K/V transposed twice, Q once)
Opening this as an issue first, per the *AI-assisted and agentic contributions*
section of CONTRIBUTING.md — I have a fix and a repro ready and would like a
maintainer's ack on the scope before opening the PR.
## What's wrong
`_cudnn_attention_forward_op` saves the attention inputs **after** transposing them
to `(B, H, S, D)`:
```python
# src/diffusers/models/attention_dispatch.py:934-937
query = query.transpose(1, 2).contiguous()
key = key.transpose(1, 2).contiguous()
value = value.transpose(1, 2).contiguous()
tensors_to_save += (query, key, value)
```
`_cudnn_attention_backward_op` then transposes `key` and `value` a **second** time,
while leaving `query` alone:
```python
# src/diffusers/models/attention_dispatch.py:971-974
query, key, value, out, lse, ... = ctx.saved_tensors # already (B, H, S, D)
grad_out = grad_out.transpose(1, 2).contiguous() # correct: grad_out arrives BSHD
key = key.transpose(1, 2).contiguous() # -> (B, S, H, D)
value = value.transpose(1, 2).contiguous() # -> (B, S, H, D)
```
so `torch.ops.aten._scaled_dot_product_cudnn_attention_backward` receives Q as
`(B, H, S, D)` and K/V as `(B, S, H, D)`. The `grad_out` transpose is correct and
should stay — the forward transposes `out` back to BSHD at :963.
## Impact
Only the context-parallel `_native_cudnn` path is affected (the simple path at
:3607 goes through `F.scaled_dot_product_attention`, which has its own autograd).
Forward is unaffected, so inference is fine and no forward-only test can catch it —
but any training/fine-tuning run on this backend optimises against garbage gradients.
## Repro
H100, bf16, `B=2 S=16 H=16 D=64`. `num_heads == seq_len` is deliberate, so a shape
check cannot mask the layout error. Errors are `max_abs` vs an eager fp32 reference:
| variant | out | dQ | dK | dV |
|---|---|---|---|---|
| all three passed as saved (BHSD) | 0.0078 | 0.0078 | 0.0078 | 0.0156 |
| current code | 0.0078 | **7.8984** | **7.9727** | **3.9785** |
Reference gradient magnitudes are 1.859 / 2.266 / 2.750 — the errors exceed the
values being approximated, so the gradients are unrelated, not merely imprecise.
## Proposed fix
Drop the two `transpose(1, 2)` calls on `key`/`value` in the backward op; the saved
tensors are already in the layout the ATen op expects.
**Split out:** the CP path also forwards a `torch.bool` mask straight into the ATen op's
additive `attn_bias` slot, where it is consumed as an additive `0`/`1`. That is tracked
separately in #14342 so this issue can be closed cleanly by #14341, which fixes the
K/V transpose only.
## Context
Found during an audit of cuDNN SDPA integrations across the ecosystem, run by the
NVIDIA cuDNN team. I have the fix, a gradient regression test, and the self-review
report ready to attach to the PR once you've had a chance to weigh in on scope.
Contributor guide
Research direction
Start in src/diffusers/models/attention_dispatch.py at the cited forward lines 934-937 and backward lines 971-974, then compare the saved tensor layouts with the ATen cuDNN backward call. Reproduce the H100 bf16 case and add or run the gradient regression test; done means K/V gradients match the eager fp32 reference without changing the separate mask issue.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100