Fix double transpose of K/V in templated attention backward

[guptaishaan]

Fixes #14338

_cudnn_attention_forward_op saves query/key/value after transposing them to (B, H, S, D),
but _cudnn_attention_backward_op transposed key and value a second time, so
aten::_scaled_dot_product_cudnn_attention_backward got Q as BHSD and K/V as BSHD. The
grad_out transpose is correct and stays, since the forward op returns out in BSHD.

_native_flash_attention_backward_op has the identical two lines against the identical
forward-op contract, so it is fixed here too. Easy to drop that hunk if you would rather split it.

Four lines deleted, no behaviour change to the forward pass.

Test

New tests/models/test_attention_dispatch.py: drives each forward/backward op pair through a
minimal autograd.Function (same contract the Templated*Attention wrappers use) and checks
all three gradients against an eager fp32 reference. num_heads == seq_len so a layout mix-up
cannot hide behind a shape check. Fails on both backends before the patch, passes after.

Verified

On 2x A40 (sm_86), torch 2.13.0+cu126, cuDNN 9.10.2, bf16:

• Standalone repro at B=2 S=16 H=16 D=64 reproduces the reporter's numbers. dQ/dK/dV max abs
  error goes from 7.90 / 7.97 / 3.98 to 0.006 / 0.008 / 0.010, against reference magnitudes of
  1.86 / 2.26 / 2.75.
• End to end: small Flux transformer, ContextParallelConfig(ulysses_degree=2),
  attention_backend("_native_flash"), parameter grads all-reduced and compared to a
  single-GPU reference. Before the patch the backward raises Number of heads in key/value must divide number of heads in query. After, the worst relative grad error is 0.0138, which is
  exactly what the plain native backend gives on the same harness.
• pytest -m "attention or context_parallel" over test_attention_dispatch.py,
  test_attention_processor.py and test_models_transformer_flux.py: 14 passed, 42 skipped.
  The 10 failures in the full Flux file (LoRA hotswap + AOT compile) are pre-existing and
  reproduce on a clean tree.
• ruff check, ruff format --check, utils/check_copies.py all clean.

Not verified

• No H100 here, so this is sm_86 only. bf16 only, no dropout, no mask, non-causal.
• flash_hub / _flash_3_hub / varlen / sage backends skipped, kernels is not installed. I
  read their backward ops and none of them have this pattern.
• Ring CP end to end. See below.

Two things left for follow-ups

1. _cudnn_attention_forward_op passes compute_log_sumexp=return_lse. The Ulysses wrapper
   forwards the caller's return_lse, which is False for ordinary model attention, so cuDNN
   never computes the lse and the backward op gets an empty logsumexp. On the end-to-end
   harness above, ulysses + _native_cudnn still gives a worst relative grad error of 328 after
   this patch, and drops to 0.0138 if I locally force compute_log_sumexp=True. So cuDNN CP
   training is still broken, for a different reason than cuDNN attention backward corrupts gradients: K/V transposed twice, Q once #14338. Torch's own SDPA dispatch
   derives that flag from query/key/value.requires_grad. Happy to send it separately.
2. The bool-mask issue @YangXu1990uiuc raised, a torch.bool attn_mask going into the ATen
   attn_bias slot without conversion to an additive -inf/0 bias. Real by inspection, not
   touched here, not tested.

Also worth noting: test_context_parallel_backward only asserts that gradients are finite,
never that they are correct, which is why this went unnoticed.

Thanks to @YangXu1990uiuc (NVIDIA cuDNN team) for the precise diagnosis and the repro numbers,
which matched what I measured on A40 to three decimal places.

[YangXu1990uiuc]

Reporter of #14338 here. You mentioned you only have Ampere, so I ran your patch on Hopper and Blackwell as well. Confirmed on all of them — and the larger shape turned up something worth adding to the PR.

Setup: torch 2.13.0+cu130, cuDNN 9.20, bf16, your exact test harness (the minimal autograd.Function driving each op pair against an eager fp32 reference). upstream = HEAD~1, patched = your two hunks.

Your test's shape, B=2 S=16 H=16 D=64

max_abs of each gradient against the fp32 reference:

