From f5e56e6720cb2872478cbbba4899d85a37fd44e2 Mon Sep 17 00:00:00 2001 From: guptaishaan Date: Thu, 30 Jul 2026 09:40:29 -0700 Subject: [PATCH 1/2] Fix double transpose of K/V in templated attention backward _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 before calling aten::_scaled_dot_product_cudnn_attention_backward. The ATen op therefore received Q as BHSD and K/V as BSHD, and returned gradients unrelated to the true ones. Only the context parallel path uses these ops, so the forward pass and plain inference were unaffected. _native_flash_attention_backward_op carried the identical two lines against the identical forward-op contract and is fixed the same way. The grad_out transpose stays: the forward op returns out in BSHD, so autograd hands back grad_out in BSHD. Adds tests/models/test_attention_dispatch.py, which drives each forward/backward op pair through a minimal autograd.Function 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. Fixes #14338 --- src/diffusers/models/attention_dispatch.py | 8 +- tests/models/test_attention_dispatch.py | 85 ++++++++++++++++++++++ 2 files changed, 89 insertions(+), 4 deletions(-) create mode 100644 tests/models/test_attention_dispatch.py diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 9414c151fd67..1a501602789e 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -969,9 +969,9 @@ def _cudnn_attention_backward_op( ): query, key, value, out, lse, cum_seq_q, cum_seq_k, philox_seed, philox_offset = ctx.saved_tensors + # Only grad_out needs to be transposed here: the saved query/key/value are the tensors the + # forward op already transposed to (B, H, S, D) before calling into cuDNN. grad_out = grad_out.transpose(1, 2).contiguous() - key = key.transpose(1, 2).contiguous() - value = value.transpose(1, 2).contiguous() # Cannot pass first 5 arguments as kwargs because: https://github.com/pytorch/pytorch/blob/d26ca5de058dbcf56ac52bb43e84dd98df2ace97/torch/_dynamo/variables/torch.py#L1341 grad_query, grad_key, grad_value = torch.ops.aten._scaled_dot_product_cudnn_attention_backward( @@ -1062,9 +1062,9 @@ def _native_flash_attention_backward_op( ): query, key, value, out, lse, cum_seq_q, cum_seq_k, philox_seed, philox_offset = ctx.saved_tensors + # Only grad_out needs to be transposed here: the saved query/key/value are the tensors the + # forward op already transposed to (B, H, S, D). grad_out = grad_out.transpose(1, 2).contiguous() - key = key.transpose(1, 2).contiguous() - value = value.transpose(1, 2).contiguous() grad_query, grad_key, grad_value = torch.ops.aten._scaled_dot_product_flash_attention_backward( grad_out, diff --git a/tests/models/test_attention_dispatch.py b/tests/models/test_attention_dispatch.py new file mode 100644 index 000000000000..7076fcf4f31e --- /dev/null +++ b/tests/models/test_attention_dispatch.py @@ -0,0 +1,85 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers.models.attention_dispatch import ( + _cudnn_attention_backward_op, + _cudnn_attention_forward_op, + _native_flash_attention_backward_op, + _native_flash_attention_forward_op, +) + +from ..testing_utils import assert_tensors_close, is_attention, require_torch_gpu + + +class _TemplatedAttentionOp(torch.autograd.Function): + """Minimal driver for a (forward_op, backward_op) pair, mirroring the context parallel wrappers.""" + + @staticmethod + def forward(ctx, query, key, value, forward_op, backward_op): + ctx.backward_op = backward_op + # The context parallel wrappers always request the lse. + out, _ = forward_op(ctx, query, key, value, return_lse=True) + return out + + @staticmethod + def backward(ctx, grad_out): + grad_query, grad_key, grad_value = ctx.backward_op(ctx, grad_out) + return grad_query, grad_key, grad_value, None, None + + +@is_attention +@require_torch_gpu +@pytest.mark.parametrize( + "forward_op,backward_op", + [ + (_cudnn_attention_forward_op, _cudnn_attention_backward_op), + (_native_flash_attention_forward_op, _native_flash_attention_backward_op), + ], + ids=["cudnn", "native_flash"], +) +def test_templated_attention_op_gradients(forward_op, backward_op): + """Gradients from the templated forward/backward op pairs must match an eager fp32 reference. + + `num_heads == seq_len` on purpose so that a (B, S, H, D) / (B, H, S, D) layout mix-up cannot be + hidden by a shape mismatch. See https://github.com/huggingface/diffusers/issues/14338. + """ + torch.manual_seed(0) + batch_size, seq_len, num_heads, head_dim = 2, 16, 16, 64 + shape = (batch_size, seq_len, num_heads, head_dim) + query, key, value, grad_out = ( + torch.randn(shape, device="cuda", dtype=torch.bfloat16, requires_grad=True) for _ in range(4) + ) + + ref_query, ref_key, ref_value = ( + x.detach().float().transpose(1, 2).requires_grad_(True) for x in (query, key, value) + ) + ref_out = torch.nn.functional.scaled_dot_product_attention(ref_query, ref_key, ref_value) + ref_out.backward(grad_out.detach().float().transpose(1, 2)) + + out = _TemplatedAttentionOp.apply(query, key, value, forward_op, backward_op) + out.backward(grad_out) + + assert_tensors_close(out.float(), ref_out.transpose(1, 2), atol=1e-2, rtol=1e-2, msg="forward output") + for name, actual, expected in ( + ("query", query.grad, ref_query.grad), + ("key", key.grad, ref_key.grad), + ("value", value.grad, ref_value.grad), + ): + assert_tensors_close( + actual.float(), expected.transpose(1, 2), atol=1e-1, rtol=1e-1, msg=f"gradient w.r.t. {name}" + ) From 0a1f670359a5cb35f353af629b539f94ad69c8e3 Mon Sep 17 00:00:00 2001 From: guptaishaan Date: Thu, 30 Jul 2026 16:20:12 -0700 Subject: [PATCH 2/2] Cover the num_heads != seq_len case in the attention op test The existing shape uses num_heads == seq_len, where the double transpose of K/V passes every shape check and silently corrupts the gradients. On any shape where num_heads != seq_len the same mix-up makes the backend reject the head count instead, so the failure is a RuntimeError rather than bad numbers: cudnn: 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: Number of heads in key/value must divide number of heads in query Parametrizes the test over a second, DiT-like shape (B=2, S=1024, H=24, D=64) so both halves of the failure mode are covered. --- tests/models/test_attention_dispatch.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/models/test_attention_dispatch.py b/tests/models/test_attention_dispatch.py index 7076fcf4f31e..7b5c6f45edae 100644 --- a/tests/models/test_attention_dispatch.py +++ b/tests/models/test_attention_dispatch.py @@ -52,14 +52,20 @@ def backward(ctx, grad_out): ], ids=["cudnn", "native_flash"], ) -def test_templated_attention_op_gradients(forward_op, backward_op): +@pytest.mark.parametrize( + "batch_size,seq_len,num_heads,head_dim", + [(2, 16, 16, 64), (2, 1024, 24, 64)], + ids=["heads_eq_seq_len", "dit_like"], +) +def test_templated_attention_op_gradients(batch_size, seq_len, num_heads, head_dim, forward_op, backward_op): """Gradients from the templated forward/backward op pairs must match an eager fp32 reference. - `num_heads == seq_len` on purpose so that a (B, S, H, D) / (B, H, S, D) layout mix-up cannot be - hidden by a shape mismatch. See https://github.com/huggingface/diffusers/issues/14338. + The first shape has `num_heads == seq_len` so that a (B, S, H, D) / (B, H, S, D) layout mix-up + cannot be hidden by a shape mismatch; it corrupts the gradients instead. The second has + `num_heads != seq_len`, where the same mix-up makes the backend reject the head count outright. + See https://github.com/huggingface/diffusers/issues/14338. """ torch.manual_seed(0) - batch_size, seq_len, num_heads, head_dim = 2, 16, 16, 64 shape = (batch_size, seq_len, num_heads, head_dim) query, key, value, grad_out = ( torch.randn(shape, device="cuda", dtype=torch.bfloat16, requires_grad=True) for _ in range(4)