cuDNN attention backward corrupts gradients: K/V transposed twice, Q once

[YangXu1990uiuc]

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):

# 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:

# 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.