GPU	op	upstream dQ/dK/dV	patched dQ/dK/dV
A100 (sm_80)	cudnn	7.8992 / 7.9726 / 3.9790	0.0059 / 0.0080 / 0.0102
A100 (sm_80)	native_flash	7.8992 / 7.9726 / 3.9790	0.0059 / 0.0080 / 0.0102
H100 (sm_90)	cudnn	7.8992 / 7.9726 / 3.9790	0.0059 / 0.0080 / 0.0102
H100 (sm_90)	native_flash	7.8992 / 7.9726 / 3.9790	0.0059 / 0.0080 / 0.0102
Blackwell (sm_100)	cudnn	7.8992 / 7.9726 / 3.9790	0.0059 / 0.0080 / 0.0102
Blackwell (sm_100)	native_flash	7.8992 / 7.9726 / 3.9790	0.0059 / 0.0080 / 0.0102

Reference magnitudes 1.86 / 2.26 / 2.75, so pre-patch the error exceeds the value being approximated. Identical across architectures because at S=16 the residual is dominated by bf16 input quantisation rather than by the kernel, so this shape does not discriminate between backends — which is fine, it is a layout test.

A DiT-realistic shape, B=2 S=1024 H=24 D=64

This is where it gets more interesting. With num_heads != seq_len, the double transpose does not silently corrupt — it hard-errors:

cudnn:        RuntimeError: cuDNN Frontend error: For group-query attention,
              number of heads for key and query must be a factor of ...
native_flash: RuntimeError: Number of heads in key/value must divide number of heads in query

on all three architectures, because K/V arrive claiming seq_len heads. After your patch:

GPU	patched dQ/dK/dV
A100 (sm_80)	0.0033 / 0.0032 / 0.0030
H100 (sm_90)	0.0031 / 0.0029 / 0.0016
Blackwell (sm_100)	0.0021 / 0.0022 / 0.0016

(Here the per-architecture spread is visible, as expected once the real kernels engage.)

Two things this suggests for the PR:

1. It might be worth saying in the description that on realistic shapes this is a crash, not silent gradient corruption. H == S is what makes it silent, and that is exactly why your test needs it — but it also explains why this survived: anyone running a real DiT through the CP path would have hit a RuntimeError, not bad outputs.
2. If you want a second test case that costs nothing, a shape with num_heads != seq_len asserts "does not raise" and covers the other half of the failure mode.

Fixing _native_flash_attention_backward_op in the same PR looks right to me — same forward-op contract, same two lines, and it reproduces identically above.

Repro script if useful: it just swaps attention_dispatch.py between the two variants and runs your harness per device. Happy to paste it or run any other shape/dtype you want on sm_80/89/90/100.

[guptaishaan]

Thanks for running this on sm_90 and sm_100, and for the larger shape. Both points taken.

Added the second case: the test is now parametrized over (B=2, S=16, H=16, D=64) and your (B=2, S=1024, H=24, D=64), against both op pairs. On A40 (sm_86, torch 2.9.0+cu126, cuDNN 9.10.2) the new case reproduces your errors exactly on the unpatched tree:

cudnn:        RuntimeError: cuDNN Frontend error: For group-query attention, number of heads
              for key and query must be a factor of number of heads for query
native_flash: RuntimeError: Number of heads in key/value must divide number of heads in query

and all 4 cases pass with the patch. Gradient max abs error on the new shape here is 0.0024 / 0.0035 / 0.0023 for both backends, reference magnitudes 0.74 / 1.13 / 0.51, in line with your A100 column. Your S=16 numbers also came out identical to mine to four decimals, which is a nice confirmation that the harness is the same one.

The test asserts on the gradients, not on the error text, since that wording comes from cuDNN and from ATen and will drift. The docstring now records which half of the failure mode each shape covers.

On your point 1: agreed, and stating it plainly here since it is the useful part for anyone finding this later. With num_heads != seq_len the double transpose is a hard RuntimeError, not silent corruption. H == S is the only regime where it passes the shape checks and returns wrong numbers, so anyone running a real DiT through the CP backward hit a crash. I will fold that into the PR description.

No need for the repro script. If you do feel like running something else, the compute_log_sumexp=return_lse problem in point 1 of the description is the one I would most like a second pair of eyes on, but that belongs in its own PR.